mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-20 00:18:55 +00:00
Merge branch 'main' into fix/grok-claude-messages-prompt-cache
Resolve const-block conflict in openai_gateway_grok_cache.go: keep #4590 client tool cache constants and #4585 xai-package limit (grokFreeRolling24hTokenLimit removed), add claudeCodeSessionHeader.
This commit is contained in:
@@ -77,8 +77,8 @@ Please read the following carefully before using this project:
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>Thanks to PPToken.org for sponsoring this project! <a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> specializes in GPT model API relay services, supporting Codex, Claude Code, OpenAI-compatible clients, and Gemini CLI integration. Top-ups are 1:1 (¥1 = $1 credit); GPT models start at 0.16x rate multiplier, with overall cost at roughly 2.2% of official pricing and first-token latency around 1 second — ideal for developers seeking low-cost, high-speed access to GPT model capabilities. Technical support: 24/7 real human responses (no bots), @tech in the group chat and get a reply within 10 minutes. Sponsor benefit: the first 200 users who register via the <a href="https://api.pptoken.org/register?promo=SUB2API">exclusive registration link</a> and enter promo code `SUB2API` can claim free Codex / Claude Code trial credits — no minimum spend, no card required.
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>Thanks to PPToken.cc for sponsoring this project! <a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> specializes in GPT model API relay services, supporting Codex, Claude Code, OpenAI-compatible clients, and Gemini CLI integration. Top-ups are 1:1 (¥1 = $1 credit); GPT models start at 0.16x rate multiplier, with overall cost at roughly 2.2% of official pricing and first-token latency around 1 second — ideal for developers seeking low-cost, high-speed access to GPT model capabilities. Technical support: 24/7 real human responses (no bots), @tech in the group chat and get a reply within 10 minutes. Sponsor benefit: the first 200 users who register via the <a href="https://api.pptoken.cc/register?promo=SUB2API">exclusive registration link</a> and enter promo code `SUB2API` can claim free Codex / Claude Code trial credits — no minimum spend, no card required.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -571,9 +571,18 @@ Additional security-related options are available in `config.yaml`:
|
||||
- `security.response_headers.enabled` to enable configurable response header filtering (disabled uses default allowlist)
|
||||
- `security.csp` to control Content-Security-Policy headers
|
||||
- `billing.circuit_breaker` to fail closed on billing errors
|
||||
- `server.trusted_proxies` to enable X-Forwarded-For parsing
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` enables legacy raw forwarded-header takeover (enabled by default for upgrade compatibility); disable it to enforce `server.trusted_proxies`, which should contain only the exact proxy CIDRs that connect directly to Sub2API
|
||||
- `security.forwarded_client_ip_headers` configures up to 16 third-party CDN client-IP header names; they are checked in order before the built-in headers only while legacy takeover is enabled
|
||||
- `turnstile.required` to require Turnstile in release mode
|
||||
|
||||
Custom client-IP headers can be set in YAML or as a comma-separated environment variable:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
Header names are validated, canonicalized, and de-duplicated. The admin security settings can update the list without a restart; new installations persist YAML/environment defaults and existing installations backfill a missing database value. When legacy takeover is disabled, all custom and built-in raw forwarding headers are ignored and Gin uses only `server.trusted_proxies`. While takeover is enabled, firewall the origin to CDN/proxy addresses and make the edge overwrite every trusted client-IP header. See [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md) for the complete migration and trust-boundary rules.
|
||||
|
||||
**⚠️ Security Warning: HTTP URL Configuration**
|
||||
|
||||
When `security.url_allowlist.enabled=false`, the system performs minimal URL validation and **allows HTTP URLs by default** (dev-friendly mode; Docker Compose deployments use the same default). For production, explicitly tighten this to HTTPS-only:
|
||||
|
||||
+12
-3
@@ -80,8 +80,8 @@
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>感谢 PPToken.org 赞助本项目! <a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> 主打 GPT 系列模型 API 中转服务,支持 Codex、Claude Code、OpenAI 兼容客户端及 Gemini CLI 等工具接入。充值 1:1,1 元=1 美元额度;GPT 模型最低 0.16 倍倍率,综合成本约为官方价格的 0.22 折,最快首字 Token 约 1 秒,适合开发者低成本、高响应速度接入 GPT 模型能力。技术支持: 7×24 小时真人响应(不是机器人),群内@技术,10 分钟内有回复 。赞助商福利:前 200 名用户通过 <a href="https://api.pptoken.org/register?promo=SUB2API">[专属注册链接]</a> 注册,输入优惠码 `SUB2API`,可领取 Codex / Claude Code 免费试用额度,无门槛、不绑卡。
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>感谢 PPToken.cc 赞助本项目! <a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> 主打 GPT 系列模型 API 中转服务,支持 Codex、Claude Code、OpenAI 兼容客户端及 Gemini CLI 等工具接入。充值 1:1,1 元=1 美元额度;GPT 模型最低 0.16 倍倍率,综合成本约为官方价格的 0.22 折,最快首字 Token 约 1 秒,适合开发者低成本、高响应速度接入 GPT 模型能力。技术支持: 7×24 小时真人响应(不是机器人),群内@技术,10 分钟内有回复 。赞助商福利:前 200 名用户通过 <a href="https://api.pptoken.cc/register?promo=SUB2API">[专属注册链接]</a> 注册,输入优惠码 `SUB2API`,可领取 Codex / Claude Code 免费试用额度,无门槛、不绑卡。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -607,9 +607,18 @@ gateway:
|
||||
- `security.response_headers.enabled` 可启用可配置响应头过滤(关闭时使用默认白名单)
|
||||
- `security.csp` 配置 Content-Security-Policy
|
||||
- `billing.circuit_breaker` 计费异常时 fail-closed
|
||||
- `server.trusted_proxies` 启用可信代理解析 X-Forwarded-For
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` 控制旧版原始转发头接管(为升级兼容默认开启);关闭后严格使用 `server.trusted_proxies`,其中只应填写直接连接 Sub2API 的精确代理 CIDR
|
||||
- `security.forwarded_client_ip_headers` 最多配置 16 个第三方 CDN 客户端 IP 请求头;仅在旧版接管开启时按顺序优先于内置请求头解析
|
||||
- `turnstile.required` 在 release 模式强制启用 Turnstile
|
||||
|
||||
自定义客户端 IP 请求头可通过 YAML 配置,也可使用逗号分隔的环境变量:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
请求头名称会经过合法性校验、规范化和大小写无关去重。管理员可在安全设置中动态更新列表,无需重启;新安装会持久化 YAML/环境变量默认值,旧安装缺少数据库字段时会自动回填。关闭旧版接管后,自定义头和内置原始转发头均被忽略,只使用 `server.trusted_proxies`。开启接管时必须限制源站仅允许 CDN/代理访问,并确保边缘代理覆盖所有受信客户端 IP 请求头。完整迁移规则和信任边界见 [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md)。
|
||||
|
||||
**网关防御纵深建议(重点)**
|
||||
|
||||
- `gateway.upstream_response_read_max_bytes`:限制非流式上游响应读取大小(默认 `8MB`),用于防止异常响应导致内存放大。
|
||||
|
||||
+12
-3
@@ -79,8 +79,8 @@
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>PPToken.org のご支援に感謝します!<a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> は GPT シリーズモデルの API 中継サービスを専門としており、Codex、Claude Code、OpenAI 互換クライアント、Gemini CLI などのツール接続をサポートしています。チャージは 1:1(1元=1ドル分のクレジット)、GPT モデルは最低 0.16 倍のレート倍率で、総合コストは公式価格の約 2.2% 、最速ファーストトークンは約1秒 — 開発者が低コスト・高速レスポンスで GPT モデル機能にアクセスするのに最適です。テクニカルサポート:24時間365日リアルな人間が対応(ボットではありません)、グループ内で @技術 すれば 10 分以内に返信。スポンサー特典:先着 200 名のユーザーが<a href="https://api.pptoken.org/register?promo=SUB2API">専用登録リンク</a>から登録し、プロモコード `SUB2API` を入力すると、Codex / Claude Code の無料トライアルクレジットを獲得できます — 最低利用額なし、カード登録不要。
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>PPToken.cc のご支援に感謝します!<a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> は GPT シリーズモデルの API 中継サービスを専門としており、Codex、Claude Code、OpenAI 互換クライアント、Gemini CLI などのツール接続をサポートしています。チャージは 1:1(1元=1ドル分のクレジット)、GPT モデルは最低 0.16 倍のレート倍率で、総合コストは公式価格の約 2.2% 、最速ファーストトークンは約1秒 — 開発者が低コスト・高速レスポンスで GPT モデル機能にアクセスするのに最適です。テクニカルサポート:24時間365日リアルな人間が対応(ボットではありません)、グループ内で @技術 すれば 10 分以内に返信。スポンサー特典:先着 200 名のユーザーが<a href="https://api.pptoken.cc/register?promo=SUB2API">専用登録リンク</a>から登録し、プロモコード `SUB2API` を入力すると、Codex / Claude Code の無料トライアルクレジットを獲得できます — 最低利用額なし、カード登録不要。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -569,9 +569,18 @@ default:
|
||||
- `security.response_headers.enabled` - 設定可能なレスポンスヘッダーフィルタリングを有効化(無効時はデフォルトの許可リストを使用)
|
||||
- `security.csp` - Content-Security-Policy ヘッダーの制御
|
||||
- `billing.circuit_breaker` - 課金エラー時にフェイルクローズ
|
||||
- `server.trusted_proxies` - X-Forwarded-For パースの有効化
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` - 従来の生転送ヘッダーによる上書きを制御(アップグレード互換性のため既定で有効)。無効にすると `server.trusted_proxies` を厳格に使用し、Sub2API に直接接続するプロキシの正確な CIDR のみを指定
|
||||
- `security.forwarded_client_ip_headers` - サードパーティ CDN のクライアント IP ヘッダーを最大 16 個指定。従来モードが有効な場合のみ、設定順で組み込みヘッダーより先に評価
|
||||
- `turnstile.required` - リリースモードでの Turnstile 必須化
|
||||
|
||||
カスタムクライアント IP ヘッダーは YAML またはカンマ区切りの環境変数で設定できます:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
ヘッダー名は検証、正規化、大小文字を区別しない重複排除が行われます。管理画面のセキュリティ設定から再起動せずに更新でき、新規インストールでは YAML/環境変数の既定値を保存し、既存環境ではデータベース値がない場合に補完します。従来モードを無効にするとカスタムおよび組み込みの生転送ヘッダーはすべて無視され、`server.trusted_proxies` のみを使用します。有効にする場合はオリジンへの接続元を CDN/プロキシに制限し、エッジで信頼する全クライアント IP ヘッダーを上書きしてください。移行規則と信頼境界の詳細は [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md) を参照してください。
|
||||
|
||||
**⚠️ セキュリティ警告: HTTP URL 設定**
|
||||
|
||||
`security.url_allowlist.enabled=false` の場合、システムは最小限の URL バリデーションのみを行い、**デフォルトで HTTP URL を許可**します(開発フレンドリーモード。Docker Compose デプロイのデフォルトも同じです)。本番環境では、以下のように明示的に HTTPS のみに制限することを推奨します:
|
||||
|
||||
@@ -200,7 +200,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
|
||||
dbDumper := repository.NewPgDumper(configConfig)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper)
|
||||
backupHandler := admin.NewBackupHandler(backupService, userService)
|
||||
imageStorageFactory := repository.ProvideImageStorageFactory()
|
||||
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
|
||||
backupHandler := admin.NewBackupHandler(backupService, userService, imageStorageSettingService)
|
||||
oAuthHandler := admin.NewOAuthHandler(oAuthService)
|
||||
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
|
||||
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
|
||||
@@ -279,11 +281,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
|
||||
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
|
||||
imageTaskStore := repository.NewImageTaskStore(redisClient)
|
||||
imageStorage, err := repository.ProvideImageStorage(configConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
imageTaskService := service.ProvideImageTaskService(imageTaskStore, imageStorage, configConfig)
|
||||
imageTaskService := service.ProvideImageTaskService(imageTaskStore, imageStorageSettingService)
|
||||
asyncImageHandler := handler.NewAsyncImageHandler(imageTaskService, openAIGatewayHandler)
|
||||
batchImageRepository := repository.NewBatchImageRepository(db)
|
||||
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
@@ -14,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
"golang.org/x/net/http/httpguts"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -255,6 +257,22 @@ func (c *ImageStorageConfig) Active() bool {
|
||||
return c.Enabled && c.IsConfigured()
|
||||
}
|
||||
|
||||
// MissingCredentialKeys 返回 IsConfigured 所缺的配置键名。
|
||||
// 用于启动日志:只说"凭证不完整"会让运维以为自己漏填了,而实际可能是值填了却没被读到。
|
||||
func (c *ImageStorageConfig) MissingCredentialKeys() []string {
|
||||
var missing []string
|
||||
if c.Bucket == "" {
|
||||
missing = append(missing, "image_storage.bucket")
|
||||
}
|
||||
if c.AccessKeyID == "" {
|
||||
missing = append(missing, "image_storage.access_key_id")
|
||||
}
|
||||
if c.SecretAccessKey == "" {
|
||||
missing = append(missing, "image_storage.secret_access_key")
|
||||
}
|
||||
return missing
|
||||
}
|
||||
|
||||
type LinuxDoConnectConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
ClientID string `mapstructure:"client_id"`
|
||||
@@ -639,17 +657,18 @@ type PricingConfig struct {
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
Mode string `mapstructure:"mode"` // debug/release
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
MaxHeaderBytes int `mapstructure:"max_header_bytes"` // 请求头最大字节数(HTTP/2 映射为 header-list 上限)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP)
|
||||
MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制
|
||||
H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
Mode string `mapstructure:"mode"` // debug/release
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
MaxHeaderBytes int `mapstructure:"max_header_bytes"` // 请求头最大字节数(HTTP/2 映射为 header-list 上限)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP)
|
||||
TrustedProxiesConfigured bool `mapstructure:"-" json:"-" yaml:"-"` // 是否显式配置了可信代理列表
|
||||
MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制
|
||||
H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置
|
||||
}
|
||||
|
||||
// H2CConfig HTTP/2 Cleartext 配置
|
||||
@@ -667,36 +686,103 @@ type CORSConfig struct {
|
||||
AllowCredentials bool `mapstructure:"allow_credentials"`
|
||||
}
|
||||
|
||||
const MaxForwardedClientIPHeaders = 16
|
||||
|
||||
type ForwardedClientIPSettings struct {
|
||||
TrustForwardedIP bool
|
||||
Headers []string
|
||||
}
|
||||
|
||||
type SecurityConfig struct {
|
||||
URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"`
|
||||
ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"`
|
||||
CSP CSPConfig `mapstructure:"csp"`
|
||||
ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"`
|
||||
ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"`
|
||||
TrustForwardedIPForAPIKeyACL bool `mapstructure:"trust_forwarded_ip_for_api_key_acl"`
|
||||
trustForwardedIPForAPIKeyACLLive *atomic.Bool `mapstructure:"-"`
|
||||
URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"`
|
||||
ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"`
|
||||
CSP CSPConfig `mapstructure:"csp"`
|
||||
ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"`
|
||||
ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"`
|
||||
// TrustForwardedIPForAPIKeyACL enables legacy raw forwarded-header takeover.
|
||||
// When disabled, server.trusted_proxies is authoritative for all client-IP consumers.
|
||||
TrustForwardedIPForAPIKeyACL bool `mapstructure:"trust_forwarded_ip_for_api_key_acl"`
|
||||
ForwardedClientIPHeaders []string `mapstructure:"forwarded_client_ip_headers" json:"forwarded_client_ip_headers" yaml:"forwarded_client_ip_headers"`
|
||||
forwardedClientIPSettingsLive *atomic.Pointer[ForwardedClientIPSettings] `mapstructure:"-" json:"-" yaml:"-"`
|
||||
}
|
||||
|
||||
func NormalizeForwardedClientIPHeaders(headers []string) ([]string, error) {
|
||||
normalized := make([]string, 0, len(headers))
|
||||
seen := make(map[string]struct{}, len(headers))
|
||||
for _, header := range headers {
|
||||
header = strings.TrimSpace(header)
|
||||
if !httpguts.ValidHeaderFieldName(header) {
|
||||
return nil, fmt.Errorf("invalid HTTP header field name %q", header)
|
||||
}
|
||||
canonical := textproto.CanonicalMIMEHeaderKey(header)
|
||||
key := strings.ToLower(canonical)
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
if len(normalized) == MaxForwardedClientIPHeaders {
|
||||
return nil, fmt.Errorf("forwarded client IP headers must contain at most %d unique names", MaxForwardedClientIPHeaders)
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
normalized = append(normalized, canonical)
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func cloneForwardedClientIPHeaders(headers []string) []string {
|
||||
if len(headers) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
return append([]string(nil), headers...)
|
||||
}
|
||||
|
||||
func (c *Config) ForwardedClientIPSettings() ForwardedClientIPSettings {
|
||||
if c == nil {
|
||||
return ForwardedClientIPSettings{Headers: []string{}}
|
||||
}
|
||||
live := c.Security.forwardedClientIPSettingsLive
|
||||
if live != nil {
|
||||
if snapshot := live.Load(); snapshot != nil {
|
||||
return ForwardedClientIPSettings{
|
||||
TrustForwardedIP: snapshot.TrustForwardedIP,
|
||||
Headers: cloneForwardedClientIPHeaders(snapshot.Headers),
|
||||
}
|
||||
}
|
||||
}
|
||||
return ForwardedClientIPSettings{
|
||||
TrustForwardedIP: c.Security.TrustForwardedIPForAPIKeyACL,
|
||||
Headers: cloneForwardedClientIPHeaders(c.Security.ForwardedClientIPHeaders),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) TrustForwardedIPForAPIKeyACL() bool {
|
||||
return c.ForwardedClientIPSettings().TrustForwardedIP
|
||||
}
|
||||
|
||||
// ForwardedClientIPTrustEnabled reports whether the legacy forwarded-header
|
||||
// compatibility mode currently overrides server.trusted_proxies.
|
||||
func (c *Config) ForwardedClientIPTrustEnabled() bool {
|
||||
return c != nil && c.TrustForwardedIPForAPIKeyACL()
|
||||
}
|
||||
|
||||
func (c *Config) SetForwardedClientIPSettings(enabled bool, headers []string) {
|
||||
if c == nil {
|
||||
return false
|
||||
return
|
||||
}
|
||||
live := c.Security.trustForwardedIPForAPIKeyACLLive
|
||||
if live == nil {
|
||||
return c.Security.TrustForwardedIPForAPIKeyACL
|
||||
headers = cloneForwardedClientIPHeaders(headers)
|
||||
if c.Security.forwardedClientIPSettingsLive == nil {
|
||||
c.Security.forwardedClientIPSettingsLive = &atomic.Pointer[ForwardedClientIPSettings]{}
|
||||
}
|
||||
return live.Load()
|
||||
c.Security.forwardedClientIPSettingsLive.Store(&ForwardedClientIPSettings{
|
||||
TrustForwardedIP: enabled,
|
||||
Headers: headers,
|
||||
})
|
||||
}
|
||||
|
||||
func (c *Config) SetTrustForwardedIPForAPIKeyACL(enabled bool) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.Security.TrustForwardedIPForAPIKeyACL = enabled
|
||||
if c.Security.trustForwardedIPForAPIKeyACLLive == nil {
|
||||
c.Security.trustForwardedIPForAPIKeyACLLive = &atomic.Bool{}
|
||||
}
|
||||
c.Security.trustForwardedIPForAPIKeyACLLive.Store(enabled)
|
||||
c.SetForwardedClientIPSettings(enabled, c.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
type URLAllowlistConfig struct {
|
||||
@@ -1564,17 +1650,31 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
}
|
||||
// 配置文件不存在时使用默认值
|
||||
}
|
||||
trustedProxiesEnv, trustedProxiesEnvConfigured := os.LookupEnv("SERVER_TRUSTED_PROXIES")
|
||||
forwardedClientIPHeadersEnv, forwardedClientIPHeadersEnvConfigured := os.LookupEnv("SECURITY_FORWARDED_CLIENT_IP_HEADERS")
|
||||
trustedProxiesConfigured := viper.InConfig("server.trusted_proxies") ||
|
||||
viper.IsSet("server.trusted_proxies") || trustedProxiesEnvConfigured
|
||||
|
||||
var cfg Config
|
||||
if err := viper.Unmarshal(&cfg); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal config error: %w", err)
|
||||
}
|
||||
if trustedProxiesEnvConfigured {
|
||||
cfg.Server.TrustedProxies = normalizeStringSlice(strings.Split(trustedProxiesEnv, ","))
|
||||
}
|
||||
if forwardedClientIPHeadersEnvConfigured {
|
||||
cfg.Security.ForwardedClientIPHeaders = normalizeStringSlice(strings.Split(forwardedClientIPHeadersEnv, ","))
|
||||
}
|
||||
cfg.Server.TrustedProxiesConfigured = trustedProxiesConfigured
|
||||
if cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs == 0 {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
|
||||
}
|
||||
if cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate == 0 {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
|
||||
}
|
||||
// Kept as a backstop: setEnvReachableDefaults now registers this key with its
|
||||
// effective default (true), so IsSet always reports true and this branch no
|
||||
// longer fires. It still guards the default if that registration is dropped.
|
||||
if !cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled && !viper.IsSet("gateway.openai_scheduler.sticky_escape_enabled") {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
|
||||
}
|
||||
@@ -1624,7 +1724,12 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
cfg.Security.ResponseHeaders.AdditionalAllowed = normalizeStringSlice(cfg.Security.ResponseHeaders.AdditionalAllowed)
|
||||
cfg.Security.ResponseHeaders.ForceRemove = normalizeStringSlice(cfg.Security.ResponseHeaders.ForceRemove)
|
||||
cfg.Security.CSP.Policy = strings.TrimSpace(cfg.Security.CSP.Policy)
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(cfg.Security.TrustForwardedIPForAPIKeyACL)
|
||||
forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(cfg.Security.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("security.forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
cfg.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders
|
||||
cfg.SetForwardedClientIPSettings(cfg.Security.TrustForwardedIPForAPIKeyACL, forwardedClientIPHeaders)
|
||||
cfg.Log.Level = strings.ToLower(strings.TrimSpace(cfg.Log.Level))
|
||||
cfg.Log.Format = strings.ToLower(strings.TrimSpace(cfg.Log.Format))
|
||||
cfg.Log.ServiceName = strings.TrimSpace(cfg.Log.ServiceName)
|
||||
@@ -1714,7 +1819,6 @@ func setDefaults() {
|
||||
viper.SetDefault("server.read_header_timeout", 10) // 10秒读取请求头
|
||||
viper.SetDefault("server.max_header_bytes", 64*1024)
|
||||
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
|
||||
viper.SetDefault("server.trusted_proxies", []string{})
|
||||
viper.SetDefault("server.max_request_body_size", int64(256*1024*1024))
|
||||
// H2C 默认配置
|
||||
viper.SetDefault("server.h2c.enabled", false)
|
||||
@@ -1771,7 +1875,7 @@ func setDefaults() {
|
||||
viper.SetDefault("security.csp.enabled", true)
|
||||
viper.SetDefault("security.csp.policy", DefaultCSPPolicy)
|
||||
viper.SetDefault("security.proxy_probe.insecure_skip_verify", false)
|
||||
viper.SetDefault("security.trust_forwarded_ip_for_api_key_acl", false)
|
||||
viper.SetDefault("security.trust_forwarded_ip_for_api_key_acl", true)
|
||||
|
||||
// Security - disable direct fallback on proxy error
|
||||
viper.SetDefault("security.proxy_fallback.allow_direct_on_error", false)
|
||||
@@ -1940,6 +2044,15 @@ func setDefaults() {
|
||||
viper.SetDefault("image_storage.force_path_style", false)
|
||||
viper.SetDefault("image_storage.presign_expiry_hours", 24)
|
||||
viper.SetDefault("image_storage.max_download_bytes", 33554432)
|
||||
// Registered with empty defaults so AutomaticEnv can reach them: viper only
|
||||
// decodes keys present in AllKeys(), so a credential that is supplied purely
|
||||
// via IMAGE_STORAGE_* and never appears in config.yaml would be dropped and
|
||||
// silently disable the whole async image feature.
|
||||
viper.SetDefault("image_storage.endpoint", "")
|
||||
viper.SetDefault("image_storage.bucket", "")
|
||||
viper.SetDefault("image_storage.access_key_id", "")
|
||||
viper.SetDefault("image_storage.secret_access_key", "")
|
||||
viper.SetDefault("image_storage.public_base_url", "")
|
||||
|
||||
// Ops (vNext)
|
||||
viper.SetDefault("ops.enabled", true)
|
||||
@@ -2226,9 +2339,97 @@ func setDefaults() {
|
||||
viper.SetDefault("subscription_maintenance.worker_count", 2)
|
||||
viper.SetDefault("subscription_maintenance.queue_size", 1024)
|
||||
|
||||
setEnvReachableDefaults()
|
||||
}
|
||||
|
||||
// setEnvReachableDefaults registers zero-valued defaults for keys that are
|
||||
// documented in deploy/config.example.yaml but had no default of their own.
|
||||
//
|
||||
// viper.Unmarshal only decodes the keys returned by AllKeys(), which unions
|
||||
// SetDefault keys, config-file keys and explicitly bound BindEnv keys.
|
||||
// AutomaticEnv can override a key already in that union, but it never adds one,
|
||||
// and the viper_bind_struct escape hatch is compiled out (we build with
|
||||
// -tags embed). So a key that lives only in the example file was unreachable by
|
||||
// environment variable: the value was read from the process environment and
|
||||
// then silently dropped. Deployments driven purely by env — which is what
|
||||
// deploy/docker-compose.yml does — got the zero value with no warning.
|
||||
//
|
||||
// The values below are deliberately zero rather than the documented example
|
||||
// values: an absent key already unmarshalled to the zero value, so registering
|
||||
// zero keeps behavior identical while making the key addressable from the
|
||||
// environment. Any subsystem that wants a richer default still applies it after
|
||||
// unmarshal, exactly as before.
|
||||
func setEnvReachableDefaults() {
|
||||
viper.SetDefault("gateway.forced_codex_instructions_template_file", "")
|
||||
viper.SetDefault("gateway.session_idle_timeout_minutes", 0)
|
||||
viper.SetDefault("gateway.user_message_queue.mode", "")
|
||||
viper.SetDefault("update.proxy_url", "")
|
||||
|
||||
// sticky_escape_enabled is the one exception to the zero-value rule: its
|
||||
// effective default is true, applied post-unmarshal via a viper.IsSet guard.
|
||||
// Registering false would make IsSet always report true and permanently
|
||||
// disable sticky escape, so register the effective default instead. An
|
||||
// explicit false in config or env still wins.
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_enabled", true)
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_error_rate", 0.0)
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_ttft_ms", 0)
|
||||
|
||||
// server.trusted_proxies and security.forwarded_client_ip_headers are the
|
||||
// other exception: load() distinguishes explicit configuration from absence
|
||||
// (issue #4600), and viper.IsSet also reports registered defaults, so a
|
||||
// SetDefault would make trusted proxies look permanently configured. Both
|
||||
// environment variables are parsed by hand in load() via os.LookupEnv and
|
||||
// were never silently dropped; binding them here records that reachability
|
||||
// where AllKeys() — and the env-reachability guard — can see it, without
|
||||
// affecting IsSet while the variables are absent. BindEnv only errors when
|
||||
// called without arguments.
|
||||
_ = viper.BindEnv("server.trusted_proxies", "SERVER_TRUSTED_PROXIES")
|
||||
_ = viper.BindEnv("security.forwarded_client_ip_headers", "SECURITY_FORWARDED_CLIENT_IP_HEADERS")
|
||||
|
||||
// Third-party login providers. These carry client secrets and are exactly
|
||||
// the settings an operator expects to inject via the environment, but every
|
||||
// key here was previously unreachable that way.
|
||||
for _, provider := range []string{"github_oauth", "google_oauth"} {
|
||||
viper.SetDefault(provider+".enabled", false)
|
||||
viper.SetDefault(provider+".client_id", "")
|
||||
viper.SetDefault(provider+".client_secret", "")
|
||||
viper.SetDefault(provider+".authorize_url", "")
|
||||
viper.SetDefault(provider+".token_url", "")
|
||||
viper.SetDefault(provider+".userinfo_url", "")
|
||||
viper.SetDefault(provider+".emails_url", "")
|
||||
viper.SetDefault(provider+".scopes", "")
|
||||
viper.SetDefault(provider+".redirect_url", "")
|
||||
viper.SetDefault(provider+".frontend_redirect_url", "")
|
||||
}
|
||||
|
||||
viper.SetDefault("dingtalk_connect.client_id", "")
|
||||
viper.SetDefault("dingtalk_connect.client_secret", "")
|
||||
viper.SetDefault("dingtalk_connect.internal_corp_id", "")
|
||||
viper.SetDefault("dingtalk_connect.redirect_url", "")
|
||||
viper.SetDefault("dingtalk_connect.bypass_registration", false)
|
||||
viper.SetDefault("dingtalk_connect.username_attribute_key", "")
|
||||
viper.SetDefault("dingtalk_connect.enable_attribute_matching", false)
|
||||
viper.SetDefault("dingtalk_connect.enable_attribute_sync", false)
|
||||
viper.SetDefault("dingtalk_connect.attribute_sync_fields", []string{})
|
||||
viper.SetDefault("dingtalk_connect.attribute_sync_overwrite_policy", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name_attr_name", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_dept", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_dept_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_dept_attr_name", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email_attr_name", "")
|
||||
}
|
||||
|
||||
func (c *Config) Validate() error {
|
||||
forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(c.Security.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return fmt.Errorf("security.forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
c.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders
|
||||
c.SetForwardedClientIPSettings(c.Security.TrustForwardedIPForAPIKeyACL, forwardedClientIPHeaders)
|
||||
if c.Server.ReadHeaderTimeout < 1 || c.Server.ReadHeaderTimeout > 60 {
|
||||
return fmt.Errorf("server.read_header_timeout must be between 1 and 60 seconds")
|
||||
}
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -41,12 +43,207 @@ func TestLoadHTTPIngressSafetyDefaults(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 10, cfg.Server.ReadHeaderTimeout)
|
||||
require.Equal(t, 64*1024, cfg.Server.MaxHeaderBytes)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.False(t, cfg.Server.TrustedProxiesConfigured)
|
||||
require.True(t, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
require.Equal(t, int64(32*1024*1024), cfg.Gateway.TextMaxBodySize)
|
||||
require.True(t, cfg.APIKeyAuth.InvalidAbuse.Enabled)
|
||||
require.Equal(t, 120, cfg.APIKeyAuth.InvalidAbuse.Threshold)
|
||||
require.Equal(t, 16384, cfg.APIKeyAuth.InvalidAbuse.Capacity)
|
||||
}
|
||||
|
||||
func TestNormalizeForwardedClientIPHeaders(t *testing.T) {
|
||||
headers, err := NormalizeForwardedClientIPHeaders([]string{
|
||||
" x-cdn-client-ip ",
|
||||
"X-CDN-CLIENT-IP",
|
||||
"true-client-ip",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"X-Cdn-Client-Ip", "True-Client-Ip"}, headers)
|
||||
|
||||
_, err = NormalizeForwardedClientIPHeaders([]string{"X Invalid"})
|
||||
require.ErrorContains(t, err, "invalid HTTP header field name")
|
||||
}
|
||||
|
||||
func TestNormalizeForwardedClientIPHeadersLimit(t *testing.T) {
|
||||
headers := make([]string, 0, MaxForwardedClientIPHeaders+1)
|
||||
for i := 0; i <= MaxForwardedClientIPHeaders; i++ {
|
||||
headers = append(headers, fmt.Sprintf("X-CDN-IP-%d", i))
|
||||
}
|
||||
|
||||
_, err := NormalizeForwardedClientIPHeaders(headers)
|
||||
require.ErrorContains(t, err, "at most 16 unique names")
|
||||
}
|
||||
|
||||
func TestLoadForwardedClientIPHeadersNormalizesAndSnapshots(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{" x-cdn-ip ", "X-CDN-IP", "true-client-ip"})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
snapshot := cfg.ForwardedClientIPSettings()
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, snapshot.Headers)
|
||||
|
||||
snapshot.Headers[0] = "X-Mutated"
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestForwardedClientIPSettingsConcurrentPublication(t *testing.T) {
|
||||
cfg := &Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Public-A"})
|
||||
|
||||
const iterations = 2000
|
||||
start := make(chan struct{})
|
||||
errCh := make(chan error, 8)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, settings := range []ForwardedClientIPSettings{
|
||||
{TrustForwardedIP: true, Headers: []string{"X-Public-A"}},
|
||||
{TrustForwardedIP: false, Headers: []string{"X-Public-B"}},
|
||||
} {
|
||||
settings := settings
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for i := 0; i < iterations; i++ {
|
||||
cfg.SetForwardedClientIPSettings(settings.TrustForwardedIP, settings.Headers)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
for i := 0; i < cap(errCh); i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for j := 0; j < iterations; j++ {
|
||||
snapshot := cfg.ForwardedClientIPSettings()
|
||||
validA := snapshot.TrustForwardedIP && len(snapshot.Headers) == 1 && snapshot.Headers[0] == "X-Public-A"
|
||||
validB := !snapshot.TrustForwardedIP && len(snapshot.Headers) == 1 && snapshot.Headers[0] == "X-Public-B"
|
||||
if !validA && !validB {
|
||||
errCh <- fmt.Errorf("observed inconsistent forwarded IP settings: %+v", snapshot)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadForwardedClientIPHeadersFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", " x-cdn-ip , X-CDN-IP, true-client-ip ")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyForwardedClientIPHeadersFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{"X-Yaml-IP"})
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", "")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestLoadRejectsInvalidForwardedClientIPHeaderFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", "X-Valid-IP, X Invalid")
|
||||
|
||||
_, err := Load()
|
||||
require.ErrorContains(t, err, "security.forwarded_client_ip_headers")
|
||||
}
|
||||
|
||||
func TestLoadRejectsInvalidForwardedClientIPHeader(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{"X Invalid"})
|
||||
|
||||
_, err := Load()
|
||||
require.ErrorContains(t, err, "security.forwarded_client_ip_headers")
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyTrustedProxiesEnablesConfiguredMode(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("server.trusted_proxies", []string{})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadExplicitTrustedProxiesEnablesConfiguredMode(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("server.trusted_proxies", []string{"127.0.0.1/32"})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"127.0.0.1/32"}, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadTrustedProxiesFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SERVER_TRUSTED_PROXIES", "127.0.0.1/32, ::1/128")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"127.0.0.1/32", "::1/128"}, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyTrustedProxiesFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SERVER_TRUSTED_PROXIES", "")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadTrustedProxiesPresenceFromYAML(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
yaml string
|
||||
want []string
|
||||
configured bool
|
||||
}{
|
||||
{name: "absent", yaml: "server:\n mode: debug\n", configured: false},
|
||||
{name: "explicit empty", yaml: "server:\n trusted_proxies: []\n", want: []string{}, configured: true},
|
||||
{
|
||||
name: "populated",
|
||||
yaml: "server:\n trusted_proxies:\n - 127.0.0.1/32\n - ::1/128\n",
|
||||
want: []string{"127.0.0.1/32", "::1/128"},
|
||||
configured: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
configDir := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(configDir, "config.yaml"), []byte(test.yaml), 0o600))
|
||||
t.Setenv("DATA_DIR", configDir)
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, test.want, cfg.Server.TrustedProxies)
|
||||
require.Equal(t, test.configured, cfg.Server.TrustedProxiesConfigured)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
|
||||
viper.Reset()
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
//go:build unit
|
||||
|
||||
package config
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"sort"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
// collectMapstructureKeys walks a config struct and returns every dotted key
|
||||
// viper would need in order to populate it.
|
||||
func collectMapstructureKeys(t reflect.Type, prefix string, out map[string]string) {
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
field := t.Field(i)
|
||||
if field.PkgPath != "" {
|
||||
continue // unexported
|
||||
}
|
||||
tag := field.Tag.Get("mapstructure")
|
||||
name, _, _ := strings.Cut(tag, ",")
|
||||
if name == "-" {
|
||||
continue
|
||||
}
|
||||
if name == "" {
|
||||
name = strings.ToLower(field.Name)
|
||||
}
|
||||
key := name
|
||||
if prefix != "" {
|
||||
key = prefix + "." + name
|
||||
}
|
||||
|
||||
ft := field.Type
|
||||
for ft.Kind() == reflect.Ptr {
|
||||
ft = ft.Elem()
|
||||
}
|
||||
if ft.Kind() == reflect.Struct {
|
||||
collectMapstructureKeys(ft, key, out)
|
||||
continue
|
||||
}
|
||||
if ft.Kind() == reflect.Map {
|
||||
// A map cannot be expressed in a single environment variable, so it
|
||||
// is out of scope here — such settings need a config file either way.
|
||||
continue
|
||||
}
|
||||
out[strings.ToLower(key)] = ft.String()
|
||||
}
|
||||
}
|
||||
|
||||
// TestConfigKeysAreEnvReachable is the systemic guard behind the image_storage
|
||||
// bug: viper.Unmarshal only decodes keys returned by AllKeys(), which unions
|
||||
// SetDefault keys, config-file keys and explicit BindEnv keys. AutomaticEnv can
|
||||
// override a key already in that union but never introduces one, and the
|
||||
// viper_bind_struct escape hatch is compiled out (we build with -tags embed).
|
||||
//
|
||||
// So a Config field with no registered default is unreachable by environment
|
||||
// variable whenever the deployment has no config.yaml containing it — the
|
||||
// operator sets the variable, the loader discards it, and the feature behaves
|
||||
// as if it were never configured. That is exactly how image_storage credentials
|
||||
// were lost, silently disabling async image tasks for env-driven deployments.
|
||||
//
|
||||
// When this fails, register a zero-valued default in setEnvReachableDefaults
|
||||
// for each reported key.
|
||||
func TestConfigKeysAreEnvReachable(t *testing.T) {
|
||||
bound := map[string]string{}
|
||||
collectMapstructureKeys(reflect.TypeOf(Config{}), "", bound)
|
||||
|
||||
viper.Reset()
|
||||
t.Cleanup(viper.Reset)
|
||||
setDefaults()
|
||||
registered := map[string]struct{}{}
|
||||
for _, key := range viper.AllKeys() {
|
||||
registered[key] = struct{}{}
|
||||
}
|
||||
|
||||
var unreachable []string
|
||||
for key, kind := range bound {
|
||||
if _, ok := registered[key]; !ok {
|
||||
unreachable = append(unreachable, key+" ("+kind+")")
|
||||
}
|
||||
}
|
||||
sort.Strings(unreachable)
|
||||
|
||||
if len(unreachable) > 0 {
|
||||
t.Fatalf("%d config keys have no default registered, so their environment variables are silently ignored:\n %s",
|
||||
len(unreachable), strings.Join(unreachable, "\n "))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//go:build unit
|
||||
|
||||
package config
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestLoadImageStorageFromEnv guards against a viper trap that silently disabled
|
||||
// asynchronous image tasks for every environment-variable-only deployment.
|
||||
//
|
||||
// viper only decodes keys returned by AllKeys(), which unions SetDefault keys,
|
||||
// config-file keys and explicit BindEnv keys. AutomaticEnv can override a key
|
||||
// that is already in that list, but it never introduces a new one. Credentials
|
||||
// such as image_storage.bucket therefore need an (empty) default registered, or
|
||||
// IMAGE_STORAGE_BUCKET is dropped on the floor and Active() stays false while
|
||||
// image_storage.enabled reads true — the endpoints 404 with no useful signal.
|
||||
func TestLoadImageStorageFromEnv(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("IMAGE_STORAGE_ENABLED", "true")
|
||||
t.Setenv("IMAGE_STORAGE_ENDPOINT", "https://acct.r2.cloudflarestorage.com")
|
||||
t.Setenv("IMAGE_STORAGE_BUCKET", "my-images")
|
||||
t.Setenv("IMAGE_STORAGE_ACCESS_KEY_ID", "ak")
|
||||
t.Setenv("IMAGE_STORAGE_SECRET_ACCESS_KEY", "sk")
|
||||
t.Setenv("IMAGE_STORAGE_PUBLIC_BASE_URL", "https://cdn.example.com")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.True(t, cfg.ImageStorage.Enabled)
|
||||
require.Equal(t, "https://acct.r2.cloudflarestorage.com", cfg.ImageStorage.Endpoint)
|
||||
require.Equal(t, "my-images", cfg.ImageStorage.Bucket)
|
||||
require.Equal(t, "ak", cfg.ImageStorage.AccessKeyID)
|
||||
require.Equal(t, "sk", cfg.ImageStorage.SecretAccessKey)
|
||||
require.Equal(t, "https://cdn.example.com", cfg.ImageStorage.PublicBaseURL)
|
||||
|
||||
require.True(t, cfg.ImageStorage.IsConfigured())
|
||||
require.True(t, cfg.ImageStorage.Active(), "async image tasks must be active when every credential is supplied via env")
|
||||
}
|
||||
@@ -10,12 +10,14 @@ import (
|
||||
type BackupHandler struct {
|
||||
backupService *service.BackupService
|
||||
userService *service.UserService
|
||||
imageStorage *service.ImageStorageSettingService
|
||||
}
|
||||
|
||||
func NewBackupHandler(backupService *service.BackupService, userService *service.UserService) *BackupHandler {
|
||||
func NewBackupHandler(backupService *service.BackupService, userService *service.UserService, imageStorage *service.ImageStorageSettingService) *BackupHandler {
|
||||
return &BackupHandler{
|
||||
backupService: backupService,
|
||||
userService: userService,
|
||||
imageStorage: imageStorage,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,3 +205,48 @@ func (h *BackupHandler) RestoreBackup(c *gin.Context) {
|
||||
}
|
||||
response.Accepted(c, record)
|
||||
}
|
||||
|
||||
// ─── 异步生图对象存储配置 ───
|
||||
//
|
||||
// 与备份共用一套 S3 客户端构造,因此放在同一个页面下:勾选"复用备份 S3"即可直接
|
||||
// 借用备份已配置的端点与密钥,只用不同的前缀区分对象(备份走 backups/,图片走 images/)。
|
||||
|
||||
func (h *BackupHandler) GetImageStorageConfig(c *gin.Context) {
|
||||
ctx := c.Request.Context()
|
||||
cfg, err := h.imageStorage.Get(ctx)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{
|
||||
"config": cfg,
|
||||
"secret_configured": h.imageStorage.SecretConfigured(ctx),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *BackupHandler) UpdateImageStorageConfig(c *gin.Context) {
|
||||
var req service.ImageStorageSettings
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
cfg, err := h.imageStorage.Update(c.Request.Context(), req)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, cfg)
|
||||
}
|
||||
|
||||
func (h *BackupHandler) TestImageStorageConnection(c *gin.Context) {
|
||||
var req service.ImageStorageSettings
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := h.imageStorage.TestConnection(c.Request.Context(), req); err != nil {
|
||||
response.Success(c, gin.H{"ok": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{"ok": true, "message": "connection successful"})
|
||||
}
|
||||
|
||||
@@ -152,6 +152,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
TurnstileSiteKey: settings.TurnstileSiteKey,
|
||||
TurnstileSecretKeyConfigured: settings.TurnstileSecretKeyConfigured,
|
||||
APIKeyACLTrustForwardedIP: settings.APIKeyACLTrustForwardedIP,
|
||||
ForwardedClientIPHeaders: settings.ForwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: settings.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: settings.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecretConfigured: settings.LinuxDoConnectClientSecretConfigured,
|
||||
|
||||
@@ -107,6 +107,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.APIKeyACLTrustForwardedIP != after.APIKeyACLTrustForwardedIP {
|
||||
changed = append(changed, "api_key_acl_trust_forwarded_ip")
|
||||
}
|
||||
if !equalStringSlice(before.ForwardedClientIPHeaders, after.ForwardedClientIPHeaders) {
|
||||
changed = append(changed, "forwarded_client_ip_headers")
|
||||
}
|
||||
if before.LinuxDoConnectEnabled != after.LinuxDoConnectEnabled {
|
||||
changed = append(changed, "linuxdo_connect_enabled")
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -155,3 +156,55 @@ func TestUpdateSettingsOmittedSecuritySwitchesKeepDisabled(t *testing.T) {
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyStepUpEnabled])
|
||||
require.Equal(t, "false", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
func TestUpdateSettingsForwardedClientIPHeadersOmittedPreservesAndEmptyClears(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyForwardedClientIPHeaders: `["X-Cdn-Ip","True-Client-Ip"]`,
|
||||
})
|
||||
|
||||
preserved := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
require.Equal(t, http.StatusOK, preserved.Code)
|
||||
require.JSONEq(t, `["X-Cdn-Ip","True-Client-Ip"]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
require.Contains(t, preserved.Body.String(), `"forwarded_client_ip_headers":["X-Cdn-Ip","True-Client-Ip"]`)
|
||||
|
||||
cleared := doUpdateSettings(t, h, map[string]any{"forwarded_client_ip_headers": []string{}}, nil)
|
||||
require.Equal(t, http.StatusOK, cleared.Code)
|
||||
require.JSONEq(t, `[]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
require.Contains(t, cleared.Body.String(), `"forwarded_client_ip_headers":[]`)
|
||||
}
|
||||
|
||||
func TestUpdateSettingsMalformedForwardedClientIPHeadersRemainFailClosedWhenOmitted(t *testing.T) {
|
||||
cfg := &config.Config{Default: config.DefaultConfig{UserConcurrency: 5}}
|
||||
repo := &settingHandlerRepoStub{values: map[string]string{
|
||||
service.SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
service.SettingKeyForwardedClientIPHeaders: `{"not":"an array"}`,
|
||||
}}
|
||||
svc := service.NewSettingService(repo, cfg)
|
||||
require.ErrorContains(t, svc.LoadForwardedClientIPSettings(context.Background()), "load forwarded client ip headers")
|
||||
require.False(t, cfg.ForwardedClientIPSettings().TrustForwardedIP)
|
||||
h := NewSettingHandler(svc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyAPIKeyACLTrustForwardedIP])
|
||||
require.JSONEq(t, `[]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Empty(t, runtimeSettings.Headers)
|
||||
require.Contains(t, rec.Body.String(), `"api_key_acl_trust_forwarded_ip":false`)
|
||||
require.Contains(t, rec.Body.String(), `"forwarded_client_ip_headers":[]`)
|
||||
}
|
||||
|
||||
func TestUpdateSettingsRejectsInvalidForwardedClientIPHeader(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyForwardedClientIPHeaders: `["X-Existing-IP"]`,
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{
|
||||
"forwarded_client_ip_headers": []string{"X Invalid"},
|
||||
}, nil)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.JSONEq(t, `["X-Existing-IP"]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
}
|
||||
|
||||
@@ -51,7 +51,8 @@ type UpdateSettingsRequest struct {
|
||||
TurnstileSecretKey string `json:"turnstile_secret_key"`
|
||||
|
||||
// API Key IP 访问控制设置
|
||||
APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
ForwardedClientIPHeaders *[]string `json:"forwarded_client_ip_headers"`
|
||||
|
||||
// LinuxDo Connect OAuth 登录
|
||||
LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"`
|
||||
@@ -400,6 +401,10 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
if req.StepUpEnabled != nil {
|
||||
stepUpEnabled = *req.StepUpEnabled
|
||||
}
|
||||
forwardedClientIPHeaders := append([]string(nil), previousSettings.ForwardedClientIPHeaders...)
|
||||
if req.ForwardedClientIPHeaders != nil {
|
||||
forwardedClientIPHeaders = append([]string(nil), (*req.ForwardedClientIPHeaders)...)
|
||||
}
|
||||
|
||||
// 开启敏感操作 step-up 门控属自锁风险操作:仅允许本人已启用 TOTP 的管理员会话开启,
|
||||
// 否则开启后操作者立即被挡在所有敏感操作之外。仅在 false→true 的开启瞬间校验,
|
||||
@@ -1269,6 +1274,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.APIKeyACLTrustForwardedIP
|
||||
}(),
|
||||
ForwardedClientIPHeaders: forwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: req.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: req.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecret: req.LinuxDoConnectClientSecret,
|
||||
@@ -1796,6 +1802,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
TurnstileSiteKey: updatedSettings.TurnstileSiteKey,
|
||||
TurnstileSecretKeyConfigured: updatedSettings.TurnstileSecretKeyConfigured,
|
||||
APIKeyACLTrustForwardedIP: updatedSettings.APIKeyACLTrustForwardedIP,
|
||||
ForwardedClientIPHeaders: updatedSettings.ForwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: updatedSettings.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: updatedSettings.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecretConfigured: updatedSettings.LinuxDoConnectClientSecretConfigured,
|
||||
|
||||
@@ -22,6 +22,27 @@ type SystemHandler struct {
|
||||
lockSvc *service.SystemOperationLockService
|
||||
}
|
||||
|
||||
// systemUpdateTimeout bounds a full in-place update or rollback: the release
|
||||
// manifest fetch plus a large binary download over slow links. It must stay
|
||||
// above the GitHub download client timeout (10 minutes) so the download owns
|
||||
// its own deadline.
|
||||
const systemUpdateTimeout = 15 * time.Minute
|
||||
|
||||
// systemUpdateContext detaches a long-running update/rollback from the HTTP
|
||||
// request lifetime. Browsers and reverse proxies commonly abort idle requests
|
||||
// after 30-60s (axios default, nginx proxy_read_timeout), which canceled
|
||||
// c.Request.Context() mid-download and killed the update with
|
||||
// "download failed: context canceled" (#4504). The swap keeps running after a
|
||||
// client disconnect; a later retry then hits the system operation lock or
|
||||
// reports "Already up to date".
|
||||
func systemUpdateContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||||
base := context.Background()
|
||||
if ctx != nil {
|
||||
base = context.WithoutCancel(ctx)
|
||||
}
|
||||
return context.WithTimeout(base, systemUpdateTimeout)
|
||||
}
|
||||
|
||||
type systemUpdateService interface {
|
||||
CheckUpdate(ctx context.Context, force bool) (*service.UpdateInfo, error)
|
||||
PerformUpdate(ctx context.Context) error
|
||||
@@ -75,9 +96,12 @@ func (h *SystemHandler) PerformUpdate(c *gin.Context) {
|
||||
release(releaseReason, succeeded)
|
||||
}()
|
||||
|
||||
if err := h.updateSvc.PerformUpdate(ctx); err != nil {
|
||||
updateCtx, cancel := systemUpdateContext(ctx)
|
||||
defer cancel()
|
||||
|
||||
if err := h.updateSvc.PerformUpdate(updateCtx); err != nil {
|
||||
if errors.Is(err, service.ErrNoUpdateAvailable) {
|
||||
info, checkErr := h.updateSvc.CheckUpdate(ctx, false)
|
||||
info, checkErr := h.updateSvc.CheckUpdate(updateCtx, false)
|
||||
if checkErr != nil {
|
||||
releaseReason = "SYSTEM_UPDATE_FAILED"
|
||||
return nil, checkErr
|
||||
@@ -152,7 +176,10 @@ func (h *SystemHandler) Rollback(c *gin.Context) {
|
||||
}()
|
||||
|
||||
if targetVersion != "" {
|
||||
err = h.updateSvc.RollbackToVersion(ctx, targetVersion)
|
||||
// 指定版本回退同样要下载完整二进制,与更新一样和请求生命周期解耦。
|
||||
rollbackCtx, cancel := systemUpdateContext(ctx)
|
||||
defer cancel()
|
||||
err = h.updateSvc.RollbackToVersion(rollbackCtx, targetVersion)
|
||||
} else {
|
||||
err = h.updateSvc.Rollback()
|
||||
}
|
||||
|
||||
@@ -18,18 +18,22 @@ import (
|
||||
)
|
||||
|
||||
type systemHandlerUpdateServiceStub struct {
|
||||
performErr error
|
||||
updateInfo *service.UpdateInfo
|
||||
checkErr error
|
||||
checkForces []bool
|
||||
performCall int
|
||||
rollbackCall int
|
||||
rollbackToCall int
|
||||
rollbackToVersions []string
|
||||
rollbackToErr error
|
||||
rollbackVersions []service.RollbackVersion
|
||||
rollbackVersionsErr error
|
||||
rollbackVersionsCall int
|
||||
performErr error
|
||||
updateInfo *service.UpdateInfo
|
||||
checkErr error
|
||||
checkForces []bool
|
||||
performCall int
|
||||
performCtxErr error
|
||||
performHasDeadline bool
|
||||
rollbackCall int
|
||||
rollbackToCall int
|
||||
rollbackToCtxErr error
|
||||
rollbackToHasDeadline bool
|
||||
rollbackToVersions []string
|
||||
rollbackToErr error
|
||||
rollbackVersions []service.RollbackVersion
|
||||
rollbackVersionsErr error
|
||||
rollbackVersionsCall int
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) CheckUpdate(_ context.Context, force bool) (*service.UpdateInfo, error) {
|
||||
@@ -37,8 +41,10 @@ func (s *systemHandlerUpdateServiceStub) CheckUpdate(_ context.Context, force bo
|
||||
return s.updateInfo, s.checkErr
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) PerformUpdate(context.Context) error {
|
||||
func (s *systemHandlerUpdateServiceStub) PerformUpdate(ctx context.Context) error {
|
||||
s.performCall++
|
||||
s.performCtxErr = ctx.Err()
|
||||
_, s.performHasDeadline = ctx.Deadline()
|
||||
return s.performErr
|
||||
}
|
||||
|
||||
@@ -52,8 +58,10 @@ func (s *systemHandlerUpdateServiceStub) ListRollbackVersions(context.Context) (
|
||||
return s.rollbackVersions, s.rollbackVersionsErr
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) RollbackToVersion(_ context.Context, version string) error {
|
||||
func (s *systemHandlerUpdateServiceStub) RollbackToVersion(ctx context.Context, version string) error {
|
||||
s.rollbackToCall++
|
||||
s.rollbackToCtxErr = ctx.Err()
|
||||
_, s.rollbackToHasDeadline = ctx.Deadline()
|
||||
s.rollbackToVersions = append(s.rollbackToVersions, version)
|
||||
return s.rollbackToErr
|
||||
}
|
||||
@@ -165,6 +173,55 @@ func TestSystemHandlerPerformUpdateFailureStillReturnsInternalError(t *testing.T
|
||||
require.Equal(t, "internal error", body.Message)
|
||||
}
|
||||
|
||||
// TestSystemHandlerPerformUpdateSurvivesClientDisconnect reproduces #4504:
|
||||
// the browser or a reverse proxy (axios 30s default, nginx proxy_read_timeout
|
||||
// 60s) aborts the long-running update request and cancels the request
|
||||
// context. The download must keep running on a detached, bounded context
|
||||
// instead of dying with "download failed: context canceled".
|
||||
func TestSystemHandlerPerformUpdateSurvivesClientDisconnect(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
router := newSystemHandlerTestRouter(t, updateSvc, repo)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/update", nil)
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
req = req.WithContext(canceledCtx)
|
||||
req.Header.Set("Idempotency-Key", "disconnected-update")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, 1, updateSvc.performCall)
|
||||
require.NoError(t, updateSvc.performCtxErr,
|
||||
"update must not observe the canceled request context")
|
||||
require.True(t, updateSvc.performHasDeadline,
|
||||
"detached update context must still be bounded by a deadline")
|
||||
requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded)
|
||||
}
|
||||
|
||||
func TestSystemHandlerRollbackToVersionSurvivesClientDisconnect(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
router := newSystemHandlerTestRouter(t, updateSvc, repo)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/rollback",
|
||||
strings.NewReader(`{"version":"0.1.146"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
req = req.WithContext(canceledCtx)
|
||||
req.Header.Set("Idempotency-Key", "disconnected-rollback")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, 1, updateSvc.rollbackToCall)
|
||||
require.NoError(t, updateSvc.rollbackToCtxErr,
|
||||
"versioned rollback must not observe the canceled request context")
|
||||
require.True(t, updateSvc.rollbackToHasDeadline,
|
||||
"detached rollback context must still be bounded by a deadline")
|
||||
requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded)
|
||||
}
|
||||
|
||||
func TestSystemHandlerRollbackWithoutBodyUsesLegacyBackup(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
|
||||
@@ -52,10 +52,11 @@ type SystemSettings struct {
|
||||
SMTPFromName string `json:"smtp_from_name"`
|
||||
SMTPUseTLS bool `json:"smtp_use_tls"`
|
||||
|
||||
TurnstileEnabled bool `json:"turnstile_enabled"`
|
||||
TurnstileSiteKey string `json:"turnstile_site_key"`
|
||||
TurnstileSecretKeyConfigured bool `json:"turnstile_secret_key_configured"`
|
||||
APIKeyACLTrustForwardedIP bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
TurnstileEnabled bool `json:"turnstile_enabled"`
|
||||
TurnstileSiteKey string `json:"turnstile_site_key"`
|
||||
TurnstileSecretKeyConfigured bool `json:"turnstile_secret_key_configured"`
|
||||
APIKeyACLTrustForwardedIP bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
ForwardedClientIPHeaders []string `json:"forwarded_client_ip_headers"`
|
||||
|
||||
LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"`
|
||||
LinuxDoConnectClientID string `json:"linuxdo_connect_client_id"`
|
||||
|
||||
@@ -83,8 +83,9 @@ func (s *FailoverState) HandleFailoverError(
|
||||
return FailoverExhausted
|
||||
}
|
||||
|
||||
// 缓存计费判断
|
||||
if needForceCacheBilling(s.hasBoundSession, failoverErr) {
|
||||
// 同账号重试不算切换账号,粘性会话仅在实际切换时强制缓存计费。
|
||||
sameAccountRetry := failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit
|
||||
if needForceCacheBilling(s.hasBoundSession, failoverErr, sameAccountRetry) {
|
||||
s.ForceCacheBilling = true
|
||||
}
|
||||
|
||||
@@ -174,9 +175,9 @@ func (s *FailoverState) HandleSelectionExhausted(ctx context.Context) FailoverAc
|
||||
}
|
||||
|
||||
// needForceCacheBilling 判断 failover 时是否需要强制缓存计费。
|
||||
// 粘性会话切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费。
|
||||
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError) bool {
|
||||
return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
||||
// 粘性会话实际切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费。
|
||||
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError, sameAccountRetry bool) bool {
|
||||
return (hasBoundSession && !sameAccountRetry) || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
||||
}
|
||||
|
||||
// failoverClientGone 判断下游客户端是否已断开(请求 context 已取消)。
|
||||
|
||||
@@ -271,7 +271,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
t.Run("hasBoundSession为true时设置ForceCacheBilling", func(t *testing.T) {
|
||||
t.Run("hasBoundSession为true且实际切换时设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true) // hasBoundSession=true
|
||||
err := newTestFailoverErr(500, false, false)
|
||||
@@ -280,6 +280,32 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
})
|
||||
|
||||
t.Run("同账号重试时仅凭hasBoundSession不设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, false)
|
||||
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.False(t, fs.ForceCacheBilling)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("同账号重试耗尽并实际切换时设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, false)
|
||||
|
||||
for i := 0; i < maxSameAccountRetries; i++ {
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
require.False(t, fs.ForceCacheBilling)
|
||||
}
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
require.Equal(t, 1, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("failoverErr.ForceCacheBilling为true时设置", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
@@ -289,6 +315,17 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
})
|
||||
|
||||
t.Run("同账号重试保留显式ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, true)
|
||||
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("两者均为false时不设置", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
@@ -633,13 +670,14 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
|
||||
for i := 0; i < maxSameAccountRetries; i++ {
|
||||
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
require.False(t, fs.ForceCacheBilling, "同账号重试期间不应仅因绑定会话强制缓存计费")
|
||||
}
|
||||
require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling")
|
||||
|
||||
// 2. 账号 100 超过重试上限 → TempUnschedule + 切换
|
||||
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
require.Equal(t, 1, fs.SwitchCount)
|
||||
require.True(t, fs.ForceCacheBilling, "实际切换账号时应设置 ForceCacheBilling")
|
||||
require.Len(t, mock.calls, 1)
|
||||
|
||||
// 3. 账号 200 遇到不可重试错误 → 直接切换
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type toggleSettingRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) Get(context.Context, string) (*service.Setting, error) { return nil, nil }
|
||||
func (r *toggleSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.values[key], nil
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *toggleSettingRepo) SetMultiple(context.Context, map[string]string) error { return nil }
|
||||
func (r *toggleSettingRepo) GetAll(context.Context) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *toggleSettingRepo) Delete(context.Context, string) error { return nil }
|
||||
|
||||
type passthroughEncryptor struct{}
|
||||
|
||||
func (passthroughEncryptor) Encrypt(plaintext string) (string, error) { return plaintext, nil }
|
||||
func (passthroughEncryptor) Decrypt(ciphertext string) (string, error) { return ciphertext, nil }
|
||||
|
||||
type noopImageStorage struct{}
|
||||
|
||||
func (noopImageStorage) Save(context.Context, string, string, []byte) (string, error) {
|
||||
return "https://cdn.example.test/object.png", nil
|
||||
}
|
||||
|
||||
// TestAsyncImageEnablesWithoutRestart drives the actual HTTP path for the bug behind
|
||||
// #4458 and #4542: with object storage unconfigured the async endpoint 404s, and the
|
||||
// only way to turn it on used to be editing config.yaml and restarting the container.
|
||||
// Flipping the admin setting must flip the endpoint over in the same process.
|
||||
func TestAsyncImageEnablesWithoutRestart(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
repo := &toggleSettingRepo{values: map[string]string{}}
|
||||
backup := service.NewBackupService(repo, &config.Config{}, passthroughEncryptor{}, nil, nil)
|
||||
factory := func(context.Context, *config.ImageStorageConfig) (service.ImageStorage, error) {
|
||||
return noopImageStorage{}, nil
|
||||
}
|
||||
settings := service.NewImageStorageSettingService(repo, passthroughEncryptor{}, backup, factory, config.ImageStorageConfig{})
|
||||
|
||||
store := &asyncImageMemoryStore{tasks: make(map[string]*service.ImageTaskRecord)}
|
||||
tasks := service.NewImageTaskServiceWithResolver(store, settings.Resolver(), time.Hour, time.Minute)
|
||||
|
||||
h := &AsyncImageHandler{tasks: tasks}
|
||||
h.execute = func(_ string, c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"created": 1, "data": []gin.H{{"url": "https://upstream.test/i.png"}}})
|
||||
}
|
||||
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
groupID := int64(3)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 9, UserID: 7, GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, AllowImageGeneration: true},
|
||||
})
|
||||
c.Next()
|
||||
})
|
||||
router.POST("/v1/images/generations/async", h.Submit)
|
||||
router.GET("/v1/images/tasks/:task_id", h.Get)
|
||||
|
||||
submit := func() *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async",
|
||||
strings.NewReader(`{"model":"gpt-image-1","prompt":"a lighthouse"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
rec := submit()
|
||||
require.Equal(t, http.StatusNotFound, rec.Code, "disabled until an admin configures object storage")
|
||||
require.Contains(t, rec.Body.String(), "async image tasks are not enabled")
|
||||
|
||||
// The admin saves the setting — no restart, same process.
|
||||
_, err := settings.Update(context.Background(), service.ImageStorageSettings{
|
||||
Enabled: true, Bucket: "my-images",
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com", AccessKeyID: "ak", SecretAccessKey: "sk",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
rec = submit()
|
||||
require.Equal(t, http.StatusAccepted, rec.Code, "the endpoint must go live as soon as the setting is saved")
|
||||
|
||||
var accepted struct {
|
||||
TaskID string `json:"task_id"`
|
||||
PollURL string `json:"poll_url"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &accepted))
|
||||
require.NotEmpty(t, accepted.TaskID)
|
||||
|
||||
// Turning the feature back off must not strand a task that was already accepted.
|
||||
_, err = settings.Update(context.Background(), service.ImageStorageSettings{Enabled: false})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, http.StatusNotFound, submit().Code, "new submissions are refused again")
|
||||
|
||||
pollRec := httptest.NewRecorder()
|
||||
router.ServeHTTP(pollRec, httptest.NewRequest(http.MethodGet, accepted.PollURL, nil))
|
||||
require.Equal(t, http.StatusOK, pollRec.Code, "an already-accepted task stays pollable after the switch is turned off")
|
||||
}
|
||||
@@ -39,6 +39,13 @@ func (h *AsyncImageHandler) enabled() bool {
|
||||
return h != nil && h.tasks != nil && h.tasks.Enabled()
|
||||
}
|
||||
|
||||
// pollable reports whether task lookups can be served. It is deliberately weaker
|
||||
// than enabled(): results already written to Redis stay readable after the
|
||||
// feature is switched off, so an in-flight task is never stranded.
|
||||
func (h *AsyncImageHandler) pollable() bool {
|
||||
return h != nil && h.tasks != nil && h.tasks.Pollable()
|
||||
}
|
||||
|
||||
// Submit accepts the same payload as the synchronous Images endpoint and
|
||||
// returns before the upstream image generation begins.
|
||||
func (h *AsyncImageHandler) Submit(c *gin.Context) {
|
||||
@@ -155,7 +162,11 @@ func (h *AsyncImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, apiKe
|
||||
}
|
||||
|
||||
func (h *AsyncImageHandler) Get(c *gin.Context) {
|
||||
if !h.enabled() {
|
||||
// Polling deliberately does not require the feature to be enabled, only that
|
||||
// the task store is reachable. Turning the switch off in the admin UI must not
|
||||
// strand tasks that were already accepted — their results are still in Redis
|
||||
// and their submitters are still polling.
|
||||
if !h.pollable() {
|
||||
imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "async image tasks are not enabled")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -53,6 +53,8 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
|
||||
return
|
||||
}
|
||||
// 让 ops 错误日志携带实际选中的上游账号,便于定位失效账号(#4544)。
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
|
||||
if err != nil {
|
||||
|
||||
@@ -13,6 +13,48 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GrokCountTokens handles Anthropic-compatible count_tokens requests locally.
|
||||
// The route middleware already authenticates the API key and resolves the
|
||||
// group; this handler intentionally does not select an account or check billing.
|
||||
func (h *OpenAIGatewayHandler) GrokCountTokens(c *gin.Context) {
|
||||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
h.anthropicErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return
|
||||
}
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||||
return
|
||||
}
|
||||
if len(body) == 0 {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||||
return
|
||||
}
|
||||
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
|
||||
if err != nil {
|
||||
logRequestBodyParseFailure(requestLogger(c, "handler.openai_gateway.grok_count_tokens"), body, err)
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
if parsedReq.Model == "" {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
|
||||
estimated, err := service.EstimateGrokCountTokens(parsedReq.Body.Bytes())
|
||||
if err != nil {
|
||||
requestLogger(c, "handler.openai_gateway.grok_count_tokens").Warn("grok_count_tokens.local_estimate_failed", zap.Error(err))
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, parsedReq.Model, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false)))
|
||||
c.JSON(http.StatusOK, gin.H{"input_tokens": estimated})
|
||||
}
|
||||
|
||||
// CountTokens handles Anthropic-compatible POST /v1/messages/count_tokens for OpenAI groups.
|
||||
// It validates billing and routes to an OpenAI token-count bridge without taking concurrency slots
|
||||
// or recording usage.
|
||||
|
||||
@@ -124,6 +124,13 @@ func openAICompatibleRequestPlatform(apiKey *service.APIKey) string {
|
||||
return service.PlatformOpenAI
|
||||
}
|
||||
|
||||
func openAIResponsesRequiredCapability(imageIntent bool, platform string) service.OpenAIEndpointCapability {
|
||||
if imageIntent && platform == service.PlatformOpenAI {
|
||||
return service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
return service.OpenAIEndpointCapabilityChatCompletions
|
||||
}
|
||||
|
||||
func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
|
||||
if apiKey == nil || apiKey.Group == nil {
|
||||
return true
|
||||
@@ -363,12 +370,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 生图意图的 /v1/responses 请求必须调度到确实支持 Responses API 的账号,否则
|
||||
// 会在 forward 阶段被静默降级为无法生图的 Chat Completions 直转(#4417)。
|
||||
// 仅对 OpenAI 平台生效:Grok 生图走独立的 forwardGrokResponses 路径,不应被过滤。
|
||||
// 使用 IsExplicitImageGenerationIntent 排除被动 image_gen namespace 声明,
|
||||
// 避免 Codex 的被动工具目录使 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, body) && requestPlatform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
// 复用前置权限与并发阶段在未修改 body 上确认的显式生图意图,避免大 tools 请求重复扫描。
|
||||
// 该判断已排除 Codex 被动 image_gen namespace,避免 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := openAIResponsesRequiredCapability(imageIntent, requestPlatform)
|
||||
|
||||
for {
|
||||
// Streaming Forward intentionally detaches the upstream request so usage can
|
||||
|
||||
@@ -138,6 +138,39 @@ func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
|
||||
}))
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesRequiredCapability(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
imageIntent bool
|
||||
platform string
|
||||
want service.OpenAIEndpointCapability
|
||||
}{
|
||||
{
|
||||
name: "OpenAI explicit image intent requires Responses",
|
||||
imageIntent: true,
|
||||
platform: service.PlatformOpenAI,
|
||||
want: service.OpenAIEndpointCapabilityResponses,
|
||||
},
|
||||
{
|
||||
name: "Grok explicit image intent keeps chat capability",
|
||||
imageIntent: true,
|
||||
platform: service.PlatformGrok,
|
||||
want: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
},
|
||||
{
|
||||
name: "non-image intent keeps chat capability",
|
||||
platform: service.PlatformOpenAI,
|
||||
want: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, openAIResponsesRequiredCapability(tt.imageIntent, tt.platform))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}`)
|
||||
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
var openAIResponsesImageIntentRoutingBenchmarkSink service.OpenAIEndpointCapability
|
||||
|
||||
func BenchmarkOpenAIResponsesImageIntentRouting_LargeToolsBody(b *testing.B) {
|
||||
body := buildLargeOpenAIResponsesToolsBody(32 << 20)
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body) {
|
||||
b.Fatal("large tools body must not have explicit image intent")
|
||||
}
|
||||
platform := service.PlatformOpenAI
|
||||
|
||||
b.Run("reuse_once", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(body)))
|
||||
b.ReportAllocs()
|
||||
for range b.N {
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body)
|
||||
openAIResponsesImageIntentRoutingBenchmarkSink = openAIResponsesRequiredCapability(imageIntent, platform)
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("rescan_twice", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(body)))
|
||||
b.ReportAllocs()
|
||||
for range b.N {
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body)
|
||||
// 对照优化前路径:路由阶段会再次扫描同一份未修改的 body。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body) && platform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
if imageIntent && requiredCapability != service.OpenAIEndpointCapabilityResponses {
|
||||
b.Fatal("explicit image intent must require Responses")
|
||||
}
|
||||
openAIResponsesImageIntentRoutingBenchmarkSink = requiredCapability
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func buildLargeOpenAIResponsesToolsBody(targetBytes int) []byte {
|
||||
var builder strings.Builder
|
||||
builder.Grow(targetBytes + 256)
|
||||
_, _ = builder.WriteString(`{"model":"gpt-5.4","tools":[{"type":"function","name":"search","description":"`)
|
||||
_, _ = builder.WriteString(strings.Repeat("x", targetBytes))
|
||||
_, _ = builder.WriteString(`"}],"tool_choice":"auto","input":"write code"}`)
|
||||
return []byte(builder.String())
|
||||
}
|
||||
@@ -232,7 +232,7 @@ func TestResponsesToAnthropic_TextOnly(t *testing.T) {
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "resp_123", anth.ID)
|
||||
assert.Equal(t, "claude-opus-4-6", anth.Model)
|
||||
assert.Equal(t, "end_turn", anth.StopReason)
|
||||
assert.Equal(t, "end_turn", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 1)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
assert.Equal(t, "Hello there!", anth.Content[0].Text)
|
||||
@@ -311,7 +311,7 @@ func TestResponsesToAnthropic_ToolUse(t *testing.T) {
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "tool_use", anth.StopReason)
|
||||
assert.Equal(t, "tool_use", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 2)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
assert.Equal(t, "tool_use", anth.Content[1].Type)
|
||||
@@ -342,7 +342,7 @@ func TestResponsesToAnthropic_ToolUseStopReasonDoesNotDependOnLastBlock(t *testi
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "tool_use", anth.StopReason)
|
||||
assert.Equal(t, "tool_use", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 2)
|
||||
assert.Equal(t, "tool_use", anth.Content[0].Type)
|
||||
assert.Equal(t, "text", anth.Content[1].Type)
|
||||
@@ -483,7 +483,7 @@ func TestResponsesToAnthropic_Incomplete(t *testing.T) {
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "max_tokens", anth.StopReason)
|
||||
assert.Equal(t, "max_tokens", AnthropicStopReasonString(anth.StopReason))
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_EmptyOutput(t *testing.T) {
|
||||
@@ -1081,7 +1081,7 @@ func TestResponsesToAnthropic_Failed(t *testing.T) {
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
// Failed status defaults to "end_turn" stop reason
|
||||
assert.Equal(t, "end_turn", anth.StopReason)
|
||||
assert.Equal(t, "end_turn", AnthropicStopReasonString(anth.StopReason))
|
||||
// Should have at least an empty text block
|
||||
require.Len(t, anth.Content, 1)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
@@ -1695,7 +1695,7 @@ func TestAnthropicToResponsesResponse_CacheTokensUseOpenAIInputSemantics(t *test
|
||||
Content: []AnthropicContentBlock{
|
||||
{Type: "text", Text: "ok"},
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 3318,
|
||||
OutputTokens: 123,
|
||||
@@ -1721,7 +1721,7 @@ func TestAnthropicToResponsesResponse_NoCacheTokens(t *testing.T) {
|
||||
Content: []AnthropicContentBlock{
|
||||
{Type: "text", Text: "ok"},
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
@@ -1818,3 +1818,25 @@ func TestAnthropicEventToResponses_CacheTokensFromMessageDelta(t *testing.T) {
|
||||
require.NotNil(t, completed.Response.Usage.InputTokensDetails)
|
||||
assert.Equal(t, 11, completed.Response.Usage.InputTokensDetails.CachedTokens)
|
||||
}
|
||||
|
||||
func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
state := NewResponsesEventToAnthropicState()
|
||||
state.Model = "grok-4.5"
|
||||
events := ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.created",
|
||||
Response: &ResponsesResponse{
|
||||
ID: "resp_test",
|
||||
Model: "grok-4.5",
|
||||
},
|
||||
}, state)
|
||||
require.Len(t, events, 1)
|
||||
require.Equal(t, "message_start", events[0].Type)
|
||||
require.NotNil(t, events[0].Message)
|
||||
require.Nil(t, events[0].Message.StopReason, "message_start must use null stop_reason, not empty string")
|
||||
|
||||
sse, err := ResponsesAnthropicEventToSSE(events[0])
|
||||
require.NoError(t, err)
|
||||
// Official Anthropic wire: "stop_reason":null
|
||||
require.Contains(t, sse, `"stop_reason":null`)
|
||||
require.NotContains(t, sse, `"stop_reason":""`)
|
||||
}
|
||||
|
||||
@@ -89,7 +89,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
out.Output = outputs
|
||||
|
||||
// Map stop_reason → status
|
||||
out.Status = anthropicStopReasonToResponsesStatus(resp.StopReason, resp.Content)
|
||||
out.Status = anthropicStopReasonToResponsesStatus(AnthropicStopReasonString(resp.StopReason), resp.Content)
|
||||
if out.Status == "incomplete" {
|
||||
out.IncompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
|
||||
}
|
||||
|
||||
@@ -390,7 +390,7 @@ func ChatCompletionsResponseToAnthropic(resp *ChatCompletionsResponse, model str
|
||||
if len(resp.Choices) > 0 {
|
||||
choice := resp.Choices[0]
|
||||
out.Content = chatMessageToAnthropicBlocks(choice.Message)
|
||||
out.StopReason = chatFinishReasonToAnthropicStopReason(choice.FinishReason, out.Content)
|
||||
out.StopReason = AnthropicStopReasonPtr(chatFinishReasonToAnthropicStopReason(choice.FinishReason, out.Content))
|
||||
// "length" → "max_tokens" is handled by chatFinishReasonToAnthropicStopReason;
|
||||
// Anthropic conveys max-tokens via stop_reason only, no incomplete_details field.
|
||||
}
|
||||
@@ -404,9 +404,9 @@ func ChatCompletionsResponseToAnthropic(resp *ChatCompletionsResponse, model str
|
||||
}
|
||||
// Empty choices / nil response never enter the choices branch above; the
|
||||
// double-conversion path still reports a completed turn ("end_turn"), and
|
||||
// stop_reason "" is invalid for strict Anthropic clients.
|
||||
if out.StopReason == "" {
|
||||
out.StopReason = chatFinishReasonToAnthropicStopReason("", out.Content)
|
||||
// stop_reason null/"" is invalid for completed non-stream responses.
|
||||
if AnthropicStopReasonString(out.StopReason) == "" {
|
||||
out.StopReason = AnthropicStopReasonPtr(chatFinishReasonToAnthropicStopReason("", out.Content))
|
||||
}
|
||||
// The double-conversion path generates a response id when the upstream
|
||||
// omits one (ChatCompletionsResponseToResponses); clients treat it as required.
|
||||
@@ -701,12 +701,13 @@ func ensureCCAnthropicMessageStart(state *ChatCompletionsToAnthropicStreamState)
|
||||
return []AnthropicStreamEvent{{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
Usage: AnthropicUsage{InputTokens: 0, OutputTokens: 0},
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
StopReason: nil, // JSON null; never ""
|
||||
Usage: AnthropicUsage{InputTokens: 0, OutputTokens: 0},
|
||||
},
|
||||
}}
|
||||
}
|
||||
|
||||
@@ -314,7 +314,7 @@ func TestChatCompletionsResponseToAnthropic_TextOnly(t *testing.T) {
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "hello world", out.Content[0].Text)
|
||||
require.Equal(t, "end_turn", out.StopReason)
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason))
|
||||
require.Equal(t, 5, out.Usage.InputTokens)
|
||||
require.Equal(t, 2, out.Usage.OutputTokens)
|
||||
}
|
||||
@@ -346,7 +346,7 @@ func TestChatCompletionsResponseToAnthropic_ToolUse(t *testing.T) {
|
||||
require.Equal(t, "call_1", out.Content[0].ID)
|
||||
require.Equal(t, "get_weather", out.Content[0].Name)
|
||||
require.Equal(t, `{"city":"SF"}`, string(out.Content[0].Input))
|
||||
require.Equal(t, "tool_use", out.StopReason)
|
||||
require.Equal(t, "tool_use", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_ReasoningOnlyFallback(t *testing.T) {
|
||||
@@ -384,7 +384,7 @@ func TestChatCompletionsResponseToAnthropic_FinishReasonLength(t *testing.T) {
|
||||
}
|
||||
|
||||
out := ChatCompletionsResponseToAnthropic(resp, "claude-sonnet-4-20250514")
|
||||
require.Equal(t, "max_tokens", out.StopReason)
|
||||
require.Equal(t, "max_tokens", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_EmptyChoices(t *testing.T) {
|
||||
@@ -398,7 +398,7 @@ func TestChatCompletionsResponseToAnthropic_EmptyChoices(t *testing.T) {
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "", out.Content[0].Text)
|
||||
require.Equal(t, "end_turn", out.StopReason, "empty choices must not produce an empty stop_reason")
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason), "empty choices must not produce an empty stop_reason")
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_CacheTokens(t *testing.T) {
|
||||
@@ -433,7 +433,7 @@ func TestChatCompletionsResponseToAnthropic_NilResponse(t *testing.T) {
|
||||
out := ChatCompletionsResponseToAnthropic(nil, "claude-sonnet-4-20250514")
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "end_turn", out.StopReason, "nil response must not produce an empty stop_reason")
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason), "nil response must not produce an empty stop_reason")
|
||||
require.NotEmpty(t, out.ID)
|
||||
}
|
||||
|
||||
@@ -678,7 +678,7 @@ func TestDirectBridge_NonStreamingMatchesDoubleConversion(t *testing.T) {
|
||||
double := ResponsesToAnthropic(responsesResp, "claude-sonnet-4-20250514")
|
||||
|
||||
// Compare key fields
|
||||
require.Equal(t, direct.StopReason, double.StopReason)
|
||||
require.Equal(t, AnthropicStopReasonString(direct.StopReason), AnthropicStopReasonString(double.StopReason))
|
||||
require.Equal(t, direct.Model, double.Model)
|
||||
require.Len(t, direct.Content, len(double.Content))
|
||||
for i := range direct.Content {
|
||||
@@ -1061,8 +1061,8 @@ func TestDirectBridge_NonStreamingMatchesDoubleConversion_EmptyChoices(t *testin
|
||||
responsesResp := ChatCompletionsResponseToResponses(resp, "claude-sonnet-4-20250514", nil, false, nil)
|
||||
double := ResponsesToAnthropic(responsesResp, "claude-sonnet-4-20250514")
|
||||
|
||||
require.Equal(t, double.StopReason, direct.StopReason)
|
||||
require.Equal(t, "end_turn", direct.StopReason)
|
||||
require.Equal(t, AnthropicStopReasonString(double.StopReason), AnthropicStopReasonString(direct.StopReason))
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(direct.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_ContentFilterWithToolUse(t *testing.T) {
|
||||
@@ -1086,5 +1086,5 @@ func TestChatCompletionsResponseToAnthropic_ContentFilterWithToolUse(t *testing.
|
||||
}
|
||||
|
||||
out := ChatCompletionsResponseToAnthropic(resp, "claude-sonnet-4-20250514")
|
||||
require.Equal(t, "tool_use", out.StopReason)
|
||||
require.Equal(t, "tool_use", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
@@ -88,7 +88,7 @@ func TestAnthropicToResponsesResponse_CacheCreation(t *testing.T) {
|
||||
CacheReadInputTokens: 4,
|
||||
CacheCreationInputTokens: 6,
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
}
|
||||
|
||||
out := AnthropicToResponsesResponse(&resp)
|
||||
|
||||
@@ -86,7 +86,7 @@ func ResponsesToAnthropic(resp *ResponsesResponse, model string) *AnthropicRespo
|
||||
}
|
||||
out.Content = blocks
|
||||
|
||||
out.StopReason = responsesStatusToAnthropicStopReason(resp.Status, resp.IncompleteDetails, blocks)
|
||||
out.StopReason = AnthropicStopReasonPtr(responsesStatusToAnthropicStopReason(resp.Status, resp.IncompleteDetails, blocks))
|
||||
|
||||
if resp.Usage != nil {
|
||||
out.Usage = anthropicUsageFromResponsesUsage(resp.Usage)
|
||||
@@ -308,14 +308,19 @@ func resToAnthHandleCreated(evt *ResponsesStreamEvent, state *ResponsesEventToAn
|
||||
}
|
||||
state.MessageStartSent = true
|
||||
|
||||
// Official Anthropic message_start uses stop_reason: null and usage with
|
||||
// input_tokens when known. We leave StopReason nil (JSON null) and usage
|
||||
// zeros until response.completed; never emit stop_reason:"" which breaks
|
||||
// strict clients' turn-finalization / session usage accounting.
|
||||
return []AnthropicStreamEvent{{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
StopReason: nil,
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 0,
|
||||
OutputTokens: 0,
|
||||
|
||||
@@ -124,17 +124,34 @@ type AnthropicCacheControl struct {
|
||||
}
|
||||
|
||||
// AnthropicResponse is the non-streaming response from POST /v1/messages.
|
||||
//
|
||||
// StopReason is a pointer so streaming message_start can emit JSON null
|
||||
// (official Anthropic wire format). A plain string zero-value would marshal as
|
||||
// "" which strict clients treat as invalid mid-stream state.
|
||||
type AnthropicResponse struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // "message"
|
||||
Role string `json:"role"` // "assistant"
|
||||
Content []AnthropicContentBlock `json:"content"`
|
||||
Model string `json:"model"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
StopReason *string `json:"stop_reason"`
|
||||
StopSequence *string `json:"stop_sequence,omitempty"`
|
||||
Usage AnthropicUsage `json:"usage"`
|
||||
}
|
||||
|
||||
// AnthropicStopReasonPtr returns a non-nil pointer to s for final stop reasons.
|
||||
func AnthropicStopReasonPtr(s string) *string {
|
||||
return &s
|
||||
}
|
||||
|
||||
// AnthropicStopReasonString returns the stop reason value, or "" when unset/null.
|
||||
func AnthropicStopReasonString(p *string) string {
|
||||
if p == nil {
|
||||
return ""
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
// AnthropicUsage holds token counts in Anthropic format.
|
||||
type AnthropicUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
|
||||
@@ -8,10 +8,138 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetClientIP resolves a client address only through Gin's configured trusted
|
||||
// proxy chain. Forwarding headers from a direct or untrusted peer are ignored.
|
||||
const forwardedIPSettingsKey = "sub2api.forwarded_ip_settings"
|
||||
|
||||
type forwardedIPSettings struct {
|
||||
trustForwarded bool
|
||||
headers []string
|
||||
}
|
||||
|
||||
// SetForwardedIPSettings snapshots the forwarded-IP mode and custom header list
|
||||
// for this request.
|
||||
func SetForwardedIPSettings(c *gin.Context, enabled bool, headers []string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.Set(forwardedIPSettingsKey, forwardedIPSettings{
|
||||
trustForwarded: enabled,
|
||||
headers: append([]string(nil), headers...),
|
||||
})
|
||||
}
|
||||
|
||||
// SetLegacyForwardedIPTrust records whether raw forwarding headers override
|
||||
// Gin's server.trusted_proxies chain for this request.
|
||||
func SetLegacyForwardedIPTrust(c *gin.Context, enabled bool) {
|
||||
SetForwardedIPSettings(c, enabled, nil)
|
||||
}
|
||||
|
||||
func requestForwardedIPSettings(c *gin.Context) (forwardedIPSettings, bool) {
|
||||
if c == nil {
|
||||
return forwardedIPSettings{}, false
|
||||
}
|
||||
value, ok := c.Get(forwardedIPSettingsKey)
|
||||
if !ok {
|
||||
return forwardedIPSettings{}, false
|
||||
}
|
||||
settings, ok := value.(forwardedIPSettings)
|
||||
return settings, ok
|
||||
}
|
||||
|
||||
func requestUsesLegacyForwardedIPTrust(c *gin.Context) bool {
|
||||
settings, ok := requestForwardedIPSettings(c)
|
||||
return !ok || settings.trustForwarded
|
||||
}
|
||||
|
||||
// GetClientIP resolves the client address using the legacy forwarding-header
|
||||
// precedence used before the trusted-proxy hardening. It remains the
|
||||
// compatibility path for request metadata and usage/error logs; security-
|
||||
// sensitive callers must use GetTrustedClientIP or GetSecurityClientIP.
|
||||
func GetClientIP(c *gin.Context) string {
|
||||
return GetTrustedClientIP(c)
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
if !requestUsesLegacyForwardedIPTrust(c) {
|
||||
return GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
settings, _ := requestForwardedIPSettings(c)
|
||||
customIP, customFallback := resolveCustomForwardedClientIP(c, settings.headers)
|
||||
if customIP != "" {
|
||||
return customIP
|
||||
}
|
||||
|
||||
// Preserve the historical precedence used by existing reverse-proxy
|
||||
// deployments, while skipping an internal proxy address when a public XFF
|
||||
// value is available. This covers Docker/Nginx setups that accidentally
|
||||
// write the bridge address into X-Real-IP.
|
||||
legacyIP, legacyFallback := resolveLegacyForwardedHeaderIP(c)
|
||||
if legacyIP != "" {
|
||||
return legacyIP
|
||||
}
|
||||
if customFallback != "" {
|
||||
return customFallback
|
||||
}
|
||||
if legacyFallback != "" {
|
||||
return legacyFallback
|
||||
}
|
||||
return normalizeIP(c.ClientIP())
|
||||
}
|
||||
|
||||
func resolveCustomForwardedClientIP(c *gin.Context, headers []string) (string, string) {
|
||||
if c == nil {
|
||||
return "", ""
|
||||
}
|
||||
var fallback string
|
||||
for _, header := range headers {
|
||||
for _, value := range c.Request.Header.Values(header) {
|
||||
for _, candidate := range strings.Split(value, ",") {
|
||||
parsed := net.ParseIP(strings.TrimSpace(candidate))
|
||||
if parsed == nil {
|
||||
continue
|
||||
}
|
||||
normalized := parsed.String()
|
||||
if isPrivateIP(normalized) {
|
||||
if fallback == "" {
|
||||
fallback = normalized
|
||||
}
|
||||
continue
|
||||
}
|
||||
return normalized, fallback
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", fallback
|
||||
}
|
||||
|
||||
func resolveLegacyForwardedHeaderIP(c *gin.Context) (string, string) {
|
||||
var fallback string
|
||||
if forwarded := normalizeIP(c.GetHeader("CF-Connecting-IP")); forwarded != "" {
|
||||
fallback = forwarded
|
||||
if !isPrivateIP(forwarded) {
|
||||
return forwarded, fallback
|
||||
}
|
||||
}
|
||||
if realIP := normalizeIP(c.GetHeader("X-Real-IP")); realIP != "" {
|
||||
if fallback == "" {
|
||||
fallback = realIP
|
||||
}
|
||||
if !isPrivateIP(realIP) {
|
||||
return realIP, fallback
|
||||
}
|
||||
}
|
||||
if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
|
||||
ips := strings.Split(xff, ",")
|
||||
for _, candidate := range ips {
|
||||
candidate = strings.TrimSpace(candidate)
|
||||
if candidate != "" && !isPrivateIP(candidate) {
|
||||
return normalizeIP(candidate), fallback
|
||||
}
|
||||
}
|
||||
if fallback == "" && len(ips) > 0 {
|
||||
fallback = normalizeIP(strings.TrimSpace(ips[0]))
|
||||
}
|
||||
}
|
||||
return "", fallback
|
||||
}
|
||||
|
||||
// GetTrustedClientIP 从 Gin 的可信代理解析链提取客户端 IP。
|
||||
@@ -24,10 +152,17 @@ func GetTrustedClientIP(c *gin.Context) string {
|
||||
return normalizeIP(c.ClientIP())
|
||||
}
|
||||
|
||||
// GetSecurityClientIP returns the address resolved through Gin's configured
|
||||
// trusted-proxy chain. The legacy toggle is retained for configuration/API
|
||||
// compatibility, but never makes raw forwarding headers trustworthy by itself.
|
||||
func GetSecurityClientIP(c *gin.Context, _ bool) string {
|
||||
// GetSecurityClientIP returns the address used by security-sensitive paths.
|
||||
// When legacy forwarded-IP trust is enabled, raw forwarding headers take over
|
||||
// client-IP resolution. When disabled, Gin's server.trusted_proxies chain is
|
||||
// authoritative.
|
||||
func GetSecurityClientIP(c *gin.Context, trustForwarded bool) string {
|
||||
if requestSettings, ok := requestForwardedIPSettings(c); ok {
|
||||
trustForwarded = requestSettings.trustForwarded
|
||||
}
|
||||
if trustForwarded {
|
||||
return GetClientIP(c)
|
||||
}
|
||||
return GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
@@ -41,6 +176,27 @@ func normalizeIP(ip string) string {
|
||||
return ip
|
||||
}
|
||||
|
||||
// privateNets contains the private/loopback ranges skipped while selecting a
|
||||
// public address from a legacy X-Forwarded-For chain.
|
||||
var privateNets []*net.IPNet
|
||||
|
||||
func init() {
|
||||
for _, cidr := range []string{
|
||||
"10.0.0.0/8",
|
||||
"172.16.0.0/12",
|
||||
"192.168.0.0/16",
|
||||
"127.0.0.0/8",
|
||||
"::1/128",
|
||||
"fc00::/7",
|
||||
} {
|
||||
_, block, err := net.ParseCIDR(cidr)
|
||||
if err != nil {
|
||||
panic("invalid CIDR: " + cidr)
|
||||
}
|
||||
privateNets = append(privateNets, block)
|
||||
}
|
||||
}
|
||||
|
||||
// CompiledIPRules 表示预编译的 IP 匹配规则。
|
||||
// PatternCount 记录原始规则数量,用于保留“规则存在但全无效”时的行为语义。
|
||||
type CompiledIPRules struct {
|
||||
@@ -96,6 +252,19 @@ func matchesCompiledRules(parsedIP net.IP, rules *CompiledIPRules) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func isPrivateIP(ipStr string) bool {
|
||||
ip := net.ParseIP(ipStr)
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
for _, block := range privateNets {
|
||||
if block.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// MatchesPattern 检查 IP 是否匹配指定的模式(支持单个 IP 或 CIDR)。
|
||||
// pattern 可以是:
|
||||
// - 单个 IP: "192.168.1.100"
|
||||
|
||||
@@ -32,6 +32,26 @@ func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) {
|
||||
require.Equal(t, "9.9.9.9", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetClientIPPreservesLegacyDockerForwardedHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetClientIP(c))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "192.168.32.1:12345"
|
||||
req.Header.Set("X-Forwarded-For", "10.0.0.2, 203.0.113.42")
|
||||
req.Header.Set("X-Real-IP", "192.168.32.1")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, "203.0.113.42", w.Body.String())
|
||||
}
|
||||
|
||||
func TestCheckIPRestrictionWithCompiledRules(t *testing.T) {
|
||||
whitelist := CompileIPRules([]string{"10.0.0.0/8", "192.168.1.2"})
|
||||
blacklist := CompileIPRules([]string{"10.1.1.1"})
|
||||
@@ -53,41 +73,125 @@ func TestCheckIPRestrictionWithCompiledRules_InvalidWhitelistStillDenies(t *test
|
||||
require.Equal(t, "access denied", reason)
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPNeverTrustsHeadersFromUntrustedPeer(t *testing.T) {
|
||||
func TestGetSecurityClientIPSwitchEnabledUsesLegacyHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, tc := range []struct {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetSecurityClientIP(c, true))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPCustomHeaderPrecedenceAndFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
trustForwarded bool
|
||||
trustForward bool
|
||||
headers []string
|
||||
requestHeaders map[string]string
|
||||
want string
|
||||
}{
|
||||
{name: "legacy toggle disabled", trustForwarded: false, want: "9.9.9.9"},
|
||||
{name: "legacy toggle enabled", trustForwarded: true, want: "9.9.9.9"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
{
|
||||
name: "configured order precedes built-ins",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-First", "X-CDN-Second"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-First": "198.51.100.10",
|
||||
"X-CDN-Second": "203.0.113.20",
|
||||
"CF-Connecting-IP": "8.8.8.8",
|
||||
},
|
||||
want: "198.51.100.10",
|
||||
},
|
||||
{
|
||||
name: "comma candidates skip invalid and private values",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-First", "X-CDN-Second"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-First": "not-an-ip, 10.0.0.8",
|
||||
"X-CDN-Second": "also-bad, 203.0.113.9",
|
||||
},
|
||||
want: "203.0.113.9",
|
||||
},
|
||||
{
|
||||
name: "legacy public header wins over custom private fallback",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "10.0.0.8",
|
||||
"X-Real-IP": "1.2.3.4",
|
||||
},
|
||||
want: "1.2.3.4",
|
||||
},
|
||||
{
|
||||
name: "custom private fallback retains configured precedence",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "10.0.0.8",
|
||||
"X-Real-IP": "192.168.1.4",
|
||||
},
|
||||
want: "10.0.0.8",
|
||||
},
|
||||
{
|
||||
name: "invalid custom value continues to built-ins",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "1.2.3.4:443",
|
||||
"CF-Connecting-IP": "4.4.4.4",
|
||||
},
|
||||
want: "4.4.4.4",
|
||||
},
|
||||
{
|
||||
name: "disabled mode ignores custom and legacy headers",
|
||||
trustForward: false,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "1.2.3.4",
|
||||
"X-Real-IP": "4.4.4.4",
|
||||
},
|
||||
want: "9.9.9.9",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetSecurityClientIP(c, tc.trustForwarded))
|
||||
SetForwardedIPSettings(c, test.trustForward, test.headers)
|
||||
c.String(200, GetSecurityClientIP(c, !test.trustForward))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
for name, value := range test.requestHeaders {
|
||||
req.Header.Set(name, value)
|
||||
}
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, tc.want, w.Body.String())
|
||||
require.Equal(t, test.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
func TestGetSecurityClientIPSwitchDisabledUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies([]string{"9.9.9.9"}))
|
||||
r.GET("/t", func(c *gin.Context) { c.String(200, GetSecurityClientIP(c, true)) })
|
||||
r.GET("/t", func(c *gin.Context) { c.String(200, GetSecurityClientIP(c, false)) })
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
@@ -97,3 +201,76 @@ func TestGetSecurityClientIPUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetClientIPSwitchDisabledUsesTrustedProxyChain(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
SetLegacyForwardedIPTrust(c, false)
|
||||
c.String(200, GetClientIP(c))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "9.9.9.9", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPRequestSnapshotCopiesCustomHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
headers := []string{"X-Original-IP"}
|
||||
SetForwardedIPSettings(c, true, headers)
|
||||
headers[0] = "X-Mutated-IP"
|
||||
c.String(200, GetSecurityClientIP(c, false))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Original-IP", "1.2.3.4")
|
||||
req.Header.Set("X-Mutated-IP", "4.4.4.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPRequestSnapshotOverridesLiveFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
requestTrust bool
|
||||
fallbackTrust bool
|
||||
want string
|
||||
}{
|
||||
{name: "captured secure mode wins", requestTrust: false, fallbackTrust: true, want: "9.9.9.9"},
|
||||
{name: "captured compatibility mode wins", requestTrust: true, fallbackTrust: false, want: "1.2.3.4"},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
SetLegacyForwardedIPTrust(c, test.requestTrust)
|
||||
c.String(200, GetSecurityClientIP(c, test.fallbackTrust))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, test.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,18 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
const GrokFreeRolling24hTokenLimit int64 = 1_000_000
|
||||
|
||||
var grokFreeRolling24hTokenLimits = map[int64]struct{}{
|
||||
GrokFreeRolling24hTokenLimit: {},
|
||||
2_000_000: {}, // Legacy Free limit observed before July 2026.
|
||||
}
|
||||
|
||||
func IsGrokFreeRolling24hTokenLimit(limit int64) bool {
|
||||
_, ok := grokFreeRolling24hTokenLimits[limit]
|
||||
return ok
|
||||
}
|
||||
|
||||
type QuotaWindow struct {
|
||||
Limit *int64 `json:"limit,omitempty"`
|
||||
Remaining *int64 `json:"remaining,omitempty"`
|
||||
|
||||
@@ -61,3 +61,11 @@ func TestObserveQuotaHeadersRecordsNoHeaderProbe(t *testing.T) {
|
||||
require.Nil(t, snapshot.Requests)
|
||||
require.Nil(t, snapshot.Tokens)
|
||||
}
|
||||
|
||||
func TestIsGrokFreeRolling24hTokenLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.True(t, IsGrokFreeRolling24hTokenLimit(GrokFreeRolling24hTokenLimit))
|
||||
require.True(t, IsGrokFreeRolling24hTokenLimit(2_000_000), "legacy snapshots remain classifiable")
|
||||
require.False(t, IsGrokFreeRolling24hTokenLimit(3_000_000))
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -18,6 +19,7 @@ import (
|
||||
type githubReleaseClient struct {
|
||||
httpClient *http.Client
|
||||
downloadHTTPClient *http.Client
|
||||
updateGitHubToken string
|
||||
}
|
||||
|
||||
type githubReleaseClientError struct {
|
||||
@@ -43,6 +45,8 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi
|
||||
}
|
||||
sharedClient = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
apiClient := cloneHTTPClient(sharedClient)
|
||||
apiClient.CheckRedirect = githubAPICheckRedirect(apiClient.CheckRedirect)
|
||||
|
||||
// 下载客户端需要更长的超时时间
|
||||
downloadClient, err := httpclient.GetClient(httpclient.Options{
|
||||
@@ -56,13 +60,50 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi
|
||||
}
|
||||
downloadClient = &http.Client{Timeout: 10 * time.Minute}
|
||||
}
|
||||
downloadClient = cloneHTTPClient(downloadClient)
|
||||
|
||||
return &githubReleaseClient{
|
||||
httpClient: sharedClient,
|
||||
httpClient: apiClient,
|
||||
downloadHTTPClient: downloadClient,
|
||||
updateGitHubToken: os.Getenv("UPDATE_GITHUB_TOKEN"),
|
||||
}
|
||||
}
|
||||
|
||||
func cloneHTTPClient(client *http.Client) *http.Client {
|
||||
cloned := *client
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func isGitHubAPIURL(url *url.URL) bool {
|
||||
return url != nil && strings.EqualFold(url.Scheme, "https") && url.User == nil &&
|
||||
strings.EqualFold(url.Host, "api.github.com")
|
||||
}
|
||||
|
||||
func githubAPICheckRedirect(previous func(*http.Request, []*http.Request) error) func(*http.Request, []*http.Request) error {
|
||||
return func(req *http.Request, via []*http.Request) error {
|
||||
if !isGitHubAPIURL(req.URL) {
|
||||
req.Header.Del("Authorization")
|
||||
}
|
||||
if previous != nil {
|
||||
return previous(req, via)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *githubReleaseClient) newAPIRequest(ctx context.Context, url string) (*http.Request, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
if c.updateGitHubToken != "" && isGitHubAPIURL(req.URL) {
|
||||
req.Header.Set("Authorization", "Bearer "+c.updateGitHubToken)
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (c *githubReleaseClientError) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
|
||||
return nil, c.err
|
||||
}
|
||||
@@ -82,12 +123,10 @@ func (c *githubReleaseClientError) FetchChecksumFile(ctx context.Context, url st
|
||||
func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
|
||||
url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
req, err := c.newAPIRequest(ctx, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -116,12 +155,10 @@ func (c *githubReleaseClient) FetchRecentReleases(ctx context.Context, repo stri
|
||||
}
|
||||
url := fmt.Sprintf("https://api.github.com/repos/%s/releases?per_page=%d", repo, perPage)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
req, err := c.newAPIRequest(ctx, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -3,6 +3,7 @@ package repository
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
@@ -47,6 +48,97 @@ func newTestGitHubReleaseClient() *githubReleaseClient {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientAPIRequestAuthorization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
wantAuth string
|
||||
}{
|
||||
{name: "exact HTTPS authority", url: "https://api.github.com/repos/test/repo", wantAuth: "Bearer update-secret"},
|
||||
{name: "HTTP", url: "http://api.github.com/repos/test/repo"},
|
||||
{name: "subdomain", url: "https://sub.api.github.com/repos/test/repo"},
|
||||
{name: "userinfo", url: "https://user@api.github.com/repos/test/repo"},
|
||||
{name: "explicit default port", url: "https://api.github.com:443/repos/test/repo"},
|
||||
{name: "custom port", url: "https://api.github.com:8443/repos/test/repo"},
|
||||
{name: "different host", url: "https://github.com/test/repo"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := newTestGitHubReleaseClient()
|
||||
client.updateGitHubToken = "update-secret"
|
||||
req, err := client.newAPIRequest(context.Background(), tt.url)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.wantAuth, req.Header.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
|
||||
client := newTestGitHubReleaseClient()
|
||||
req, err := client.newAPIRequest(context.Background(), "https://api.github.com/repos/test/repo")
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientRedirectAuthorization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
wantAuth string
|
||||
}{
|
||||
{name: "same HTTPS authority", url: "https://api.github.com/redirected", wantAuth: "Bearer update-secret"},
|
||||
{name: "HTTP", url: "http://api.github.com/redirected"},
|
||||
{name: "subdomain", url: "https://sub.api.github.com/redirected"},
|
||||
{name: "userinfo", url: "https://user@api.github.com/redirected"},
|
||||
{name: "custom port", url: "https://api.github.com:8443/redirected"},
|
||||
{name: "different host", url: "https://example.com/redirected"},
|
||||
}
|
||||
|
||||
checkRedirect := githubAPICheckRedirect(nil)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodGet, tt.url, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer update-secret")
|
||||
|
||||
require.NoError(t, checkRedirect(req, nil))
|
||||
require.Equal(t, tt.wantAuth, req.Header.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientDoesNotAuthorizeDownloads(t *testing.T) {
|
||||
client := newTestGitHubReleaseClient()
|
||||
client.updateGitHubToken = "update-secret"
|
||||
|
||||
var headers []http.Header
|
||||
transport := githubReleaseRoundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
headers = append(headers, req.Header.Clone())
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader("checksum")),
|
||||
Request: req,
|
||||
}, nil
|
||||
})
|
||||
client.httpClient.Transport = transport
|
||||
client.downloadHTTPClient.Transport = transport
|
||||
|
||||
dest := filepath.Join(t.TempDir(), "asset")
|
||||
require.NoError(t, client.DownloadFile(context.Background(), "https://objects.githubusercontent.com/asset", dest, 100))
|
||||
_, err := client.FetchChecksumFile(context.Background(), "https://github.com/test/repo/releases/download/v1/checksums.txt")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, headers, 2)
|
||||
for _, header := range headers {
|
||||
require.Empty(t, header.Get("Authorization"))
|
||||
}
|
||||
}
|
||||
|
||||
type githubReleaseRoundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f githubReleaseRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func (s *GitHubReleaseServiceSuite) SetupTest() {
|
||||
s.tempDir = s.T().TempDir()
|
||||
}
|
||||
|
||||
@@ -142,7 +142,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewS3BackupStoreFactory,
|
||||
|
||||
// Image storage (async image task result offload)
|
||||
ProvideImageStorage,
|
||||
ProvideImageStorageFactory,
|
||||
|
||||
// HTTP service ports (DI Strategy A: return interface directly)
|
||||
NewTurnstileVerifier,
|
||||
@@ -175,17 +175,14 @@ func ProvideEnt(cfg *config.Config) (*ent.Client, error) {
|
||||
return client, err
|
||||
}
|
||||
|
||||
// ProvideImageStorage 提供异步图片任务结果转存所用的对象存储实现。
|
||||
// 仅当开关打开且 S3 凭证齐全时返回具体实现,否则返回 nil(功能整体禁用)。
|
||||
func ProvideImageStorage(cfg *config.Config) (service.ImageStorage, error) {
|
||||
if !cfg.ImageStorage.Active() {
|
||||
return nil, nil
|
||||
// ProvideImageStorageFactory 提供按需构造对象存储客户端的工厂。
|
||||
//
|
||||
// 这里返回工厂而不是实例:异步生图的开关与凭证可以在后台随时改动,客户端必须能在
|
||||
// 设置保存后重建,而不是在启动时定死一份。
|
||||
func ProvideImageStorageFactory() service.ImageStorageFactory {
|
||||
return func(ctx context.Context, cfg *config.ImageStorageConfig) (service.ImageStorage, error) {
|
||||
return NewS3ImageStorage(ctx, cfg)
|
||||
}
|
||||
store, err := NewS3ImageStorage(context.Background(), &cfg.ImageStorage)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return store, nil
|
||||
}
|
||||
|
||||
// ProvideSQLDB 从 Ent 客户端提取底层的 *sql.DB 连接。
|
||||
|
||||
@@ -791,6 +791,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"site_subtitle": "Subtitle",
|
||||
"api_base_url": "https://api.example.com",
|
||||
"api_key_acl_trust_forwarded_ip": false,
|
||||
"forwarded_client_ip_headers": [],
|
||||
"contact_info": "support",
|
||||
"doc_url": "https://docs.example.com",
|
||||
"auth_source_default_email_balance": 0,
|
||||
@@ -1102,6 +1103,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"site_subtitle": "Subscription to API Conversion Platform",
|
||||
"api_base_url": "",
|
||||
"api_key_acl_trust_forwarded_ip": false,
|
||||
"forwarded_client_ip_headers": [],
|
||||
"contact_info": "",
|
||||
"doc_url": "",
|
||||
"home_content": "",
|
||||
|
||||
@@ -47,18 +47,7 @@ func ProvideRouter(
|
||||
|
||||
r := gin.New()
|
||||
r.Use(middleware2.Recovery())
|
||||
if len(cfg.Server.TrustedProxies) > 0 {
|
||||
if err := r.SetTrustedProxies(cfg.Server.TrustedProxies); err != nil {
|
||||
log.Printf("Failed to set trusted proxies: %v", err)
|
||||
}
|
||||
} else {
|
||||
if err := r.SetTrustedProxies(nil); err != nil {
|
||||
log.Printf("Failed to disable trusted proxies: %v", err)
|
||||
}
|
||||
if cfg.Server.Mode == "release" {
|
||||
log.Printf("Warning: server.trusted_proxies is empty in release mode; client IP trust chain is disabled")
|
||||
}
|
||||
}
|
||||
configureTrustedProxies(r, cfg.Server)
|
||||
|
||||
// Wire up websearch Manager builder so it initializes on startup and rebuilds on config save.
|
||||
settingService.SetWebSearchManagerBuilder(context.Background(), func(cfg *service.WebSearchEmulationConfig, proxyURLs map[int64]string) {
|
||||
@@ -99,6 +88,25 @@ func ProvideRouter(
|
||||
return SetupRouter(r, handlers, jwtAuth, adminAuth, apiKeyAuth, auditLog, stepUpAuth, apiKeyService, subscriptionService, opsService, settingService, cfg, redisClient)
|
||||
}
|
||||
|
||||
func configureTrustedProxies(r *gin.Engine, cfg config.ServerConfig) {
|
||||
if cfg.TrustedProxiesConfigured {
|
||||
if err := r.SetTrustedProxies(cfg.TrustedProxies); err != nil {
|
||||
log.Printf("Failed to set trusted proxies: %v", err)
|
||||
_ = r.SetTrustedProxies(nil)
|
||||
}
|
||||
if len(cfg.TrustedProxies) == 0 && cfg.Mode == "release" {
|
||||
log.Printf("Warning: server.trusted_proxies is explicitly empty; forwarded client IP trust is disabled")
|
||||
}
|
||||
} else {
|
||||
if err := r.SetTrustedProxies(nil); err != nil {
|
||||
log.Printf("Failed to disable trusted proxies: %v", err)
|
||||
}
|
||||
if cfg.Mode == "release" {
|
||||
log.Printf("Warning: server.trusted_proxies is not configured; disabling the forwarded-IP compatibility switch will use direct peer addresses only")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ProvideHTTPServer 提供 HTTP 服务器
|
||||
func ProvideHTTPServer(cfg *config.Config, router *gin.Engine) *http.Server {
|
||||
httpHandler := http.Handler(router)
|
||||
|
||||
@@ -55,6 +55,56 @@ func TestProvideHTTPServerEnablesBoundedH2C(t *testing.T) {
|
||||
require.True(t, srv.Protocols.HTTP1())
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxies(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
cfg config.ServerConfig
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "configured proxy resolves forwarded client",
|
||||
cfg: config.ServerConfig{
|
||||
TrustedProxies: []string{"9.9.9.9/32"},
|
||||
TrustedProxiesConfigured: true,
|
||||
},
|
||||
want: "1.2.3.4",
|
||||
},
|
||||
{
|
||||
name: "explicit empty list ignores forwarded client",
|
||||
cfg: config.ServerConfig{
|
||||
TrustedProxiesConfigured: true,
|
||||
},
|
||||
want: "9.9.9.9",
|
||||
},
|
||||
{
|
||||
name: "invalid proxy list fails closed",
|
||||
cfg: config.ServerConfig{
|
||||
TrustedProxies: []string{"not-a-cidr"},
|
||||
TrustedProxiesConfigured: true,
|
||||
},
|
||||
want: "9.9.9.9",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
configureTrustedProxies(r, tc.cfg)
|
||||
r.GET("/t", func(c *gin.Context) { c.String(http.StatusOK, c.ClientIP()) })
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, tc.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPServerRejectsOversizedHTTP1Header(t *testing.T) {
|
||||
r := gin.New()
|
||||
r.GET("/", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
|
||||
@@ -863,7 +863,7 @@ func TestRequireGroupAssignmentMarksUngroupedKeyBusinessLimited(t *testing.T) {
|
||||
require.Equal(t, service.OpsClientBusinessLimitedReasonAPIKeyGroupUnassigned, businessLimitedReason)
|
||||
}
|
||||
|
||||
func TestAPIKeyAuthIPRestrictionDoesNotTrustForwardedClientIPByDefault(t *testing.T) {
|
||||
func TestAPIKeyAuthIPRestrictionUsesTrustedPathWhenSwitchDisabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
user := &service.User{
|
||||
@@ -893,6 +893,7 @@ func TestAPIKeyAuthIPRestrictionDoesNotTrustForwardedClientIPByDefault(t *testin
|
||||
}
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(false)
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies(nil))
|
||||
@@ -1003,7 +1004,7 @@ func TestAPIKeyAuthIPRestrictionUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
}
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(true)
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(false)
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"}))
|
||||
@@ -1054,7 +1055,7 @@ func TestAPIKeyAuthIPRestrictionUsesForwardedClientIPInDenialWhenTrusted(t *test
|
||||
}
|
||||
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(true)
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(false)
|
||||
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies([]string{"9.9.9.9"}))
|
||||
|
||||
@@ -13,14 +13,16 @@ import (
|
||||
// SessionBindingContext 全局中间件:将请求的客户端 IP 与 User-Agent 注入
|
||||
// request context,供 token 签发路径(登录 / 刷新 / OAuth 回调)读取并写入会话绑定,
|
||||
// 同时作为审计日志、会话绑定校验的统一客户端 IP 来源。
|
||||
// IP 取值与 API Key IP 限制共用 Gin trusted_proxies 解析链;旧设置开关
|
||||
// 仅为配置兼容保留,不能单独使直连请求的转发头变为可信。
|
||||
// IP 取值与 API Key IP 限制共用转发 IP 开关:开启时旧版原始转发头逻辑
|
||||
// 接管解析,关闭时使用 Gin 的 server.trusted_proxies 可信代理链。
|
||||
func SessionBindingContext(cfg *config.Config) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
forwardedIPSettings := cfg.ForwardedClientIPSettings()
|
||||
ip.SetForwardedIPSettings(c, forwardedIPSettings.TrustForwardedIP, forwardedIPSettings.Headers)
|
||||
userAgent := normalizePersistentText(c.Request.UserAgent(), maxPersistentUserAgentBytes)
|
||||
c.Request.Header.Set("User-Agent", userAgent)
|
||||
binding := &service.SessionBinding{
|
||||
IP: ip.GetSecurityClientIP(c, cfg.TrustForwardedIPForAPIKeyACL()),
|
||||
IP: ip.GetSecurityClientIP(c, forwardedIPSettings.TrustForwardedIP),
|
||||
UserAgent: userAgent,
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithSessionBinding(c.Request.Context(), binding))
|
||||
@@ -29,8 +31,7 @@ func SessionBindingContext(cfg *config.Config) gin.HandlerFunc {
|
||||
}
|
||||
|
||||
// requestSessionBinding 返回当前请求的会话指纹,优先取 SessionBindingContext
|
||||
// 注入的解析结果(保证与 token 签发路径取值一致);注入缺失时按 trusted_proxies
|
||||
// 链回退兜底(等价于开关关闭时的行为)。
|
||||
// 注入的解析结果(保证与 token 签发路径取值一致);注入缺失时使用安全回退。
|
||||
func requestSessionBinding(c *gin.Context) *service.SessionBinding {
|
||||
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil {
|
||||
return binding
|
||||
@@ -42,7 +43,7 @@ func requestSessionBinding(c *gin.Context) *service.SessionBinding {
|
||||
}
|
||||
|
||||
// SecurityClientIP 返回当前请求用于安全敏感记录(审计日志等)的客户端 IP。
|
||||
// 与会话绑定、API Key IP 限制共用同一套「信任反代传递的客户端 IP」开关语义。
|
||||
// 与会话绑定、API Key IP 限制共用同一套客户端 IP 来源。
|
||||
func SecurityClientIP(c *gin.Context) string {
|
||||
if binding := service.SessionBindingFromContext(c.Request.Context()); binding != nil &&
|
||||
strings.TrimSpace(binding.IP) != "" {
|
||||
|
||||
@@ -8,29 +8,32 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSessionBindingContextDoesNotTrustHeadersWithoutTrustedProxy(t *testing.T) {
|
||||
func TestSessionBindingContextFollowsForwardedIPSwitch(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
trustForwarded bool
|
||||
trustedProxies []string
|
||||
wantIP string
|
||||
}{
|
||||
{name: "trust disabled records proxy address", trustForwarded: false, wantIP: "127.0.0.1"},
|
||||
{name: "legacy trust toggle cannot bypass trusted proxies", trustForwarded: true, wantIP: "127.0.0.1"},
|
||||
{name: "enabled switch takes over raw headers", trustForwarded: true, wantIP: "1.2.3.4"},
|
||||
{name: "disabled switch ignores untrusted headers", trustForwarded: false, wantIP: "127.0.0.1"},
|
||||
{name: "disabled switch uses configured Gin proxy", trustForwarded: false, trustedProxies: []string{"127.0.0.1"}, wantIP: "1.2.3.4"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(tc.trustForwarded)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
require.NoError(t, r.SetTrustedProxies(tc.trustedProxies))
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
binding := service.SessionBindingFromContext(c.Request.Context())
|
||||
@@ -53,6 +56,39 @@ func TestSessionBindingContextDoesNotTrustHeadersWithoutTrustedProxy(t *testing.
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBindingContextSnapshotsForwardedModeAndHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Initial-IP"})
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.Use(SessionBindingContext(cfg))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
binding := service.SessionBindingFromContext(c.Request.Context())
|
||||
require.NotNil(t, binding)
|
||||
require.Equal(t, "1.2.3.4", binding.IP)
|
||||
|
||||
cfg.SetForwardedClientIPSettings(false, []string{"X-Changed-IP"})
|
||||
require.Equal(t, "1.2.3.4", ip.GetSecurityClientIP(c, false))
|
||||
c.Status(200)
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Initial-IP", "1.2.3.4")
|
||||
req.Header.Set("X-Changed-IP", "4.4.4.4")
|
||||
req.Header.Set("X-Real-IP", "8.8.8.8")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Changed-IP"}, runtimeSettings.Headers)
|
||||
}
|
||||
|
||||
func TestSessionBindingContextBoundsPersistedUserAgent(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
r := gin.New()
|
||||
|
||||
@@ -56,7 +56,7 @@ func SetupRouter(
|
||||
// 应用中间件
|
||||
r.Use(middleware2.RequestLogger())
|
||||
// 将客户端 IP + UA 注入 request context,供 token 签发/会话绑定/审计日志统一读取。
|
||||
// IP 取值与 API Key IP 限制共用 server.trusted_proxies 信任链。
|
||||
// 解析模式按请求快照:兼容开关开启时信任原始转发头,关闭时使用 server.trusted_proxies。
|
||||
r.Use(middleware2.SessionBindingContext(cfg))
|
||||
r.Use(middleware2.Logger())
|
||||
r.Use(middleware2.CORS(cfg.CORS))
|
||||
|
||||
@@ -579,6 +579,12 @@ func registerBackupRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAut
|
||||
backup.PUT("/s3-config", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.UpdateS3Config)
|
||||
backup.POST("/s3-config/test", h.Admin.Backup.TestS3Connection)
|
||||
|
||||
// 异步生图对象存储配置(与备份共用 S3 客户端,可直接复用备份凭证)
|
||||
backup.GET("/image-storage", h.Admin.Backup.GetImageStorageConfig)
|
||||
// 同 S3 配置:改写对象存储目标可将生成内容导向外部账号——要求 step-up 2FA
|
||||
backup.PUT("/image-storage", gin.HandlerFunc(stepUpAuth), h.Admin.Backup.UpdateImageStorageConfig)
|
||||
backup.POST("/image-storage/test", h.Admin.Backup.TestImageStorageConnection)
|
||||
|
||||
// 定时备份配置
|
||||
backup.GET("/schedule", h.Admin.Backup.GetSchedule)
|
||||
backup.PUT("/schedule", h.Admin.Backup.UpdateSchedule)
|
||||
|
||||
@@ -43,6 +43,16 @@ func RegisterGatewayRoutes(
|
||||
isOpenAIGatewayPlatform := func(c *gin.Context) bool {
|
||||
return getGroupPlatform(c) == service.PlatformOpenAI
|
||||
}
|
||||
countTokensHandler := func(c *gin.Context) {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI:
|
||||
h.OpenAIGateway.CountTokens(c)
|
||||
case service.PlatformGrok:
|
||||
h.OpenAIGateway.GrokCountTokens(c)
|
||||
default:
|
||||
h.Gateway.CountTokens(c)
|
||||
}
|
||||
}
|
||||
modelsHandler := func(c *gin.Context) {
|
||||
if isOpenAIGatewayPlatform(c) && c.Query("client_version") != "" {
|
||||
h.OpenAIGateway.CodexModels(c)
|
||||
@@ -139,26 +149,9 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.Messages(c)
|
||||
})
|
||||
// /v1/messages/count_tokens: OpenAI uses Anthropic-compat bridge; other
|
||||
// OpenAI-compatible platforms keep the prior unsupported response.
|
||||
gateway.POST("/messages/count_tokens", func(c *gin.Context) {
|
||||
if isOpenAIGatewayPlatform(c) {
|
||||
h.OpenAIGateway.CountTokens(c)
|
||||
return
|
||||
}
|
||||
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{
|
||||
"type": "not_found_error",
|
||||
"message": "Token counting is not supported for this platform",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
h.Gateway.CountTokens(c)
|
||||
})
|
||||
// /v1/messages/count_tokens: OpenAI bridges upstream, Grok estimates
|
||||
// locally, and Anthropic-compatible platforms retain their existing path.
|
||||
gateway.POST("/messages/count_tokens", countTokensHandler)
|
||||
// Codex CLI / Codex app refresh their model picker from the provider's
|
||||
// /models endpoint with a client_version query and expect the ChatGPT
|
||||
// Codex manifest format; other clients keep the OpenAI-style list.
|
||||
@@ -256,6 +249,7 @@ func RegisterGatewayRoutes(
|
||||
h.OpenAIGateway.ResponsesWebSocket(c)
|
||||
})
|
||||
r.GET("/models", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, modelsHandler)
|
||||
r.POST("/messages/count_tokens", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, countTokensHandler)
|
||||
codexDirect := r.Group("/backend-api/codex")
|
||||
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic)
|
||||
{
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -15,6 +16,10 @@ import (
|
||||
)
|
||||
|
||||
func newGatewayRoutesTestRouter(platform ...string) *gin.Engine {
|
||||
return newGatewayRoutesTestRouterWithConfig(&config.Config{}, platform...)
|
||||
}
|
||||
|
||||
func newGatewayRoutesTestRouterWithConfig(cfg *config.Config, platform ...string) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
|
||||
@@ -42,7 +47,7 @@ func newGatewayRoutesTestRouter(platform ...string) *gin.Engine {
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
&config.Config{},
|
||||
cfg,
|
||||
)
|
||||
|
||||
return router
|
||||
@@ -224,13 +229,22 @@ func TestGatewayRoutesGrokAllowsCLICompatibilityEntrypoints(t *testing.T) {
|
||||
require.NotContains(t, w.Body.String(), "not supported for Grok groups")
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", strings.NewReader(`{"model":"grok","messages":[{"role":"user","content":"hi"}]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
countTokensRouter := newGatewayRoutesTestRouterWithConfig(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxBodySize: 1024 * 1024},
|
||||
}, service.PlatformGrok)
|
||||
for _, path := range []string{"/v1/messages/count_tokens", "/messages/count_tokens"} {
|
||||
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok","messages":[{"role":"user","content":"hi"}]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusNotFound, w.Code)
|
||||
require.Contains(t, w.Body.String(), "Token counting is not supported for this platform")
|
||||
countTokensRouter.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code, "path=%s", path)
|
||||
var response struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response), "path=%s", path)
|
||||
require.Positive(t, response.InputTokens, "path=%s", path)
|
||||
}
|
||||
|
||||
for _, path := range []string{
|
||||
"/v1/responses",
|
||||
|
||||
@@ -711,7 +711,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
return s.sendErrorAndEnd(c, "Grok token provider not configured")
|
||||
}
|
||||
var err error
|
||||
authToken, err = s.grokTokenProvider.GetAccessToken(ctx, account)
|
||||
// 手动测试不走生产调度资格门:关闭调度、限流/过载/临时冷却中的账号
|
||||
// 也应能被管理员探测(#4598),与 Codex/OpenAI 测试行为一致。
|
||||
authToken, err = s.grokTokenProvider.GetAccessTokenForManualTest(ctx, account)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get Grok access token: %s", err.Error()))
|
||||
}
|
||||
|
||||
@@ -207,6 +207,7 @@ type UsageInfo struct {
|
||||
GrokLastQuotaProbeAt string `json:"grok_last_quota_probe_at,omitempty"`
|
||||
GrokLastHeadersSeenAt string `json:"grok_last_headers_seen_at,omitempty"`
|
||||
GrokLastStatusCode int `json:"grok_last_status_code,omitempty"`
|
||||
GrokFreeTokenLimit int64 `json:"grok_free_token_limit,omitempty"`
|
||||
GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"`
|
||||
GrokLocalUsage24h *WindowStats `json:"grok_local_usage_24h,omitempty"`
|
||||
GrokLocalUsage7d *WindowStats `json:"grok_local_usage_7d,omitempty"`
|
||||
|
||||
@@ -576,6 +576,14 @@ var proxyQualityTargets = []proxyQualityTarget{
|
||||
http.StatusOK: {},
|
||||
},
|
||||
},
|
||||
{
|
||||
Target: "grok",
|
||||
URL: "https://api.x.ai/v1/models",
|
||||
Method: http.MethodGet,
|
||||
AllowedStatuses: map[int]struct{}{
|
||||
http.StatusUnauthorized: {},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
const (
|
||||
|
||||
@@ -93,3 +93,44 @@ func TestRunProxyQualityTarget_AllowedStatusPassForUnauthorized(t *testing.T) {
|
||||
require.Equal(t, http.StatusUnauthorized, item.HTTPStatus)
|
||||
require.Contains(t, item.Message, "目标可达")
|
||||
}
|
||||
|
||||
func TestProxyQualityTargets_IncludesGrok(t *testing.T) {
|
||||
var grokTarget *proxyQualityTarget
|
||||
for i := range proxyQualityTargets {
|
||||
if proxyQualityTargets[i].Target == "grok" {
|
||||
grokTarget = &proxyQualityTargets[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
require.NotNil(t, grokTarget)
|
||||
require.Equal(t, "https://api.x.ai/v1/models", grokTarget.URL)
|
||||
require.Equal(t, http.MethodGet, grokTarget.Method)
|
||||
require.Contains(t, grokTarget.AllowedStatuses, http.StatusUnauthorized)
|
||||
}
|
||||
|
||||
func TestRunProxyQualityTarget_GrokUnauthorizedPasses(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
t.Errorf("unexpected method: %s", r.Method)
|
||||
}
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
target := proxyQualityTarget{
|
||||
Target: "grok",
|
||||
URL: server.URL,
|
||||
Method: http.MethodGet,
|
||||
AllowedStatuses: map[int]struct{}{
|
||||
http.StatusUnauthorized: {},
|
||||
},
|
||||
}
|
||||
|
||||
item := runProxyQualityTarget(context.Background(), server.Client(), target)
|
||||
require.Equal(t, "grok", item.Target)
|
||||
require.Equal(t, "pass", item.Status)
|
||||
require.Equal(t, http.StatusUnauthorized, item.HTTPStatus)
|
||||
require.Contains(t, item.Message, "目标可达")
|
||||
}
|
||||
|
||||
@@ -165,6 +165,8 @@ const (
|
||||
|
||||
// API Key IP 访问控制设置
|
||||
SettingKeyAPIKeyACLTrustForwardedIP = "api_key_acl_trust_forwarded_ip" // API Key IP 白/黑名单是否信任转发 IP
|
||||
SettingKeyForwardedClientIPHeaders = "forwarded_client_ip_headers" // 自定义 CDN 客户端 IP 请求头(JSON 数组)
|
||||
settingKeyForwardedClientIPModeV2 = "forwarded_client_ip_mode_v2_migrated"
|
||||
|
||||
// TOTP 双因素认证设置
|
||||
SettingKeyTotpEnabled = "totp_enabled" // 是否启用 TOTP 2FA 功能
|
||||
|
||||
@@ -270,7 +270,7 @@ func (s *GatewayService) handleCCBufferedFromAnthropic(
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
if event.Delta != nil && event.Delta.StopReason != "" && finalResp != nil {
|
||||
finalResp.StopReason = event.Delta.StopReason
|
||||
finalResp.StopReason = apicompat.AnthropicStopReasonPtr(event.Delta.StopReason)
|
||||
}
|
||||
}
|
||||
if event.Type == "content_block_start" && event.ContentBlock != nil && finalResp != nil {
|
||||
|
||||
@@ -285,7 +285,7 @@ func (s *GatewayService) handleResponsesBufferedStreamingResponse(
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
if event.Delta != nil && event.Delta.StopReason != "" && finalResp != nil {
|
||||
finalResp.StopReason = event.Delta.StopReason
|
||||
finalResp.StopReason = apicompat.AnthropicStopReasonPtr(event.Delta.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -173,7 +173,23 @@ func openAIStreamEventIsTerminal(data string) bool {
|
||||
if trimmed == "[DONE]" {
|
||||
return true
|
||||
}
|
||||
switch gjson.Get(trimmed, "type").String() {
|
||||
return openAIStreamEventTypeIsTerminal(gjson.Get(trimmed, "type").String())
|
||||
}
|
||||
|
||||
// openAIStreamEventIsTerminalWithType 复用已提取的 type,避免 SSE 热路径重复扫描 JSON。
|
||||
func openAIStreamEventIsTerminalWithType(data, eventType string) bool {
|
||||
trimmed := strings.TrimSpace(data)
|
||||
if trimmed == "" {
|
||||
return false
|
||||
}
|
||||
if trimmed == "[DONE]" {
|
||||
return true
|
||||
}
|
||||
return openAIStreamEventTypeIsTerminal(eventType)
|
||||
}
|
||||
|
||||
func openAIStreamEventTypeIsTerminal(eventType string) bool {
|
||||
switch eventType {
|
||||
case "response.completed", "response.done", "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
|
||||
return true
|
||||
default:
|
||||
|
||||
@@ -556,12 +556,13 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom
|
||||
if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{
|
||||
Type: "message_start",
|
||||
Message: &apicompat.AnthropicResponse{
|
||||
ID: messageID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Model: originalModel,
|
||||
Content: []apicompat.AnthropicContentBlock{},
|
||||
Usage: apicompat.AnthropicUsage{},
|
||||
ID: messageID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Model: originalModel,
|
||||
Content: []apicompat.AnthropicContentBlock{},
|
||||
StopReason: nil, // JSON null
|
||||
Usage: apicompat.AnthropicUsage{},
|
||||
},
|
||||
}) {
|
||||
return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil
|
||||
|
||||
@@ -21,8 +21,9 @@ func NewGrokQuotaFetcher() *GrokQuotaFetcher {
|
||||
func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo {
|
||||
now := time.Now()
|
||||
usage := &UsageInfo{
|
||||
Source: "passive",
|
||||
UpdatedAt: &now,
|
||||
Source: "passive",
|
||||
UpdatedAt: &now,
|
||||
GrokFreeTokenLimit: xai.GrokFreeRolling24hTokenLimit,
|
||||
}
|
||||
if account == nil {
|
||||
usage.ErrorCode = "quota_unknown"
|
||||
|
||||
@@ -19,6 +19,7 @@ func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) {
|
||||
|
||||
usage := NewGrokQuotaFetcher().BuildUsageInfo(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth})
|
||||
require.Equal(t, "passive", usage.Source)
|
||||
require.Equal(t, xai.GrokFreeRolling24hTokenLimit, usage.GrokFreeTokenLimit)
|
||||
require.Equal(t, "quota_unknown", usage.ErrorCode)
|
||||
require.Contains(t, usage.Error, "unknown until billing is probed")
|
||||
}
|
||||
|
||||
@@ -23,6 +23,8 @@ const (
|
||||
grokQuotaProbeInput = "hi"
|
||||
grokQuotaDefaultModel = grokDefaultResponsesModel
|
||||
grokBillingExtraKey = "grok_billing_snapshot"
|
||||
grokBillingMaxAttempts = 2
|
||||
grokBillingRetryDelay = 100 * time.Millisecond
|
||||
)
|
||||
|
||||
type GrokQuotaProbeResult struct {
|
||||
@@ -356,33 +358,63 @@ func (s *GrokQuotaService) fetchBilling(
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil)
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err)
|
||||
}
|
||||
xai.ApplyCLIBillingHeaders(req, token)
|
||||
// billing 探测与真实转发保持同一套账号级请求头覆写。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2))
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
for attempt := 0; attempt < grokBillingMaxAttempts; attempt++ {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil)
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err)
|
||||
}
|
||||
xai.ApplyCLIBillingHeaders(req, token)
|
||||
// billing 探测与真实转发保持同一套账号级请求头覆写。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
resp, requestErr := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2))
|
||||
|
||||
bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
return nil, resp.StatusCode, nil
|
||||
statusCode := 0
|
||||
var bodyBytes []byte
|
||||
if requestErr == nil {
|
||||
statusCode = resp.StatusCode
|
||||
bodyBytes, _ = io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
|
||||
shouldRetry := requestErr != nil || isRetryableGrokBillingStatus(statusCode)
|
||||
if shouldRetry && attempt+1 < grokBillingMaxAttempts {
|
||||
timer := time.NewTimer(grokBillingRetryDelay)
|
||||
select {
|
||||
case <-timer.C:
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", ctx.Err())
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if requestErr != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", requestErr)
|
||||
}
|
||||
if statusCode == http.StatusTooManyRequests {
|
||||
return nil, statusCode, nil
|
||||
}
|
||||
if statusCode >= 400 {
|
||||
bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240)
|
||||
slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", statusCode, "body", bodyText)
|
||||
return nil, statusCode, infraerrors.Newf(mapUpstreamStatus(statusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", statusCode, bodyText)
|
||||
}
|
||||
payload, err := xai.ParseBillingPayload(bodyBytes)
|
||||
if err != nil {
|
||||
return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err)
|
||||
}
|
||||
return xai.BuildBillingSummary(payload.Config), statusCode, nil
|
||||
}
|
||||
if resp.StatusCode >= 400 {
|
||||
bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240)
|
||||
slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", resp.StatusCode, "body", bodyText)
|
||||
return nil, resp.StatusCode, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", resp.StatusCode, bodyText)
|
||||
return nil, 0, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed")
|
||||
}
|
||||
|
||||
func isRetryableGrokBillingStatus(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
payload, err := xai.ParseBillingPayload(bodyBytes)
|
||||
if err != nil {
|
||||
return nil, resp.StatusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err)
|
||||
}
|
||||
return xai.BuildBillingSummary(payload.Config), resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func mergeGrokBillingProbeErrors(weeklyStatus, monthlyStatus int, weeklyErr, monthlyErr error) error {
|
||||
|
||||
@@ -5,6 +5,7 @@ package service
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -127,6 +128,44 @@ type grokHybridUpstream struct {
|
||||
billingHeaders http.Header
|
||||
}
|
||||
|
||||
type grokQuotaUpstreamStep struct {
|
||||
status int
|
||||
body string
|
||||
err error
|
||||
}
|
||||
|
||||
type grokQuotaSequenceUpstream struct {
|
||||
httpUpstreamRecorder
|
||||
mu sync.Mutex
|
||||
steps []grokQuotaUpstreamStep
|
||||
requests []*http.Request
|
||||
}
|
||||
|
||||
func (u *grokQuotaSequenceUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
u.requests = append(u.requests, req)
|
||||
index := len(u.requests) - 1
|
||||
if index >= len(u.steps) {
|
||||
return nil, errors.New("unexpected upstream request")
|
||||
}
|
||||
step := u.steps[index]
|
||||
if step.err != nil {
|
||||
return nil, step.err
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: step.status,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(step.body)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *grokQuotaSequenceUpstream) snapshotRequests() []*http.Request {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]*http.Request(nil), u.requests...)
|
||||
}
|
||||
|
||||
func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
var body []byte
|
||||
if req != nil && req.Body != nil {
|
||||
@@ -228,6 +267,139 @@ func healthyGrokQuotaOAuthAccount(id int64) *Account {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceFetchBillingRetries502ThenSucceeds(t *testing.T) {
|
||||
account := healthyGrokQuotaOAuthAccount(401)
|
||||
upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{
|
||||
{status: http.StatusBadGateway, body: `The origin web server returned an invalid or incomplete response to Cloudflare.`},
|
||||
{status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"},"creditUsagePercent":12}}`},
|
||||
}}
|
||||
svc := &GrokQuotaService{httpUpstream: upstream}
|
||||
|
||||
summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, status)
|
||||
require.NotNil(t, summary)
|
||||
require.NotNil(t, summary.UsagePercent)
|
||||
require.Equal(t, 12.0, *summary.UsagePercent)
|
||||
requests := upstream.snapshotRequests()
|
||||
require.Len(t, requests, 2)
|
||||
require.Equal(t, http.MethodGet, requests[0].Method)
|
||||
require.Equal(t, "/v1/billing", requests[0].URL.Path)
|
||||
require.Equal(t, "format=credits", requests[0].URL.RawQuery)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceFetchBillingRetriesTransportErrorThenSucceeds(t *testing.T) {
|
||||
account := healthyGrokQuotaOAuthAccount(402)
|
||||
upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{
|
||||
{err: errors.New("temporary transport failure")},
|
||||
{status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"},"creditUsagePercent":8}}`},
|
||||
}}
|
||||
svc := &GrokQuotaService{httpUpstream: upstream}
|
||||
|
||||
summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, status)
|
||||
require.NotNil(t, summary)
|
||||
require.Len(t, upstream.snapshotRequests(), 2)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceFetchBillingStopsAfterSingleTransientRetry(t *testing.T) {
|
||||
account := healthyGrokQuotaOAuthAccount(403)
|
||||
upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{
|
||||
{status: http.StatusBadGateway, body: `cloudflare failure`},
|
||||
{status: http.StatusBadGateway, body: `cloudflare failure`},
|
||||
{status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"}}}`},
|
||||
}}
|
||||
svc := &GrokQuotaService{httpUpstream: upstream}
|
||||
|
||||
summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, summary)
|
||||
require.Equal(t, http.StatusBadGateway, status)
|
||||
require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err))
|
||||
require.Contains(t, infraerrors.Message(err), "billing returned 502: cloudflare failure")
|
||||
require.Len(t, upstream.snapshotRequests(), 2)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceFetchBillingDoesNotRetryNonTransientStatuses(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
status int
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "unauthorized", status: http.StatusUnauthorized, wantErr: true},
|
||||
{name: "forbidden", status: http.StatusForbidden, wantErr: true},
|
||||
{name: "rate limited", status: http.StatusTooManyRequests, wantErr: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := healthyGrokQuotaOAuthAccount(404)
|
||||
upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{
|
||||
{status: tt.status, body: `{"error":{"message":"rejected"}}`},
|
||||
{status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"}}}`},
|
||||
}}
|
||||
svc := &GrokQuotaService{httpUpstream: upstream}
|
||||
|
||||
summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true)
|
||||
|
||||
if tt.wantErr {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Nil(t, summary)
|
||||
require.Equal(t, tt.status, status)
|
||||
require.Len(t, upstream.snapshotRequests(), 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsRetryableGrokBillingStatus(t *testing.T) {
|
||||
tests := []struct {
|
||||
status int
|
||||
want bool
|
||||
}{
|
||||
{status: http.StatusBadGateway, want: true},
|
||||
{status: http.StatusServiceUnavailable, want: true},
|
||||
{status: http.StatusGatewayTimeout, want: true},
|
||||
{status: http.StatusUnauthorized, want: false},
|
||||
{status: http.StatusForbidden, want: false},
|
||||
{status: http.StatusTooManyRequests, want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(strconv.Itoa(tt.status), func(t *testing.T) {
|
||||
require.Equal(t, tt.want, isRetryableGrokBillingStatus(tt.status))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageDoesNotRetryResponsesPost(t *testing.T) {
|
||||
account := healthyGrokQuotaOAuthAccount(405)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{
|
||||
{status: http.StatusBadGateway, body: `cloudflare failure`},
|
||||
{status: http.StatusOK, body: `{"id":"unexpected_retry"}`},
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), account.ID)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err))
|
||||
requests := upstream.snapshotRequests()
|
||||
require.Len(t, requests, 1)
|
||||
require.Equal(t, http.MethodPost, requests[0].Method)
|
||||
require.Equal(t, "/v1/responses", requests[0].URL.Path)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -167,6 +167,77 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
// GetAccessTokenForManualTest returns an access token for an admin-initiated
|
||||
// "test connection" probe. Unlike GetAccessToken it does not apply the
|
||||
// request-path scheduling eligibility gate (manual Schedulable switch,
|
||||
// rate-limit / overload / temp-unschedulable cooldowns): a manual test exists
|
||||
// precisely to check accounts in those states, matching how Codex/OpenAI
|
||||
// account tests read credentials regardless of scheduling state (#4598).
|
||||
//
|
||||
// Credential integrity still applies: the configured-proxy-missing check, the
|
||||
// shared refresh lock protocol, and the refresh API's own account re-read.
|
||||
// Credential rotation for non-active (disabled/error) accounts remains
|
||||
// blocked inside RefreshIfNeeded; their still-valid tokens are probed as-is.
|
||||
func (p *GrokTokenProvider) GetAccessTokenForManualTest(ctx context.Context, account *Account) (string, error) {
|
||||
if account == nil {
|
||||
return "", errors.New("account is nil")
|
||||
}
|
||||
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth {
|
||||
return "", errors.New("not a grok oauth account")
|
||||
}
|
||||
if account.ProxyID != nil && account.Proxy == nil {
|
||||
return "", errGrokOAuthConfiguredProxyMiss
|
||||
}
|
||||
if strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
|
||||
return "", errGrokOAuthRefreshTokenMissing
|
||||
}
|
||||
|
||||
accessToken := strings.TrimSpace(account.GetGrokAccessToken())
|
||||
expiresAt := account.GetCredentialAsTime("expires_at")
|
||||
tokenValid := accessToken != "" && expiresAt != nil && time.Now().Before(*expiresAt)
|
||||
if accessToken != "" && expiresAt != nil && time.Until(*expiresAt) > grokTokenRefreshSkew {
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
if p.refreshAPI == nil || p.executor == nil {
|
||||
if tokenValid {
|
||||
return accessToken, nil
|
||||
}
|
||||
return "", errGrokOAuthRefreshNotConfigured
|
||||
}
|
||||
|
||||
// Deliberately not marked as a request-path refresh: the request path
|
||||
// re-applies scheduling eligibility inside RefreshIfNeeded, which is
|
||||
// exactly what a manual test must bypass.
|
||||
refreshCtx, cancel := context.WithTimeout(ctx, grokRequestRefreshTimeout)
|
||||
defer cancel()
|
||||
result, err := p.refreshAPI.RefreshIfNeeded(refreshCtx, account, p.executor, grokTokenRefreshSkew)
|
||||
if err != nil {
|
||||
if tokenValid {
|
||||
return accessToken, nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if result != nil && result.LockHeld {
|
||||
if tokenValid {
|
||||
return accessToken, nil
|
||||
}
|
||||
return "", errors.New("token refresh is already in progress on another worker; retry in a few seconds")
|
||||
}
|
||||
if result != nil && result.Account != nil {
|
||||
account = result.Account
|
||||
}
|
||||
|
||||
accessToken = strings.TrimSpace(account.GetGrokAccessToken())
|
||||
if accessToken == "" {
|
||||
return "", errGrokOAuthAccessTokenMissing
|
||||
}
|
||||
if latestExpiry := account.GetCredentialAsTime("expires_at"); latestExpiry != nil && !time.Now().Before(*latestExpiry) {
|
||||
return "", errGrokOAuthAccessTokenExpired
|
||||
}
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
func (p *GrokTokenProvider) waitForRefreshedToken(ctx context.Context, account *Account, cacheKey string) (string, error) {
|
||||
waitCtx, cancel := context.WithTimeout(ctx, grokRefreshLockWaitTimeout)
|
||||
defer cancel()
|
||||
|
||||
@@ -287,6 +287,150 @@ func TestGrokTokenProviderRejectsStaleDBTokenWithoutExpiry(t *testing.T) {
|
||||
require.Empty(t, token)
|
||||
}
|
||||
|
||||
// TestGrokTokenProviderManualTestBypassesSchedulingGate reproduces #4598:
|
||||
// admins must be able to run "test connection" against accounts that the
|
||||
// scheduler currently excludes (manual switch off, rate limited, overloaded,
|
||||
// temporarily cooled down). The production request path keeps rejecting them.
|
||||
func TestGrokTokenProviderManualTestBypassesSchedulingGate(t *testing.T) {
|
||||
future := time.Now().Add(time.Hour)
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(*Account)
|
||||
}{
|
||||
{name: "not schedulable", mutate: func(account *Account) { account.Schedulable = false }},
|
||||
{name: "temporarily unschedulable", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }},
|
||||
{name: "rate limited", mutate: func(account *Account) { account.RateLimitResetAt = &future }},
|
||||
{name: "overloaded", mutate: func(account *Account) { account.OverloadUntil = &future }},
|
||||
{name: "disabled by error", mutate: func(account *Account) { account.Status = StatusError }},
|
||||
}
|
||||
|
||||
for index, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(int64(120 + index))
|
||||
account.Credentials["access_token"] = "still-valid-token"
|
||||
account.Credentials["expires_at"] = time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339)
|
||||
tt.mutate(account)
|
||||
provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, &grokTokenCacheForProviderTest{})
|
||||
|
||||
// Production request path keeps excluding this account.
|
||||
_, requestErr := provider.GetAccessToken(context.Background(), account)
|
||||
require.ErrorIs(t, requestErr, errOAuthRefreshAccountStateChanged)
|
||||
|
||||
// Manual test path returns the valid credential for probing.
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "still-valid-token", token)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderManualTestRefreshesExpiredTokenWhileUnschedulable(t *testing.T) {
|
||||
t.Setenv(xai.EnvBaseURL, xai.DefaultCLIBaseURL)
|
||||
|
||||
expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
|
||||
account := &Account{
|
||||
ID: 130,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: false,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "expired-access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": expiredAt,
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
"client_id": "client-id",
|
||||
},
|
||||
}
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{130: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: true}
|
||||
oauthSvc := NewGrokOAuthService(nil, &grokOAuthClientStub{
|
||||
refreshResponse: &xai.TokenResponse{
|
||||
AccessToken: "manual-test-refreshed-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresIn: 3600,
|
||||
},
|
||||
})
|
||||
defer oauthSvc.Stop()
|
||||
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
|
||||
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "manual-test-refreshed-token", token)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderManualTestFallsBackToValidTokenOnRefreshFailure(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 131,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: false,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "near-expiry-token",
|
||||
"refresh_token": "refresh-token",
|
||||
// Inside the refresh window but not expired yet.
|
||||
"expires_at": time.Now().Add(10 * time.Minute).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{131: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: true}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
|
||||
err: errors.New("upstream refresh unavailable"),
|
||||
})
|
||||
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "near-expiry-token", token)
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderManualTestReportsRefreshFailureWhenTokenExpired(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(132)
|
||||
account.Schedulable = false
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: true}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
|
||||
err: errors.New("invalid_client: client credentials rejected"),
|
||||
})
|
||||
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, token)
|
||||
require.Contains(t, err.Error(), "invalid_client")
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderManualTestLockHeldWithExpiredTokenReturnsSpecificError(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(133)
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: false}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
|
||||
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, token)
|
||||
require.Contains(t, err.Error(), "refresh is already in progress")
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderManualTestRequiresRefreshToken(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(134)
|
||||
delete(account.Credentials, "refresh_token")
|
||||
provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, &grokTokenCacheForProviderTest{})
|
||||
|
||||
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
|
||||
require.ErrorIs(t, err, errGrokOAuthRefreshTokenMissing)
|
||||
require.Empty(t, token)
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderRejectsIneligibleSelectedAccountBeforeWarmCache(t *testing.T) {
|
||||
future := time.Now().Add(time.Hour)
|
||||
tests := []struct {
|
||||
|
||||
@@ -0,0 +1,342 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const settingKeyImageStorageConfig = "image_storage_config"
|
||||
|
||||
// ErrImageStorageIncomplete 表示开关已打开但凭证不全,无法启用异步生图。
|
||||
var ErrImageStorageIncomplete = errors.New("image storage is enabled but bucket/access_key_id/secret_access_key are incomplete")
|
||||
|
||||
// ImageStorageFactory 由 repository 层提供,把配置变成一个可用的对象存储实现。
|
||||
// 与 BackupObjectStoreFactory 同样的注入方式,避免 service 反向依赖 repository。
|
||||
type ImageStorageFactory func(ctx context.Context, cfg *config.ImageStorageConfig) (ImageStorage, error)
|
||||
|
||||
// ImageStorageSettings 是后台可编辑的异步生图对象存储配置。
|
||||
//
|
||||
// ReuseBackupS3 为真时不保存自己的凭证,直接借用数据库备份已配置的 S3 端点与密钥,
|
||||
// 只用自己的 Bucket/Prefix 区分对象;这样"数据走 backups/、图片走 images/"无需重复配置。
|
||||
type ImageStorageSettings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
ReuseBackupS3 bool `json:"reuse_backup_s3"`
|
||||
|
||||
Bucket string `json:"bucket"` // 留空且复用备份时,沿用备份桶
|
||||
Prefix string `json:"prefix"`
|
||||
PublicBaseURL string `json:"public_base_url"`
|
||||
PresignExpiry int `json:"presign_expiry_hours"`
|
||||
MaxDownloadBytes int64 `json:"max_download_bytes"`
|
||||
|
||||
// 以下仅在 ReuseBackupS3 为假时使用
|
||||
Endpoint string `json:"endpoint"`
|
||||
Region string `json:"region"`
|
||||
AccessKeyID string `json:"access_key_id"`
|
||||
SecretAccessKey string `json:"secret_access_key,omitempty"` //nolint:revive // field name follows AWS convention
|
||||
ForcePathStyle bool `json:"force_path_style"`
|
||||
}
|
||||
|
||||
// ImageStorageSettingService 读写后台设置,并把结果解析成一个可直接使用的 uploader。
|
||||
//
|
||||
// 解析结果带缓存:网关每次请求都要判断功能是否开启,不能每次都查库。保存设置时调用
|
||||
// Invalidate 清缓存,下一次请求即重建客户端——这是"后台开关立即生效、无需重启"的实现。
|
||||
type ImageStorageSettingService struct {
|
||||
settingRepo SettingRepository
|
||||
encryptor SecretEncryptor
|
||||
backup *BackupService
|
||||
factory ImageStorageFactory
|
||||
|
||||
// fallback 是 config.yaml 里的配置。后台从未保存过设置时沿用它,
|
||||
// 保证升级前已用配置文件开启该功能的部署不被打断。
|
||||
fallback config.ImageStorageConfig
|
||||
|
||||
mu sync.Mutex
|
||||
resolved bool
|
||||
uploader *ImageResultUploader
|
||||
enabled bool
|
||||
}
|
||||
|
||||
func NewImageStorageSettingService(
|
||||
settingRepo SettingRepository,
|
||||
encryptor SecretEncryptor,
|
||||
backup *BackupService,
|
||||
factory ImageStorageFactory,
|
||||
fallback config.ImageStorageConfig,
|
||||
) *ImageStorageSettingService {
|
||||
return &ImageStorageSettingService{
|
||||
settingRepo: settingRepo,
|
||||
encryptor: encryptor,
|
||||
backup: backup,
|
||||
factory: factory,
|
||||
fallback: fallback,
|
||||
}
|
||||
}
|
||||
|
||||
// Resolver 返回可注入 ImageTaskService 的解析函数。
|
||||
func (s *ImageStorageSettingService) Resolver() ImageStorageResolver {
|
||||
return func() (*ImageResultUploader, bool) {
|
||||
return s.resolve()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ImageStorageSettingService) resolve() (*ImageResultUploader, bool) {
|
||||
if s == nil {
|
||||
return nil, false
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.resolved {
|
||||
return s.uploader, s.enabled
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
s.resolved = true
|
||||
s.uploader, s.enabled = nil, false
|
||||
|
||||
cfg, err := s.effectiveConfig(ctx)
|
||||
if err != nil {
|
||||
logger.L().Warn("image_storage.settings_load_failed; async image tasks stay disabled", zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
if !cfg.Enabled {
|
||||
return nil, false
|
||||
}
|
||||
if !cfg.IsConfigured() {
|
||||
logger.L().Warn("image_storage is enabled but not fully configured; async image tasks are disabled",
|
||||
zap.Strings("missing_keys", cfg.MissingCredentialKeys()))
|
||||
return nil, false
|
||||
}
|
||||
|
||||
storage, err := s.factory(ctx, cfg)
|
||||
if err != nil {
|
||||
logger.L().Error("image_storage.client_build_failed; async image tasks stay disabled", zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
s.uploader = NewImageResultUploader(storage, cfg.Prefix, cfg.MaxDownloadByte, nil)
|
||||
s.enabled = true
|
||||
return s.uploader, true
|
||||
}
|
||||
|
||||
// Invalidate 丢弃缓存,使下一次请求按最新设置重新解析。
|
||||
func (s *ImageStorageSettingService) Invalidate() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.resolved = false
|
||||
s.uploader = nil
|
||||
s.enabled = false
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Get 返回后台设置(SecretAccessKey 已脱敏)。从未保存过时返回 config.yaml 的等价值。
|
||||
func (s *ImageStorageSettingService) Get(ctx context.Context) (*ImageStorageSettings, error) {
|
||||
settings, err := s.load(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if settings == nil {
|
||||
settings = settingsFromConfig(s.fallback)
|
||||
}
|
||||
settings.SecretAccessKey = ""
|
||||
return settings, nil
|
||||
}
|
||||
|
||||
// SecretConfigured 供前端展示"已配置"占位符。
|
||||
func (s *ImageStorageSettingService) SecretConfigured(ctx context.Context) bool {
|
||||
settings, err := s.load(ctx)
|
||||
if err != nil || settings == nil {
|
||||
return s.fallback.SecretAccessKey != ""
|
||||
}
|
||||
if settings.ReuseBackupS3 {
|
||||
cfg, err := s.backupCredentials(ctx)
|
||||
return err == nil && cfg != nil && cfg.SecretAccessKey != ""
|
||||
}
|
||||
return settings.SecretAccessKey != ""
|
||||
}
|
||||
|
||||
// Update 保存设置并立即生效。SecretAccessKey 留空表示沿用已保存的值。
|
||||
func (s *ImageStorageSettingService) Update(ctx context.Context, in ImageStorageSettings) (*ImageStorageSettings, error) {
|
||||
normalizeImageStorageSettings(&in)
|
||||
|
||||
if in.ReuseBackupS3 {
|
||||
// 复用备份凭证时不落自己的密钥,避免同一份密钥在库里存两份。
|
||||
in.Endpoint, in.Region, in.AccessKeyID, in.SecretAccessKey = "", "", "", ""
|
||||
in.ForcePathStyle = false
|
||||
} else if in.SecretAccessKey == "" {
|
||||
if old, err := s.load(ctx); err == nil && old != nil {
|
||||
in.SecretAccessKey = old.SecretAccessKey
|
||||
}
|
||||
} else {
|
||||
encrypted, err := s.encryptor.Encrypt(in.SecretAccessKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encrypt secret: %w", err)
|
||||
}
|
||||
in.SecretAccessKey = encrypted
|
||||
}
|
||||
|
||||
data, err := json.Marshal(in)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal image storage settings: %w", err)
|
||||
}
|
||||
if err := s.settingRepo.Set(ctx, settingKeyImageStorageConfig, string(data)); err != nil {
|
||||
return nil, fmt.Errorf("save image storage settings: %w", err)
|
||||
}
|
||||
s.Invalidate()
|
||||
|
||||
in.SecretAccessKey = ""
|
||||
return &in, nil
|
||||
}
|
||||
|
||||
// TestConnection 用给定设置试建一次客户端,用于后台的"测试连接"按钮。
|
||||
// 与 Update 一样支持留空 SecretAccessKey 表示沿用已保存的值。
|
||||
func (s *ImageStorageSettingService) TestConnection(ctx context.Context, in ImageStorageSettings) error {
|
||||
normalizeImageStorageSettings(&in)
|
||||
if !in.ReuseBackupS3 && in.SecretAccessKey == "" {
|
||||
old, err := s.load(ctx)
|
||||
if err == nil && old != nil {
|
||||
in.SecretAccessKey = old.SecretAccessKey
|
||||
}
|
||||
}
|
||||
cfg, err := s.toImageStorageConfig(ctx, &in)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !cfg.IsConfigured() {
|
||||
return ErrImageStorageIncomplete
|
||||
}
|
||||
if _, err := s.factory(ctx, cfg); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// effectiveConfig 把后台设置(或 config.yaml 回落)解析成运行时配置。
|
||||
func (s *ImageStorageSettingService) effectiveConfig(ctx context.Context) (*config.ImageStorageConfig, error) {
|
||||
settings, err := s.load(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if settings == nil {
|
||||
fallback := s.fallback
|
||||
return &fallback, nil
|
||||
}
|
||||
return s.toImageStorageConfig(ctx, settings)
|
||||
}
|
||||
|
||||
func (s *ImageStorageSettingService) toImageStorageConfig(ctx context.Context, in *ImageStorageSettings) (*config.ImageStorageConfig, error) {
|
||||
cfg := &config.ImageStorageConfig{
|
||||
Enabled: in.Enabled,
|
||||
Bucket: in.Bucket,
|
||||
Prefix: in.Prefix,
|
||||
PublicBaseURL: in.PublicBaseURL,
|
||||
PresignExpiry: in.PresignExpiry,
|
||||
MaxDownloadByte: in.MaxDownloadBytes,
|
||||
Endpoint: in.Endpoint,
|
||||
Region: in.Region,
|
||||
AccessKeyID: in.AccessKeyID,
|
||||
SecretAccessKey: in.SecretAccessKey,
|
||||
ForcePathStyle: in.ForcePathStyle,
|
||||
}
|
||||
|
||||
if in.ReuseBackupS3 {
|
||||
backupCfg, err := s.backupCredentials(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if backupCfg == nil {
|
||||
return nil, errors.New("image storage is set to reuse the backup S3 configuration, but no backup S3 configuration exists")
|
||||
}
|
||||
cfg.Endpoint = backupCfg.Endpoint
|
||||
cfg.Region = backupCfg.Region
|
||||
cfg.AccessKeyID = backupCfg.AccessKeyID
|
||||
cfg.SecretAccessKey = backupCfg.SecretAccessKey
|
||||
cfg.ForcePathStyle = backupCfg.ForcePathStyle
|
||||
if cfg.Bucket == "" {
|
||||
cfg.Bucket = backupCfg.Bucket
|
||||
}
|
||||
} else if cfg.SecretAccessKey != "" {
|
||||
decrypted, err := s.encryptor.Decrypt(cfg.SecretAccessKey)
|
||||
if err != nil {
|
||||
// 兼容未加密的旧数据,与备份配置的处理保持一致。
|
||||
logger.L().Warn("image_storage secret decrypt failed; treating the stored value as plaintext", zap.Error(err))
|
||||
} else {
|
||||
cfg.SecretAccessKey = decrypted
|
||||
}
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// backupCredentials 取备份已配置的 S3 凭证(已解密)。
|
||||
func (s *ImageStorageSettingService) backupCredentials(ctx context.Context) (*BackupS3Config, error) {
|
||||
if s.backup == nil {
|
||||
return nil, errors.New("backup service is unavailable")
|
||||
}
|
||||
return s.backup.loadS3Config(ctx)
|
||||
}
|
||||
|
||||
// load 读出后台设置;从未保存过时返回 nil。
|
||||
func (s *ImageStorageSettingService) load(ctx context.Context) (*ImageStorageSettings, error) {
|
||||
if s.settingRepo == nil {
|
||||
return nil, nil //nolint:nilnil // no repository means no stored settings
|
||||
}
|
||||
raw, err := s.settingRepo.GetValue(ctx, settingKeyImageStorageConfig)
|
||||
if err != nil || strings.TrimSpace(raw) == "" {
|
||||
return nil, nil //nolint:nilnil // never configured is a valid state
|
||||
}
|
||||
var settings ImageStorageSettings
|
||||
if err := json.Unmarshal([]byte(raw), &settings); err != nil {
|
||||
return nil, fmt.Errorf("parse image storage settings: %w", err)
|
||||
}
|
||||
return &settings, nil
|
||||
}
|
||||
|
||||
func settingsFromConfig(cfg config.ImageStorageConfig) *ImageStorageSettings {
|
||||
return &ImageStorageSettings{
|
||||
Enabled: cfg.Enabled,
|
||||
Bucket: cfg.Bucket,
|
||||
Prefix: cfg.Prefix,
|
||||
PublicBaseURL: cfg.PublicBaseURL,
|
||||
PresignExpiry: cfg.PresignExpiry,
|
||||
MaxDownloadBytes: cfg.MaxDownloadByte,
|
||||
Endpoint: cfg.Endpoint,
|
||||
Region: cfg.Region,
|
||||
AccessKeyID: cfg.AccessKeyID,
|
||||
SecretAccessKey: cfg.SecretAccessKey,
|
||||
ForcePathStyle: cfg.ForcePathStyle,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeImageStorageSettings(in *ImageStorageSettings) {
|
||||
in.Bucket = strings.TrimSpace(in.Bucket)
|
||||
in.Endpoint = strings.TrimSpace(in.Endpoint)
|
||||
in.Region = strings.TrimSpace(in.Region)
|
||||
in.AccessKeyID = strings.TrimSpace(in.AccessKeyID)
|
||||
in.SecretAccessKey = strings.TrimSpace(in.SecretAccessKey)
|
||||
in.PublicBaseURL = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(in.PublicBaseURL), "/"))
|
||||
|
||||
in.Prefix = strings.TrimSpace(in.Prefix)
|
||||
if in.Prefix == "" {
|
||||
in.Prefix = "images/"
|
||||
}
|
||||
if !strings.HasSuffix(in.Prefix, "/") {
|
||||
in.Prefix += "/"
|
||||
}
|
||||
if in.Region == "" {
|
||||
in.Region = "auto"
|
||||
}
|
||||
if in.PresignExpiry <= 0 {
|
||||
in.PresignExpiry = 24
|
||||
}
|
||||
if in.MaxDownloadBytes <= 0 {
|
||||
in.MaxDownloadBytes = defaultImageMaxDownloadBytes
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type stubSettingRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func newStubSettingRepo() *stubSettingRepo {
|
||||
return &stubSettingRepo{values: map[string]string{}}
|
||||
}
|
||||
|
||||
func (r *stubSettingRepo) Get(context.Context, string) (*Setting, error) { return nil, nil }
|
||||
func (r *stubSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.values[key], nil
|
||||
}
|
||||
|
||||
func (r *stubSettingRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
func (r *stubSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *stubSettingRepo) SetMultiple(context.Context, map[string]string) error { return nil }
|
||||
func (r *stubSettingRepo) GetAll(context.Context) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *stubSettingRepo) Delete(context.Context, string) error { return nil }
|
||||
|
||||
// reversibleEncryptor stands in for AES: prefixed so a test can tell ciphertext
|
||||
// from plaintext, and so decrypting a plaintext value fails like the real one.
|
||||
type reversibleEncryptor struct{}
|
||||
|
||||
func (reversibleEncryptor) Encrypt(plaintext string) (string, error) {
|
||||
return "enc:" + plaintext, nil
|
||||
}
|
||||
|
||||
func (reversibleEncryptor) Decrypt(ciphertext string) (string, error) {
|
||||
rest, ok := strings.CutPrefix(ciphertext, "enc:")
|
||||
if !ok {
|
||||
return "", errors.New("not encrypted")
|
||||
}
|
||||
return rest, nil
|
||||
}
|
||||
|
||||
type recordingStorage struct{ saved []string }
|
||||
|
||||
func (s *recordingStorage) Save(_ context.Context, key, _ string, _ []byte) (string, error) {
|
||||
s.saved = append(s.saved, key)
|
||||
return "https://cdn.example.com/" + key, nil
|
||||
}
|
||||
|
||||
func newImageStorageFixture(t *testing.T, fallback config.ImageStorageConfig) (*ImageStorageSettingService, *stubSettingRepo, *[]config.ImageStorageConfig) {
|
||||
t.Helper()
|
||||
repo := newStubSettingRepo()
|
||||
encryptor := reversibleEncryptor{}
|
||||
backup := NewBackupService(repo, &config.Config{}, encryptor, nil, nil)
|
||||
|
||||
var built []config.ImageStorageConfig
|
||||
factory := func(_ context.Context, cfg *config.ImageStorageConfig) (ImageStorage, error) {
|
||||
built = append(built, *cfg)
|
||||
return &recordingStorage{}, nil
|
||||
}
|
||||
return NewImageStorageSettingService(repo, encryptor, backup, factory, fallback), repo, &built
|
||||
}
|
||||
|
||||
func seedBackupS3(t *testing.T, repo *stubSettingRepo, cfg BackupS3Config) {
|
||||
t.Helper()
|
||||
cfg.SecretAccessKey = "enc:" + cfg.SecretAccessKey
|
||||
data, err := json.Marshal(cfg)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, repo.Set(context.Background(), settingKeyBackupS3Config, string(data)))
|
||||
}
|
||||
|
||||
// The admin switch must take effect without a restart: that is the entire point
|
||||
// of moving image_storage out of config.yaml (#4542).
|
||||
func TestImageStorageSettingsToggleTakesEffectWithoutRestart(t *testing.T) {
|
||||
svc, repo, built := newImageStorageFixture(t, config.ImageStorageConfig{})
|
||||
ctx := context.Background()
|
||||
seedBackupS3(t, repo, BackupS3Config{
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com", Region: "auto",
|
||||
Bucket: "backup-bucket", AccessKeyID: "ak", SecretAccessKey: "sk",
|
||||
Prefix: "backups/",
|
||||
})
|
||||
|
||||
uploader, enabled := svc.resolve()
|
||||
require.False(t, enabled, "disabled until an admin turns it on")
|
||||
require.Nil(t, uploader)
|
||||
|
||||
_, err := svc.Update(ctx, ImageStorageSettings{Enabled: true, ReuseBackupS3: true})
|
||||
require.NoError(t, err)
|
||||
|
||||
uploader, enabled = svc.resolve()
|
||||
require.True(t, enabled, "saving the setting must enable the feature immediately")
|
||||
require.NotNil(t, uploader)
|
||||
|
||||
_, err = svc.Update(ctx, ImageStorageSettings{Enabled: false, ReuseBackupS3: true})
|
||||
require.NoError(t, err)
|
||||
_, enabled = svc.resolve()
|
||||
require.False(t, enabled, "turning it back off must also apply immediately")
|
||||
|
||||
require.Len(t, *built, 1, "the S3 client is built only when the feature is on")
|
||||
}
|
||||
|
||||
func TestImageStorageSettingsReuseBackupCredentials(t *testing.T) {
|
||||
svc, repo, built := newImageStorageFixture(t, config.ImageStorageConfig{})
|
||||
ctx := context.Background()
|
||||
seedBackupS3(t, repo, BackupS3Config{
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com", Region: "wnam",
|
||||
Bucket: "backup-bucket", AccessKeyID: "backup-ak", SecretAccessKey: "backup-sk",
|
||||
Prefix: "backups/", ForcePathStyle: true,
|
||||
})
|
||||
|
||||
_, err := svc.Update(ctx, ImageStorageSettings{Enabled: true, ReuseBackupS3: true, Prefix: "images"})
|
||||
require.NoError(t, err)
|
||||
_, enabled := svc.resolve()
|
||||
require.True(t, enabled)
|
||||
|
||||
require.Len(t, *built, 1)
|
||||
got := (*built)[0]
|
||||
require.Equal(t, "https://acct.r2.cloudflarestorage.com", got.Endpoint)
|
||||
require.Equal(t, "wnam", got.Region)
|
||||
require.Equal(t, "backup-ak", got.AccessKeyID)
|
||||
require.Equal(t, "backup-sk", got.SecretAccessKey, "the backup secret must be decrypted before use")
|
||||
require.True(t, got.ForcePathStyle)
|
||||
require.Equal(t, "backup-bucket", got.Bucket, "an empty bucket falls back to the backup bucket")
|
||||
require.Equal(t, "images/", got.Prefix, "images stay under their own prefix so they never collide with backups/")
|
||||
|
||||
// Reusing must not duplicate the secret into a second row.
|
||||
raw, err := repo.GetValue(ctx, settingKeyImageStorageConfig)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, raw, "backup-sk")
|
||||
require.NotContains(t, raw, "enc:")
|
||||
}
|
||||
|
||||
func TestImageStorageSettingsOwnCredentialsAreEncryptedAndMasked(t *testing.T) {
|
||||
svc, repo, built := newImageStorageFixture(t, config.ImageStorageConfig{})
|
||||
ctx := context.Background()
|
||||
|
||||
saved, err := svc.Update(ctx, ImageStorageSettings{
|
||||
Enabled: true, Bucket: "my-images",
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com",
|
||||
AccessKeyID: "ak", SecretAccessKey: "super-secret",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, saved.SecretAccessKey, "the response must never echo the secret back")
|
||||
|
||||
raw, err := repo.GetValue(ctx, settingKeyImageStorageConfig)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, raw, `"secret_access_key":"super-secret"`, "the secret must be encrypted at rest")
|
||||
require.Contains(t, raw, "enc:super-secret")
|
||||
|
||||
fetched, err := svc.Get(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, fetched.SecretAccessKey)
|
||||
require.True(t, svc.SecretConfigured(ctx))
|
||||
|
||||
_, enabled := svc.resolve()
|
||||
require.True(t, enabled)
|
||||
require.Equal(t, "super-secret", (*built)[0].SecretAccessKey, "the stored secret must be decrypted before use")
|
||||
|
||||
// An update that omits the secret keeps the stored one rather than wiping it.
|
||||
_, err = svc.Update(ctx, ImageStorageSettings{
|
||||
Enabled: true, Bucket: "my-images",
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com", AccessKeyID: "ak",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
svc.resolve()
|
||||
require.Equal(t, "super-secret", (*built)[1].SecretAccessKey)
|
||||
}
|
||||
|
||||
func TestImageStorageSettingsIncompleteStaysDisabled(t *testing.T) {
|
||||
svc, _, built := newImageStorageFixture(t, config.ImageStorageConfig{})
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := svc.Update(ctx, ImageStorageSettings{Enabled: true, Bucket: "my-images"})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, enabled := svc.resolve()
|
||||
require.False(t, enabled, "missing credentials must not enable the feature")
|
||||
require.Empty(t, *built, "no client is built from an incomplete configuration")
|
||||
}
|
||||
|
||||
// Deployments that already enabled the feature through config.yaml must keep
|
||||
// working after the setting moves into the database.
|
||||
func TestImageStorageSettingsFallBackToConfigFile(t *testing.T) {
|
||||
svc, _, built := newImageStorageFixture(t, config.ImageStorageConfig{
|
||||
Enabled: true, Endpoint: "https://acct.r2.cloudflarestorage.com", Region: "auto",
|
||||
Bucket: "yaml-bucket", AccessKeyID: "yaml-ak", SecretAccessKey: "yaml-sk",
|
||||
Prefix: "images/", MaxDownloadByte: 1024,
|
||||
})
|
||||
|
||||
_, enabled := svc.resolve()
|
||||
require.True(t, enabled, "config.yaml still enables the feature when nothing is stored yet")
|
||||
require.Equal(t, "yaml-bucket", (*built)[0].Bucket)
|
||||
|
||||
fetched, err := svc.Get(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.True(t, fetched.Enabled)
|
||||
require.Equal(t, "yaml-bucket", fetched.Bucket)
|
||||
require.Empty(t, fetched.SecretAccessKey)
|
||||
}
|
||||
@@ -69,10 +69,17 @@ type ImageTaskStore interface {
|
||||
Get(ctx context.Context, id string) (*ImageTaskRecord, error)
|
||||
}
|
||||
|
||||
// ImageStorageResolver reports the currently effective object-storage binding.
|
||||
// It exists so the async image feature can be switched on and off from the admin
|
||||
// UI without a restart: the wiring below is fixed at startup, but the answer to
|
||||
// "is object storage configured right now" is re-read (and cached) per call.
|
||||
type ImageStorageResolver func() (uploader *ImageResultUploader, enabled bool)
|
||||
|
||||
type ImageTaskService struct {
|
||||
store ImageTaskStore
|
||||
uploader *ImageResultUploader
|
||||
enabled bool
|
||||
resolve ImageStorageResolver
|
||||
ttl time.Duration
|
||||
executionTimeout time.Duration
|
||||
}
|
||||
@@ -100,10 +107,40 @@ func NewImageTaskServiceWithUploader(store ImageTaskStore, uploader *ImageResult
|
||||
return s
|
||||
}
|
||||
|
||||
// NewImageTaskServiceWithResolver 构造一个由 resolver 决定启用状态的服务:
|
||||
// 开关与凭证来自后台设置,保存后立即生效,无需重启。
|
||||
func NewImageTaskServiceWithResolver(store ImageTaskStore, resolve ImageStorageResolver, ttl, executionTimeout time.Duration) *ImageTaskService {
|
||||
s := NewImageTaskServiceWithOptions(store, ttl, executionTimeout)
|
||||
s.resolve = resolve
|
||||
return s
|
||||
}
|
||||
|
||||
// current 返回当前生效的 uploader 与启用状态。
|
||||
// 注入了 resolver 时以 resolver 为准(后台设置可热切换),否则回落到构造时固定的值。
|
||||
func (s *ImageTaskService) current() (*ImageResultUploader, bool) {
|
||||
if s == nil {
|
||||
return nil, false
|
||||
}
|
||||
if s.resolve != nil {
|
||||
return s.resolve()
|
||||
}
|
||||
return s.uploader, s.enabled
|
||||
}
|
||||
|
||||
// Enabled 表示异步图片任务功能是否可用(总开关 + 凭证齐全)。
|
||||
// 关闭时 handler 直接返回 404,不创建任务、不写 Redis。
|
||||
func (s *ImageTaskService) Enabled() bool {
|
||||
return s != nil && s.enabled && s.store != nil
|
||||
if s == nil || s.store == nil {
|
||||
return false
|
||||
}
|
||||
_, enabled := s.current()
|
||||
return enabled
|
||||
}
|
||||
|
||||
// Pollable 表示已创建的任务能否被查询。
|
||||
// 比 Enabled 弱:只要 store 可用即可,从而在功能被关掉后仍能取回进行中的任务结果。
|
||||
func (s *ImageTaskService) Pollable() bool {
|
||||
return s != nil && s.store != nil
|
||||
}
|
||||
|
||||
func (s *ImageTaskService) ExecutionTimeout() time.Duration {
|
||||
@@ -154,8 +191,8 @@ func (s *ImageTaskService) Complete(ctx context.Context, id string, statusCode i
|
||||
if !json.Valid(result) {
|
||||
return s.Fail(ctx, id, http.StatusBadGateway, imageTaskErrorJSON("api_error", "upstream returned a non-JSON image response"))
|
||||
}
|
||||
if s.uploader != nil {
|
||||
rewritten, err := s.uploader.Rewrite(ctx, id, result)
|
||||
if uploader, _ := s.current(); uploader != nil {
|
||||
rewritten, err := uploader.Rewrite(ctx, id, result)
|
||||
if err != nil {
|
||||
// 转存失败不回退存 base64,避免大 blob 撑爆 Redis:直接把任务标记为失败。
|
||||
logger.L().Error("image_task.offload_failed", zap.String("task_id", id), zap.Error(err))
|
||||
|
||||
@@ -51,6 +51,31 @@ var (
|
||||
notificationEmailPlaceholderPattern = regexp.MustCompile(`{{\s*([a-zA-Z][a-zA-Z0-9_]*)\s*}}`)
|
||||
notificationEmailLocales = []string{notificationEmailDefaultLocale, notificationEmailLocaleChinese}
|
||||
notificationEmailCommonPlaceholders = []string{"site_name", "recipient_name", "recipient_email"}
|
||||
// Keep summary values separate so admins can rearrange or omit individual metrics in the template.
|
||||
notificationEmailOpsSummaryPlaceholders = []string{
|
||||
"report_summary_display",
|
||||
"report_total_requests",
|
||||
"report_success_count",
|
||||
"report_sla_error_count",
|
||||
"report_business_limited_count",
|
||||
"report_sla",
|
||||
"report_error_rate",
|
||||
"report_upstream_error_rate",
|
||||
"report_upstream_error_count_excl_429_529",
|
||||
"report_upstream_429_count",
|
||||
"report_upstream_529_count",
|
||||
"report_latency_p50",
|
||||
"report_latency_p99",
|
||||
"report_ttft_p50",
|
||||
"report_ttft_p99",
|
||||
"report_tokens",
|
||||
"report_qps_current",
|
||||
"report_qps_peak",
|
||||
"report_qps_avg",
|
||||
"report_tps_current",
|
||||
"report_tps_peak",
|
||||
"report_tps_avg",
|
||||
}
|
||||
)
|
||||
|
||||
type NotificationEmailService struct {
|
||||
@@ -504,6 +529,34 @@ func (s *NotificationEmailService) runtimeVariables(ctx context.Context, event,
|
||||
for key, value := range input.Variables {
|
||||
variables[key] = value
|
||||
}
|
||||
if event == NotificationEmailEventOpsScheduledReport {
|
||||
// Scheduled reports may be sent by integrations that only provide report_html.
|
||||
// Do not let preview sample values appear in a live email in that case.
|
||||
if _, ok := input.Variables["report_html"]; !ok {
|
||||
variables["report_html"] = ""
|
||||
}
|
||||
if _, ok := input.Variables["report_detail_display"]; !ok {
|
||||
// Keep legacy/custom templates useful when they only render report_html.
|
||||
variables["report_detail_display"] = "block"
|
||||
}
|
||||
hasSummaryValues := false
|
||||
for _, placeholder := range notificationEmailOpsSummaryPlaceholders {
|
||||
if _, ok := input.Variables[placeholder]; ok {
|
||||
if placeholder != "report_summary_display" {
|
||||
hasSummaryValues = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
variables[placeholder] = "-"
|
||||
}
|
||||
if _, ok := input.Variables["report_summary_display"]; !ok {
|
||||
if hasSummaryValues {
|
||||
variables["report_summary_display"] = "block"
|
||||
} else {
|
||||
variables["report_summary_display"] = "none"
|
||||
}
|
||||
}
|
||||
}
|
||||
variables["site_name"] = s.siteName(ctx)
|
||||
variables["recipient_email"] = input.RecipientEmail
|
||||
if strings.TrimSpace(input.RecipientName) != "" {
|
||||
@@ -844,7 +897,7 @@ func isSafeNotificationEmailURL(raw string) bool {
|
||||
|
||||
func notificationEmailSampleVariables(locale string) map[string]string {
|
||||
if normalizeNotificationLocale(locale) == notificationEmailLocaleChinese {
|
||||
return map[string]string{
|
||||
variables := map[string]string{
|
||||
"site_name": defaultSiteName,
|
||||
"recipient_name": "张三",
|
||||
"recipient_email": "user@example.com",
|
||||
@@ -885,12 +938,14 @@ func notificationEmailSampleVariables(locale string) map[string]string {
|
||||
"alert_description": "最近 10 分钟错误率超过阈值",
|
||||
"report_name": "日报",
|
||||
"report_type": "daily_summary",
|
||||
"report_start_time": "2026-05-19 12:00",
|
||||
"report_end_time": "2026-05-20 12:00",
|
||||
"report_html": "<h2>日报</h2><p>请求量:1024</p>",
|
||||
"report_start_time": "2026-07-18T01:00:26Z",
|
||||
"report_end_time": "2026-07-19T01:00:26Z",
|
||||
"report_html": "<h2>日报</h2><p>请求量:2,374</p>",
|
||||
}
|
||||
addNotificationEmailOpsSummarySampleVariables(variables)
|
||||
return variables
|
||||
}
|
||||
return map[string]string{
|
||||
variables := map[string]string{
|
||||
"site_name": defaultSiteName,
|
||||
"recipient_name": "Alex",
|
||||
"recipient_email": "user@example.com",
|
||||
@@ -931,10 +986,38 @@ func notificationEmailSampleVariables(locale string) map[string]string {
|
||||
"alert_description": "Error rate exceeded threshold in the last 10 minutes.",
|
||||
"report_name": "Daily summary",
|
||||
"report_type": "daily_summary",
|
||||
"report_start_time": "2026-05-19 12:00",
|
||||
"report_end_time": "2026-05-20 12:00",
|
||||
"report_html": "<h2>Daily summary</h2><p>Requests: 1024</p>",
|
||||
"report_start_time": "2026-07-18T01:00:26Z",
|
||||
"report_end_time": "2026-07-19T01:00:26Z",
|
||||
"report_html": "<h2>Daily summary</h2><p>Requests: 2,374</p>",
|
||||
}
|
||||
addNotificationEmailOpsSummarySampleVariables(variables)
|
||||
return variables
|
||||
}
|
||||
|
||||
func addNotificationEmailOpsSummarySampleVariables(variables map[string]string) {
|
||||
variables["report_summary_display"] = "block"
|
||||
variables["report_detail_display"] = "none"
|
||||
variables["report_total_requests"] = "2,374"
|
||||
variables["report_success_count"] = "1,451"
|
||||
variables["report_sla_error_count"] = "2"
|
||||
variables["report_business_limited_count"] = "921"
|
||||
variables["report_sla"] = "99.86%"
|
||||
variables["report_error_rate"] = "0.14%"
|
||||
variables["report_upstream_error_rate"] = "0.28%"
|
||||
variables["report_upstream_error_count_excl_429_529"] = "4"
|
||||
variables["report_upstream_429_count"] = "0"
|
||||
variables["report_upstream_529_count"] = "0"
|
||||
variables["report_latency_p50"] = "8,231 ms"
|
||||
variables["report_latency_p99"] = "151,260 ms"
|
||||
variables["report_ttft_p50"] = "1,674 ms"
|
||||
variables["report_ttft_p99"] = "11,222 ms"
|
||||
variables["report_tokens"] = "121,550,190"
|
||||
variables["report_qps_current"] = "0.0"
|
||||
variables["report_qps_peak"] = "1.2"
|
||||
variables["report_qps_avg"] = "0.0"
|
||||
variables["report_tps_current"] = "0.0"
|
||||
variables["report_tps_peak"] = "133421.2"
|
||||
variables["report_tps_avg"] = "1406.8"
|
||||
}
|
||||
|
||||
var notificationEmailEventOrder = []string{
|
||||
@@ -1061,8 +1144,13 @@ var notificationEmailEventDefinitions = map[string]NotificationEmailEventInfo{
|
||||
Description: "Sent to configured operations recipients for scheduled daily/weekly/error/account-health reports.",
|
||||
Category: "ops",
|
||||
Optional: false,
|
||||
Placeholders: append(append([]string{}, notificationEmailCommonPlaceholders...),
|
||||
"report_name", "report_type", "report_start_time", "report_end_time", "report_html"),
|
||||
Placeholders: append(
|
||||
append(
|
||||
append([]string{}, notificationEmailCommonPlaceholders...),
|
||||
"report_name", "report_type", "report_start_time", "report_end_time",
|
||||
),
|
||||
append(append([]string{}, notificationEmailOpsSummaryPlaceholders...), "report_detail_display", "report_html")...,
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -1294,11 +1382,11 @@ var notificationEmailOfficialTemplates = map[string]map[string]notificationEmail
|
||||
HTML: notificationEmailCard("#ef4444", "Cyber-security policy notice", `
|
||||
<p>Hello {{recipient_name}},</p>
|
||||
<p>Your request was blocked by the upstream provider's cyber-security policy.</p>
|
||||
<table style="width:100%;border-collapse:collapse;">
|
||||
<tr><td>Triggered at</td><td>{{triggered_at}}</td></tr>
|
||||
<tr><td>Model</td><td>{{model}}</td></tr>
|
||||
<tr><td>Group</td><td>{{group_name}}</td></tr>
|
||||
<tr><td>Upstream message</td><td>{{upstream_message}}</td></tr>
|
||||
<table style="width:100%;border-collapse:collapse;table-layout:fixed;">
|
||||
<tr><td style="width:128px;vertical-align:top;">Triggered at</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{triggered_at}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">Model</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{model}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">Group</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{group_name}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">Upstream message</td><td style="overflow-wrap:anywhere;word-break:break-all;white-space:pre-wrap;">{{upstream_message}}</td></tr>
|
||||
</table>
|
||||
<p>If you believe this is a mistake, try rephrasing your request, or apply for authorized security access.</p>`),
|
||||
},
|
||||
@@ -1307,11 +1395,11 @@ var notificationEmailOfficialTemplates = map[string]map[string]notificationEmail
|
||||
HTML: notificationEmailCard("#ef4444", "网络安全策略拦截提醒", `
|
||||
<p>{{recipient_name}},您好:</p>
|
||||
<p>您的请求被上游服务商的网络安全策略(cyber policy)拦截。</p>
|
||||
<table style="width:100%;border-collapse:collapse;">
|
||||
<tr><td>触发时间</td><td>{{triggered_at}}</td></tr>
|
||||
<tr><td>模型</td><td>{{model}}</td></tr>
|
||||
<tr><td>所属分组</td><td>{{group_name}}</td></tr>
|
||||
<tr><td>上游说明</td><td>{{upstream_message}}</td></tr>
|
||||
<table style="width:100%;border-collapse:collapse;table-layout:fixed;">
|
||||
<tr><td style="width:128px;vertical-align:top;">触发时间</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{triggered_at}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">模型</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{model}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">所属分组</td><td style="overflow-wrap:anywhere;word-break:break-word;">{{group_name}}</td></tr>
|
||||
<tr><td style="width:128px;vertical-align:top;">上游说明</td><td style="overflow-wrap:anywhere;word-break:break-all;white-space:pre-wrap;">{{upstream_message}}</td></tr>
|
||||
</table>
|
||||
<p>如认为系误判,可调整请求措辞后重试,或申请获得授权的安全访问权限。</p>`),
|
||||
},
|
||||
@@ -1341,23 +1429,207 @@ var notificationEmailOfficialTemplates = map[string]map[string]notificationEmail
|
||||
NotificationEmailEventOpsScheduledReport: {
|
||||
notificationEmailDefaultLocale: {
|
||||
Subject: "[Ops Report] {{report_name}}",
|
||||
HTML: notificationEmailCard("#0891b2", "Ops report", `
|
||||
<p><strong>Report</strong>: {{report_name}}</p>
|
||||
<p><strong>Type</strong>: {{report_type}}</p>
|
||||
<p><strong>Range</strong>: {{report_start_time}} - {{report_end_time}}</p>
|
||||
<div>{{report_html}}</div>`),
|
||||
HTML: notificationEmailOpsScheduledReportTemplate(notificationEmailDefaultLocale),
|
||||
},
|
||||
notificationEmailLocaleChinese: {
|
||||
Subject: "[运维报表] {{report_name}}",
|
||||
HTML: notificationEmailCard("#0891b2", "运维报表", `
|
||||
<p><strong>报表</strong>:{{report_name}}</p>
|
||||
<p><strong>类型</strong>:{{report_type}}</p>
|
||||
<p><strong>时间范围</strong>:{{report_start_time}} - {{report_end_time}}</p>
|
||||
<div>{{report_html}}</div>`),
|
||||
HTML: notificationEmailOpsScheduledReportTemplate(notificationEmailLocaleChinese),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
func notificationEmailOpsScheduledReportTemplate(locale string) string {
|
||||
if normalizeNotificationLocale(locale) == notificationEmailLocaleChinese {
|
||||
return `<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<style>
|
||||
body { margin: 0; padding: 24px 12px; background: #f4f6f8; color: #1f2937; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", "PingFang SC", "Microsoft YaHei", sans-serif; }
|
||||
.container { width: 100%; max-width: 680px; margin: 0 auto; background: #ffffff; border: 1px solid #dfe7ea; border-radius: 8px; overflow: hidden; }
|
||||
.header { padding: 28px 32px 24px; background: #0f766e; color: #ffffff; }
|
||||
.eyebrow { margin: 0 0 8px; color: #ccfbf1; font-size: 12px; font-weight: 700; letter-spacing: 0; text-transform: uppercase; }
|
||||
h1 { margin: 0; font-size: 26px; line-height: 1.3; }
|
||||
.header p { margin: 8px 0 0; color: #e6fffb; font-size: 14px; }
|
||||
.content { padding: 28px 32px 32px; }
|
||||
.meta { width: 100%; margin: 0 0 20px; border-collapse: collapse; background: #f8fafc; border: 1px solid #e2e8f0; }
|
||||
.meta td { padding: 10px 12px; border-bottom: 1px solid #e2e8f0; font-size: 13px; vertical-align: top; }
|
||||
.meta tr:last-child td { border-bottom: 0; }
|
||||
.meta-label { width: 112px; color: #64748b; font-weight: 600; }
|
||||
.section-title { margin: 28px 0 12px; color: #0f172a; font-size: 16px; line-height: 1.4; }
|
||||
.metric-grid { width: 100%; border-collapse: separate; border-spacing: 8px; margin: -8px; }
|
||||
.metric-cell { width: 50%; padding: 14px 16px; border: 1px solid #e2e8f0; background: #ffffff; vertical-align: top; }
|
||||
.metric-label { display: block; color: #64748b; font-size: 12px; line-height: 1.4; }
|
||||
.metric-value { display: block; margin-top: 6px; color: #0f172a; font-size: 20px; font-weight: 700; line-height: 1.2; }
|
||||
.metric-value.good { color: #15803d; }
|
||||
.metric-value.alert { color: #b91c1c; }
|
||||
.detail { width: 100%; border-collapse: collapse; }
|
||||
.detail td { padding: 10px 12px; border-bottom: 1px solid #e2e8f0; font-size: 13px; }
|
||||
.detail td:first-child { width: 56%; color: #475569; }
|
||||
.detail td:last-child { color: #0f172a; font-weight: 600; text-align: right; }
|
||||
.report-detail { margin-top: 28px; }
|
||||
.report-detail:empty { display: none; }
|
||||
.footer { padding: 18px 32px; background: #f8fafc; border-top: 1px solid #e2e8f0; color: #64748b; font-size: 12px; line-height: 1.6; }
|
||||
@media only screen and (max-width: 620px) {
|
||||
body { padding: 0; }
|
||||
.container { border: 0; border-radius: 0; }
|
||||
.header, .content, .footer { padding-left: 20px; padding-right: 20px; }
|
||||
.metric-grid, .metric-grid tbody, .metric-grid tr, .metric-cell { display: block; width: 100% !important; box-sizing: border-box; }
|
||||
.metric-cell { margin: 8px 0; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<div class="header">
|
||||
<p class="eyebrow">运维报表</p>
|
||||
<h1>{{report_name}}</h1>
|
||||
<p>{{site_name}} 的运行概览</p>
|
||||
</div>
|
||||
<div class="content">
|
||||
<table class="meta" role="presentation">
|
||||
<tr><td class="meta-label">报表</td><td>{{report_name}}</td></tr>
|
||||
<tr><td class="meta-label">类型</td><td>{{report_type}}</td></tr>
|
||||
<tr><td class="meta-label">统计周期</td><td>{{report_start_time}} 至 {{report_end_time}} (UTC)</td></tr>
|
||||
</table>
|
||||
|
||||
<div style="display: {{report_summary_display}};">
|
||||
<h2 class="section-title">请求概览</h2>
|
||||
<table class="metric-grid" role="presentation"><tr>
|
||||
<td class="metric-cell"><span class="metric-label">总请求数</span><span class="metric-value">{{report_total_requests}}</span></td>
|
||||
<td class="metric-cell"><span class="metric-label">成功请求</span><span class="metric-value good">{{report_success_count}}</span></td>
|
||||
</tr><tr>
|
||||
<td class="metric-cell"><span class="metric-label">SLA 错误</span><span class="metric-value alert">{{report_sla_error_count}}</span></td>
|
||||
<td class="metric-cell"><span class="metric-label">业务限流</span><span class="metric-value">{{report_business_limited_count}}</span></td>
|
||||
</tr></table>
|
||||
|
||||
<h2 class="section-title">可靠性</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>SLA</td><td>{{report_sla}}</td></tr>
|
||||
<tr><td>错误率</td><td>{{report_error_rate}}</td></tr>
|
||||
<tr><td>上游错误率(不含 429 / 529)</td><td>{{report_upstream_error_rate}}</td></tr>
|
||||
<tr><td>上游错误(不含 429 / 529)</td><td>{{report_upstream_error_count_excl_429_529}}</td></tr>
|
||||
<tr><td>上游 429 / 529</td><td>{{report_upstream_429_count}} / {{report_upstream_529_count}}</td></tr>
|
||||
</table>
|
||||
|
||||
<h2 class="section-title">延迟表现</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>请求延迟 p50 / p99</td><td>{{report_latency_p50}} / {{report_latency_p99}}</td></tr>
|
||||
<tr><td>首 Token 时间 p50 / p99</td><td>{{report_ttft_p50}} / {{report_ttft_p99}}</td></tr>
|
||||
</table>
|
||||
|
||||
<h2 class="section-title">吞吐量</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>Token 消耗</td><td>{{report_tokens}}</td></tr>
|
||||
<tr><td>QPS(当前 / 峰值 / 平均)</td><td>{{report_qps_current}} / {{report_qps_peak}} / {{report_qps_avg}}</td></tr>
|
||||
<tr><td>TPS(当前 / 峰值 / 平均)</td><td>{{report_tps_current}} / {{report_tps_peak}} / {{report_tps_avg}}</td></tr>
|
||||
</table>
|
||||
|
||||
</div>
|
||||
<div class="report-detail" style="display: {{report_detail_display}};">{{report_html}}</div>
|
||||
</div>
|
||||
<div class="footer">此邮件由 {{site_name}} 自动发送,请勿直接回复。</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>`
|
||||
}
|
||||
|
||||
return `<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<style>
|
||||
body { margin: 0; padding: 24px 12px; background: #f4f6f8; color: #1f2937; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
|
||||
.container { width: 100%; max-width: 680px; margin: 0 auto; background: #ffffff; border: 1px solid #dfe7ea; border-radius: 8px; overflow: hidden; }
|
||||
.header { padding: 28px 32px 24px; background: #0f766e; color: #ffffff; }
|
||||
.eyebrow { margin: 0 0 8px; color: #ccfbf1; font-size: 12px; font-weight: 700; letter-spacing: 0; text-transform: uppercase; }
|
||||
h1 { margin: 0; font-size: 26px; line-height: 1.3; }
|
||||
.header p { margin: 8px 0 0; color: #e6fffb; font-size: 14px; }
|
||||
.content { padding: 28px 32px 32px; }
|
||||
.meta { width: 100%; margin: 0 0 20px; border-collapse: collapse; background: #f8fafc; border: 1px solid #e2e8f0; }
|
||||
.meta td { padding: 10px 12px; border-bottom: 1px solid #e2e8f0; font-size: 13px; vertical-align: top; }
|
||||
.meta tr:last-child td { border-bottom: 0; }
|
||||
.meta-label { width: 112px; color: #64748b; font-weight: 600; }
|
||||
.section-title { margin: 28px 0 12px; color: #0f172a; font-size: 16px; line-height: 1.4; }
|
||||
.metric-grid { width: 100%; border-collapse: separate; border-spacing: 8px; margin: -8px; }
|
||||
.metric-cell { width: 50%; padding: 14px 16px; border: 1px solid #e2e8f0; background: #ffffff; vertical-align: top; }
|
||||
.metric-label { display: block; color: #64748b; font-size: 12px; line-height: 1.4; }
|
||||
.metric-value { display: block; margin-top: 6px; color: #0f172a; font-size: 20px; font-weight: 700; line-height: 1.2; }
|
||||
.metric-value.good { color: #15803d; }
|
||||
.metric-value.alert { color: #b91c1c; }
|
||||
.detail { width: 100%; border-collapse: collapse; }
|
||||
.detail td { padding: 10px 12px; border-bottom: 1px solid #e2e8f0; font-size: 13px; }
|
||||
.detail td:first-child { width: 56%; color: #475569; }
|
||||
.detail td:last-child { color: #0f172a; font-weight: 600; text-align: right; }
|
||||
.report-detail { margin-top: 28px; }
|
||||
.report-detail:empty { display: none; }
|
||||
.footer { padding: 18px 32px; background: #f8fafc; border-top: 1px solid #e2e8f0; color: #64748b; font-size: 12px; line-height: 1.6; }
|
||||
@media only screen and (max-width: 620px) {
|
||||
body { padding: 0; }
|
||||
.container { border: 0; border-radius: 0; }
|
||||
.header, .content, .footer { padding-left: 20px; padding-right: 20px; }
|
||||
.metric-grid, .metric-grid tbody, .metric-grid tr, .metric-cell { display: block; width: 100% !important; box-sizing: border-box; }
|
||||
.metric-cell { margin: 8px 0; }
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<div class="header">
|
||||
<p class="eyebrow">Operations report</p>
|
||||
<h1>{{report_name}}</h1>
|
||||
<p>{{site_name}} runtime overview</p>
|
||||
</div>
|
||||
<div class="content">
|
||||
<table class="meta" role="presentation">
|
||||
<tr><td class="meta-label">Report</td><td>{{report_name}}</td></tr>
|
||||
<tr><td class="meta-label">Type</td><td>{{report_type}}</td></tr>
|
||||
<tr><td class="meta-label">Reporting period</td><td>{{report_start_time}} to {{report_end_time}} (UTC)</td></tr>
|
||||
</table>
|
||||
|
||||
<div style="display: {{report_summary_display}};">
|
||||
<h2 class="section-title">Request Overview</h2>
|
||||
<table class="metric-grid" role="presentation"><tr>
|
||||
<td class="metric-cell"><span class="metric-label">Total Requests</span><span class="metric-value">{{report_total_requests}}</span></td>
|
||||
<td class="metric-cell"><span class="metric-label">Successful Requests</span><span class="metric-value good">{{report_success_count}}</span></td>
|
||||
</tr><tr>
|
||||
<td class="metric-cell"><span class="metric-label">SLA Errors</span><span class="metric-value alert">{{report_sla_error_count}}</span></td>
|
||||
<td class="metric-cell"><span class="metric-label">Business Limited</span><span class="metric-value">{{report_business_limited_count}}</span></td>
|
||||
</tr></table>
|
||||
|
||||
<h2 class="section-title">Reliability</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>SLA</td><td>{{report_sla}}</td></tr>
|
||||
<tr><td>Error Rate</td><td>{{report_error_rate}}</td></tr>
|
||||
<tr><td>Upstream Error Rate (excluding 429 / 529)</td><td>{{report_upstream_error_rate}}</td></tr>
|
||||
<tr><td>Upstream Errors (excluding 429 / 529)</td><td>{{report_upstream_error_count_excl_429_529}}</td></tr>
|
||||
<tr><td>Upstream 429 / 529</td><td>{{report_upstream_429_count}} / {{report_upstream_529_count}}</td></tr>
|
||||
</table>
|
||||
|
||||
<h2 class="section-title">Latency</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>Request Latency p50 / p99</td><td>{{report_latency_p50}} / {{report_latency_p99}}</td></tr>
|
||||
<tr><td>Time to First Token p50 / p99</td><td>{{report_ttft_p50}} / {{report_ttft_p99}}</td></tr>
|
||||
</table>
|
||||
|
||||
<h2 class="section-title">Throughput</h2>
|
||||
<table class="detail" role="presentation">
|
||||
<tr><td>Tokens Consumed</td><td>{{report_tokens}}</td></tr>
|
||||
<tr><td>QPS (current / peak / average)</td><td>{{report_qps_current}} / {{report_qps_peak}} / {{report_qps_avg}}</td></tr>
|
||||
<tr><td>TPS (current / peak / average)</td><td>{{report_tps_current}} / {{report_tps_peak}} / {{report_tps_avg}}</td></tr>
|
||||
</table>
|
||||
|
||||
</div>
|
||||
<div class="report-detail" style="display: {{report_detail_display}};">{{report_html}}</div>
|
||||
</div>
|
||||
<div class="footer">This email was sent automatically by {{site_name}}. Please do not reply directly.</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>`
|
||||
}
|
||||
|
||||
func notificationEmailCard(accent, title, content string) string {
|
||||
return `<!DOCTYPE html>
|
||||
<html>
|
||||
|
||||
@@ -158,6 +158,104 @@ func TestNotificationEmailAdditionalEventsAreListedAndPreviewable(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCyberPolicyNoticeTemplateWrapsLongUpstreamMessages(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc := NewNotificationEmailService(newNotificationEmailMemorySettingRepo(), nil)
|
||||
longMessage := strings.Repeat("0123456789abcdef", 256)
|
||||
|
||||
for _, locale := range []string{"en", "zh"} {
|
||||
preview, err := svc.PreviewTemplate(ctx, NotificationEmailPreviewInput{
|
||||
Event: NotificationEmailEventCyberPolicyNotice,
|
||||
Locale: locale,
|
||||
Variables: map[string]string{
|
||||
"upstream_message": longMessage,
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, preview.HTML, "table-layout:fixed")
|
||||
require.Contains(t, preview.HTML, "overflow-wrap:anywhere")
|
||||
require.Contains(t, preview.HTML, "word-break:break-all")
|
||||
require.NotContains(t, preview.HTML, "{{upstream_message}}")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsScheduledReportTemplateExposesEditableSummaryMetrics(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc := NewNotificationEmailService(newNotificationEmailMemorySettingRepo(), nil)
|
||||
|
||||
requiredPlaceholders := []string{
|
||||
"report_summary_display",
|
||||
"report_detail_display",
|
||||
"report_total_requests",
|
||||
"report_success_count",
|
||||
"report_sla_error_count",
|
||||
"report_business_limited_count",
|
||||
"report_sla",
|
||||
"report_error_rate",
|
||||
"report_upstream_error_rate",
|
||||
"report_upstream_error_count_excl_429_529",
|
||||
"report_upstream_429_count",
|
||||
"report_upstream_529_count",
|
||||
"report_latency_p50",
|
||||
"report_latency_p99",
|
||||
"report_ttft_p50",
|
||||
"report_ttft_p99",
|
||||
"report_tokens",
|
||||
"report_qps_current",
|
||||
"report_qps_peak",
|
||||
"report_qps_avg",
|
||||
"report_tps_current",
|
||||
"report_tps_peak",
|
||||
"report_tps_avg",
|
||||
}
|
||||
|
||||
for _, locale := range []string{"en", "zh"} {
|
||||
tmpl, err := svc.GetTemplate(ctx, NotificationEmailEventOpsScheduledReport, locale)
|
||||
require.NoError(t, err)
|
||||
for _, placeholder := range requiredPlaceholders {
|
||||
require.Contains(t, tmpl.Placeholders, placeholder)
|
||||
require.Contains(t, tmpl.HTML, "{{"+placeholder+"}}")
|
||||
}
|
||||
|
||||
preview, err := svc.PreviewTemplate(ctx, NotificationEmailPreviewInput{
|
||||
Event: NotificationEmailEventOpsScheduledReport,
|
||||
Locale: locale,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, preview.HTML, "2,374")
|
||||
require.Contains(t, preview.HTML, "99.86%")
|
||||
require.Contains(t, preview.HTML, "151,260 ms")
|
||||
require.Contains(t, preview.HTML, `style="display: none;"`)
|
||||
require.NotContains(t, preview.HTML, "{{report_total_requests}}")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsScheduledReportRuntimeVariablesDoNotLeakPreviewSamples(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc := NewNotificationEmailService(newNotificationEmailMemorySettingRepo(), nil)
|
||||
|
||||
variables := svc.runtimeVariables(ctx, NotificationEmailEventOpsScheduledReport, "en", NotificationEmailSendInput{})
|
||||
require.Equal(t, "none", variables["report_summary_display"])
|
||||
require.Equal(t, "block", variables["report_detail_display"])
|
||||
require.Empty(t, variables["report_html"])
|
||||
for _, placeholder := range notificationEmailOpsSummaryPlaceholders {
|
||||
if placeholder == "report_summary_display" {
|
||||
continue
|
||||
}
|
||||
require.Equal(t, "-", variables[placeholder])
|
||||
}
|
||||
|
||||
rendered, err := renderNotificationEmail(
|
||||
NotificationEmailEventOpsScheduledReport,
|
||||
"Report",
|
||||
`<div style="display: {{report_detail_display}};">{{report_html}}</div>`,
|
||||
variables,
|
||||
nil,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, rendered.HTML, "<h2>Daily summary</h2>")
|
||||
}
|
||||
|
||||
func TestNotificationEmailRawHTMLVariablesAreTrustedOnlyForHTMLPlaceholders(t *testing.T) {
|
||||
require.True(t, notificationEmailRawHTMLAllowed(NotificationEmailEventOpsScheduledReport, "report_html"))
|
||||
require.False(t, notificationEmailRawHTMLAllowed(NotificationEmailEventOpsScheduledReport, "recipient_name"))
|
||||
@@ -455,9 +553,11 @@ func TestNotificationEmailMemorySettingRepoSatisfiesInterface(t *testing.T) {
|
||||
}
|
||||
|
||||
type notificationEmailTestSMTPServer struct {
|
||||
listener net.Listener
|
||||
wg sync.WaitGroup
|
||||
messages atomic.Int64
|
||||
listener net.Listener
|
||||
wg sync.WaitGroup
|
||||
messages atomic.Int64
|
||||
messageMu sync.Mutex
|
||||
messageBodies []string
|
||||
}
|
||||
|
||||
func startNotificationEmailTestSMTPServer(t *testing.T) *notificationEmailTestSMTPServer {
|
||||
@@ -489,6 +589,15 @@ func (s *notificationEmailTestSMTPServer) messageCount() int64 {
|
||||
return s.messages.Load()
|
||||
}
|
||||
|
||||
func (s *notificationEmailTestSMTPServer) lastMessage() string {
|
||||
s.messageMu.Lock()
|
||||
defer s.messageMu.Unlock()
|
||||
if len(s.messageBodies) == 0 {
|
||||
return ""
|
||||
}
|
||||
return s.messageBodies[len(s.messageBodies)-1]
|
||||
}
|
||||
|
||||
func (s *notificationEmailTestSMTPServer) close() {
|
||||
_ = s.listener.Close()
|
||||
s.wg.Wait()
|
||||
@@ -547,6 +656,7 @@ func (s *notificationEmailTestSMTPServer) handleConn(conn net.Conn) {
|
||||
if !writeLine("354 End data with <CR><LF>.<CR><LF>") {
|
||||
return
|
||||
}
|
||||
var message strings.Builder
|
||||
for {
|
||||
dataLine, err := rw.ReadString('\n')
|
||||
if err != nil {
|
||||
@@ -555,7 +665,11 @@ func (s *notificationEmailTestSMTPServer) handleConn(conn net.Conn) {
|
||||
if strings.TrimRight(dataLine, "\r\n") == "." {
|
||||
break
|
||||
}
|
||||
_, _ = message.WriteString(dataLine)
|
||||
}
|
||||
s.messageMu.Lock()
|
||||
s.messageBodies = append(s.messageBodies, message.String())
|
||||
s.messageMu.Unlock()
|
||||
s.messages.Add(1)
|
||||
if !writeLine("250 2.0.0 OK") {
|
||||
return
|
||||
|
||||
@@ -47,6 +47,7 @@ type codexModelsManifestUpstreamError struct {
|
||||
err error
|
||||
retryable bool
|
||||
statusCode int
|
||||
headers http.Header
|
||||
body []byte
|
||||
}
|
||||
|
||||
@@ -56,7 +57,13 @@ func (e *codexModelsManifestUpstreamError) Unwrap() error { return e.err }
|
||||
|
||||
// IsRetryableCodexModelsManifestError reports whether another selected account
|
||||
// may succeed without changing the request. Configuration and upstream 4xx
|
||||
// responses, except 429, are intentionally not retried.
|
||||
// responses, except 429 and ChatGPT-backend 401, are intentionally not
|
||||
// retried. A manifest 401 from the ChatGPT Codex backend reflects the selected
|
||||
// OAuth account's upstream token rather than the client request (the client's
|
||||
// own API key was already validated locally), so a different account may still
|
||||
// serve the manifest. Custom API key upstreams keep the old no-failover 401
|
||||
// behavior because their /models auth semantics are not authoritative for the
|
||||
// account.
|
||||
func IsRetryableCodexModelsManifestError(err error) bool {
|
||||
var upstreamErr *codexModelsManifestUpstreamError
|
||||
return errors.As(err, &upstreamErr) && upstreamErr.retryable
|
||||
@@ -322,6 +329,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
|
||||
}
|
||||
manifest, fetchErr := s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch)
|
||||
if !credAccount.IsOpenAIAgentIdentity() || !isAgentIdentityTaskInvalidCodexModelsError(fetchErr) {
|
||||
s.handleCodexModelsManifestAccountAuthError(ctx, account, credAccount, fetchErr)
|
||||
return manifest, fetchErr
|
||||
}
|
||||
expectedTaskID := strings.TrimSpace(credAccount.GetCredential("task_id"))
|
||||
@@ -349,6 +357,37 @@ func isAgentIdentityTaskInvalidCodexModelsError(err error) bool {
|
||||
isAgentIdentityTaskInvalidHTTPResponse(upstreamErr.statusCode, upstreamErr.body)
|
||||
}
|
||||
|
||||
// handleCodexModelsManifestAccountAuthError feeds manifest 401s from the
|
||||
// ChatGPT Codex backend into the shared upstream-error state machinery
|
||||
// (token cache invalidation, temp-unschedulable cooldown, or permanent
|
||||
// disable for token_revoked/token_invalidated). Without this, an account
|
||||
// whose OAuth token was revoked upstream stays active and schedulable and
|
||||
// keeps being selected for every subsequent /models request (#4544).
|
||||
//
|
||||
// Scope is deliberately limited to plain OAuth accounts: the manifest
|
||||
// endpoint authenticates with the same token as /responses forwarding, so a
|
||||
// 401 is authoritative for the account. Agent Identity accounts are excluded
|
||||
// because their 401s can be task-scoped and have a dedicated recovery flow,
|
||||
// and API key manifests come from custom upstreams whose /models auth may
|
||||
// diverge from their chat endpoints.
|
||||
func (s *OpenAIGatewayService) handleCodexModelsManifestAccountAuthError(ctx context.Context, account, credAccount *Account, err error) {
|
||||
if s == nil || account == nil || err == nil {
|
||||
return
|
||||
}
|
||||
if credAccount == nil || !credAccount.IsOpenAIOAuth() || credAccount.IsOpenAIAgentIdentity() {
|
||||
return
|
||||
}
|
||||
var upstreamErr *codexModelsManifestUpstreamError
|
||||
if !errors.As(err, &upstreamErr) || upstreamErr.statusCode != http.StatusUnauthorized {
|
||||
return
|
||||
}
|
||||
headers := upstreamErr.headers
|
||||
if headers == nil {
|
||||
headers = http.Header{}
|
||||
}
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.statusCode, headers, upstreamErr.body)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) fetchCachedAPIKeyCodexModelsManifest(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
@@ -450,8 +489,10 @@ func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Cont
|
||||
return nil, &codexModelsManifestUpstreamError{
|
||||
err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message),
|
||||
statusCode: resp.StatusCode,
|
||||
headers: resp.Header.Clone(),
|
||||
body: body,
|
||||
retryable: resp.StatusCode == http.StatusTooManyRequests ||
|
||||
retryable: (resp.StatusCode == http.StatusUnauthorized && !request.useAPIKeyUpstream) ||
|
||||
resp.StatusCode == http.StatusTooManyRequests ||
|
||||
(resp.StatusCode >= http.StatusInternalServerError && resp.StatusCode < 600),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1094,6 +1094,147 @@ func TestFetchCodexModelsManifestAPIKeyRejectsBaseURLFragment(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// codexModelsAccountStateRepo records account state transitions triggered by
|
||||
// manifest upstream errors (#4544).
|
||||
type codexModelsAccountStateRepo struct {
|
||||
AccountRepository
|
||||
mu sync.Mutex
|
||||
setErrorCalls int
|
||||
lastErrorMsg string
|
||||
setTempUnschedCalls int
|
||||
lastTempReason string
|
||||
}
|
||||
|
||||
func (r *codexModelsAccountStateRepo) SetError(_ context.Context, _ int64, errorMsg string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.setErrorCalls++
|
||||
r.lastErrorMsg = errorMsg
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *codexModelsAccountStateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.setTempUnschedCalls++
|
||||
r.lastTempReason = reason
|
||||
return nil
|
||||
}
|
||||
|
||||
func newCodexModels401TestService(repo AccountRepository) *OpenAIGatewayService {
|
||||
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
s := &OpenAIGatewayService{rateLimitService: rateLimitService}
|
||||
rateLimitService.SetAccountRuntimeBlocker(s)
|
||||
return s
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestOAuth401MarksAccountUnschedulable(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"detail":{"message":"invalid token"}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
original := chatgptCodexModelsURL
|
||||
chatgptCodexModelsURL = server.URL
|
||||
defer func() { chatgptCodexModelsURL = original }()
|
||||
|
||||
repo := &codexModelsAccountStateRepo{}
|
||||
s := newCodexModels401TestService(repo)
|
||||
account := newCodexModelsTestAccount()
|
||||
account.Credentials["refresh_token"] = "test-refresh-token"
|
||||
|
||||
_, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
|
||||
require.Error(t, err)
|
||||
require.True(t, IsRetryableCodexModelsManifestError(err), "manifest 401 should allow account failover")
|
||||
require.Equal(t, 1, repo.setTempUnschedCalls, "OAuth 401 should temp-unschedule the account")
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.True(t, s.isOpenAIAccountRuntimeBlocked(account), "account should be runtime-blocked after manifest 401")
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestOAuth401TokenRevokedDisablesAccount(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"error":{"code":"token_revoked","message":"token has been revoked"}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
original := chatgptCodexModelsURL
|
||||
chatgptCodexModelsURL = server.URL
|
||||
defer func() { chatgptCodexModelsURL = original }()
|
||||
|
||||
repo := &codexModelsAccountStateRepo{}
|
||||
s := newCodexModels401TestService(repo)
|
||||
account := newCodexModelsTestAccount()
|
||||
account.Credentials["refresh_token"] = "test-refresh-token"
|
||||
|
||||
_, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
|
||||
require.Error(t, err)
|
||||
require.True(t, IsRetryableCodexModelsManifestError(err))
|
||||
require.Equal(t, 1, repo.setErrorCalls, "revoked token should permanently disable the account")
|
||||
require.Contains(t, repo.lastErrorMsg, "Token revoked")
|
||||
require.Equal(t, 0, repo.setTempUnschedCalls)
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestAgentIdentity401DoesNotDisableAccount(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 6,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": key.taskID,
|
||||
"chatgpt_account_id": "acc-agent-401",
|
||||
},
|
||||
}
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"detail":"some non-task 401"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
original := chatgptCodexModelsURL
|
||||
chatgptCodexModelsURL = server.URL
|
||||
defer func() { chatgptCodexModelsURL = original }()
|
||||
|
||||
repo := &codexModelsAccountStateRepo{}
|
||||
s := newCodexModels401TestService(repo)
|
||||
|
||||
_, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
|
||||
require.Error(t, err)
|
||||
require.Equal(t, 0, repo.setErrorCalls, "agent identity 401s must not disable the account")
|
||||
require.Equal(t, 0, repo.setTempUnschedCalls)
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestAPIKey401KeepsNoFailoverAndNoDisable(t *testing.T) {
|
||||
upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Status: "401 Unauthorized",
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":"invalid api key"}`)),
|
||||
}, nil
|
||||
}}
|
||||
|
||||
repo := &codexModelsAccountStateRepo{}
|
||||
s := newCodexModelsAPIKeyTestService(upstream)
|
||||
s.rateLimitService = NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
|
||||
_, err := s.FetchCodexModelsManifest(
|
||||
context.Background(),
|
||||
newCodexModelsAPIKeyTestAccount("https://upstream.example"),
|
||||
"0.144.0",
|
||||
"",
|
||||
)
|
||||
require.Error(t, err)
|
||||
require.False(t, IsRetryableCodexModelsManifestError(err), "custom upstream manifest 401 keeps the no-failover behavior")
|
||||
require.Equal(t, 0, repo.setErrorCalls, "custom upstream manifest 401 must not disable the account")
|
||||
require.Equal(t, 0, repo.setTempUnschedCalls)
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestAPIKeyUpstreamError(t *testing.T) {
|
||||
upstream := &codexModelsHTTPUpstreamStub{do: func(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
@@ -77,10 +79,39 @@ type codexTransformResult struct {
|
||||
}
|
||||
|
||||
type codexOAuthTransformOptions struct {
|
||||
IsCodexCLI bool
|
||||
IsCompact bool
|
||||
SkipDefaultInstructions bool
|
||||
PreserveToolCallIDs bool
|
||||
IsCodexCLI bool
|
||||
IsCompact bool
|
||||
SkipDefaultInstructions bool
|
||||
PreserveToolCallIDs bool
|
||||
OmitPromotedSystemMessagesFromInput bool
|
||||
}
|
||||
|
||||
const (
|
||||
codexCallIDMaxLength = 64
|
||||
codexCallIDPrefix = "fc_"
|
||||
)
|
||||
|
||||
func normalizeCodexCallID(id string) string {
|
||||
candidate := id
|
||||
switch {
|
||||
case id == "":
|
||||
return ""
|
||||
case strings.HasPrefix(id, "fc"):
|
||||
case strings.HasPrefix(id, "call_"):
|
||||
candidate = codexCallIDPrefix + strings.TrimPrefix(id, "call_")
|
||||
default:
|
||||
candidate = codexCallIDPrefix + id
|
||||
}
|
||||
if len(candidate) <= codexCallIDMaxLength {
|
||||
return candidate
|
||||
}
|
||||
return compactCodexCallID(candidate)
|
||||
}
|
||||
|
||||
func compactCodexCallID(id string) string {
|
||||
digest := sha256.Sum256([]byte("sub2api:codex-call-id:v1:" + id))
|
||||
encoded := hex.EncodeToString(digest[:])
|
||||
return codexCallIDPrefix + encoded[:codexCallIDMaxLength-len(codexCallIDPrefix)]
|
||||
}
|
||||
|
||||
const codexImageGenerationFunctionToolName = "image_gen.imagegen"
|
||||
@@ -219,9 +250,11 @@ func applyCodexOAuthTransformWithOptions(reqBody map[string]any, opts codexOAuth
|
||||
}
|
||||
|
||||
// ChatGPT internal Codex endpoint does not accept role:"system".
|
||||
// Keep the guidance in input as developer for Responses JSON mode, and
|
||||
// also mirror it into instructions because Codex OAuth requires it.
|
||||
if extractSystemMessagesFromInput(reqBody) {
|
||||
// Mirror its text into instructions because Codex OAuth requires it. Some
|
||||
// callers must also keep the guidance in input as developer (notably
|
||||
// Responses JSON object mode), while Chat Completions compatibility can
|
||||
// omit text-only messages after promoting them losslessly.
|
||||
if extractSystemMessagesFromInput(reqBody, opts.OmitPromotedSystemMessagesFromInput) {
|
||||
result.Modified = true
|
||||
}
|
||||
|
||||
@@ -1097,31 +1130,46 @@ func extractTextFromContent(content any) string {
|
||||
}
|
||||
}
|
||||
|
||||
// extractSystemMessagesFromInput scans input for role=="system", maps those
|
||||
// items to developer, and mirrors their text into reqBody["instructions"].
|
||||
// It preserves the input items so Responses JSON mode can still see JSON
|
||||
// instructions in input messages.
|
||||
func extractSystemMessagesFromInput(reqBody map[string]any) bool {
|
||||
// extractSystemMessagesFromInput scans input for role=="system" and mirrors
|
||||
// their text into reqBody["instructions"]. By default it maps those items to
|
||||
// developer so Responses JSON mode can still see JSON instructions in input.
|
||||
// When omitPromoted is true, text-only items are removed after their content is
|
||||
// losslessly promoted; mixed or malformed content is retained as developer.
|
||||
func extractSystemMessagesFromInput(reqBody map[string]any, omitPromoted bool) bool {
|
||||
input, ok := reqBody["input"].([]any)
|
||||
if !ok || len(input) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
var systemTexts []string
|
||||
filteredInput := make([]any, 0, len(input))
|
||||
modified := false
|
||||
for _, item := range input {
|
||||
m, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
if !ok || m["role"] != "system" {
|
||||
filteredInput = append(filteredInput, item)
|
||||
continue
|
||||
}
|
||||
if role, _ := m["role"].(string); role != "system" {
|
||||
continue
|
||||
|
||||
if omitPromoted {
|
||||
if losslessText, lossless := extractLosslessTextFromContent(m["content"]); lossless {
|
||||
if losslessText != "" {
|
||||
systemTexts = append(systemTexts, losslessText)
|
||||
}
|
||||
modified = true
|
||||
continue
|
||||
}
|
||||
}
|
||||
m["role"] = "developer"
|
||||
modified = true
|
||||
|
||||
if text := extractTextFromContent(m["content"]); text != "" {
|
||||
systemTexts = append(systemTexts, text)
|
||||
}
|
||||
m["role"] = "developer"
|
||||
filteredInput = append(filteredInput, item)
|
||||
modified = true
|
||||
}
|
||||
if omitPromoted && len(filteredInput) != len(input) {
|
||||
reqBody["input"] = filteredInput
|
||||
}
|
||||
|
||||
if len(systemTexts) == 0 {
|
||||
@@ -1137,6 +1185,35 @@ func extractSystemMessagesFromInput(reqBody map[string]any) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// extractLosslessTextFromContent returns text only when the entire content can
|
||||
// be represented by an instructions string without dropping non-text parts.
|
||||
func extractLosslessTextFromContent(content any) (string, bool) {
|
||||
switch v := content.(type) {
|
||||
case string:
|
||||
return v, true
|
||||
case []any:
|
||||
var b strings.Builder
|
||||
for _, part := range v {
|
||||
m, ok := part.(map[string]any)
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
typeName, ok := m["type"].(string)
|
||||
if !ok || (typeName != "text" && typeName != "input_text" && typeName != "output_text") {
|
||||
return "", false
|
||||
}
|
||||
text, ok := m["text"].(string)
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
_, _ = b.WriteString(text)
|
||||
}
|
||||
return b.String(), true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func extractPromptLikeInstructionsFromInput(reqBody map[string]any) string {
|
||||
input, ok := reqBody["input"].([]any)
|
||||
if !ok || len(input) == 0 {
|
||||
@@ -1334,13 +1411,7 @@ func filterCodexInputWithOptions(input []any, opts codexInputFilterOptions) []an
|
||||
if opts.PreserveCallIDs {
|
||||
return id
|
||||
}
|
||||
if id == "" || strings.HasPrefix(id, "fc") {
|
||||
return id
|
||||
}
|
||||
if strings.HasPrefix(id, "call_") {
|
||||
return "fc_" + strings.TrimPrefix(id, "call_")
|
||||
}
|
||||
return "fc_" + id
|
||||
return normalizeCodexCallID(id)
|
||||
}
|
||||
|
||||
if typ == "item_reference" {
|
||||
|
||||
@@ -127,6 +127,70 @@ func TestApplyCodexOAuthTransform_ToolContinuationNormalizesToolReferenceIDsOnly
|
||||
require.Equal(t, "fc_1", second["call_id"])
|
||||
}
|
||||
|
||||
func TestApplyCodexOAuthTransform_BoundsLongCallIDsAndPreservesPairing(t *testing.T) {
|
||||
suffix := strings.Repeat("z", 62)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
callID string
|
||||
outputCallID string
|
||||
}{
|
||||
{name: "non-native boundary id", callID: "call-" + strings.Repeat("x", 59), outputCallID: "call-" + strings.Repeat("x", 59)},
|
||||
{name: "overlong fc id", callID: "fc_" + strings.Repeat("y", 62), outputCallID: "fc_" + strings.Repeat("y", 62)},
|
||||
{name: "equivalent prefixes", callID: "call_" + suffix, outputCallID: "fc_" + suffix},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
reqBody := map[string]any{
|
||||
"model": "gpt-5.2",
|
||||
"input": []any{
|
||||
map[string]any{"type": "function_call", "call_id": tc.callID, "name": "shell"},
|
||||
map[string]any{"type": "function_call_output", "call_id": tc.outputCallID, "output": "done"},
|
||||
},
|
||||
}
|
||||
|
||||
applyCodexOAuthTransform(reqBody, false, false)
|
||||
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
call, ok := input[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
output, ok := input[1].(map[string]any)
|
||||
require.True(t, ok)
|
||||
|
||||
fixedCallID, ok := call["call_id"].(string)
|
||||
require.True(t, ok)
|
||||
require.LessOrEqual(t, len(fixedCallID), codexCallIDMaxLength)
|
||||
require.True(t, strings.HasPrefix(fixedCallID, codexCallIDPrefix))
|
||||
require.Equal(t, fixedCallID, output["call_id"])
|
||||
require.Equal(t, fixedCallID, normalizeCodexCallID(tc.callID))
|
||||
require.Equal(t, fixedCallID, normalizeCodexCallID(tc.outputCallID))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyCodexOAuthTransform_PreservesLongCallIDsWhenRequested(t *testing.T) {
|
||||
callID := "call-" + strings.Repeat("x", 70)
|
||||
reqBody := map[string]any{
|
||||
"model": "gpt-5.2",
|
||||
"input": []any{
|
||||
map[string]any{"type": "function_call", "call_id": callID, "name": "shell"},
|
||||
map[string]any{"type": "function_call_output", "call_id": callID, "output": "done"},
|
||||
},
|
||||
}
|
||||
|
||||
applyCodexOAuthTransformWithOptions(reqBody, codexOAuthTransformOptions{
|
||||
PreserveToolCallIDs: true,
|
||||
})
|
||||
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
call, ok := input[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
output, ok := input[1].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, callID, call["call_id"])
|
||||
require.Equal(t, callID, output["call_id"])
|
||||
}
|
||||
|
||||
func TestApplyCodexOAuthTransform_ToolSearchOutputPreservesCallID(t *testing.T) {
|
||||
reqBody := map[string]any{
|
||||
"model": "gpt-5.2",
|
||||
@@ -1336,7 +1400,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.False(t, result)
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
@@ -1352,7 +1416,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.True(t, result)
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
@@ -1378,7 +1442,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
},
|
||||
},
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.True(t, result)
|
||||
require.Equal(t, "Be helpful.", reqBody["instructions"])
|
||||
input, ok := reqBody["input"].([]any)
|
||||
@@ -1401,7 +1465,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
map[string]any{"role": "user", "content": "hi"},
|
||||
},
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.True(t, result)
|
||||
require.Equal(t, "First.\n\nSecond.", reqBody["instructions"])
|
||||
input, ok := reqBody["input"].([]any)
|
||||
@@ -1427,7 +1491,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
map[string]any{"role": "assistant", "content": "Hi there"},
|
||||
},
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.True(t, result)
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
@@ -1452,7 +1516,7 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
},
|
||||
"instructions": "Existing instructions.",
|
||||
}
|
||||
result := extractSystemMessagesFromInput(reqBody)
|
||||
result := extractSystemMessagesFromInput(reqBody, false)
|
||||
require.True(t, result)
|
||||
require.Equal(t, "Extracted.\n\nExisting instructions.", reqBody["instructions"])
|
||||
input, ok := reqBody["input"].([]any)
|
||||
@@ -1461,6 +1525,62 @@ func TestExtractSystemMessagesFromInput(t *testing.T) {
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "developer", msg["role"])
|
||||
})
|
||||
|
||||
t.Run("omit losslessly promoted text-only messages", func(t *testing.T) {
|
||||
reqBody := map[string]any{
|
||||
"input": []any{
|
||||
map[string]any{"role": "system", "content": "First."},
|
||||
map[string]any{
|
||||
"role": "system",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "Second "},
|
||||
map[string]any{"type": "input_text", "text": "and "},
|
||||
map[string]any{"type": "output_text", "text": "third."},
|
||||
},
|
||||
},
|
||||
map[string]any{"role": "user", "content": "hi"},
|
||||
},
|
||||
"instructions": "Existing.",
|
||||
}
|
||||
|
||||
result := extractSystemMessagesFromInput(reqBody, true)
|
||||
|
||||
require.True(t, result)
|
||||
require.Equal(t, "First.\n\nSecond and third.\n\nExisting.", reqBody["instructions"])
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, input, 1)
|
||||
user, ok := input[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "user", user["role"])
|
||||
})
|
||||
|
||||
t.Run("omit keeps mixed system content as developer", func(t *testing.T) {
|
||||
reqBody := map[string]any{
|
||||
"input": []any{
|
||||
map[string]any{
|
||||
"role": "system",
|
||||
"content": []any{
|
||||
map[string]any{"type": "input_text", "text": "Inspect this image."},
|
||||
map[string]any{"type": "input_image", "image_url": "https://example.com/image.png"},
|
||||
},
|
||||
},
|
||||
map[string]any{"role": "user", "content": "hi"},
|
||||
},
|
||||
}
|
||||
|
||||
result := extractSystemMessagesFromInput(reqBody, true)
|
||||
|
||||
require.True(t, result)
|
||||
require.Equal(t, "Inspect this image.", reqBody["instructions"])
|
||||
input, ok := reqBody["input"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, input, 2)
|
||||
developer, ok := input[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "developer", developer["role"])
|
||||
require.Len(t, developer["content"], 2)
|
||||
})
|
||||
}
|
||||
|
||||
// TestApplyCodexOAuthTransform_StripsPromptCacheRetention is a regression
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -1721,6 +1722,43 @@ func TestForwardAsAnthropic_BufferedTerminalWithoutUpstreamCloseReturns(t *testi
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleAnthropicBufferedStreamingResponse_OverridesUpstreamContentType(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`data: {"type":"response.completed","response":{"id":"resp_buffered_json","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":15,"output_tokens":6,"total_tokens":21,"input_tokens_details":{"cached_tokens":5}}}}` + "\n\n",
|
||||
)),
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: cfg,
|
||||
responseHeaderFilter: compileResponseHeaderFilter(cfg),
|
||||
}
|
||||
|
||||
result, err := svc.handleAnthropicBufferedStreamingResponse(
|
||||
resp, c, &Account{}, "claude-sonnet-4-5", "gpt-5.4", "gpt-5.4", time.Now(),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "application/json; charset=utf-8", rec.Header().Get("Content-Type"))
|
||||
require.NotContains(t, rec.Header().Get("Content-Type"), "text/event-stream")
|
||||
|
||||
var message apicompat.AnthropicResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &message))
|
||||
require.Equal(t, "message", message.Type)
|
||||
require.Equal(t, "resp_buffered_json", message.ID)
|
||||
require.Equal(t, 10, message.Usage.InputTokens)
|
||||
require.Equal(t, 6, message.Usage.OutputTokens)
|
||||
require.Equal(t, 5, message.Usage.CacheReadInputTokens)
|
||||
}
|
||||
|
||||
func TestForwardAsAnthropic_BufferedEventNamedTerminalWithoutUpstreamCloseReturns(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -196,8 +196,10 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
|
||||
if err := json.Unmarshal(responsesBody, &reqBody); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal for codex transform: %w", err)
|
||||
}
|
||||
isJSONObjectFormat := strings.EqualFold(strings.TrimSpace(gjson.GetBytes(responsesBody, "text.format.type").String()), "json_object")
|
||||
codexResult := applyCodexOAuthTransformWithOptions(reqBody, codexOAuthTransformOptions{
|
||||
SkipDefaultInstructions: !isResponsesShape,
|
||||
SkipDefaultInstructions: !isResponsesShape,
|
||||
OmitPromotedSystemMessagesFromInput: !isResponsesShape && !isJSONObjectFormat,
|
||||
})
|
||||
if !isResponsesShape {
|
||||
ensureCodexOAuthInstructionsField(reqBody)
|
||||
|
||||
@@ -279,6 +279,79 @@ func TestForwardAsChatCompletions_OAuthDoesNotInjectDefaultInstructions(t *testi
|
||||
require.NotContains(t, string(upstream.lastBody), "Communicate with the user by streaming thinking")
|
||||
}
|
||||
|
||||
func forwardOAuthChatCompletionsForUpstreamBody(t *testing.T, body []byte) []byte {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_system_promotion"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{
|
||||
ID: 4,
|
||||
Name: "openai-oauth",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "oauth-token",
|
||||
"chatgpt_account_id": "chatgpt-acc",
|
||||
},
|
||||
}
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4")
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.NotEmpty(t, upstream.lastBody)
|
||||
return upstream.lastBody
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletions_OAuthPromotesSystemMessageWithoutDuplication(t *testing.T) {
|
||||
const systemPrompt = "Unique system prefix for token accounting."
|
||||
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"` + systemPrompt + `"},{"role":"user","content":"hello"}],"stream":false}`)
|
||||
|
||||
upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
|
||||
|
||||
require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
|
||||
require.Equal(t, int64(1), gjson.GetBytes(upstreamBody, "input.#").Int())
|
||||
require.Equal(t, "user", gjson.GetBytes(upstreamBody, "input.0.role").String())
|
||||
require.Equal(t, 1, strings.Count(string(upstreamBody), systemPrompt))
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletions_OAuthJsonObjectKeepsSystemMessageInInput(t *testing.T) {
|
||||
const systemPrompt = "Return JSON only."
|
||||
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"` + systemPrompt + `"},{"role":"user","content":"symbol data"}],"response_format":{"type":" JSON_OBJECT "},"stream":false}`)
|
||||
|
||||
upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
|
||||
|
||||
require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
|
||||
require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.#").Int())
|
||||
require.Equal(t, "developer", gjson.GetBytes(upstreamBody, "input.0.role").String())
|
||||
require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "input.0.content").String())
|
||||
require.Equal(t, 2, strings.Count(string(upstreamBody), systemPrompt))
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletions_OAuthKeepsMixedSystemContentInInput(t *testing.T) {
|
||||
const systemPrompt = "Inspect this reference image."
|
||||
body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":[{"type":"text","text":"` + systemPrompt + `"},{"type":"image_url","image_url":{"url":"https://example.com/reference.png"}}]},{"role":"user","content":"hello"}],"stream":false}`)
|
||||
|
||||
upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
|
||||
|
||||
require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
|
||||
require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.#").Int())
|
||||
require.Equal(t, "developer", gjson.GetBytes(upstreamBody, "input.0.role").String())
|
||||
require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.0.content.#").Int())
|
||||
require.Equal(t, "input_image", gjson.GetBytes(upstreamBody, "input.0.content.1.type").String())
|
||||
require.Equal(t, "https://example.com/reference.png", gjson.GetBytes(upstreamBody, "input.0.content.1.image_url").String())
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletions_ClientDisconnectDrainsUpstreamUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -39,6 +39,39 @@ type openAIInputTokensCountPrepared struct {
|
||||
UpstreamModel string
|
||||
}
|
||||
|
||||
// EstimateGrokCountTokens estimates an Anthropic-compatible count_tokens request
|
||||
// locally. Grok does not expose a compatible token-counting endpoint, so this
|
||||
// path deliberately avoids account selection, credentials, and upstream calls.
|
||||
func EstimateGrokCountTokens(body []byte) (int, error) {
|
||||
var anthropicReq apicompat.AnthropicRequest
|
||||
if err := json.Unmarshal(body, &anthropicReq); err != nil {
|
||||
return 0, fmt.Errorf("parse anthropic count_tokens request: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(anthropicReq.Model) == "" {
|
||||
return 0, fmt.Errorf("parse anthropic count_tokens request: model is required")
|
||||
}
|
||||
|
||||
responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("convert anthropic request to responses: %w", err)
|
||||
}
|
||||
|
||||
estimated, err := estimateOpenAIInputTokens(openAIInputTokensCountRequest{
|
||||
Model: anthropicReq.Model,
|
||||
Instructions: responsesReq.Instructions,
|
||||
Input: responsesReq.Input,
|
||||
Tools: responsesReq.Tools,
|
||||
ToolChoice: responsesReq.ToolChoice,
|
||||
})
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("estimate grok input tokens: %w", err)
|
||||
}
|
||||
if estimated < openAIInputTokensFallbackMinimum {
|
||||
estimated = openAIInputTokensFallbackMinimum
|
||||
}
|
||||
return estimated, nil
|
||||
}
|
||||
|
||||
// ForwardCountTokensAsAnthropic bridges Anthropic /v1/messages/count_tokens to
|
||||
// OpenAI POST /v1/responses/input_tokens and returns Anthropic-compatible output.
|
||||
func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic(
|
||||
|
||||
@@ -236,6 +236,51 @@ func TestEstimateOpenAIInputTokens_RequestSamples(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateGrokCountTokens_AnthropicRequests(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{
|
||||
name: "simple message",
|
||||
body: `{"model":"grok-4","messages":[{"role":"user","content":"hello world"}]}`,
|
||||
},
|
||||
{
|
||||
name: "system blocks and tools",
|
||||
body: `{
|
||||
"model":"grok-4",
|
||||
"system":[{"type":"text","text":"You are helpful."}],
|
||||
"messages":[{"role":"user","content":[{"type":"text","text":"look up the weather"}]}],
|
||||
"tools":[{"name":"lookup_weather","description":"Look up weather","input_schema":{"type":"object","properties":{"city":{"type":"string"}}}}],
|
||||
"tool_choice":{"type":"auto"}
|
||||
}`,
|
||||
},
|
||||
{
|
||||
name: "empty conversation uses positive minimum",
|
||||
body: `{"model":"grok-4","messages":[]}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range cases {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := EstimateGrokCountTokens([]byte(tt.body))
|
||||
require.NoError(t, err)
|
||||
require.Positive(t, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateGrokCountTokens_RejectsInvalidRequests(t *testing.T) {
|
||||
for _, body := range []string{
|
||||
`{`,
|
||||
`{"messages":[{"role":"user","content":"hello"}]}`,
|
||||
`{"model":"grok-4","messages":[{"role":"user","content":{"unexpected":true}}]}`,
|
||||
} {
|
||||
_, err := EstimateGrokCountTokens([]byte(body))
|
||||
require.Error(t, err, "body=%s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIInputTokensEncodingForModel(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
|
||||
@@ -52,18 +52,21 @@ func (s *OpenAIGatewayService) forwardGrokResponses(
|
||||
if isGrokImageGenerationModel(upstreamModel) {
|
||||
return nil, fmt.Errorf("model %s is an image model and is not available on the Responses endpoint; use /v1/images/generations instead", upstreamModel)
|
||||
}
|
||||
cacheIdentity := resolveGrokCacheIdentity(c, body, "", upstreamModel)
|
||||
patchedBody, err := patchGrokResponsesBody(body, upstreamModel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Derive the identity from the request xAI will actually see. This makes
|
||||
// Codex Responses Lite additional_tools part of the stable tool prefix.
|
||||
cacheIdentity := resolveGrokCacheIdentity(c, patchedBody, "", upstreamModel)
|
||||
mixedCacheIntentBody := append([]byte(nil), patchedBody...)
|
||||
patchedBody, err = applyGrokResponsesCacheIdentity(patchedBody, body, cacheIdentity, account.IsGrokOAuth())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
|
||||
}
|
||||
// Free OAuth + client function tools: reuse Messages mixed-tools cache route
|
||||
// (append web_search/x_search so xAI does not force non-cacheable build-free).
|
||||
patchedBody, err = applyGrokFreeMessagesFunctionToolCacheRoute(patchedBody, body, account, cacheIdentity)
|
||||
patchedBody, err = applyGrokFreeRequestToolCacheRoute(c, patchedBody, mixedCacheIntentBody, account, cacheIdentity)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
|
||||
}
|
||||
@@ -484,10 +487,9 @@ func deleteJSONFields(value any, fields map[string]struct{}) bool {
|
||||
}
|
||||
|
||||
// additional_tools is a Codex/Responses Lite private input carrier. xAI's
|
||||
// Responses schema accepts ordinary message/function-call input items but
|
||||
// rejects this carrier before inference with a ModelInput deserialization
|
||||
// error. Top-level supported tools remain available through the separate
|
||||
// sanitizeGrokResponsesTools path.
|
||||
// Responses schema rejects the carrier itself, but accepts supported tools at
|
||||
// the top level. Preserve top-level order, append newly discovered tools in
|
||||
// carrier order, then let sanitizeGrokResponsesTools filter unsupported types.
|
||||
func sanitizeGrokResponsesInput(body []byte) ([]byte, error) {
|
||||
if !bytes.Contains(body, []byte(`"additional_tools"`)) {
|
||||
return body, nil
|
||||
@@ -499,8 +501,36 @@ func sanitizeGrokResponsesInput(body []byte) ([]byte, error) {
|
||||
|
||||
rawItems := input.Array()
|
||||
filtered := make([]json.RawMessage, 0, len(rawItems))
|
||||
topLevelTools := gjson.GetBytes(body, "tools")
|
||||
mergedTools := make([]json.RawMessage, 0)
|
||||
seenTools := make(map[string]struct{})
|
||||
appendTool := func(tool gjson.Result) bool {
|
||||
key := grokResponsesToolDedupKey(tool)
|
||||
if _, exists := seenTools[key]; exists {
|
||||
return false
|
||||
}
|
||||
seenTools[key] = struct{}{}
|
||||
mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
|
||||
return true
|
||||
}
|
||||
if topLevelTools.IsArray() {
|
||||
for _, tool := range topLevelTools.Array() {
|
||||
seenTools[grokResponsesToolDedupKey(tool)] = struct{}{}
|
||||
mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
|
||||
}
|
||||
}
|
||||
|
||||
promoted := false
|
||||
for _, item := range rawItems {
|
||||
if strings.TrimSpace(item.Get("type").String()) == "additional_tools" {
|
||||
tools := item.Get("tools")
|
||||
if tools.IsArray() {
|
||||
for _, tool := range tools.Array() {
|
||||
if appendTool(tool) {
|
||||
promoted = true
|
||||
}
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, json.RawMessage(item.Raw))
|
||||
@@ -512,7 +542,30 @@ func sanitizeGrokResponsesInput(body []byte) ([]byte, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sjson.SetRawBytes(body, "input", encoded)
|
||||
body, err = sjson.SetRawBytes(body, "input", encoded)
|
||||
if err != nil || !promoted {
|
||||
return body, err
|
||||
}
|
||||
encodedTools, err := json.Marshal(mergedTools)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sjson.SetRawBytes(body, "tools", encodedTools)
|
||||
}
|
||||
|
||||
func grokResponsesToolDedupKey(tool gjson.Result) string {
|
||||
toolType := strings.TrimSpace(tool.Get("type").String())
|
||||
if toolType != "" {
|
||||
if name := strings.TrimSpace(tool.Get("name").String()); name != "" {
|
||||
return "type:" + toolType + "\x00name:" + name
|
||||
}
|
||||
if toolType == "mcp" {
|
||||
if label := strings.TrimSpace(tool.Get("server_label").String()); label != "" {
|
||||
return "type:mcp\x00server_label:" + label
|
||||
}
|
||||
}
|
||||
}
|
||||
return "json:" + normalizeCompatSeedJSON(json.RawMessage(tool.Raw))
|
||||
}
|
||||
|
||||
// sanitizeGrokReasoningNullContent 删除 reasoning 项中的 "content": null。
|
||||
|
||||
@@ -7,17 +7,19 @@ import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
const (
|
||||
grokConversationIDHeader = "X-Grok-Conv-Id"
|
||||
claudeCodeSessionHeader = "X-Claude-Code-Session-Id"
|
||||
grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]`
|
||||
grokFreeCacheDisabledToolChoice = "none"
|
||||
grokFreeRolling24hTokenLimit = int64(2_000_000)
|
||||
grokConversationIDHeader = "X-Grok-Conv-Id"
|
||||
claudeCodeSessionHeader = "X-Claude-Code-Session-Id"
|
||||
grokClientToolCacheOptInHeader = "X-Sub2API-Grok-Client-Tool-Cache"
|
||||
grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]`
|
||||
grokFreeCacheDisabledToolChoice = "none"
|
||||
grokClientToolCacheOptInExtraKey = "grok_client_tool_cache_enabled"
|
||||
)
|
||||
|
||||
// Claude Code metadata.user_id often ends with _session_<uuid>.
|
||||
@@ -167,7 +169,7 @@ func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity str
|
||||
// Inspect the pre-sanitization source. patchGrokResponsesBody may remove an
|
||||
// unsupported client tool and its tool_choice; that must not turn an
|
||||
// explicit client tool intent into an eligible native-tool request.
|
||||
if gjson.GetBytes(intentSourceBody, "tools").Exists() || gjson.GetBytes(intentSourceBody, "tool_choice").Exists() {
|
||||
if hasGrokResponsesToolIntent(intentSourceBody) {
|
||||
return out, nil
|
||||
}
|
||||
out, err = sjson.SetRawBytes(out, "tools", []byte(grokFreeCacheNativeToolsJSON))
|
||||
@@ -177,12 +179,113 @@ func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity str
|
||||
return sjson.SetBytes(out, "tool_choice", grokFreeCacheDisabledToolChoice)
|
||||
}
|
||||
|
||||
func hasGrokResponsesToolIntent(body []byte) bool {
|
||||
if gjson.GetBytes(body, "tools").Exists() || gjson.GetBytes(body, "tool_choice").Exists() {
|
||||
return true
|
||||
}
|
||||
input := gjson.GetBytes(body, "input")
|
||||
if !input.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, item := range input.Array() {
|
||||
if strings.TrimSpace(item.Get("type").String()) != "additional_tools" {
|
||||
continue
|
||||
}
|
||||
tools := item.Get("tools")
|
||||
if !tools.Exists() || !tools.IsArray() || len(tools.Array()) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// applyGrokFreeMessagesFunctionToolCacheRoute enables xAI's cache-capable
|
||||
// mixed-tools route only for the Anthropic Messages bridge and only when the
|
||||
// selected account is known to be Free. Native tools become eligible under
|
||||
// auto selection, so callers must not apply this policy to paid accounts or
|
||||
// other ingress protocols implicitly.
|
||||
// mixed-tools route only for known Free accounts. Pure client tools default to
|
||||
// the cache-capable route so an intermediate sub2api does not need to preserve
|
||||
// client-specific opt-in headers. Operators can explicitly disable this per
|
||||
// account when native search tools would change the desired behavior (#4486).
|
||||
func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) {
|
||||
allowPureClientTools, _ := grokClientToolCacheAccountPolicy(account)
|
||||
return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, true)
|
||||
}
|
||||
|
||||
// applyGrokFreeRequestToolCacheRoute also accepts a request-scoped opt-in. The
|
||||
// sub2api header is consumed locally because buildGrokResponsesRequest only
|
||||
// forwards the explicitly supported OpenAI-Beta header from downstream.
|
||||
func applyGrokFreeRequestToolCacheRoute(c *gin.Context, body, intentSourceBody []byte, account *Account, cacheIdentity string) ([]byte, error) {
|
||||
allowPureClientTools, accountPolicyExplicit := grokClientToolCacheAccountPolicy(account)
|
||||
requestOptOut := false
|
||||
if c != nil {
|
||||
switch strings.ToLower(strings.TrimSpace(c.GetHeader(grokClientToolCacheOptInHeader))) {
|
||||
case "1", "true", "yes", "on", "prefer-cache":
|
||||
allowPureClientTools = true
|
||||
case "0", "false", "no", "off":
|
||||
allowPureClientTools = false
|
||||
requestOptOut = true
|
||||
}
|
||||
}
|
||||
if !allowPureClientTools && !accountPolicyExplicit && !requestOptOut && isGrokClaudeDesktopResponsesCacheRequest(c) {
|
||||
allowPureClientTools = true
|
||||
}
|
||||
// A function merely named web_search/x_search is still a client function.
|
||||
// Known Free OAuth accounts use the cache route by default; a request-scoped
|
||||
// opt-in may override an account opt-out, while an explicit request opt-out
|
||||
// always wins. The legacy Claude fingerprint remains only as a compatibility
|
||||
// fallback when no account policy has been recorded (#4486).
|
||||
return applyGrokFreeToolCacheRoute(body, intentSourceBody, account, cacheIdentity, allowPureClientTools, allowPureClientTools)
|
||||
}
|
||||
|
||||
// grokClientToolCacheAccountPolicy is intentionally strict for configured
|
||||
// values: only a JSON boolean is accepted. A missing key defaults on solely for
|
||||
// accounts positively identified as Grok Free OAuth; paid, API-key, and unknown
|
||||
// accounts remain fail-closed.
|
||||
func grokClientToolCacheAccountPolicy(account *Account) (enabled, explicit bool) {
|
||||
if !isKnownGrokFreeAccount(account) {
|
||||
return false, false
|
||||
}
|
||||
if account.Extra == nil {
|
||||
return true, false
|
||||
}
|
||||
value, exists := account.Extra[grokClientToolCacheOptInExtraKey]
|
||||
if !exists {
|
||||
return true, false
|
||||
}
|
||||
enabled, valid := value.(bool)
|
||||
if !valid {
|
||||
return false, true
|
||||
}
|
||||
return enabled, true
|
||||
}
|
||||
|
||||
// isGrokClaudeDesktopResponsesCacheRequest recognizes the strict wire
|
||||
// fingerprint emitted when Claude Desktop's local agent is translated by
|
||||
// CC Switch into an OpenAI Responses request. Requiring every independent
|
||||
// signal prevents a generic Claude-compatible client (or the Chat bridge)
|
||||
// from silently opting into the mixed native/client tool route.
|
||||
func isGrokClaudeDesktopResponsesCacheRequest(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil || isOpenAIResponsesCompactPath(c) {
|
||||
return false
|
||||
}
|
||||
path := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
|
||||
if !strings.HasSuffix(path, "/responses") {
|
||||
return false
|
||||
}
|
||||
|
||||
if !claudeCodeUAPattern.MatchString(strings.TrimSpace(c.GetHeader("User-Agent"))) {
|
||||
return false
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.GetHeader("X-App"))) {
|
||||
case "cli", "cli-bg":
|
||||
default:
|
||||
return false
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(c.GetHeader("anthropic-client-platform")), "desktop_app") {
|
||||
return false
|
||||
}
|
||||
return strings.TrimSpace(c.GetHeader("X-Claude-Code-Session-Id")) != ""
|
||||
}
|
||||
|
||||
func applyGrokFreeToolCacheRoute(body, intentSourceBody []byte, account *Account, cacheIdentity string, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) {
|
||||
if strings.TrimSpace(cacheIdentity) == "" || !isKnownGrokFreeAccount(account) {
|
||||
return body, nil
|
||||
}
|
||||
@@ -191,7 +294,12 @@ func applyGrokFreeMessagesFunctionToolCacheRoute(body, intentSourceBody []byte,
|
||||
if !isGrokFreeCacheFunctionToolIntent(intentTools, intentToolChoice) {
|
||||
return body, nil
|
||||
}
|
||||
return appendMissingGrokFreeCacheNativeTools(body)
|
||||
if intentToolChoice.Type == gjson.String && strings.TrimSpace(intentToolChoice.String()) == grokFreeCacheDisabledToolChoice {
|
||||
// Adding native cache-routing tools cannot change behavior when the
|
||||
// client has explicitly disabled all tool execution.
|
||||
return appendGrokFreeCacheNativeToolsWithPolicy(body, true, false)
|
||||
}
|
||||
return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, allowFunctionSearch)
|
||||
}
|
||||
|
||||
func isKnownGrokFreeAccount(account *Account) bool {
|
||||
@@ -232,7 +340,7 @@ func isKnownGrokFreeAccount(account *Account) bool {
|
||||
}
|
||||
}
|
||||
if snapshot.Tokens != nil && snapshot.Tokens.Limit != nil &&
|
||||
*snapshot.Tokens.Limit == grokFreeRolling24hTokenLimit {
|
||||
xai.IsGrokFreeRolling24hTokenLimit(*snapshot.Tokens.Limit) {
|
||||
inferredFreeSignal = true
|
||||
}
|
||||
}
|
||||
@@ -276,22 +384,44 @@ func isGrokFreeCacheFunctionToolIntent(tools, toolChoice gjson.Result) bool {
|
||||
return false
|
||||
}
|
||||
for _, tool := range items {
|
||||
if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) != "function" {
|
||||
if !tool.IsObject() {
|
||||
return false
|
||||
}
|
||||
// Responses function declarations keep name at the top level. Reject
|
||||
// Chat Completions' nested function shape and incomplete declarations.
|
||||
if strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() {
|
||||
toolType := strings.TrimSpace(tool.Get("type").String())
|
||||
if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok {
|
||||
return false
|
||||
}
|
||||
if toolType == "function" {
|
||||
// Responses function declarations keep name at the top level. Reject
|
||||
// Chat Completions' nested function shape and incomplete declarations.
|
||||
if strings.TrimSpace(tool.Get("name").String()) == "" || tool.Get("function").Exists() {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
if !toolChoice.Exists() {
|
||||
return true
|
||||
}
|
||||
return toolChoice.Type == gjson.String && strings.TrimSpace(toolChoice.String()) == "auto"
|
||||
if toolChoice.Type != gjson.String {
|
||||
return false
|
||||
}
|
||||
switch strings.TrimSpace(toolChoice.String()) {
|
||||
case "auto", grokFreeCacheDisabledToolChoice:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
|
||||
return appendGrokFreeCacheNativeTools(body, false)
|
||||
}
|
||||
|
||||
func appendGrokFreeCacheNativeTools(body []byte, allowPureClientTools bool) ([]byte, error) {
|
||||
return appendGrokFreeCacheNativeToolsWithPolicy(body, allowPureClientTools, true)
|
||||
}
|
||||
|
||||
func appendGrokFreeCacheNativeToolsWithPolicy(body []byte, allowPureClientTools, allowFunctionSearch bool) ([]byte, error) {
|
||||
tools := gjson.GetBytes(body, "tools")
|
||||
if !tools.Exists() || !tools.IsArray() {
|
||||
return body, nil
|
||||
@@ -301,9 +431,19 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
|
||||
if len(items) == 0 {
|
||||
return body, nil
|
||||
}
|
||||
hasNativeSearch := false
|
||||
for _, tool := range items {
|
||||
switch strings.TrimSpace(tool.Get("type").String()) {
|
||||
case "web_search", "x_search":
|
||||
hasNativeSearch = true
|
||||
}
|
||||
}
|
||||
if !allowPureClientTools && !allowFunctionSearch && !hasNativeSearch {
|
||||
return body, nil
|
||||
}
|
||||
merged := make([]json.RawMessage, 0, len(items)+2)
|
||||
present := make(map[string]bool, 2)
|
||||
hasFunction := false
|
||||
hasCompanionTool := false
|
||||
for _, tool := range items {
|
||||
toolType := strings.TrimSpace(tool.Get("type").String())
|
||||
switch toolType {
|
||||
@@ -314,7 +454,7 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
|
||||
}
|
||||
// Grok Build may declare search as function tools. Convert to native
|
||||
// entries so Free OAuth stays cache-capable without duplicate names.
|
||||
if name == "web_search" || name == "x_search" {
|
||||
if (name == "web_search" || name == "x_search") && allowFunctionSearch {
|
||||
if present[name] {
|
||||
continue
|
||||
}
|
||||
@@ -324,9 +464,17 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
|
||||
}
|
||||
merged = append(merged, raw)
|
||||
present[name] = true
|
||||
if allowPureClientTools {
|
||||
hasCompanionTool = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
hasFunction = true
|
||||
if name == "web_search" || name == "x_search" {
|
||||
// Keep the client function intact and avoid adding a same-named
|
||||
// native tool unless conversion was explicitly enabled.
|
||||
present[name] = true
|
||||
}
|
||||
hasCompanionTool = true
|
||||
merged = append(merged, json.RawMessage(tool.Raw))
|
||||
case "web_search", "x_search":
|
||||
if present[toolType] {
|
||||
@@ -335,17 +483,21 @@ func appendMissingGrokFreeCacheNativeTools(body []byte) ([]byte, error) {
|
||||
merged = append(merged, json.RawMessage(tool.Raw))
|
||||
present[toolType] = true
|
||||
default:
|
||||
return body, nil
|
||||
if _, ok := grokResponsesSupportedToolTypes[toolType]; !ok {
|
||||
return body, nil
|
||||
}
|
||||
hasCompanionTool = true
|
||||
merged = append(merged, json.RawMessage(tool.Raw))
|
||||
}
|
||||
}
|
||||
if !hasFunction {
|
||||
if !hasCompanionTool {
|
||||
return body, nil
|
||||
}
|
||||
// Only complement missing native search tools when the request already contains
|
||||
// at least one search tool (native or function-form). Pure client function tools
|
||||
// (e.g. view_image) must not trigger injection to avoid biasing model tool
|
||||
// selection (#4486).
|
||||
if !present["web_search"] && !present["x_search"] {
|
||||
if !allowPureClientTools && !present["web_search"] && !present["x_search"] {
|
||||
return body, nil
|
||||
}
|
||||
for _, toolType := range []string{"web_search", "x_search"} {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -219,7 +220,7 @@ func TestApplyGrokCacheIdentityWritesResponsesBodyAndHeader(t *testing.T) {
|
||||
require.False(t, gjson.GetBytes(unscopedBody, "tool_choice").Exists())
|
||||
}
|
||||
|
||||
func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing.T) {
|
||||
func TestGrokFreeMessagesClientToolCacheDefaultsOnForKnownFree(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(901, "access-token")
|
||||
account.Credentials["subscription_tier"] = " FREE "
|
||||
tests := []struct {
|
||||
@@ -233,7 +234,6 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing.
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Pure client function tools without search → no native injection (#4486).
|
||||
intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup","description":"look up a value","parameters":{"type":"object"}},{"type":"function","name":"save","parameters":{"type":"object"}}]` + tt.toolChoiceJSON + `}`)
|
||||
body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
|
||||
require.NoError(t, err)
|
||||
@@ -242,16 +242,46 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsToResponseFunctions(t *testing.
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String())
|
||||
tools := gjson.GetBytes(body, "tools").Array()
|
||||
require.Len(t, tools, 2, "pure client functions should not get native search injected")
|
||||
require.Len(t, tools, 4)
|
||||
require.Equal(t, "function", tools[0].Get("type").String())
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.Equal(t, "function", tools[1].Get("type").String())
|
||||
require.Equal(t, "save", tools[1].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[3].Get("type").String())
|
||||
require.Equal(t, tt.wantChoice, gjson.GetBytes(body, "tool_choice").Exists())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokFreeMessagesClientToolCacheDefaultsWithMissingAccountSetting(t *testing.T) {
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
tests := []struct {
|
||||
name string
|
||||
extra map[string]any
|
||||
}{
|
||||
{name: "nil extra"},
|
||||
{name: "empty extra", extra: map[string]any{}},
|
||||
{name: "unrelated extra", extra: map[string]any{"other_setting": true}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90101, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
account.Extra = tt.extra
|
||||
|
||||
patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(901, "access-token")
|
||||
account.Credentials["subscription_tier"] = " FREE "
|
||||
@@ -271,6 +301,330 @@ func TestApplyGrokCacheIdentityAppendsNativeToolsWhenSearchPresent(t *testing.T)
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheAccountOptIn(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9011, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true}
|
||||
intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"read_file","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
|
||||
require.NoError(t, err)
|
||||
body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := gjson.GetBytes(body, "tools").Array()
|
||||
require.Len(t, tools, 4)
|
||||
require.Equal(t, "view_image", tools[0].Get("name").String())
|
||||
require.Equal(t, "read_file", tools[1].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[3].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeMessagesClientToolCacheAccountOptOut(t *testing.T) {
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
tests := []struct {
|
||||
name string
|
||||
value any
|
||||
}{
|
||||
{name: "explicit false", value: false},
|
||||
{name: "string false is malformed", value: "false"},
|
||||
{name: "numeric value is malformed", value: 1},
|
||||
{name: "null value is malformed", value: nil},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90111, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: tt.value}
|
||||
|
||||
patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(body), string(patched))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheRequestOptInOverridesAccountOptOut(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9014, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: false}
|
||||
c := newGrokCacheTestContext(9014)
|
||||
c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
|
||||
intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
|
||||
require.NoError(t, err)
|
||||
body, err = applyGrokFreeRequestToolCacheRoute(c, body, intentBody, account, "isolated-id")
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := gjson.GetBytes(body, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "view_image", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeChatRequestClientToolCacheDefaultsOnWithoutClientFingerprint(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90140, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
c := newGrokCacheTestContext(90140)
|
||||
c.Request.URL.Path = "/v1/chat/completions"
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "view_image", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheClaudeDesktopResponsesAutoOptIn(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90141, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}},{"type":"function","name":"Edit","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
for _, xApp := range []string{"cli", "cli-bg"} {
|
||||
t.Run(xApp, func(t *testing.T) {
|
||||
c := newGrokCacheTestContext(90141)
|
||||
// The desktop marker text is intentionally not required; CC Switch and
|
||||
// Claude Desktop may change the descriptive User-Agent suffix.
|
||||
c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, future-desktop, agent-sdk/0.3.215)")
|
||||
c.Request.Header.Set("X-App", xApp)
|
||||
c.Request.Header.Set("anthropic-client-platform", "desktop_app")
|
||||
c.Request.Header.Set("X-Claude-Code-Session-Id", "desktop-session-1")
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 4)
|
||||
require.Equal(t, "Read", tools[0].Get("name").String())
|
||||
require.Equal(t, "Edit", tools[1].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[3].Get("type").String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheClaudeDesktopFingerprintRequiresAllSignals(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90142, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: false}
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
ua string
|
||||
xApp string
|
||||
platform string
|
||||
session string
|
||||
}{
|
||||
{
|
||||
name: "chat path",
|
||||
path: "/v1/chat/completions",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "cli",
|
||||
platform: "desktop_app",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "compact responses path",
|
||||
path: "/v1/responses/compact",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "cli",
|
||||
platform: "desktop_app",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "non claude cli user agent",
|
||||
path: "/v1/responses",
|
||||
ua: "Mozilla/5.0 (claude-desktop-3p)",
|
||||
xApp: "cli",
|
||||
platform: "desktop_app",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "missing x app",
|
||||
path: "/v1/responses",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
platform: "desktop_app",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "wrong x app",
|
||||
path: "/v1/responses",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "desktop",
|
||||
platform: "desktop_app",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "missing client platform",
|
||||
path: "/v1/responses",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "cli",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "wrong client platform",
|
||||
path: "/v1/responses",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "cli",
|
||||
platform: "web",
|
||||
session: "desktop-session-1",
|
||||
},
|
||||
{
|
||||
name: "missing session header",
|
||||
path: "/v1/responses",
|
||||
ua: "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)",
|
||||
xApp: "cli",
|
||||
platform: "desktop_app",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := newGrokCacheTestContext(90142)
|
||||
c.Request.URL.Path = tt.path
|
||||
if tt.ua != "" {
|
||||
c.Request.Header.Set("User-Agent", tt.ua)
|
||||
}
|
||||
if tt.xApp != "" {
|
||||
c.Request.Header.Set("X-App", tt.xApp)
|
||||
}
|
||||
if tt.platform != "" {
|
||||
c.Request.Header.Set("anthropic-client-platform", tt.platform)
|
||||
}
|
||||
if tt.session != "" {
|
||||
c.Request.Header.Set("X-Claude-Code-Session-Id", tt.session)
|
||||
}
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(body), string(patched))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheExplicitRequestOptOut(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90144, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
for _, value := range []string{"0", "false", "no", "off"} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
c := newGrokCacheTestContext(90144)
|
||||
c.Request.URL.Path = "/v1/chat/completions"
|
||||
c.Request.Header.Set(grokClientToolCacheOptInHeader, value)
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(body), string(patched))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokFreeClientToolCacheClaudeDesktopAutoOptInDoesNotOverridePaidTier(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(90143, "access-token")
|
||||
account.Credentials["subscription_tier"] = "supergrok"
|
||||
c := newGrokCacheTestContext(90143)
|
||||
c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)")
|
||||
c.Request.Header.Set("X-App", "cli")
|
||||
c.Request.Header.Set("anthropic-client-platform", "desktop_app")
|
||||
c.Request.Header.Set("X-Claude-Code-Session-Id", "desktop-session-1")
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"Read","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(body), string(patched))
|
||||
}
|
||||
|
||||
func TestGrokFreeRequestClientSearchFunctionUsesDefaultAccountPolicy(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9015, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
c := newGrokCacheTestContext(9015)
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"function","name":"web_search","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "view_image", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeRequestToolChoiceNoneUsesSafeCacheRoute(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9016, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
c := newGrokCacheTestContext(9016)
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"none"}`)
|
||||
|
||||
patched, err := applyGrokFreeRequestToolCacheRoute(c, body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "view_image", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "none", gjson.GetBytes(patched, "tool_choice").String())
|
||||
}
|
||||
|
||||
func TestApplyGrokCacheIdentityRecognizesResponsesLiteAdditionalTools(t *testing.T) {
|
||||
intentBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"hello"}]}`)
|
||||
patchedBody := []byte(`{"model":"grok-4.5","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"type":"message","role":"user","content":"hello"}]}`)
|
||||
|
||||
patched, err := applyGrokResponsesCacheIdentity(patchedBody, intentBody, "isolated-id", true)
|
||||
|
||||
require.NoError(t, err)
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 1)
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
|
||||
require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String())
|
||||
}
|
||||
|
||||
func TestGrokFreeCacheRoutePreservesMixedSupportedToolsWithSearchIntent(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9012, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}},{"type":"shell"},{"type":"web_search"}],"tool_choice":"auto"}`)
|
||||
|
||||
body, err := applyGrokResponsesCacheIdentity(intentBody, intentBody, "isolated-id", true)
|
||||
require.NoError(t, err)
|
||||
body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, intentBody, account, "isolated-id")
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := gjson.GetBytes(body, "tools").Array()
|
||||
require.Len(t, tools, 4)
|
||||
require.Equal(t, "function", tools[0].Get("type").String())
|
||||
require.Equal(t, "shell", tools[1].Get("type").String())
|
||||
require.Equal(t, "web_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[3].Get("type").String())
|
||||
}
|
||||
|
||||
func TestGrokClientToolCacheOptInDoesNotOverridePaidTier(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(9013, "access-token")
|
||||
account.Credentials["subscription_tier"] = "supergrok"
|
||||
account.Extra = map[string]any{grokClientToolCacheOptInExtraKey: true}
|
||||
body := []byte(`{"model":"grok","tools":[{"type":"function","name":"view_image","parameters":{"type":"object"}}],"tool_choice":"auto"}`)
|
||||
|
||||
patched, err := applyGrokFreeMessagesFunctionToolCacheRoute(body, body, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, string(body), string(patched))
|
||||
}
|
||||
|
||||
func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) {
|
||||
account := healthyGrokOAuthGatewayTestAccount(902, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
@@ -282,7 +636,7 @@ func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) {
|
||||
{name: "missing tools", patchedBody: `{"model":"grok-4.5"}`},
|
||||
{name: "empty tools", patchedBody: `{"model":"grok-4.5","tools":[]}`},
|
||||
{name: "native tools only", patchedBody: `{"model":"grok-4.5","tools":[{"type":"web_search"}]}`},
|
||||
{name: "unexpected patched tool", patchedBody: `{"model":"grok-4.5","tools":[{"type":"function","name":"lookup"},{"type":"mcp","name":"server"}]}`},
|
||||
{name: "unexpected patched tool", patchedBody: `{"model":"grok-4.5","tools":[{"type":"function","name":"lookup"},{"type":"namespace","name":"server"}]}`},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -302,8 +656,7 @@ func TestApplyGrokCacheIdentityRequiresPatchedFunctionTools(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.T) {
|
||||
// Include web_search as function to trigger native tool injection (pure client
|
||||
// functions no longer trigger injection after #4486 fix).
|
||||
// Include web_search as a function to also cover its conversion to a native tool.
|
||||
intentBody := []byte(`{"model":"grok","tools":[{"type":"function","name":"lookup"},{"type":"function","name":"web_search"}],"tool_choice":"auto"}`)
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -347,7 +700,19 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.
|
||||
a := healthyGrokOAuthGatewayTestAccount(9112, "access-token")
|
||||
a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{
|
||||
"headers_observed": true,
|
||||
"tokens": map[string]any{"limit": grokFreeRolling24hTokenLimit},
|
||||
"tokens": map[string]any{"limit": xai.GrokFreeRolling24hTokenLimit},
|
||||
}}
|
||||
return a
|
||||
}(),
|
||||
wantMix: true,
|
||||
},
|
||||
{
|
||||
name: "legacy free rolling token quota",
|
||||
account: func() *Account {
|
||||
a := healthyGrokOAuthGatewayTestAccount(9113, "access-token")
|
||||
a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{
|
||||
"headers_observed": true,
|
||||
"tokens": map[string]any{"limit": int64(2_000_000)},
|
||||
}}
|
||||
return a
|
||||
}(),
|
||||
@@ -369,7 +734,7 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing.
|
||||
grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": http.StatusOK},
|
||||
grokQuotaSnapshotExtraKey: map[string]any{
|
||||
"headers_observed": true,
|
||||
"tokens": map[string]any{"limit": grokFreeRolling24hTokenLimit},
|
||||
"tokens": map[string]any{"limit": int64(2_000_000)},
|
||||
},
|
||||
}
|
||||
return a
|
||||
|
||||
@@ -22,17 +22,23 @@ const (
|
||||
var grokChatResponsesBridgeTopLevelFields = map[string]struct{}{
|
||||
"model": {},
|
||||
"messages": {},
|
||||
"instructions": {},
|
||||
"stream": {},
|
||||
"stream_options": {},
|
||||
"max_tokens": {},
|
||||
"max_completion_tokens": {},
|
||||
"temperature": {},
|
||||
"top_p": {},
|
||||
"stop": {},
|
||||
"reasoning_effort": {},
|
||||
"prompt_cache_key": {},
|
||||
"tools": {},
|
||||
"tool_choice": {},
|
||||
"functions": {},
|
||||
"function_call": {},
|
||||
"parallel_tool_calls": {},
|
||||
"response_format": {},
|
||||
"service_tier": {},
|
||||
}
|
||||
|
||||
// grokChatResponsesBridgeEligibility deliberately accepts only request shapes
|
||||
@@ -45,18 +51,49 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) {
|
||||
return false, "invalid_json"
|
||||
}
|
||||
|
||||
// These fields have no effect when explicitly set to JSON null. Accepting
|
||||
// that common SDK representation keeps the request on the bridge path,
|
||||
// while non-null values remain unsupported because the Responses converter
|
||||
// cannot preserve their Chat Completions semantics.
|
||||
for _, field := range []string{"stop", "reasoning_effort"} {
|
||||
if _, exists := root[field]; exists {
|
||||
if raw, exists := root[field]; exists && !grokChatJSONNull(raw) {
|
||||
return false, "unsupported_" + field
|
||||
}
|
||||
}
|
||||
for _, field := range []string{"tools", "functions"} {
|
||||
if raw, exists := root[field]; exists && !grokChatNullOrEmptyArray(raw) {
|
||||
return false, "unsupported_" + field
|
||||
if raw, exists := root["instructions"]; exists {
|
||||
var instructions string
|
||||
if !grokChatJSONNull(raw) && json.Unmarshal(raw, &instructions) != nil {
|
||||
return false, "invalid_instructions"
|
||||
}
|
||||
}
|
||||
if raw, exists := root["tool_choice"]; exists && !grokChatNullOrNone(raw) {
|
||||
return false, "unsupported_tool_choice"
|
||||
if raw, exists := root["response_format"]; exists {
|
||||
var responseFormat map[string]json.RawMessage
|
||||
if !grokChatJSONNull(raw) && (json.Unmarshal(raw, &responseFormat) != nil || responseFormat == nil) {
|
||||
return false, "invalid_response_format"
|
||||
}
|
||||
}
|
||||
if raw, exists := root["service_tier"]; exists {
|
||||
var serviceTier string
|
||||
if !grokChatJSONNull(raw) && json.Unmarshal(raw, &serviceTier) != nil {
|
||||
return false, "invalid_service_tier"
|
||||
}
|
||||
}
|
||||
if raw, exists := root["tools"]; exists {
|
||||
if ok, reason := grokChatFunctionDeclarationsBridgeable(raw); !ok {
|
||||
return false, reason
|
||||
}
|
||||
}
|
||||
if raw, exists := root["functions"]; exists && !grokChatNullOrEmptyArray(raw) {
|
||||
return false, "unsupported_functions"
|
||||
}
|
||||
if raw, exists := root["tool_choice"]; exists {
|
||||
if ok, reason := grokChatToolChoiceBridgeable(raw); !ok {
|
||||
return false, reason
|
||||
}
|
||||
var choice string
|
||||
if json.Unmarshal(raw, &choice) == nil && choice == "required" && !grokChatHasFunctionDeclarations(root) {
|
||||
return false, "required_tool_choice_without_tools"
|
||||
}
|
||||
}
|
||||
if raw, exists := root["function_call"]; exists && !grokChatNullOrNone(raw) {
|
||||
return false, "unsupported_function_call"
|
||||
@@ -78,6 +115,12 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) {
|
||||
return false, "invalid_stream"
|
||||
}
|
||||
}
|
||||
if raw, ok := root["parallel_tool_calls"]; ok {
|
||||
var parallelToolCalls *bool
|
||||
if json.Unmarshal(raw, ¶llelToolCalls) != nil || parallelToolCalls == nil {
|
||||
return false, "invalid_parallel_tool_calls"
|
||||
}
|
||||
}
|
||||
if raw, ok := root["stream_options"]; ok {
|
||||
var options map[string]json.RawMessage
|
||||
if json.Unmarshal(raw, &options) != nil || options == nil {
|
||||
@@ -128,42 +171,256 @@ func grokChatResponsesBridgeEligibility(body []byte) (bool, string) {
|
||||
return false, "invalid_messages"
|
||||
}
|
||||
for _, message := range messages {
|
||||
for field := range message {
|
||||
if field != "role" && field != "content" {
|
||||
return false, "unsafe_message_field_" + field
|
||||
}
|
||||
}
|
||||
var role string
|
||||
if raw, exists := message["role"]; !exists || json.Unmarshal(raw, &role) != nil {
|
||||
return false, "invalid_message_role"
|
||||
}
|
||||
switch role {
|
||||
case "system", "user", "assistant":
|
||||
case "system", "user":
|
||||
if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content"); !ok {
|
||||
return false, reason
|
||||
}
|
||||
raw, exists := message["content"]
|
||||
if !exists {
|
||||
return false, "non_text_message_content"
|
||||
}
|
||||
if ok, reason := grokChatRequiredMessageContentBridgeable(raw); !ok {
|
||||
return false, reason
|
||||
}
|
||||
case "assistant":
|
||||
if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "reasoning_content", "tool_calls"); !ok {
|
||||
return false, reason
|
||||
}
|
||||
reasoningContent := ""
|
||||
if raw, exists := message["reasoning_content"]; exists {
|
||||
if !grokChatJSONNull(raw) && json.Unmarshal(raw, &reasoningContent) != nil {
|
||||
return false, "invalid_reasoning_content"
|
||||
}
|
||||
}
|
||||
hasReasoningContent := strings.TrimSpace(reasoningContent) != ""
|
||||
toolCallCount := 0
|
||||
if raw, exists := message["tool_calls"]; exists {
|
||||
var reason string
|
||||
toolCallCount, reason = grokChatAssistantToolCallsBridgeable(raw)
|
||||
if reason != "" {
|
||||
return false, reason
|
||||
}
|
||||
}
|
||||
raw, hasContent := message["content"]
|
||||
if !hasContent || strings.TrimSpace(string(raw)) == "null" {
|
||||
if toolCallCount == 0 && !hasReasoningContent {
|
||||
return false, "non_text_message_content"
|
||||
}
|
||||
continue
|
||||
}
|
||||
var content string
|
||||
if json.Unmarshal(raw, &content) == nil {
|
||||
if strings.TrimSpace(content) == "" && toolCallCount == 0 && !hasReasoningContent {
|
||||
return false, "empty_message_content"
|
||||
}
|
||||
continue
|
||||
}
|
||||
if ok, reason := grokChatStructuredContentBridgeable(raw); !ok {
|
||||
// The converter can still emit a standalone reasoning part when
|
||||
// an otherwise empty content array accompanies reasoning_content.
|
||||
// Do not broaden this exception to unsupported/malformed parts.
|
||||
if !hasReasoningContent || reason != "empty_message_content" {
|
||||
return false, reason
|
||||
}
|
||||
}
|
||||
case "tool":
|
||||
if ok, reason := grokChatMessageFieldsBridgeable(message, "role", "content", "tool_call_id"); !ok {
|
||||
return false, reason
|
||||
}
|
||||
var callID string
|
||||
if raw, exists := message["tool_call_id"]; !exists || json.Unmarshal(raw, &callID) != nil || strings.TrimSpace(callID) == "" {
|
||||
return false, "invalid_tool_call_id"
|
||||
}
|
||||
var output string
|
||||
if raw, exists := message["content"]; !exists || json.Unmarshal(raw, &output) != nil || output == "" {
|
||||
return false, "invalid_tool_message_content"
|
||||
}
|
||||
default:
|
||||
return false, "unsupported_message_role_" + role
|
||||
}
|
||||
raw, exists := message["content"]
|
||||
if !exists {
|
||||
return false, "non_text_message_content"
|
||||
}
|
||||
var content string
|
||||
if json.Unmarshal(raw, &content) == nil {
|
||||
if strings.TrimSpace(content) == "" {
|
||||
return false, "empty_message_content"
|
||||
}
|
||||
continue
|
||||
}
|
||||
// Structured content: only allow arrays whose parts are text or
|
||||
// image_url. These are losslessly convertible to Responses input_text/
|
||||
// input_image parts, so the bridge preserves Chat Completions semantics.
|
||||
if ok, reason := grokChatStructuredContentBridgeable(raw); !ok {
|
||||
return false, reason
|
||||
}
|
||||
}
|
||||
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func grokChatFunctionDeclarationsBridgeable(raw json.RawMessage) (bool, string) {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return true, ""
|
||||
}
|
||||
var declarations []json.RawMessage
|
||||
if json.Unmarshal(raw, &declarations) != nil {
|
||||
return false, "invalid_tools"
|
||||
}
|
||||
for _, declaration := range declarations {
|
||||
var tool map[string]json.RawMessage
|
||||
if json.Unmarshal(declaration, &tool) != nil || tool == nil {
|
||||
return false, "invalid_tool"
|
||||
}
|
||||
for field := range tool {
|
||||
if field != "type" && field != "function" {
|
||||
return false, "unsafe_tool_field_" + field
|
||||
}
|
||||
}
|
||||
var toolType string
|
||||
if rawType, exists := tool["type"]; !exists || json.Unmarshal(rawType, &toolType) != nil || toolType != "function" {
|
||||
return false, "unsupported_tool_type"
|
||||
}
|
||||
functionRaw, exists := tool["function"]
|
||||
if !exists {
|
||||
return false, "invalid_tool_function"
|
||||
}
|
||||
|
||||
var function map[string]json.RawMessage
|
||||
if json.Unmarshal(functionRaw, &function) != nil || function == nil {
|
||||
return false, "invalid_tool_function"
|
||||
}
|
||||
for field := range function {
|
||||
switch field {
|
||||
case "name", "description", "parameters", "strict":
|
||||
default:
|
||||
return false, "unsafe_tool_function_field_" + field
|
||||
}
|
||||
}
|
||||
var name string
|
||||
if rawName, exists := function["name"]; !exists || json.Unmarshal(rawName, &name) != nil || strings.TrimSpace(name) == "" {
|
||||
return false, "invalid_tool_function_name"
|
||||
}
|
||||
if rawDescription, exists := function["description"]; exists {
|
||||
var description string
|
||||
if json.Unmarshal(rawDescription, &description) != nil {
|
||||
return false, "invalid_tool_function_description"
|
||||
}
|
||||
}
|
||||
var parameters map[string]json.RawMessage
|
||||
if rawParameters, exists := function["parameters"]; !exists || json.Unmarshal(rawParameters, ¶meters) != nil || parameters == nil {
|
||||
return false, "invalid_tool_function_parameters"
|
||||
}
|
||||
if rawStrict, exists := function["strict"]; exists {
|
||||
var strict bool
|
||||
if json.Unmarshal(rawStrict, &strict) != nil {
|
||||
return false, "invalid_tool_function_strict"
|
||||
}
|
||||
}
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func grokChatToolChoiceBridgeable(raw json.RawMessage) (bool, string) {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return true, ""
|
||||
}
|
||||
var choice string
|
||||
if json.Unmarshal(raw, &choice) != nil {
|
||||
return false, "unsupported_tool_choice"
|
||||
}
|
||||
switch choice {
|
||||
case "auto", "none", "required":
|
||||
return true, ""
|
||||
default:
|
||||
return false, "unsupported_tool_choice"
|
||||
}
|
||||
}
|
||||
|
||||
func grokChatHasFunctionDeclarations(root map[string]json.RawMessage) bool {
|
||||
for _, field := range []string{"tools", "functions"} {
|
||||
raw, exists := root[field]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
var declarations []json.RawMessage
|
||||
if json.Unmarshal(raw, &declarations) == nil && len(declarations) > 0 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func grokChatMessageFieldsBridgeable(message map[string]json.RawMessage, allowedFields ...string) (bool, string) {
|
||||
allowed := make(map[string]struct{}, len(allowedFields))
|
||||
for _, field := range allowedFields {
|
||||
allowed[field] = struct{}{}
|
||||
}
|
||||
for field := range message {
|
||||
if _, ok := allowed[field]; !ok {
|
||||
return false, "unsafe_message_field_" + field
|
||||
}
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func grokChatRequiredMessageContentBridgeable(raw json.RawMessage) (bool, string) {
|
||||
var content string
|
||||
if json.Unmarshal(raw, &content) == nil {
|
||||
if strings.TrimSpace(content) == "" {
|
||||
return false, "empty_message_content"
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
// Structured content: only allow arrays whose parts are text or
|
||||
// image_url. These are losslessly convertible to Responses input_text/
|
||||
// input_image parts, so the bridge preserves Chat Completions semantics.
|
||||
return grokChatStructuredContentBridgeable(raw)
|
||||
}
|
||||
|
||||
func grokChatAssistantToolCallsBridgeable(raw json.RawMessage) (int, string) {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return 0, ""
|
||||
}
|
||||
var calls []map[string]json.RawMessage
|
||||
if json.Unmarshal(raw, &calls) != nil {
|
||||
return 0, "invalid_tool_calls"
|
||||
}
|
||||
for _, call := range calls {
|
||||
if call == nil {
|
||||
return 0, "invalid_tool_call"
|
||||
}
|
||||
for field := range call {
|
||||
switch field {
|
||||
case "id", "type", "function", "index":
|
||||
default:
|
||||
return 0, "unsafe_tool_call_field_" + field
|
||||
}
|
||||
}
|
||||
if rawIndex, exists := call["index"]; exists {
|
||||
var index *int
|
||||
if json.Unmarshal(rawIndex, &index) != nil || (index != nil && *index < 0) {
|
||||
return 0, "invalid_tool_call_index"
|
||||
}
|
||||
}
|
||||
var callID string
|
||||
if rawID, exists := call["id"]; !exists || json.Unmarshal(rawID, &callID) != nil || strings.TrimSpace(callID) == "" {
|
||||
return 0, "invalid_tool_call_id"
|
||||
}
|
||||
var callType string
|
||||
if rawType, exists := call["type"]; !exists || json.Unmarshal(rawType, &callType) != nil || callType != "function" {
|
||||
return 0, "unsupported_tool_call_type"
|
||||
}
|
||||
var function map[string]json.RawMessage
|
||||
if rawFunction, exists := call["function"]; !exists || json.Unmarshal(rawFunction, &function) != nil || function == nil {
|
||||
return 0, "invalid_tool_call_function"
|
||||
}
|
||||
for field := range function {
|
||||
if field != "name" && field != "arguments" {
|
||||
return 0, "unsafe_tool_call_function_field_" + field
|
||||
}
|
||||
}
|
||||
var name string
|
||||
if rawName, exists := function["name"]; !exists || json.Unmarshal(rawName, &name) != nil || strings.TrimSpace(name) == "" {
|
||||
return 0, "invalid_tool_call_function_name"
|
||||
}
|
||||
var arguments string
|
||||
if rawArguments, exists := function["arguments"]; !exists || json.Unmarshal(rawArguments, &arguments) != nil || !json.Valid([]byte(arguments)) {
|
||||
return 0, "invalid_tool_call_arguments"
|
||||
}
|
||||
}
|
||||
return len(calls), ""
|
||||
}
|
||||
|
||||
func grokChatStructuredContentBridgeable(raw json.RawMessage) (bool, string) {
|
||||
var parts []map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &parts); err != nil {
|
||||
@@ -199,14 +456,6 @@ func grokChatStructuredContentBridgeable(raw json.RawMessage) (bool, string) {
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func grokChatNullOrEmptyArray(raw json.RawMessage) bool {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return true
|
||||
}
|
||||
var values []json.RawMessage
|
||||
return json.Unmarshal(raw, &values) == nil && len(values) == 0
|
||||
}
|
||||
|
||||
func grokChatNullOrNone(raw json.RawMessage) bool {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return true
|
||||
@@ -215,14 +464,34 @@ func grokChatNullOrNone(raw json.RawMessage) bool {
|
||||
return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "none")
|
||||
}
|
||||
|
||||
func grokChatCacheIntentBody(body []byte) ([]byte, error) {
|
||||
func grokChatJSONNull(raw json.RawMessage) bool {
|
||||
return strings.TrimSpace(string(raw)) == "null"
|
||||
}
|
||||
|
||||
func grokChatNullOrEmptyArray(raw json.RawMessage) bool {
|
||||
if strings.TrimSpace(string(raw)) == "null" {
|
||||
return true
|
||||
}
|
||||
var values []json.RawMessage
|
||||
return json.Unmarshal(raw, &values) == nil && len(values) == 0
|
||||
}
|
||||
|
||||
func grokChatResponsesCacheIntentBody(body []byte) ([]byte, error) {
|
||||
// An empty Chat tools array is omitted by the Responses converter. In that
|
||||
// case auto/none is also a semantic no-op and must not suppress the normal
|
||||
// tool-free cache route. Non-empty converted tools are always kept intact.
|
||||
if gjson.GetBytes(body, "tools").Exists() {
|
||||
return append([]byte(nil), body...), nil
|
||||
}
|
||||
choice := gjson.GetBytes(body, "tool_choice")
|
||||
if !choice.Exists() || choice.Type != gjson.String || (choice.String() != "auto" && choice.String() != "none") {
|
||||
return append([]byte(nil), body...), nil
|
||||
}
|
||||
var root map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &root); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, field := range []string{"tools", "tool_choice", "functions", "function_call"} {
|
||||
delete(root, field)
|
||||
}
|
||||
delete(root, "tool_choice")
|
||||
return json.Marshal(root)
|
||||
}
|
||||
|
||||
@@ -268,6 +537,10 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
|
||||
}
|
||||
responsesReq.Model = upstreamModel
|
||||
responsesReq.Stream = true
|
||||
// Keep Chat and native Responses paths aligned for OpenAI-compatible
|
||||
// service_tier aliases (for example, "fast" -> "priority"). Unknown
|
||||
// values are omitted by the shared normalizer instead of reaching xAI.
|
||||
normalizeResponsesRequestServiceTier(responsesReq)
|
||||
// These fields are useful to Codex but are not needed by the Grok CLI
|
||||
// protocol. Keep the bridge request as close as possible to native Grok.
|
||||
responsesReq.Include = nil
|
||||
@@ -277,18 +550,25 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal grok responses bridge request: %w", err)
|
||||
}
|
||||
// Preserve the converted Responses intent before Grok capability
|
||||
// sanitization. Cache routing must see the actual client function tools,
|
||||
// not the nested Chat Completions declarations and not a tool-free copy.
|
||||
intentBody, err := grokChatResponsesCacheIntentBody(responsesBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("normalize grok responses bridge cache intent: %w", err)
|
||||
}
|
||||
responsesBody, err = patchGrokResponsesBody(responsesBody, upstreamModel)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("patch grok responses bridge request: %w", err)
|
||||
}
|
||||
intentBody, err := grokChatCacheIntentBody(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("normalize grok responses bridge tool intent: %w", err)
|
||||
}
|
||||
responsesBody, err = applyGrokResponsesCacheIdentity(responsesBody, intentBody, cacheIdentity, true)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("apply grok responses bridge cache identity: %w", err)
|
||||
}
|
||||
responsesBody, err = applyGrokFreeRequestToolCacheRoute(c, responsesBody, intentBody, account, cacheIdentity)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("apply grok responses bridge function-tool cache route: %w", err)
|
||||
}
|
||||
|
||||
updatedBody, policyErr := s.applyOpenAIFastPolicyToBody(ctx, account, upstreamModel, responsesBody)
|
||||
if policyErr != nil {
|
||||
|
||||
@@ -39,6 +39,31 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) {
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stream_options":{"include_usage":true},"max_completion_tokens":256,"temperature":0.2,"top_p":0.9,"prompt_cache_key":"session","tools":[],"functions":null,"tool_choice":"none"}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "Trae compatible fields",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"return json"}],"instructions":"Be concise","response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "Trae nullable SDK fields",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"instructions":null,"response_format":null,"service_tier":null,"stop":null,"reasoning_effort":null}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "invalid instructions fall back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"instructions":{"text":"invalid"}}`,
|
||||
reason: "invalid_instructions",
|
||||
},
|
||||
{
|
||||
name: "invalid response format falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"response_format":"json_object"}`,
|
||||
reason: "invalid_response_format",
|
||||
},
|
||||
{
|
||||
name: "invalid service tier falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"service_tier":1}`,
|
||||
reason: "invalid_service_tier",
|
||||
},
|
||||
{
|
||||
name: "stop falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stop":"done"}`,
|
||||
@@ -75,15 +100,80 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) {
|
||||
reason: "empty_message_content",
|
||||
},
|
||||
{
|
||||
name: "function tools fall back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`,
|
||||
reason: "unsupported_tools",
|
||||
name: "function tools bridge",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"},"strict":false}}]}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "automatic tool choice falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`,
|
||||
name: "legacy functions fall back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"functions":[{"name":"lookup","parameters":{"type":"object"}}]}`,
|
||||
reason: "unsupported_functions",
|
||||
},
|
||||
{
|
||||
name: "automatic tool choice bridges",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "required tool choice without tools falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"required"}`,
|
||||
reason: "required_tool_choice_without_tools",
|
||||
},
|
||||
{
|
||||
name: "tool history bridges",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"summarize"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":"auto","parallel_tool_calls":true}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "Trae reasoning and indexed tool history bridges",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","content":null,"reasoning_content":"I should call lookup","tool_calls":[{"index":0,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"summarize"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}]}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "reasoning only assistant history bridges",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","reasoning_content":"Prior reasoning"},{"role":"user","content":"continue"}]}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "invalid reasoning content falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","reasoning_content":{"text":"invalid"}},{"role":"user","content":"continue"}]}`,
|
||||
reason: "invalid_reasoning_content",
|
||||
},
|
||||
{
|
||||
name: "negative tool call index falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"index":-1,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{}"}}]}]}`,
|
||||
reason: "invalid_tool_call_index",
|
||||
},
|
||||
{
|
||||
name: "unknown tool type falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"web_search","function":{"name":"lookup","parameters":{"type":"object"}}}]}`,
|
||||
reason: "unsupported_tool_type",
|
||||
},
|
||||
{
|
||||
name: "missing tool schema falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`,
|
||||
reason: "invalid_tool_function_parameters",
|
||||
},
|
||||
{
|
||||
name: "named tool choice falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup","parameters":{"type":"object"}}}],"tool_choice":{"type":"function","function":{"name":"lookup"}}}`,
|
||||
reason: "unsupported_tool_choice",
|
||||
},
|
||||
{
|
||||
name: "invalid tool call arguments fall back",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{"}}]}]}`,
|
||||
reason: "invalid_tool_call_arguments",
|
||||
},
|
||||
{
|
||||
name: "tool result without call id falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"tool","content":"ok"}]}`,
|
||||
reason: "invalid_tool_call_id",
|
||||
},
|
||||
{
|
||||
name: "non boolean parallel tool calls falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"parallel_tool_calls":"true"}`,
|
||||
reason: "invalid_parallel_tool_calls",
|
||||
},
|
||||
{
|
||||
name: "reasoning effort falls back because conversion adds summary",
|
||||
body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`,
|
||||
@@ -100,9 +190,9 @@ func TestGrokChatResponsesBridgeEligibility(t *testing.T) {
|
||||
reason: "empty_message_content",
|
||||
},
|
||||
{
|
||||
name: "tool history falls back",
|
||||
name: "empty tool history falls back",
|
||||
body: `{"model":"grok","messages":[{"role":"assistant","content":"","tool_calls":[]}]}`,
|
||||
reason: "unsafe_message_field_tool_calls",
|
||||
reason: "empty_message_content",
|
||||
},
|
||||
{
|
||||
name: "unknown field falls back",
|
||||
@@ -183,6 +273,116 @@ func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing.
|
||||
require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey])
|
||||
}
|
||||
|
||||
func TestForwardGrokChatViaResponsesTraeToolHistoryKeepsCacheRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
firstTurnBody := []byte(`{"model":"grok","messages":[{"role":"system","content":"Be concise"},{"role":"user","content":"Find alpha"}],"stream":false,"prompt_cache_key":"trae-session","tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]},"strict":false}}],"tool_choice":"auto","parallel_tool_calls":true}`)
|
||||
body := []byte(`{"model":"grok","messages":[{"role":"system","content":"Be concise"},{"role":"user","content":"Find alpha"},{"role":"assistant","content":null,"tool_calls":[{"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"Summarize"}],"stream":false,"prompt_cache_key":"trae-session","tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]},"strict":false}}],"tool_choice":"auto","parallel_tool_calls":true}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
|
||||
c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
|
||||
c.Set("api_key", &APIKey{ID: 7151})
|
||||
|
||||
account := grokChatBridgeTestAccount(715)
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_trae", 8192)}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
firstTurnIdentity := resolveGrokCacheIdentity(c, firstTurnBody, "", "grok-4.5")
|
||||
extendedTurnIdentity := resolveGrokCacheIdentity(c, body, "", "grok-4.5")
|
||||
require.NotEmpty(t, firstTurnIdentity)
|
||||
require.Equal(t, firstTurnIdentity, extendedTurnIdentity)
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint)
|
||||
require.Equal(t, extendedTurnIdentity, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
|
||||
require.Equal(t, extendedTurnIdentity, upstream.lastReq.Header.Get(grokConversationIDHeader))
|
||||
|
||||
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "function", tools[0].Get("type").String())
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.Equal(t, "Lookup a value", tools[0].Get("description").String())
|
||||
require.Equal(t, "string", tools[0].Get("parameters.properties.key.type").String())
|
||||
require.True(t, tools[0].Get("strict").Exists())
|
||||
require.False(t, tools[0].Get("strict").Bool())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
require.Equal(t, "auto", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, "parallel_tool_calls").Bool())
|
||||
|
||||
require.Equal(t, "function_call", gjson.GetBytes(upstream.lastBody, "input.2.type").String())
|
||||
require.Equal(t, "call_lookup", gjson.GetBytes(upstream.lastBody, "input.2.call_id").String())
|
||||
require.Equal(t, "lookup", gjson.GetBytes(upstream.lastBody, "input.2.name").String())
|
||||
require.Equal(t, `{"key":"alpha"}`, gjson.GetBytes(upstream.lastBody, "input.2.arguments").String())
|
||||
require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.3.type").String())
|
||||
require.Equal(t, "call_lookup", gjson.GetBytes(upstream.lastBody, "input.3.call_id").String())
|
||||
require.Equal(t, `{"value":"ok"}`, gjson.GetBytes(upstream.lastBody, "input.3.output").String())
|
||||
}
|
||||
|
||||
func TestForwardGrokChatViaResponsesTraeCompatibilityFieldsKeepCacheRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
firstTurnBody := []byte(`{"model":"grok","messages":[{"role":"user","content":"Find alpha"}],"instructions":"Return concise JSON","stream":false,"response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null,"tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}}],"tool_choice":"auto","parallel_tool_calls":true}`)
|
||||
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"Find alpha"},{"role":"assistant","content":null,"reasoning_content":"I should use lookup","tool_calls":[{"index":0,"id":"call_lookup","type":"function","function":{"name":"lookup","arguments":"{\"key\":\"alpha\"}"}}]},{"role":"tool","tool_call_id":"call_lookup","content":"{\"value\":\"ok\"}"},{"role":"user","content":"Summarize"}],"instructions":"Return concise JSON","stream":false,"response_format":{"type":"json_object"},"service_tier":"fast","stop":null,"reasoning_effort":null,"tools":[{"type":"function","function":{"name":"lookup","description":"Lookup a value","parameters":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}}}],"tool_choice":"auto","parallel_tool_calls":true}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
|
||||
c.Set("api_key", &APIKey{ID: 7161})
|
||||
|
||||
account := grokChatBridgeTestAccount(716)
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_trae_compat", 12288)}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
firstTurnIdentity := resolveGrokCacheIdentity(c, firstTurnBody, "", "grok-4.5")
|
||||
extendedTurnIdentity := resolveGrokCacheIdentity(c, body, "", "grok-4.5")
|
||||
require.NotEmpty(t, firstTurnIdentity)
|
||||
require.Equal(t, firstTurnIdentity, extendedTurnIdentity)
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint)
|
||||
require.Equal(t, 12288, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, extendedTurnIdentity, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
|
||||
require.Equal(t, extendedTurnIdentity, upstream.lastReq.Header.Get(grokConversationIDHeader))
|
||||
require.Equal(t, "Return concise JSON", gjson.GetBytes(upstream.lastBody, "instructions").String())
|
||||
require.Equal(t, "json_object", gjson.GetBytes(upstream.lastBody, "text.format.type").String())
|
||||
require.Equal(t, "priority", gjson.GetBytes(upstream.lastBody, "service_tier").String())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "stop").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "reasoning").Exists())
|
||||
require.Contains(t, gjson.GetBytes(upstream.lastBody, "input.1.content.0.text").String(), "<thinking>I should use lookup</thinking>")
|
||||
require.Equal(t, "function_call", gjson.GetBytes(upstream.lastBody, "input.2.type").String())
|
||||
require.Equal(t, "function_call_output", gjson.GetBytes(upstream.lastBody, "input.3.type").String())
|
||||
|
||||
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
require.Equal(t, int64(12288), gjson.Get(recorder.Body.String(), "usage.prompt_tokens_details.cached_tokens").Int())
|
||||
}
|
||||
|
||||
func TestForwardGrokChatViaResponsesStreamingPropagatesCachedUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -206,18 +207,27 @@ func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *test
|
||||
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
|
||||
}
|
||||
|
||||
func TestPatchGrokResponsesBodyDropsCodexAdditionalToolsInputItems(t *testing.T) {
|
||||
func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := []byte(`{
|
||||
"model": "grok",
|
||||
"tools": [
|
||||
{"type": "function", "name": "existing", "description": "top-level wins"},
|
||||
{"type": "web_search"}
|
||||
],
|
||||
"tool_choice": "auto",
|
||||
"input": [
|
||||
{
|
||||
"type": "additional_tools",
|
||||
"role": "developer",
|
||||
"tools": [
|
||||
{"type": "namespace", "name": "image_gen"},
|
||||
{"type": "function", "name": "wait"}
|
||||
{"type": "function", "name": "existing", "description": "duplicate carrier definition"},
|
||||
{"type": "function", "name": "wait"},
|
||||
{"type": "web_search"},
|
||||
{"type": "shell"},
|
||||
{"type": "custom", "name": "apply_patch"},
|
||||
{"type": "namespace", "name": "collaboration"}
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -239,12 +249,235 @@ func TestPatchGrokResponsesBodyDropsCodexAdditionalToolsInputItems(t *testing.T)
|
||||
require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String())
|
||||
require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array()))
|
||||
require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists())
|
||||
tools := gjson.GetBytes(patched, "tools").Array()
|
||||
require.Len(t, tools, 4)
|
||||
require.Equal(t, "existing", tools[0].Get("name").String())
|
||||
require.Equal(t, "top-level wins", tools[0].Get("description").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "wait", tools[2].Get("name").String())
|
||||
require.Equal(t, "shell", tools[3].Get("type").String())
|
||||
require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists())
|
||||
require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists())
|
||||
require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String())
|
||||
require.Equal(t, "developer", gjson.GetBytes(patched, "input.0.role").String())
|
||||
require.Equal(t, "system prompt", gjson.GetBytes(patched, "input.0.content.0.text").String())
|
||||
require.Equal(t, "user", gjson.GetBytes(patched, "input.1.role").String())
|
||||
require.Equal(t, "hello", gjson.GetBytes(patched, "input.1.content.0.text").String())
|
||||
}
|
||||
|
||||
func TestForwardGrokResponsesCodexAdditionalToolsUsesMixedCacheIntent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
body := []byte(`{
|
||||
"model":"grok",
|
||||
"stream":false,
|
||||
"prompt_cache_key":"codex-session",
|
||||
"input":[
|
||||
{"type":"additional_tools","role":"developer","tools":[
|
||||
{"type":"function","name":"lookup","description":"look up a key","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}},
|
||||
{"type":"custom","name":"apply_patch"},
|
||||
{"type":"namespace","name":"collaboration"}
|
||||
]},
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}
|
||||
]
|
||||
}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
|
||||
c.Set("api_key", &APIKey{ID: 4501})
|
||||
|
||||
account := healthyGrokOAuthGatewayTestAccount(4501, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{
|
||||
"id":"resp_codex_lite","object":"response","model":"grok-4.5","status":"completed",
|
||||
"output":[],"usage":{"input_tokens":10,"output_tokens":1}
|
||||
}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "resp_codex_lite", result.ResponseID)
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists())
|
||||
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "function", tools[0].Get("type").String())
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())
|
||||
identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
|
||||
require.NotEmpty(t, identity)
|
||||
require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
|
||||
require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
|
||||
}
|
||||
|
||||
func TestForwardGrokResponsesClaudeDesktopClientToolsUseCacheRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
firstBody := []byte(`{
|
||||
"model":"grok","stream":false,"instructions":"You are Claude Desktop.",
|
||||
"tools":[
|
||||
{"type":"function","name":"Read","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"Edit","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"WebSearch","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
|
||||
],
|
||||
"input":[{"role":"user","content":[{"type":"input_text","text":"first turn"}]}]
|
||||
}`)
|
||||
secondBody := []byte(`{
|
||||
"model":"grok","stream":false,"instructions":"You are Claude Desktop.",
|
||||
"tools":[
|
||||
{"type":"function","name":"Read","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"Edit","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"WebSearch","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
|
||||
],
|
||||
"input":[
|
||||
{"role":"user","content":[{"type":"input_text","text":"first turn"}]},
|
||||
{"role":"assistant","content":[{"type":"output_text","text":"first answer"}]},
|
||||
{"role":"user","content":[{"type":"input_text","text":"second turn"}]}
|
||||
]
|
||||
}`)
|
||||
|
||||
account := healthyGrokOAuthGatewayTestAccount(4504, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{
|
||||
"id":"resp_claude_desktop_1","object":"response","model":"grok-4.5","status":"completed",
|
||||
"output":[],"usage":{"input_tokens":30000,"output_tokens":10,"input_tokens_details":{"cached_tokens":0}}
|
||||
}`)),
|
||||
},
|
||||
{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{
|
||||
"id":"resp_claude_desktop_2","object":"response","model":"grok-4.5","status":"completed",
|
||||
"output":[],"usage":{"input_tokens":30100,"output_tokens":12,"input_tokens_details":{"cached_tokens":28672}}
|
||||
}`)),
|
||||
},
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
newContext := func(body []byte) *gin.Context {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)")
|
||||
c.Request.Header.Set("X-App", "cli")
|
||||
c.Request.Header.Set("anthropic-client-platform", "desktop_app")
|
||||
c.Request.Header.Set("X-Claude-Code-Session-Id", "claude-desktop-session")
|
||||
c.Set("api_key", &APIKey{ID: 4504})
|
||||
return c
|
||||
}
|
||||
|
||||
first, err := svc.forwardGrokResponses(context.Background(), newContext(firstBody), account, firstBody, "grok", false, time.Now())
|
||||
require.NoError(t, err)
|
||||
second, err := svc.forwardGrokResponses(context.Background(), newContext(secondBody), account, secondBody, "grok", false, time.Now())
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, 0, first.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, 28672, second.Usage.CacheReadInputTokens)
|
||||
require.Len(t, upstream.bodies, 2)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
for i := range upstream.bodies {
|
||||
tools := gjson.GetBytes(upstream.bodies[i], "tools").Array()
|
||||
require.Len(t, tools, 6)
|
||||
require.Equal(t, "Read", tools[0].Get("name").String())
|
||||
require.Equal(t, "Edit", tools[1].Get("name").String())
|
||||
require.Equal(t, "WebSearch", tools[2].Get("name").String())
|
||||
require.Equal(t, "mcp__workspace__bash", tools[3].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[4].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[5].Get("type").String())
|
||||
require.False(t, gjson.GetBytes(upstream.bodies[i], "tool_choice").Exists())
|
||||
require.Empty(t, upstream.requests[i].Header.Get("X-App"))
|
||||
require.Empty(t, upstream.requests[i].Header.Get("anthropic-client-platform"))
|
||||
require.Empty(t, upstream.requests[i].Header.Get("X-Claude-Code-Session-Id"))
|
||||
}
|
||||
firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
|
||||
secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
|
||||
require.NotEmpty(t, firstIdentity)
|
||||
require.Equal(t, firstIdentity, secondIdentity)
|
||||
require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
|
||||
require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
|
||||
}
|
||||
|
||||
func TestGrokResponsesCacheIdentityIncludesPromotedCodexTools(t *testing.T) {
|
||||
c := newGrokCacheTestContext(4503)
|
||||
lookupBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
|
||||
readBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"read_file","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
|
||||
|
||||
patchedLookup, err := patchGrokResponsesBody(lookupBody, "grok-4.5")
|
||||
require.NoError(t, err)
|
||||
patchedRead, err := patchGrokResponsesBody(readBody, "grok-4.5")
|
||||
require.NoError(t, err)
|
||||
|
||||
lookupIdentity := resolveGrokCacheIdentity(c, patchedLookup, "", "grok-4.5")
|
||||
readIdentity := resolveGrokCacheIdentity(c, patchedRead, "", "grok-4.5")
|
||||
require.NotEmpty(t, lookupIdentity)
|
||||
require.NotEmpty(t, readIdentity)
|
||||
require.NotEqual(t, lookupIdentity, readIdentity)
|
||||
}
|
||||
|
||||
func TestCodexUnsupportedAdditionalToolsDoNotBecomeToolFreeCacheIntent(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model":"grok","tool_choice":"auto",
|
||||
"input":[
|
||||
{"type":"additional_tools","role":"developer","tools":[
|
||||
{"type":"custom","name":"apply_patch"},
|
||||
{"type":"namespace","name":"collaboration"}
|
||||
]},
|
||||
{"type":"message","role":"user","content":"hello"}
|
||||
]
|
||||
}`)
|
||||
patched, err := patchGrokResponsesBody(body, "grok-4.5")
|
||||
require.NoError(t, err)
|
||||
require.False(t, gjson.GetBytes(patched, "tools").Exists())
|
||||
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
|
||||
|
||||
mixedCacheIntent := patched
|
||||
patched, err = applyGrokResponsesCacheIdentity(patched, body, "isolated-id", true)
|
||||
require.NoError(t, err)
|
||||
account := healthyGrokOAuthGatewayTestAccount(4502, "access-token")
|
||||
account.Credentials["subscription_tier"] = "free"
|
||||
patched, err = applyGrokFreeRequestToolCacheRoute(nil, patched, mixedCacheIntent, account, "isolated-id")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, gjson.GetBytes(patched, "tools").Exists())
|
||||
require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
|
||||
require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String())
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
|
||||
@@ -2531,6 +2764,54 @@ func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) {
|
||||
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
||||
}
|
||||
|
||||
func TestOpenAIWSHTTPBridgeSSEErrorSideEffectsRunOncePerPlatform(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, platform := range []string{PlatformOpenAI, PlatformGrok} {
|
||||
t.Run(platform, func(t *testing.T) {
|
||||
repo := &grokQuotaAccountRepo{}
|
||||
cfg := &config.Config{}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n",
|
||||
)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: cfg,
|
||||
accountRepo: repo,
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
if platform == PlatformOpenAI {
|
||||
svc.rateLimitService = NewRateLimitService(repo, nil, cfg, nil, nil)
|
||||
}
|
||||
account := &Account{ID: 70, Platform: platform, Type: AccountTypeOAuth, Concurrency: 1}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
writes := 0
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", 1,
|
||||
func([]byte) error {
|
||||
writes++
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||||
require.Zero(t, writes)
|
||||
require.Equal(t, 1, repo.rateLimitedCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) {
|
||||
repo := &grokQuotaAccountRepo{}
|
||||
resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second)
|
||||
|
||||
@@ -586,6 +586,7 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse(
|
||||
if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
c.Header("Content-Type", "application/json; charset=utf-8")
|
||||
c.JSON(http.StatusOK, anthropicResp)
|
||||
|
||||
return &OpenAIForwardResult{
|
||||
|
||||
@@ -402,6 +402,39 @@ func TestOpenAIResponseFlush_FailedAndErrorEventsFlushAtBoundaries(t *testing.T)
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIResponseFlush_ReusedTypeKeepsSSEBytesAndTerminalSemantics(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
flushCount int
|
||||
}{
|
||||
{
|
||||
name: "whitespace around done",
|
||||
body: "data: \t[DONE] \n\n",
|
||||
flushCount: 1,
|
||||
},
|
||||
{
|
||||
name: "invalid JSON before done",
|
||||
body: "data: {\"type\":\n\ndata: [DONE]\n\n",
|
||||
flushCount: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := newOpenAIResponseFlushRecorder()
|
||||
|
||||
result, err := runOpenAIResponseFlushTest(recorder, io.NopCloser(strings.NewReader(tt.body)), config.GatewayConfig{})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
gotBody, flushes := recorder.snapshot()
|
||||
require.Equal(t, tt.body, gotBody)
|
||||
require.Len(t, flushes, tt.flushCount)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponseFlush_ClientDisconnectStillDrainsUsage(t *testing.T) {
|
||||
first := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"a\"}\n\n"
|
||||
terminal := "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":7,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":2}}}}\n\n"
|
||||
|
||||
@@ -396,10 +396,12 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
|
||||
// Extract data from SSE line (supports both "data: " and "data:" formats)
|
||||
if data, ok := extractOpenAISSEDataLine(line); ok {
|
||||
dataBytes := []byte(data)
|
||||
if openAIStreamEventIsTerminal(data) {
|
||||
eventTypeRaw := gjson.GetBytes(dataBytes, "type").String()
|
||||
eventType := strings.TrimSpace(eventTypeRaw)
|
||||
// 初始上游 data 的 type 只解析一次:原始值保持终止事件的精确匹配,规范化值供后续分支复用。
|
||||
if openAIStreamEventIsTerminalWithType(data, eventTypeRaw) {
|
||||
sawTerminalEvent = true
|
||||
}
|
||||
eventType := strings.TrimSpace(gjson.GetBytes(dataBytes, "type").String())
|
||||
if responseID == "" {
|
||||
responseID = extractOpenAIResponseIDFromJSONBytes(dataBytes)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestOpenAIStreamEventIsTerminalWithTypeMatchesExistingSemantics(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data string
|
||||
want bool
|
||||
}{
|
||||
{name: "empty", data: "", want: false},
|
||||
{name: "whitespace", data: " \t ", want: false},
|
||||
{name: "done", data: " [DONE] ", want: true},
|
||||
{name: "JSON outer whitespace", data: " \n\t {\"type\":\"response.completed\"} \r\n", want: true},
|
||||
{name: "completed", data: `{"type":"response.completed"}`, want: true},
|
||||
{name: "response done", data: `{"type":"response.done"}`, want: true},
|
||||
{name: "failed", data: `{"type":"response.failed"}`, want: true},
|
||||
{name: "incomplete", data: `{"type":"response.incomplete"}`, want: true},
|
||||
{name: "cancelled", data: `{"type":"response.cancelled"}`, want: true},
|
||||
{name: "canceled", data: `{"type":"response.canceled"}`, want: true},
|
||||
{name: "delta", data: `{"type":"response.output_text.delta"}`, want: false},
|
||||
{name: "invalid JSON", data: `{"type":`, want: false},
|
||||
{name: "terminal with trailing garbage", data: `{"type":"response.completed"} trailing`, want: true},
|
||||
{name: "nonterminal with trailing garbage", data: `{"type":"response.output_text.delta"} trailing`, want: false},
|
||||
{name: "type whitespace remains nonterminal", data: `{"type":" response.completed "}`, want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
eventType := gjson.GetBytes([]byte(tt.data), "type").String()
|
||||
got := openAIStreamEventIsTerminalWithType(tt.data, eventType)
|
||||
|
||||
require.Equal(t, tt.want, got)
|
||||
require.Equal(t, openAIStreamEventIsTerminal(tt.data), got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
benchmarkOpenAIResponseSSEEventTypeSink string
|
||||
benchmarkOpenAIResponseSSETerminalSink bool
|
||||
)
|
||||
|
||||
func BenchmarkOpenAIResponseSSETypeExtraction(b *testing.B) {
|
||||
data := `{"type":"response.output_text.delta","sequence_number":42,"delta":"streaming response benchmark payload"}`
|
||||
dataBytes := []byte(data)
|
||||
|
||||
b.Run("legacy double parse", func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(dataBytes)))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
benchmarkOpenAIResponseSSETerminalSink = openAIStreamEventIsTerminal(data)
|
||||
benchmarkOpenAIResponseSSEEventTypeSink = strings.TrimSpace(gjson.GetBytes(dataBytes, "type").String())
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("reused single parse", func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(dataBytes)))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
eventTypeRaw := gjson.GetBytes(dataBytes, "type").String()
|
||||
benchmarkOpenAIResponseSSEEventTypeSink = strings.TrimSpace(eventTypeRaw)
|
||||
benchmarkOpenAIResponseSSETerminalSink = openAIStreamEventIsTerminalWithType(data, eventTypeRaw)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -196,12 +196,13 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
releaseUpstreamCtx()
|
||||
return nil, err
|
||||
}
|
||||
grokMixedCacheIntentBody := append([]byte(nil), body...)
|
||||
body, err = applyGrokResponsesCacheIdentity(body, grokIntentSourceBody, grokCacheIdentity, account.IsGrokOAuth())
|
||||
if err != nil {
|
||||
releaseUpstreamCtx()
|
||||
return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
|
||||
}
|
||||
body, err = applyGrokFreeMessagesFunctionToolCacheRoute(body, grokIntentSourceBody, account, grokCacheIdentity)
|
||||
body, err = applyGrokFreeRequestToolCacheRoute(c, body, grokMixedCacheIntentBody, account, grokCacheIdentity)
|
||||
if err != nil {
|
||||
releaseUpstreamCtx()
|
||||
return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
|
||||
@@ -230,6 +231,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
turnStart := time.Now()
|
||||
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
if turn == 1 {
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
|
||||
}
|
||||
safeErr := sanitizeUpstreamErrorMessage(err.Error())
|
||||
_ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed"))
|
||||
return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr)
|
||||
@@ -238,16 +242,23 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes))
|
||||
if account.Platform == PlatformGrok {
|
||||
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
} else if shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody) {
|
||||
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel)
|
||||
}
|
||||
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
||||
if upstreamMsg == "" {
|
||||
upstreamMsg = http.StatusText(resp.StatusCode)
|
||||
}
|
||||
shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody)
|
||||
if account.Platform == PlatformGrok {
|
||||
s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if turn == 1 && shouldFailover {
|
||||
return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, false)
|
||||
}
|
||||
} else if turn == 1 && shouldFailover {
|
||||
return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, respBody)
|
||||
}
|
||||
if account.Platform != PlatformGrok && (shouldFailover || shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody)) {
|
||||
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel)
|
||||
}
|
||||
_ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg))
|
||||
return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
@@ -372,6 +383,31 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
}
|
||||
replayCollector.AddEvent(eventType, upstreamMessage)
|
||||
|
||||
var upstreamEventErr error
|
||||
if eventType == "error" {
|
||||
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
|
||||
errMessage := strings.TrimSpace(errMsgRaw)
|
||||
if errMessage == "" {
|
||||
errMessage = "upstream error event"
|
||||
}
|
||||
statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw)
|
||||
shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(statusCode, errMessage, upstreamMessage)
|
||||
if account.Platform == PlatformGrok {
|
||||
s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage)
|
||||
} else if shouldFailover {
|
||||
accountStatus := statusCode
|
||||
if transientStatus := openAIWSPayloadTransientStatus(upstreamMessage); transientStatus != 0 {
|
||||
accountStatus = transientStatus
|
||||
}
|
||||
canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel)
|
||||
s.handleOpenAIAccountUpstreamError(ctx, account, accountStatus, resp.Header, upstreamMessage, canonicalModel)
|
||||
}
|
||||
if turn == 1 && !wroteDownstream && shouldFailover {
|
||||
return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false)
|
||||
}
|
||||
upstreamEventErr = errors.New(errMessage)
|
||||
}
|
||||
|
||||
if !clientDisconnected {
|
||||
if err := writeClientMessage(upstreamMessage); err != nil {
|
||||
if isOpenAIWSClientDisconnectError(err) {
|
||||
@@ -396,15 +432,8 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
}
|
||||
}
|
||||
|
||||
if eventType == "error" {
|
||||
s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage)
|
||||
errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage)
|
||||
s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw)
|
||||
errMessage := strings.TrimSpace(errMsgRaw)
|
||||
if errMessage == "" {
|
||||
errMessage = "upstream error event"
|
||||
}
|
||||
return resultWithUsage(), errors.New(errMessage)
|
||||
if upstreamEventErr != nil {
|
||||
return resultWithUsage(), upstreamEventErr
|
||||
}
|
||||
if isOpenAIWSTerminalEvent(eventType) {
|
||||
upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage)
|
||||
@@ -432,12 +461,20 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return resultWithUsage(), fmt.Errorf("read upstream http bridge stream: %w", err)
|
||||
streamErr := fmt.Errorf("read upstream http bridge stream: %w", err)
|
||||
if turn == 1 && !wroteDownstream {
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, streamErr, true)
|
||||
}
|
||||
return resultWithUsage(), streamErr
|
||||
}
|
||||
if sawDone && eventCount > 0 {
|
||||
return resultWithUsage(), nil
|
||||
terminalErr := errors.New("upstream http bridge stream ended before terminal event")
|
||||
if sawDone {
|
||||
terminalErr = errors.New("upstream http bridge stream sent [DONE] before terminal event")
|
||||
}
|
||||
return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event")
|
||||
if turn == 1 && !wroteDownstream {
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, terminalErr, true)
|
||||
}
|
||||
return resultWithUsage(), terminalErr
|
||||
}
|
||||
|
||||
func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte, originalModel string) (string, error) {
|
||||
@@ -446,6 +483,10 @@ func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte
|
||||
return "", err
|
||||
}
|
||||
upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel)
|
||||
body, err = patchGrokResponsesBody(body, upstreamModel)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return resolveGrokCacheIdentity(c, body, "", upstreamModel), nil
|
||||
}
|
||||
|
||||
|
||||
@@ -62,6 +62,219 @@ func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) {
|
||||
require.True(t, svc.shouldBridgeOpenAIWSHTTP(&Account{Platform: PlatformGrok}, 1, "resp_existing"))
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnTransportErrorFailoverSafety(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
turn int
|
||||
wantFailover bool
|
||||
wantWrites int
|
||||
}{
|
||||
{name: "first_turn_fails_over_before_downstream_event", turn: 1, wantFailover: true},
|
||||
{name: "later_turn_does_not_replay_completed_turns", turn: 2, wantWrites: 1},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{err: io.EOF}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: &config.Config{},
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
account := &Account{
|
||||
ID: 8,
|
||||
Name: "api-key",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
var writes [][]byte
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", tt.turn,
|
||||
func(message []byte) error {
|
||||
writes = append(writes, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if tt.wantFailover {
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
||||
require.JSONEq(t, string(openAITransportFailoverBody), string(failoverErr.ResponseBody))
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
}
|
||||
require.Len(t, writes, tt.wantWrites)
|
||||
if tt.wantWrites > 0 {
|
||||
require.Equal(t, "error", gjson.GetBytes(writes[0], "type").String())
|
||||
require.Equal(t, int64(http.StatusBadGateway), gjson.GetBytes(writes[0], "status").Int())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnHTTPStatusFailoverSafety(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
turn int
|
||||
status int
|
||||
wantFailover bool
|
||||
wantWrites int
|
||||
}{
|
||||
{name: "first_turn_401", turn: 1, status: http.StatusUnauthorized, wantFailover: true},
|
||||
{name: "first_turn_429", turn: 1, status: http.StatusTooManyRequests, wantFailover: true},
|
||||
{name: "first_turn_500", turn: 1, status: http.StatusInternalServerError, wantFailover: true},
|
||||
{name: "later_turn_500_does_not_replay", turn: 2, status: http.StatusInternalServerError, wantWrites: 1},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: tt.status,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"type":"server_error","message":"temporary upstream failure"}}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{ID: 9, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
var writes [][]byte
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", tt.turn,
|
||||
func(message []byte) error {
|
||||
writes = append(writes, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if tt.wantFailover {
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, tt.status, failoverErr.StatusCode)
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
}
|
||||
require.Len(t, writes, tt.wantWrites)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnSSEErrorFailoverSafety(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, turn := range []int{1, 2} {
|
||||
t.Run(fmt.Sprintf("turn_%d", turn), func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n",
|
||||
)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
var writes [][]byte
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", turn,
|
||||
func(message []byte) error {
|
||||
writes = append(writes, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if turn == 1 {
|
||||
require.Nil(t, result)
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||||
require.Empty(t, writes)
|
||||
} else {
|
||||
require.NotNil(t, result)
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
require.Len(t, writes, 1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantFailover bool
|
||||
wantWrites int
|
||||
}{
|
||||
{name: "done_without_events_fails_over", body: "data: [DONE]\n\n", wantFailover: true},
|
||||
{
|
||||
name: "created_then_done_is_truncated_not_success",
|
||||
body: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_truncated\"}}\n\n" +
|
||||
"data: [DONE]\n\n",
|
||||
wantWrites: 1,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(tt.body)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{ID: 11, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
var writes [][]byte
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", 1,
|
||||
func(message []byte) error {
|
||||
writes = append(writes, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if tt.wantFailover {
|
||||
require.Nil(t, result)
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
} else {
|
||||
require.NotNil(t, result)
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
}
|
||||
require.Len(t, writes, tt.wantWrites)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -235,6 +448,78 @@ func TestProxyOpenAIWSHTTPBridgeTurnForGrokDefaultsEmptyModelTo45(t *testing.T)
|
||||
require.Len(t, events, 2)
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnPromotesCodexAdditionalToolsForMixedCache(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`data: {"type":"response.created","response":{"id":"resp_grok_codex_lite","model":"grok-4.5"}}`,
|
||||
"",
|
||||
`data: {"type":"response.completed","response":{"id":"resp_grok_codex_lite","model":"grok-4.5","usage":{"input_tokens":4,"output_tokens":1}}}`,
|
||||
"",
|
||||
}, "\n"))),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}},
|
||||
httpUpstream: upstream,
|
||||
}
|
||||
account := &Account{
|
||||
ID: 73,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
"subscription_tier": "free",
|
||||
},
|
||||
}
|
||||
payload := []byte(`{
|
||||
"type":"response.create","generate":true,"model":"grok","stream":true,
|
||||
"input":[
|
||||
{"type":"additional_tools","role":"developer","tools":[
|
||||
{"type":"function","name":"lookup","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"web_search","parameters":{"type":"object"}},
|
||||
{"type":"custom","name":"apply_patch"},
|
||||
{"type":"namespace","name":"collaboration"}
|
||||
]},
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}
|
||||
]
|
||||
}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
|
||||
var events [][]byte
|
||||
|
||||
result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "access-token", payload, len(payload),
|
||||
"grok", "", "", "", "isolated-ws-cache-id", 1,
|
||||
func(message []byte) error {
|
||||
events = append(events, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Len(t, events, 2)
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists())
|
||||
tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
|
||||
require.Len(t, tools, 3)
|
||||
require.Equal(t, "function", tools[0].Get("type").String())
|
||||
require.Equal(t, "lookup", tools[0].Get("name").String())
|
||||
require.Equal(t, "web_search", tools[1].Get("type").String())
|
||||
require.Equal(t, "x_search", tools[2].Get("type").String())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())
|
||||
require.Equal(t, "isolated-ws-cache-id", gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
|
||||
require.Equal(t, "isolated-ws-cache-id", upstream.lastReq.Header.Get(grokConversationIDHeader))
|
||||
require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
|
||||
}
|
||||
|
||||
func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -219,6 +219,11 @@ type opsScheduledReport struct {
|
||||
NextRunAt time.Time
|
||||
}
|
||||
|
||||
type opsScheduledReportContent struct {
|
||||
html string
|
||||
overview *OpsDashboardOverview
|
||||
}
|
||||
|
||||
func (s *OpsScheduledReportService) listScheduledReports(ctx context.Context, now time.Time) []*opsScheduledReport {
|
||||
if s == nil || s.opsService == nil {
|
||||
return nil
|
||||
@@ -316,11 +321,11 @@ func (s *OpsScheduledReportService) runReport(ctx context.Context, report *opsSc
|
||||
// Mark as "run" up-front so a broken SMTP config doesn't spam retries every minute.
|
||||
s.setLastRunAt(ctx, report.ReportType, now)
|
||||
|
||||
content, err := s.generateReportHTML(ctx, report, now)
|
||||
content, err := s.generateReportContent(ctx, report, now)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if strings.TrimSpace(content) == "" {
|
||||
if strings.TrimSpace(content.html) == "" {
|
||||
// Skip sending when the report decides not to emit content (e.g., digest below min count).
|
||||
return 0, nil
|
||||
}
|
||||
@@ -336,9 +341,6 @@ func (s *OpsScheduledReportService) runReport(ctx context.Context, report *opsSc
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
subject := fmt.Sprintf("[Ops Report] %s", strings.TrimSpace(report.Name))
|
||||
templateVariables := opsScheduledReportEmailVariables(report, now)
|
||||
|
||||
attempts := 0
|
||||
for _, to := range recipients {
|
||||
addr := strings.TrimSpace(to)
|
||||
@@ -346,25 +348,40 @@ func (s *OpsScheduledReportService) runReport(ctx context.Context, report *opsSc
|
||||
continue
|
||||
}
|
||||
attempts++
|
||||
locale := ""
|
||||
if s.emailService.notificationEmailService != nil {
|
||||
locale = s.emailService.notificationEmailService.ResolveRecipientLocale(ctx, 0, addr)
|
||||
templateVariables := opsScheduledReportLocalizedEmailVariables(report, now, locale)
|
||||
rawHTMLVariables := map[string]string{"report_html": content.html}
|
||||
if isOpsSummaryReport(report) {
|
||||
templateVariables = opsSummaryReportEmailVariables(report, now, content.overview, locale)
|
||||
}
|
||||
if err := s.emailService.notificationEmailService.Send(ctx, NotificationEmailSendInput{
|
||||
Event: NotificationEmailEventOpsScheduledReport,
|
||||
RecipientEmail: addr,
|
||||
RecipientName: emailRecipientName(addr),
|
||||
SourceType: "ops_scheduled_report",
|
||||
SourceID: opsScheduledReportDeliverySourceID(report),
|
||||
ReminderKey: now.UTC().Format("2006-01-02T15:04"),
|
||||
Variables: templateVariables,
|
||||
RawHTMLVariables: map[string]string{
|
||||
"report_html": content,
|
||||
},
|
||||
Event: NotificationEmailEventOpsScheduledReport,
|
||||
Locale: locale,
|
||||
RecipientEmail: addr,
|
||||
RecipientName: emailRecipientName(addr),
|
||||
SourceType: "ops_scheduled_report",
|
||||
SourceID: opsScheduledReportDeliverySourceID(report),
|
||||
ReminderKey: now.UTC().Format("2006-01-02T15:04"),
|
||||
Variables: templateVariables,
|
||||
RawHTMLVariables: rawHTMLVariables,
|
||||
}); err == nil {
|
||||
continue
|
||||
} else if !shouldFallbackNotificationEmail(err) {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if err := s.emailService.SendEmail(ctx, addr, subject, content); err != nil {
|
||||
subjectName := strings.TrimSpace(report.Name)
|
||||
if locale != "" {
|
||||
subjectName = opsScheduledReportLocalizedName(report, locale)
|
||||
}
|
||||
subjectPrefix := "[Ops Report]"
|
||||
if strings.HasPrefix(strings.ToLower(locale), "zh") {
|
||||
subjectPrefix = "[运维报表]"
|
||||
}
|
||||
subject := fmt.Sprintf("%s %s", subjectPrefix, subjectName)
|
||||
if err := s.emailService.SendEmail(ctx, addr, subject, content.html); err != nil {
|
||||
// Ignore per-recipient failures; continue best-effort.
|
||||
continue
|
||||
}
|
||||
@@ -412,12 +429,138 @@ func opsScheduledReportEmailVariables(report *opsScheduledReport, now time.Time)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpsScheduledReportService) generateReportHTML(ctx context.Context, report *opsScheduledReport, now time.Time) (string, error) {
|
||||
func opsScheduledReportLocalizedEmailVariables(report *opsScheduledReport, now time.Time, locale string) map[string]string {
|
||||
variables := opsScheduledReportEmailVariables(report, now)
|
||||
variables["report_html"] = ""
|
||||
variables["report_detail_display"] = "block"
|
||||
for _, placeholder := range notificationEmailOpsSummaryPlaceholders {
|
||||
variables[placeholder] = "-"
|
||||
}
|
||||
variables["report_summary_display"] = "none"
|
||||
if name := opsScheduledReportLocalizedName(report, locale); name != "" {
|
||||
variables["report_name"] = name
|
||||
}
|
||||
return variables
|
||||
}
|
||||
|
||||
func opsScheduledReportLocalizedName(report *opsScheduledReport, locale string) string {
|
||||
if report == nil {
|
||||
return "Ops report"
|
||||
}
|
||||
chinese := strings.HasPrefix(strings.ToLower(strings.TrimSpace(locale)), "zh")
|
||||
switch strings.TrimSpace(report.ReportType) {
|
||||
case "daily_summary":
|
||||
if chinese {
|
||||
return "日报"
|
||||
}
|
||||
return "Daily summary"
|
||||
case "weekly_summary":
|
||||
if chinese {
|
||||
return "周报"
|
||||
}
|
||||
return "Weekly summary"
|
||||
case "error_digest":
|
||||
if chinese {
|
||||
return "错误摘要"
|
||||
}
|
||||
return "Error digest"
|
||||
case "account_health":
|
||||
if chinese {
|
||||
return "账号健康"
|
||||
}
|
||||
return "Account health"
|
||||
default:
|
||||
return strings.TrimSpace(report.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func isOpsSummaryReport(report *opsScheduledReport) bool {
|
||||
if report == nil {
|
||||
return false
|
||||
}
|
||||
switch strings.TrimSpace(report.ReportType) {
|
||||
case "daily_summary", "weekly_summary":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func opsSummaryReportEmailVariables(report *opsScheduledReport, now time.Time, overview *OpsDashboardOverview, locale string) map[string]string {
|
||||
variables := opsScheduledReportLocalizedEmailVariables(report, now, locale)
|
||||
variables["report_detail_display"] = "none"
|
||||
if overview == nil {
|
||||
for _, placeholder := range notificationEmailOpsSummaryPlaceholders {
|
||||
if placeholder == "report_summary_display" {
|
||||
continue
|
||||
}
|
||||
variables[placeholder] = "-"
|
||||
}
|
||||
variables["report_summary_display"] = "block"
|
||||
return variables
|
||||
}
|
||||
variables["report_summary_display"] = "block"
|
||||
|
||||
variables["report_total_requests"] = formatOpsReportInteger(overview.RequestCountTotal)
|
||||
variables["report_success_count"] = formatOpsReportInteger(overview.SuccessCount)
|
||||
variables["report_sla_error_count"] = formatOpsReportInteger(overview.ErrorCountSLA)
|
||||
variables["report_business_limited_count"] = formatOpsReportInteger(overview.BusinessLimitedCount)
|
||||
variables["report_sla"] = fmt.Sprintf("%.2f%%", overview.SLA*100)
|
||||
variables["report_error_rate"] = fmt.Sprintf("%.2f%%", overview.ErrorRate*100)
|
||||
variables["report_upstream_error_rate"] = fmt.Sprintf("%.2f%%", overview.UpstreamErrorRate*100)
|
||||
variables["report_upstream_error_count_excl_429_529"] = formatOpsReportInteger(overview.UpstreamErrorCountExcl429529)
|
||||
variables["report_upstream_429_count"] = formatOpsReportInteger(overview.Upstream429Count)
|
||||
variables["report_upstream_529_count"] = formatOpsReportInteger(overview.Upstream529Count)
|
||||
variables["report_latency_p50"] = formatOpsReportMilliseconds(overview.Duration.P50)
|
||||
variables["report_latency_p99"] = formatOpsReportMilliseconds(overview.Duration.P99)
|
||||
variables["report_ttft_p50"] = formatOpsReportMilliseconds(overview.TTFT.P50)
|
||||
variables["report_ttft_p99"] = formatOpsReportMilliseconds(overview.TTFT.P99)
|
||||
variables["report_tokens"] = formatOpsReportInteger(overview.TokenConsumed)
|
||||
variables["report_qps_current"] = fmt.Sprintf("%.1f", overview.QPS.Current)
|
||||
variables["report_qps_peak"] = fmt.Sprintf("%.1f", overview.QPS.Peak)
|
||||
variables["report_qps_avg"] = fmt.Sprintf("%.1f", overview.QPS.Avg)
|
||||
variables["report_tps_current"] = fmt.Sprintf("%.1f", overview.TPS.Current)
|
||||
variables["report_tps_peak"] = fmt.Sprintf("%.1f", overview.TPS.Peak)
|
||||
variables["report_tps_avg"] = fmt.Sprintf("%.1f", overview.TPS.Avg)
|
||||
return variables
|
||||
}
|
||||
|
||||
func formatOpsReportInteger(value int64) string {
|
||||
raw := strconv.FormatInt(value, 10)
|
||||
start := 0
|
||||
if strings.HasPrefix(raw, "-") {
|
||||
start = 1
|
||||
}
|
||||
if len(raw)-start <= 3 {
|
||||
return raw
|
||||
}
|
||||
|
||||
var builder strings.Builder
|
||||
builder.Grow(len(raw) + (len(raw)-start-1)/3)
|
||||
_, _ = builder.WriteString(raw[:start])
|
||||
digitLen := len(raw) - start
|
||||
for offset := 0; offset < digitLen; offset++ {
|
||||
if offset > 0 && (digitLen-offset)%3 == 0 {
|
||||
_ = builder.WriteByte(',')
|
||||
}
|
||||
_ = builder.WriteByte(raw[start+offset])
|
||||
}
|
||||
return builder.String()
|
||||
}
|
||||
|
||||
func formatOpsReportMilliseconds(value *int) string {
|
||||
if value == nil {
|
||||
return "-"
|
||||
}
|
||||
return fmt.Sprintf("%s ms", formatOpsReportInteger(int64(*value)))
|
||||
}
|
||||
|
||||
func (s *OpsScheduledReportService) generateReportContent(ctx context.Context, report *opsScheduledReport, now time.Time) (opsScheduledReportContent, error) {
|
||||
if s == nil || s.opsService == nil || report == nil {
|
||||
return "", fmt.Errorf("service not initialized")
|
||||
return opsScheduledReportContent{}, fmt.Errorf("service not initialized")
|
||||
}
|
||||
if report.TimeRange <= 0 {
|
||||
return "", fmt.Errorf("invalid time range")
|
||||
return opsScheduledReportContent{}, fmt.Errorf("invalid time range")
|
||||
}
|
||||
|
||||
end := now.UTC()
|
||||
@@ -444,10 +587,13 @@ func (s *OpsScheduledReportService) generateReportHTML(ctx context.Context, repo
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
return opsScheduledReportContent{}, err
|
||||
}
|
||||
}
|
||||
return buildOpsSummaryEmailHTML(report.Name, start, end, overview), nil
|
||||
return opsScheduledReportContent{
|
||||
html: buildOpsSummaryEmailHTML(report.Name, start, end, overview),
|
||||
overview: overview,
|
||||
}, nil
|
||||
case "error_digest":
|
||||
// Lightweight digest: list recent errors (status>=400) and breakdown by type.
|
||||
startTime := start
|
||||
@@ -460,22 +606,22 @@ func (s *OpsScheduledReportService) generateReportHTML(ctx context.Context, repo
|
||||
}
|
||||
out, err := s.opsService.GetErrorLogs(ctx, filter)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return opsScheduledReportContent{}, err
|
||||
}
|
||||
if report.ErrorDigestMinCount > 0 && out != nil && out.Total < report.ErrorDigestMinCount {
|
||||
return "", nil
|
||||
return opsScheduledReportContent{}, nil
|
||||
}
|
||||
return buildOpsErrorDigestEmailHTML(report.Name, start, end, out), nil
|
||||
return opsScheduledReportContent{html: buildOpsErrorDigestEmailHTML(report.Name, start, end, out)}, nil
|
||||
case "account_health":
|
||||
// Best-effort: use account availability (not error rate yet).
|
||||
avail, err := s.opsService.GetAccountAvailability(ctx, "", nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return opsScheduledReportContent{}, err
|
||||
}
|
||||
_ = report.AccountHealthErrorRateThreshold // reserved for future per-account error rate report
|
||||
return buildOpsAccountHealthEmailHTML(report.Name, start, end, avail), nil
|
||||
return opsScheduledReportContent{html: buildOpsAccountHealthEmailHTML(report.Name, start, end, avail)}, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unknown report type: %s", report.ReportType)
|
||||
return opsScheduledReportContent{}, fmt.Errorf("unknown report type: %s", report.ReportType)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpsSummaryReportEmailVariables(t *testing.T) {
|
||||
latencyP50 := 8231
|
||||
latencyP99 := 151260
|
||||
ttftP50 := 1674
|
||||
ttftP99 := 11222
|
||||
now := time.Date(2026, time.July, 19, 1, 0, 26, 0, time.UTC)
|
||||
report := &opsScheduledReport{
|
||||
Name: "日报",
|
||||
ReportType: "daily_summary",
|
||||
TimeRange: 24 * time.Hour,
|
||||
}
|
||||
overview := &OpsDashboardOverview{
|
||||
RequestCountTotal: 2374,
|
||||
SuccessCount: 1451,
|
||||
ErrorCountSLA: 2,
|
||||
BusinessLimitedCount: 921,
|
||||
SLA: 0.9986,
|
||||
ErrorRate: 0.0014,
|
||||
UpstreamErrorRate: 0.0028,
|
||||
UpstreamErrorCountExcl429529: 4,
|
||||
Upstream429Count: 0,
|
||||
Upstream529Count: 0,
|
||||
TokenConsumed: 121550190,
|
||||
Duration: OpsPercentiles{P50: &latencyP50, P99: &latencyP99},
|
||||
TTFT: OpsPercentiles{P50: &ttftP50, P99: &ttftP99},
|
||||
QPS: OpsRateSummary{Current: 0, Peak: 1.2, Avg: 0},
|
||||
TPS: OpsRateSummary{Current: 0, Peak: 133421.2, Avg: 1406.8},
|
||||
}
|
||||
|
||||
variables := opsSummaryReportEmailVariables(report, now, overview, "en")
|
||||
require.Equal(t, "Daily summary", variables["report_name"])
|
||||
require.Equal(t, "daily_summary", variables["report_type"])
|
||||
require.Equal(t, "2026-07-18T01:00:26Z", variables["report_start_time"])
|
||||
require.Equal(t, "2026-07-19T01:00:26Z", variables["report_end_time"])
|
||||
require.Equal(t, "2,374", variables["report_total_requests"])
|
||||
require.Equal(t, "1,451", variables["report_success_count"])
|
||||
require.Equal(t, "2", variables["report_sla_error_count"])
|
||||
require.Equal(t, "921", variables["report_business_limited_count"])
|
||||
require.Equal(t, "99.86%", variables["report_sla"])
|
||||
require.Equal(t, "0.14%", variables["report_error_rate"])
|
||||
require.Equal(t, "0.28%", variables["report_upstream_error_rate"])
|
||||
require.Equal(t, "4", variables["report_upstream_error_count_excl_429_529"])
|
||||
require.Equal(t, "8,231 ms", variables["report_latency_p50"])
|
||||
require.Equal(t, "151,260 ms", variables["report_latency_p99"])
|
||||
require.Equal(t, "1,674 ms", variables["report_ttft_p50"])
|
||||
require.Equal(t, "11,222 ms", variables["report_ttft_p99"])
|
||||
require.Equal(t, "121,550,190", variables["report_tokens"])
|
||||
require.Equal(t, "1.2", variables["report_qps_peak"])
|
||||
require.Equal(t, "133421.2", variables["report_tps_peak"])
|
||||
require.Equal(t, "1406.8", variables["report_tps_avg"])
|
||||
require.Empty(t, variables["report_html"])
|
||||
require.Equal(t, "none", variables["report_detail_display"])
|
||||
|
||||
zhVariables := opsSummaryReportEmailVariables(report, now, overview, "zh-CN")
|
||||
require.Equal(t, "日报", zhVariables["report_name"])
|
||||
emptyVariables := opsSummaryReportEmailVariables(report, now, nil, "en")
|
||||
require.Equal(t, "block", emptyVariables["report_summary_display"])
|
||||
require.Equal(t, "none", emptyVariables["report_detail_display"])
|
||||
|
||||
legacyTemplate, err := renderNotificationEmail(
|
||||
NotificationEmailEventOpsScheduledReport,
|
||||
"Report",
|
||||
`<section>{{report_html}}</section>`,
|
||||
variables,
|
||||
map[string]string{"report_html": `<h2>generated summary</h2>`},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, legacyTemplate.HTML, `<h2>generated summary</h2>`)
|
||||
}
|
||||
|
||||
func TestOpsScheduledReportVariablesDoNotUsePreviewMetrics(t *testing.T) {
|
||||
now := time.Date(2026, time.July, 19, 1, 0, 26, 0, time.UTC)
|
||||
report := &opsScheduledReport{
|
||||
Name: "错误摘要",
|
||||
ReportType: "error_digest",
|
||||
TimeRange: 24 * time.Hour,
|
||||
}
|
||||
|
||||
variables := opsScheduledReportLocalizedEmailVariables(report, now, "en")
|
||||
require.Equal(t, "Error digest", variables["report_name"])
|
||||
require.Empty(t, variables["report_html"])
|
||||
require.Equal(t, "block", variables["report_detail_display"])
|
||||
for _, placeholder := range notificationEmailOpsSummaryPlaceholders {
|
||||
if placeholder == "report_summary_display" {
|
||||
require.Equal(t, "none", variables[placeholder])
|
||||
continue
|
||||
}
|
||||
require.Equal(t, "-", variables[placeholder])
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpsScheduledReportLegacyTemplateReceivesSummaryHTML(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := newNotificationEmailMemorySettingRepo()
|
||||
smtpServer := startNotificationEmailTestSMTPServer(t)
|
||||
require.NoError(t, repo.SetMultiple(ctx, smtpServer.settings()))
|
||||
|
||||
emailService := NewEmailService(repo, nil)
|
||||
notificationService := NewNotificationEmailService(repo, emailService)
|
||||
_, err := notificationService.UpdateTemplate(
|
||||
ctx,
|
||||
NotificationEmailEventOpsScheduledReport,
|
||||
"en",
|
||||
"Legacy report {{report_name}}",
|
||||
`<section data-template="legacy">{{report_html}}</section>`,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
svc := &OpsScheduledReportService{
|
||||
opsService: &OpsService{opsRepo: &opsRepoMock{}},
|
||||
emailService: emailService,
|
||||
}
|
||||
report := &opsScheduledReport{
|
||||
Name: "日报",
|
||||
ReportType: "daily_summary",
|
||||
TimeRange: 24 * time.Hour,
|
||||
Recipients: []string{"ops@example.com"},
|
||||
}
|
||||
|
||||
attempts, err := svc.runReport(ctx, report, time.Date(2026, time.July, 19, 1, 0, 26, 0, time.UTC))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, attempts)
|
||||
require.Equal(t, int64(1), smtpServer.messageCount())
|
||||
message := smtpServer.lastMessage()
|
||||
require.Contains(t, message, `<section data-template="legacy">`)
|
||||
require.Contains(t, message, `<h2>日报</h2>`)
|
||||
}
|
||||
|
||||
func TestFormatOpsReportIntegerGroupsDigits(t *testing.T) {
|
||||
require.Equal(t, "2,374", formatOpsReportInteger(2374))
|
||||
require.Equal(t, "-1,234", formatOpsReportInteger(-1234))
|
||||
require.Equal(t, "42", formatOpsReportInteger(42))
|
||||
}
|
||||
@@ -43,6 +43,14 @@ func (s *SettingService) InitializeDefaultSettings(ctx context.Context) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
forwardedClientIPHeaders := []string{}
|
||||
if s != nil && s.cfg != nil {
|
||||
forwardedClientIPHeaders = s.cfg.ForwardedClientIPSettings().Headers
|
||||
}
|
||||
forwardedClientIPHeadersJSON, err := json.Marshal(forwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal default forwarded client IP headers: %w", err)
|
||||
}
|
||||
|
||||
// 初始化默认设置
|
||||
defaults := map[string]string{
|
||||
@@ -54,7 +62,9 @@ func (s *SettingService) InitializeDefaultSettings(ctx context.Context) error {
|
||||
SettingKeyLoginAgreementMode: defaultLoginAgreementMode,
|
||||
SettingKeyLoginAgreementUpdatedAt: defaultLoginAgreementDate,
|
||||
SettingKeyLoginAgreementDocuments: loginAgreementDocumentsJSON,
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "false",
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
SettingKeyForwardedClientIPHeaders: string(forwardedClientIPHeadersJSON),
|
||||
settingKeyForwardedClientIPModeV2: "true",
|
||||
SettingKeySiteName: "Sub2API",
|
||||
SettingKeySiteLogo: "",
|
||||
SettingKeyPurchaseSubscriptionEnabled: "false",
|
||||
@@ -237,6 +247,21 @@ func (s *SettingService) InitializeDefaultSettings(ctx context.Context) error {
|
||||
return s.settingRepo.SetMultiple(ctx, defaults)
|
||||
}
|
||||
|
||||
func parseForwardedClientIPHeadersSetting(value string) ([]string, error) {
|
||||
var headers []string
|
||||
if err := json.Unmarshal([]byte(value), &headers); err != nil {
|
||||
return nil, fmt.Errorf("parse forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
if headers == nil {
|
||||
return nil, fmt.Errorf("parse forwarded_client_ip_headers: value must be a JSON array")
|
||||
}
|
||||
normalized, err := config.NormalizeForwardedClientIPHeaders(headers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
// parseSettings 解析设置到结构体
|
||||
func (s *SettingService) parseSettings(settings map[string]string) *SystemSettings {
|
||||
emailVerifyEnabled := settings[SettingKeyEmailVerifyEnabled] == "true"
|
||||
@@ -246,10 +271,24 @@ func (s *SettingService) parseSettings(settings map[string]string) *SystemSettin
|
||||
loginAgreementUpdatedAt = defaultLoginAgreementDate
|
||||
}
|
||||
apiKeyACLTrustForwardedIP := false
|
||||
forwardedClientIPHeaders := []string{}
|
||||
if s != nil && s.cfg != nil {
|
||||
runtimeSettings := s.cfg.ForwardedClientIPSettings()
|
||||
apiKeyACLTrustForwardedIP = runtimeSettings.TrustForwardedIP
|
||||
forwardedClientIPHeaders = runtimeSettings.Headers
|
||||
}
|
||||
if value, ok := settings[SettingKeyAPIKeyACLTrustForwardedIP]; ok {
|
||||
apiKeyACLTrustForwardedIP = value == "true"
|
||||
} else if s != nil && s.cfg != nil {
|
||||
apiKeyACLTrustForwardedIP = s.cfg.Security.TrustForwardedIPForAPIKeyACL
|
||||
}
|
||||
if value, ok := settings[SettingKeyForwardedClientIPHeaders]; ok {
|
||||
parsed, err := parseForwardedClientIPHeadersSetting(value)
|
||||
if err != nil {
|
||||
slog.Error("invalid persisted forwarded client IP headers; forwarded trust disabled", "error", err)
|
||||
apiKeyACLTrustForwardedIP = false
|
||||
forwardedClientIPHeaders = []string{}
|
||||
} else {
|
||||
forwardedClientIPHeaders = parsed
|
||||
}
|
||||
}
|
||||
result := &SystemSettings{
|
||||
RegistrationEnabled: settings[SettingKeyRegistrationEnabled] == "true",
|
||||
@@ -277,6 +316,7 @@ func (s *SettingService) parseSettings(settings map[string]string) *SystemSettin
|
||||
TurnstileSiteKey: settings[SettingKeyTurnstileSiteKey],
|
||||
TurnstileSecretKeyConfigured: settings[SettingKeyTurnstileSecretKey] != "",
|
||||
APIKeyACLTrustForwardedIP: apiKeyACLTrustForwardedIP,
|
||||
ForwardedClientIPHeaders: forwardedClientIPHeaders,
|
||||
SiteName: s.getStringOrDefault(settings, SettingKeySiteName, "Sub2API"),
|
||||
SiteLogo: settings[SettingKeySiteLogo],
|
||||
SiteSubtitle: s.getStringOrDefault(settings, SettingKeySiteSubtitle, "Subscription to API Conversion Platform"),
|
||||
|
||||
@@ -2,6 +2,7 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
@@ -216,21 +217,66 @@ func (s *SettingService) SetProxyRepository(repo ProxyRepository) {
|
||||
s.proxyRepo = repo
|
||||
}
|
||||
|
||||
func (s *SettingService) LoadAPIKeyACLTrustForwardedIPSetting(ctx context.Context) error {
|
||||
func (s *SettingService) LoadForwardedClientIPSettings(ctx context.Context) error {
|
||||
if s == nil || s.cfg == nil || s.settingRepo == nil {
|
||||
return nil
|
||||
}
|
||||
value, err := s.settingRepo.GetValue(ctx, SettingKeyAPIKeyACLTrustForwardedIP)
|
||||
|
||||
values, err := s.settingRepo.GetMultiple(ctx, []string{
|
||||
SettingKeyAPIKeyACLTrustForwardedIP,
|
||||
SettingKeyForwardedClientIPHeaders,
|
||||
settingKeyForwardedClientIPModeV2,
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrSettingNotFound) {
|
||||
s.cfg.SetTrustForwardedIPForAPIKeyACL(s.cfg.Security.TrustForwardedIPForAPIKeyACL)
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("get api key acl forwarded ip setting: %w", err)
|
||||
s.cfg.SetForwardedClientIPSettings(false, nil)
|
||||
return fmt.Errorf("get forwarded client ip settings: %w", err)
|
||||
}
|
||||
enabled := value == "true"
|
||||
s.cfg.SetTrustForwardedIPForAPIKeyACL(enabled)
|
||||
return nil
|
||||
|
||||
enabled := s.cfg.Security.TrustForwardedIPForAPIKeyACL
|
||||
headers := s.cfg.ForwardedClientIPSettings().Headers
|
||||
storedValue, hasStoredValue := values[SettingKeyAPIKeyACLTrustForwardedIP]
|
||||
if hasStoredValue {
|
||||
enabled = storedValue == "true"
|
||||
}
|
||||
|
||||
var headersErr error
|
||||
if storedHeaders, ok := values[SettingKeyForwardedClientIPHeaders]; ok {
|
||||
headers, headersErr = parseForwardedClientIPHeadersSetting(storedHeaders)
|
||||
if headersErr != nil {
|
||||
enabled = false
|
||||
headers = []string{}
|
||||
headersErr = fmt.Errorf("load forwarded client ip headers: %w", headersErr)
|
||||
}
|
||||
}
|
||||
|
||||
updates := make(map[string]string)
|
||||
if _, hasStoredHeaders := values[SettingKeyForwardedClientIPHeaders]; !hasStoredHeaders {
|
||||
headersJSON, marshalErr := json.Marshal(headers)
|
||||
if marshalErr != nil {
|
||||
headers = []string{}
|
||||
headersErr = errors.Join(headersErr, fmt.Errorf("marshal forwarded client ip headers: %w", marshalErr))
|
||||
headersJSON = []byte("[]")
|
||||
}
|
||||
updates[SettingKeyForwardedClientIPHeaders] = string(headersJSON)
|
||||
}
|
||||
if values[settingKeyForwardedClientIPModeV2] != "true" {
|
||||
updates[settingKeyForwardedClientIPModeV2] = "true"
|
||||
// Before this migration, new installations persisted false by default.
|
||||
// Restore compatibility only when no trusted-proxy policy was configured.
|
||||
if headersErr == nil && hasStoredValue && !enabled && !s.cfg.Server.TrustedProxiesConfigured {
|
||||
enabled = true
|
||||
updates[SettingKeyAPIKeyACLTrustForwardedIP] = "true"
|
||||
}
|
||||
}
|
||||
if len(updates) > 0 {
|
||||
if err := s.settingRepo.SetMultiple(ctx, updates); err != nil {
|
||||
s.cfg.SetForwardedClientIPSettings(enabled, headers)
|
||||
return errors.Join(headersErr, fmt.Errorf("migrate forwarded client ip setting: %w", err))
|
||||
}
|
||||
}
|
||||
|
||||
s.cfg.SetForwardedClientIPSettings(enabled, headers)
|
||||
return headersErr
|
||||
}
|
||||
|
||||
// GetAllSettings 获取所有系统设置
|
||||
|
||||
@@ -5,6 +5,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math"
|
||||
"strconv"
|
||||
"testing"
|
||||
@@ -16,7 +17,8 @@ import (
|
||||
)
|
||||
|
||||
type settingUpdateRepoStub struct {
|
||||
updates map[string]string
|
||||
updates map[string]string
|
||||
setMultipleErr error
|
||||
}
|
||||
|
||||
func (s *settingUpdateRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
|
||||
@@ -40,7 +42,7 @@ func (s *settingUpdateRepoStub) SetMultiple(ctx context.Context, settings map[st
|
||||
for k, v := range settings {
|
||||
s.updates[k] = v
|
||||
}
|
||||
return nil
|
||||
return s.setMultipleErr
|
||||
}
|
||||
|
||||
func (s *settingUpdateRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
|
||||
@@ -87,6 +89,62 @@ func (s *settingGetAllRepoStub) Delete(ctx context.Context, key string) error {
|
||||
panic("unexpected Delete call")
|
||||
}
|
||||
|
||||
type forwardedIPMigrationRepoStub struct {
|
||||
values map[string]string
|
||||
updates map[string]string
|
||||
getMultipleErr error
|
||||
setMultipleErr error
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) Get(context.Context, string) (*Setting, error) {
|
||||
panic("unexpected Get call")
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) GetValue(_ context.Context, key string) (string, error) {
|
||||
value, ok := s.values[key]
|
||||
if !ok {
|
||||
return "", ErrSettingNotFound
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) Set(context.Context, string, string) error {
|
||||
panic("unexpected Set call")
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
|
||||
if s.getMultipleErr != nil {
|
||||
return nil, s.getMultipleErr
|
||||
}
|
||||
result := make(map[string]string, len(keys))
|
||||
for _, key := range keys {
|
||||
if value, ok := s.values[key]; ok {
|
||||
result[key] = value
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) SetMultiple(_ context.Context, values map[string]string) error {
|
||||
if s.setMultipleErr != nil {
|
||||
return s.setMultipleErr
|
||||
}
|
||||
s.updates = make(map[string]string, len(values))
|
||||
for key, value := range values {
|
||||
s.values[key] = value
|
||||
s.updates[key] = value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) GetAll(context.Context) (map[string]string, error) {
|
||||
panic("unexpected GetAll call")
|
||||
}
|
||||
|
||||
func (s *forwardedIPMigrationRepoStub) Delete(context.Context, string) error {
|
||||
panic("unexpected Delete call")
|
||||
}
|
||||
|
||||
type settingAntigravityUARepoStub struct {
|
||||
values map[string]string
|
||||
}
|
||||
@@ -495,6 +553,16 @@ func TestSettingService_UpdateSettings_AntigravityUserAgentVersion(t *testing.T)
|
||||
require.Equal(t, "1.23.2", repo.updates[SettingKeyAntigravityUserAgentVersion])
|
||||
}
|
||||
|
||||
func TestSettingService_InitializeDefaultSettingsPersistsConfiguredForwardedClientIPHeaders(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{values: map[string]string{}}
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Cdn-Ip", "True-Client-Ip"})
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
require.NoError(t, svc.InitializeDefaultSettings(context.Background()))
|
||||
require.JSONEq(t, `["X-Cdn-Ip","True-Client-Ip"]`, repo.values[SettingKeyForwardedClientIPHeaders])
|
||||
}
|
||||
|
||||
func TestSettingService_UpdateSettings_APIKeyACLTrustForwardedIPRefreshesConfig(t *testing.T) {
|
||||
repo := &settingUpdateRepoStub{}
|
||||
cfg := &config.Config{}
|
||||
@@ -502,11 +570,51 @@ func TestSettingService_UpdateSettings_APIKeyACLTrustForwardedIPRefreshesConfig(
|
||||
|
||||
err := svc.UpdateSettings(context.Background(), &SystemSettings{
|
||||
APIKeyACLTrustForwardedIP: true,
|
||||
ForwardedClientIPHeaders: []string{" x-cdn-ip ", "X-CDN-IP", "true-client-ip"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "true", repo.updates[SettingKeyAPIKeyACLTrustForwardedIP])
|
||||
require.True(t, cfg.Security.TrustForwardedIPForAPIKeyACL)
|
||||
require.True(t, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
require.JSONEq(t, `["X-Cdn-Ip","True-Client-Ip"]`, repo.updates[SettingKeyForwardedClientIPHeaders])
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.True(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, runtimeSettings.Headers)
|
||||
|
||||
runtimeSettings.Headers[0] = "X-Mutated"
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestSettingService_UpdateSettings_RejectsInvalidForwardedClientIPHeadersWithoutRefreshing(t *testing.T) {
|
||||
repo := &settingUpdateRepoStub{}
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Existing-IP"})
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
err := svc.UpdateSettings(context.Background(), &SystemSettings{
|
||||
ForwardedClientIPHeaders: []string{"X Invalid"},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, repo.updates)
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.True(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Existing-IP"}, runtimeSettings.Headers)
|
||||
}
|
||||
|
||||
func TestSettingService_UpdateSettings_WriteFailureDoesNotRefreshForwardedIPRuntime(t *testing.T) {
|
||||
repo := &settingUpdateRepoStub{setMultipleErr: errors.New("database unavailable")}
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(false, []string{"X-Existing-IP"})
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
err := svc.UpdateSettings(context.Background(), &SystemSettings{
|
||||
APIKeyACLTrustForwardedIP: true,
|
||||
ForwardedClientIPHeaders: []string{"X-New-IP"},
|
||||
})
|
||||
|
||||
require.ErrorContains(t, err, "database unavailable")
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Existing-IP"}, runtimeSettings.Headers)
|
||||
}
|
||||
|
||||
func TestSettingService_ParseSettings_APIKeyACLTrustForwardedIPFallsBackToConfigWhenMissing(t *testing.T) {
|
||||
@@ -519,6 +627,194 @@ func TestSettingService_ParseSettings_APIKeyACLTrustForwardedIPFallsBackToConfig
|
||||
require.True(t, got.APIKeyACLTrustForwardedIP)
|
||||
}
|
||||
|
||||
func TestSettingService_ParseSettings_APIKeyACLTrustForwardedIPUsesStoredValue(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(true)
|
||||
svc := NewSettingService(&settingUpdateRepoStub{}, cfg)
|
||||
|
||||
got := svc.parseSettings(map[string]string{SettingKeyAPIKeyACLTrustForwardedIP: "false"})
|
||||
|
||||
require.False(t, got.APIKeyACLTrustForwardedIP)
|
||||
}
|
||||
|
||||
func TestSettingService_ParseSettings_ForwardedClientIPHeaders(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Config-IP"})
|
||||
svc := NewSettingService(&settingUpdateRepoStub{}, cfg)
|
||||
|
||||
t.Run("stored value is normalized", func(t *testing.T) {
|
||||
got := svc.parseSettings(map[string]string{
|
||||
SettingKeyForwardedClientIPHeaders: `[" x-cdn-ip ","X-CDN-IP","true-client-ip"]`,
|
||||
})
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, got.ForwardedClientIPHeaders)
|
||||
})
|
||||
|
||||
t.Run("missing value falls back to config", func(t *testing.T) {
|
||||
got := svc.parseSettings(map[string]string{})
|
||||
require.Equal(t, []string{"X-Config-IP"}, got.ForwardedClientIPHeaders)
|
||||
})
|
||||
|
||||
t.Run("malformed value disables forwarded trust", func(t *testing.T) {
|
||||
got := svc.parseSettings(map[string]string{
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
SettingKeyForwardedClientIPHeaders: `{"not":"an array"}`,
|
||||
})
|
||||
require.False(t, got.APIKeyACLTrustForwardedIP)
|
||||
require.Empty(t, got.ForwardedClientIPHeaders)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsMigration(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
values map[string]string
|
||||
trustedProxiesSet bool
|
||||
configDefault bool
|
||||
wantEnabled bool
|
||||
wantForwardedIPUpdate string
|
||||
wantMigrationMarkerSet bool
|
||||
}{
|
||||
{
|
||||
name: "missing setting follows configured default",
|
||||
values: map[string]string{},
|
||||
configDefault: true,
|
||||
wantEnabled: true,
|
||||
wantMigrationMarkerSet: true,
|
||||
},
|
||||
{
|
||||
name: "legacy false without proxy config migrates to compatibility",
|
||||
values: map[string]string{SettingKeyAPIKeyACLTrustForwardedIP: "false"},
|
||||
wantEnabled: true,
|
||||
wantForwardedIPUpdate: "true",
|
||||
wantMigrationMarkerSet: true,
|
||||
},
|
||||
{
|
||||
name: "legacy false with explicit proxy config stays secure",
|
||||
values: map[string]string{SettingKeyAPIKeyACLTrustForwardedIP: "false"},
|
||||
trustedProxiesSet: true,
|
||||
wantEnabled: false,
|
||||
wantMigrationMarkerSet: true,
|
||||
},
|
||||
{
|
||||
name: "completed migration preserves later false choice",
|
||||
values: map[string]string{
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "false",
|
||||
settingKeyForwardedClientIPModeV2: "true",
|
||||
},
|
||||
wantEnabled: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{values: test.values}
|
||||
cfg := &config.Config{Server: config.ServerConfig{TrustedProxiesConfigured: test.trustedProxiesSet}}
|
||||
cfg.Security.TrustForwardedIPForAPIKeyACL = test.configDefault
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
require.NoError(t, svc.LoadForwardedClientIPSettings(context.Background()))
|
||||
require.Equal(t, test.wantEnabled, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
require.Equal(t, test.wantForwardedIPUpdate, repo.updates[SettingKeyAPIKeyACLTrustForwardedIP])
|
||||
require.JSONEq(t, `[]`, repo.updates[SettingKeyForwardedClientIPHeaders])
|
||||
if test.wantMigrationMarkerSet {
|
||||
require.Equal(t, "true", repo.updates[settingKeyForwardedClientIPModeV2])
|
||||
} else {
|
||||
require.NotContains(t, repo.updates, settingKeyForwardedClientIPModeV2)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsLoadsHeaders(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{values: map[string]string{
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
SettingKeyForwardedClientIPHeaders: `[" x-cdn-ip ","true-client-ip"]`,
|
||||
settingKeyForwardedClientIPModeV2: "true",
|
||||
}}
|
||||
cfg := &config.Config{}
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
require.NoError(t, svc.LoadForwardedClientIPSettings(context.Background()))
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.True(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, runtimeSettings.Headers)
|
||||
require.Nil(t, repo.updates)
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsMalformedHeadersDisablesCustomTrust(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{values: map[string]string{
|
||||
SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
SettingKeyForwardedClientIPHeaders: `["X Invalid"]`,
|
||||
}}
|
||||
cfg := &config.Config{}
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
err := svc.LoadForwardedClientIPSettings(context.Background())
|
||||
|
||||
require.ErrorContains(t, err, "load forwarded client ip headers")
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Empty(t, runtimeSettings.Headers)
|
||||
require.Equal(t, "true", repo.updates[settingKeyForwardedClientIPModeV2])
|
||||
require.NotContains(t, repo.updates, SettingKeyAPIKeyACLTrustForwardedIP)
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsBackfillsConfigHeaders(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{values: map[string]string{
|
||||
settingKeyForwardedClientIPModeV2: "true",
|
||||
}}
|
||||
cfg := &config.Config{}
|
||||
cfg.SetForwardedClientIPSettings(false, []string{"X-Config-IP"})
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
require.NoError(t, svc.LoadForwardedClientIPSettings(context.Background()))
|
||||
require.JSONEq(t, `["X-Config-IP"]`, repo.updates[SettingKeyForwardedClientIPHeaders])
|
||||
require.Equal(t, []string{"X-Config-IP"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsReadFailureFailsClosed(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{
|
||||
getMultipleErr: errors.New("database unavailable"),
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(true)
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
err := svc.LoadForwardedClientIPSettings(context.Background())
|
||||
|
||||
require.ErrorContains(t, err, "get forwarded client ip settings")
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Empty(t, runtimeSettings.Headers)
|
||||
}
|
||||
|
||||
func TestSettingService_LoadForwardedClientIPSettingsWriteFailureUsesComputedMode(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
trustedProxiesSet bool
|
||||
wantEnabled bool
|
||||
}{
|
||||
{name: "compatibility migration remains effective", wantEnabled: true},
|
||||
{name: "explicit proxy policy remains secure", trustedProxiesSet: true, wantEnabled: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
repo := &forwardedIPMigrationRepoStub{
|
||||
values: map[string]string{SettingKeyAPIKeyACLTrustForwardedIP: "false"},
|
||||
setMultipleErr: errors.New("database unavailable"),
|
||||
}
|
||||
cfg := &config.Config{Server: config.ServerConfig{TrustedProxiesConfigured: test.trustedProxiesSet}}
|
||||
svc := NewSettingService(repo, cfg)
|
||||
|
||||
err := svc.LoadForwardedClientIPSettings(context.Background())
|
||||
|
||||
require.ErrorContains(t, err, "migrate forwarded client ip setting")
|
||||
require.Equal(t, test.wantEnabled, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingService_GetAntigravityUserAgentVersion_Precedence(t *testing.T) {
|
||||
t.Run("后台设置优先", func(t *testing.T) {
|
||||
svc := NewSettingService(&settingAntigravityUARepoStub{values: map[string]string{
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
)
|
||||
@@ -62,6 +63,11 @@ func (s *SettingService) buildSystemSettingsUpdates(ctx context.Context, setting
|
||||
normalizedWhitelist = []string{}
|
||||
}
|
||||
settings.RegistrationEmailSuffixWhitelist = normalizedWhitelist
|
||||
normalizedForwardedClientIPHeaders, err := config.NormalizeForwardedClientIPHeaders(settings.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return nil, infraerrors.BadRequest("INVALID_FORWARDED_CLIENT_IP_HEADERS", err.Error())
|
||||
}
|
||||
settings.ForwardedClientIPHeaders = normalizedForwardedClientIPHeaders
|
||||
alipaySource, err := normalizeVisibleMethodSettingSource("alipay", settings.PaymentVisibleMethodAlipaySource, settings.PaymentVisibleMethodAlipayEnabled)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -156,6 +162,11 @@ func (s *SettingService) buildSystemSettingsUpdates(ctx context.Context, setting
|
||||
updates[SettingKeyTurnstileSecretKey] = settings.TurnstileSecretKey
|
||||
}
|
||||
updates[SettingKeyAPIKeyACLTrustForwardedIP] = strconv.FormatBool(settings.APIKeyACLTrustForwardedIP)
|
||||
forwardedClientIPHeadersJSON, err := json.Marshal(settings.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal forwarded client IP headers: %w", err)
|
||||
}
|
||||
updates[SettingKeyForwardedClientIPHeaders] = string(forwardedClientIPHeadersJSON)
|
||||
|
||||
// LinuxDo Connect OAuth 登录
|
||||
updates[SettingKeyLinuxDoConnectEnabled] = strconv.FormatBool(settings.LinuxDoConnectEnabled)
|
||||
@@ -579,7 +590,7 @@ func (s *SettingService) refreshCachedSettings(settings *SystemSettings) {
|
||||
})
|
||||
}
|
||||
if s.cfg != nil {
|
||||
s.cfg.SetTrustForwardedIPForAPIKeyACL(settings.APIKeyACLTrustForwardedIP)
|
||||
s.cfg.SetForwardedClientIPSettings(settings.APIKeyACLTrustForwardedIP, settings.ForwardedClientIPHeaders)
|
||||
}
|
||||
// codex_cli_only 加固策略缓存:设置更新后强制下次重载(涉及 4 个键 + JSON 解析,直接置过期)。
|
||||
s.codexRestrictionPolicySF.Forget("codex_restriction_policy")
|
||||
|
||||
@@ -42,6 +42,7 @@ type SystemSettings struct {
|
||||
TurnstileSecretKey string
|
||||
TurnstileSecretKeyConfigured bool
|
||||
APIKeyACLTrustForwardedIP bool
|
||||
ForwardedClientIPHeaders []string
|
||||
|
||||
// LinuxDo Connect OAuth 登录
|
||||
LinuxDoConnectEnabled bool
|
||||
|
||||
@@ -34,7 +34,7 @@ func TestCalculateProgress_BasicFields(t *testing.T) {
|
||||
assert.Equal(t, int64(100), progress.ID)
|
||||
assert.Equal(t, "Premium", progress.GroupName)
|
||||
assert.Equal(t, sub.ExpiresAt, progress.ExpiresAt)
|
||||
assert.True(t, progress.ExpiresInDays == 29 || progress.ExpiresInDays == 30, "ExpiresInDays should be 29 or 30, got %d", progress.ExpiresInDays)
|
||||
assert.Equal(t, 30, progress.ExpiresInDays)
|
||||
assert.Nil(t, progress.Daily, "无日限额时 Daily 应为 nil")
|
||||
assert.Nil(t, progress.Weekly, "无周限额时 Weekly 应为 nil")
|
||||
assert.Nil(t, progress.Monthly, "无月限额时 Monthly 应为 nil")
|
||||
|
||||
@@ -2,6 +2,8 @@ package service
|
||||
|
||||
import "time"
|
||||
|
||||
const subscriptionDayDuration = 24 * time.Hour
|
||||
|
||||
type UserSubscription struct {
|
||||
ID int64
|
||||
UserID int64
|
||||
@@ -41,10 +43,20 @@ func (s *UserSubscription) IsExpired() bool {
|
||||
}
|
||||
|
||||
func (s *UserSubscription) DaysRemaining() int {
|
||||
if s.IsExpired() {
|
||||
return s.daysRemainingAt(time.Now())
|
||||
}
|
||||
|
||||
func (s *UserSubscription) daysRemainingAt(now time.Time) int {
|
||||
remaining := s.ExpiresAt.Sub(now)
|
||||
if remaining <= 0 {
|
||||
return 0
|
||||
}
|
||||
return int(time.Until(s.ExpiresAt).Hours() / 24)
|
||||
|
||||
days := int(remaining / subscriptionDayDuration)
|
||||
if remaining%subscriptionDayDuration != 0 {
|
||||
days++
|
||||
}
|
||||
return days
|
||||
}
|
||||
|
||||
func (s *UserSubscription) IsWindowActivated() bool {
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUserSubscriptionDaysRemainingAt(t *testing.T) {
|
||||
now := time.Date(2026, time.July, 19, 12, 0, 0, 0, time.UTC)
|
||||
tests := []struct {
|
||||
name string
|
||||
expiresAt time.Time
|
||||
want int
|
||||
}{
|
||||
{name: "expired", expiresAt: now.Add(-time.Nanosecond), want: 0},
|
||||
{name: "expires now", expiresAt: now, want: 0},
|
||||
{name: "less than one day", expiresAt: now.Add(subscriptionDayDuration - time.Nanosecond), want: 1},
|
||||
{name: "exactly one day", expiresAt: now.Add(subscriptionDayDuration), want: 1},
|
||||
{name: "over one day", expiresAt: now.Add(subscriptionDayDuration + time.Nanosecond), want: 2},
|
||||
{name: "less than two days", expiresAt: now.Add(2*subscriptionDayDuration - time.Nanosecond), want: 2},
|
||||
{name: "exactly two days", expiresAt: now.Add(2 * subscriptionDayDuration), want: 2},
|
||||
{name: "over two days", expiresAt: now.Add(2*subscriptionDayDuration + time.Nanosecond), want: 3},
|
||||
{name: "exactly seven days", expiresAt: now.Add(7 * subscriptionDayDuration), want: 7},
|
||||
{name: "over seven days", expiresAt: now.Add(7*subscriptionDayDuration + time.Nanosecond), want: 8},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sub := &UserSubscription{ExpiresAt: tt.expiresAt}
|
||||
require.Equal(t, tt.want, sub.daysRemainingAt(now))
|
||||
})
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user