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:
shaw
2026-07-20 11:39:37 +08:00
135 changed files with 7818 additions and 553 deletions
+12 -3
View File
@@ -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
View File
@@ -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
View File
@@ -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 のみに制限することを推奨します:
+4 -6
View File
@@ -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)
+232 -31
View File
@@ -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")
}
+197
View File
@@ -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()
+5 -4
View File
@@ -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"`
+6 -5
View File
@@ -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 已取消)。
+40 -2
View File
@@ -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")
}
+12 -1
View File
@@ -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,
+18 -1
View File
@@ -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"`
+176 -7
View File
@@ -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"
+190 -13
View File
@@ -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())
})
}
}
+12
View File
@@ -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"`
+8
View File
@@ -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()
}
+8 -11
View File
@@ -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": "",
+20 -12
View File
@@ -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()
+1 -1
View File
@@ -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))
+6
View File
@@ -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)
+14 -20
View File
@@ -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)
{
+21 -7
View File
@@ -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)
}
}
+17 -1
View File
@@ -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")
}
+56 -24
View File
@@ -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)
}
+40 -3
View File
@@ -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, &parallelToolCalls) != 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, &parameters) != 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 -3
View File
@@ -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"),
+56 -10
View File
@@ -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{
+12 -1
View File
@@ -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")
+14 -2
View File
@@ -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