Merge branch 'main' into fix/issues-4561-4562-4566-4582

Resolve const-block conflict in openai_gateway_grok_cache.go:
keep #4590's client tool cache constants, drop grokFreeRolling24hTokenLimit
(moved to pkg/xai as IsGrokFreeRolling24hTokenLimit with legacy 2M support).
This commit is contained in:
shaw committed 2026-07-20 11:25:11 +08:00
commit a90c18cbea
117 files changed
+7309 -469

No files matched your search

+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 {
@@ -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())
}
@@ -208,7 +208,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)
@@ -287,7 +287,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)
@@ -318,7 +318,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)
@@ -411,7 +411,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) {
@@ -995,7 +995,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)
@@ -1609,7 +1609,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,
@@ -1635,7 +1635,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,
@@ -1732,3 +1732,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)
@@ -81,7 +81,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)
@@ -297,14 +297,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
@@ -121,17 +121,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())
})
}
}
@@ -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)
@@ -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()))
}
@@ -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
@@ -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)
@@ -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)
}
@@ -374,10 +377,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
@@ -389,8 +391,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))
@@ -402,7 +432,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。
@@ -13,9 +13,11 @@ import (
)
const (
grokConversationIDHeader = "X-Grok-Conv-Id"
grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]`
grokFreeCacheDisabledToolChoice = "none"
grokConversationIDHeader = "X-Grok-Conv-Id"
grokClientToolCacheOptInHeader = "X-Sub2API-Grok-Client-Tool-Cache"
grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]`
grokFreeCacheDisabledToolChoice = "none"
grokClientToolCacheOptInExtraKey = "grok_client_tool_cache_enabled"
)
// resolveGrokCacheIdentity derives one stable, tenant-isolated routing identity
@@ -122,7 +124,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))
@@ -132,12 +134,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
}
@@ -146,7 +249,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 {
@@ -231,22 +339,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
@@ -256,9 +386,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 {
@@ -269,7 +409,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
}
@@ -279,9 +419,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] {
@@ -290,17 +438,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"} {
@@ -198,7 +198,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 {
@@ -212,7 +212,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)
@@ -221,16 +220,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 "
@@ -250,6 +279,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"
@@ -261,7 +614,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 {
@@ -281,8 +634,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
@@ -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")
@@ -2498,6 +2731,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)
@@ -512,6 +512,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))
})
}
}
+28 -13
View File
@@ -12,6 +12,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/google/wire"
"github.com/redis/go-redis/v9"
"go.uber.org/zap"
)
// BuildInfo contains build information
@@ -521,20 +522,33 @@ func ProvideAPIKeyAuthCacheInvalidator(apiKeyService *APIKeyService) APIKeyAuthC
return apiKeyService
}
// ProvideImageStorageSettingService 构造异步生图对象存储的后台设置服务。
//
// config.yaml 里的 image_storage 作为回落:后台从未保存过设置时沿用它,
// 使升级前已通过配置文件开启该功能的部署不被打断。
func ProvideImageStorageSettingService(
settingRepo SettingRepository,
encryptor SecretEncryptor,
backup *BackupService,
factory ImageStorageFactory,
cfg *config.Config,
) *ImageStorageSettingService {
if cfg.ImageStorage.Enabled && !cfg.ImageStorage.Active() {
// 列出具体缺失的键。若这些键其实已在环境变量里设过,说明它们没被读进来,
// 请确认 setDefaults 中已为其注册默认值(见 config.setEnvReachableDefaults)。
logger.L().Warn("image_storage.enabled is true in config but object storage is not fully configured; configure it in the admin UI or complete the config file",
zap.Strings("missing_keys", cfg.ImageStorage.MissingCredentialKeys()))
}
return NewImageStorageSettingService(settingRepo, encryptor, backup, factory, cfg.ImageStorage)
}
// ProvideImageTaskService 构造异步图片任务服务。
//
// 对象存储是异步图片任务的启用前提:仅当 image_storage 开关打开且凭证齐全时,
// 服务才启用,并挂上把结果转存到对象存储的 uploader;否则功能整体禁用
// 对象存储是异步图片任务的启用前提:仅当开关打开且凭证齐全时功能才可用,否则整体禁用
// (handler 返回 404,不创建任务、不写 Redis),从而避免大 base64 结果撑爆 Redis。
func ProvideImageTaskService(store ImageTaskStore, storage ImageStorage, cfg *config.Config) *ImageTaskService {
if !cfg.ImageStorage.Active() {
if cfg.ImageStorage.Enabled {
logger.L().Warn("image_storage.enabled is true but object storage is not fully configured; async image tasks are disabled")
}
return NewImageTaskService(store)
}
uploader := NewImageResultUploader(storage, cfg.ImageStorage.Prefix, cfg.ImageStorage.MaxDownloadByte, nil)
return NewImageTaskServiceWithUploader(store, uploader, defaultImageTaskTTL, defaultImageTaskExecutionTimeout)
// 启用状态由 settings 服务在运行时解析,因此后台改开关后无需重启即可生效。
func ProvideImageTaskService(store ImageTaskStore, settings *ImageStorageSettingService) *ImageTaskService {
return NewImageTaskServiceWithResolver(store, settings.Resolver(), defaultImageTaskTTL, defaultImageTaskExecutionTimeout)
}
// ProvideBackupService creates and starts BackupService
@@ -613,8 +627,8 @@ func ProvideSettingService(settingRepo SettingRepository, groupRepo GroupReposit
svc := NewSettingService(settingRepo, cfg)
svc.SetDefaultSubscriptionGroupReader(groupRepo)
svc.SetProxyRepository(proxyRepo)
if err := svc.LoadAPIKeyACLTrustForwardedIPSetting(context.Background()); err != nil {
logger.LegacyPrintf("service.setting", "Warning: load api key acl forwarded ip setting failed: %v", err)
if err := svc.LoadForwardedClientIPSettings(context.Background()); err != nil {
logger.LegacyPrintf("service.setting", "Warning: load forwarded client IP settings failed: %v", err)
}
if err := svc.MigrateOpenAIAllowClaudeCodeCodexPluginSetting(context.Background()); err != nil {
logger.LegacyPrintf("service.setting", "Warning: migrate openai allow Claude Code Codex plugin setting failed: %v", err)
@@ -680,6 +694,7 @@ var ProviderSet = wire.NewSet(
NewAdminService,
NewGatewayService,
NewOpenAIGatewayService,
ProvideImageStorageSettingService,
ProvideImageTaskService,
ProvideBatchImageModelPricingResolver,
NewBatchImagePublicService,
+4
View File
@@ -23,6 +23,10 @@ SERVER_PORT=8080
# Server mode: release or debug
SERVER_MODE=release
# Optional token used only for GitHub Release API update checks. Release asset
# downloads remain anonymous. GITHUB_TOKEN and GH_TOKEN are not used.
UPDATE_GITHUB_TOKEN=
# Return Server-Timing for authenticated requests made by the Admin web UI
ENABLE_SERVER_TIMING=false
+35 -6
View File
@@ -29,13 +29,42 @@ the application's responsibility.
## Trusted client IPs
`server.trusted_proxies` must contain only the CIDR/IP addresses that connect
directly to Sub2API, normally the local Nginx/Caddy address or the private load
balancer subnet. An empty list disables forwarded-IP trust.
`security.trust_forwarded_ip_for_api_key_acl` is enabled by default for upgrade
compatibility. While enabled, raw forwarding headers take over client-IP
resolution for logs and security-sensitive paths. Custom headers from
`security.forwarded_client_ip_headers` are checked in configured order before
the built-in `CF-Connecting-IP`, `X-Real-IP`, and `X-Forwarded-For` fallback.
Header names are case-insensitive, normalized when loaded, de-duplicated, and
limited to 16 unique valid HTTP field names. Header values must contain IP
literals; comma-separated values are supported, invalid entries are skipped,
and public addresses are preferred over private fallback addresses.
Never trust `CF-Connecting-IP`, `X-Real-IP`, or `X-Forwarded-For` merely because
the header exists. A CDN deployment must firewall the origin so only the CDN or
load balancer can reach it, and the proxy must overwrite forwarded headers.
The list can be supplied in YAML or with the comma-separated environment
variable `SECURITY_FORWARDED_CLIENT_IP_HEADERS`; an explicitly empty environment
value clears YAML values. It is also editable from the admin security settings
and updates at runtime without a restart. A request snapshots the switch and
header list together, so one request cannot mix old and new settings. Custom
headers are ignored completely when the switch is disabled. In that mode Gin's
`server.trusted_proxies` chain is authoritative: configure only the exact
CIDR/IP addresses that connect directly to Sub2API. An explicit empty list
trusts no forwarded client IPs.
On the first upgrade to this mode, a legacy `false` value is changed to `true`
only when `server.trusted_proxies` was not explicitly configured; explicit
proxy policies remain in secure mode. New installations persist the configured
custom header list during database initialization. Existing installations
backfill a missing database value from the YAML configuration. A hidden
migration marker prevents later administrator changes from being overwritten.
If settings cannot be read or the persisted custom-header list is malformed,
the process fails closed to trusted-proxy mode with no custom headers. If a
migration write fails, the computed mode remains active for the current process
and startup records a warning.
Compatibility takeover accepts forwarded headers without validating the direct
peer, including any configured custom header. Protect the origin from direct
access while it is enabled. A CDN deployment must firewall the origin so only
the CDN or load balancer can reach it, and that proxy must overwrite every
trusted client-IP header rather than append an untrusted client value.
Example for a proxy on the same host:
+1
View File
@@ -238,6 +238,7 @@ docker compose down -v
| `ADMIN_EMAIL` | No | `admin@sub2api.local` | Admin email |
| `ADMIN_PASSWORD` | No | *(auto-generated)* | Admin password |
| `TZ` | No | `Asia/Shanghai` | Timezone |
| `UPDATE_GITHUB_TOKEN` | No | *(empty)* | Token for `api.github.com` release checks only; asset downloads remain anonymous. |
| `GEMINI_OAUTH_CLIENT_ID` | No | *(builtin)* | Google OAuth client ID (Gemini OAuth). Leave empty to use the built-in Gemini CLI client. |
| `GEMINI_OAUTH_CLIENT_SECRET` | No | *(builtin)* | Google OAuth client secret (Gemini OAuth). Leave empty to use the built-in Gemini CLI client. |
| `GEMINI_OAUTH_SCOPES` | No | *(default)* | OAuth scopes (Gemini OAuth) |
+25 -3
View File
@@ -36,9 +36,15 @@ server:
# Keep-alive idle timeout in seconds.
# Keep-Alive 空闲连接超时(秒)。
idle_timeout: 120
# Trusted proxies for X-Forwarded-For parsing (CIDR/IP). Empty disables trusted proxies.
# 信任的代理地址(CIDR/IP 格式),用于解析 X-Forwarded-For 头。留空则禁用代理信任。
trusted_proxies: []
# Trusted proxies used when security.trust_forwarded_ip_for_api_key_acl is false.
# List only the exact proxy addresses that connect directly to Sub2API.
# Set [] explicitly to disable forwarded-IP trust in high-security mode.
# security.trust_forwarded_ip_for_api_key_acl=false 时使用的可信代理。
# 只填写直接连接 Sub2API 的精确代理地址;显式设置 [] 可在高安全模式下
# 禁用转发 IP 信任。
trusted_proxies:
- 127.0.0.1/32
- ::1/128
# Global max request body size in bytes (default: 256MB)
# 全局最大请求体大小(字节,默认 256MB)
# Applies to all requests, especially important for h2c first request memory protection
@@ -95,6 +101,22 @@ cors:
# 安全配置
# =============================================================================
security:
# Legacy compatibility switch. When true, raw forwarded headers take over
# server.trusted_proxies. Set false to enforce the trusted proxy chain above.
# 旧版兼容开关。开启时原始转发头会接管 server.trusted_proxies;关闭后严格
# 使用上方配置的可信代理链。示例配置采用高安全模式。
trust_forwarded_ip_for_api_key_acl: false
# Optional client-IP headers for third-party CDNs, checked in list order before
# CF-Connecting-IP, X-Real-IP, and X-Forwarded-For. Valid only while the legacy
# compatibility switch above is true; maximum 16 unique HTTP header names.
# 第三方 CDN 的可选客户端 IP 请求头,按列表顺序优先于内置请求头解析。
# 仅在上方旧版兼容开关为 true 时生效;最多配置 16 个不重复的合法 HTTP 头名。
# Environment / 环境变量: SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
forwarded_client_ip_headers: []
# Example / 示例:
# forwarded_client_ip_headers:
# - "True-Client-IP"
# - "X-CDN-Client-IP"
url_allowlist:
# Enable URL allowlist validation (disable to skip all URL checks)
# 启用 URL 白名单验证(禁用则跳过所有 URL 检查)
+8 -4
View File
@@ -114,12 +114,16 @@ services:
restart: unless-stopped
volumes:
- ./redis_data:/data:Z
# The command is one quoted script for the inner `sh -c`. Compose keeps
# the newlines inside the quoted string, so every line needs a trailing
# `\` — without it, `redis-server` on the first line runs with no flags
# at all, and the --save/--appendonly/--appendfsync lines are never read.
command: >
sh -c '
redis-server
--save 60 1
--appendonly yes
--appendfsync everysec
redis-server \
--save 60 1 \
--appendonly yes \
--appendfsync everysec \
${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}'
environment:
- TZ=${TZ:-Asia/Shanghai}
+9 -4
View File
@@ -53,6 +53,7 @@ services:
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
- UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-}
# =======================================================================
# Database Configuration (PostgreSQL)
@@ -230,12 +231,16 @@ services:
volumes:
# Local directory mapping for easy migration
- ./redis_data:/data:Z
# The command is one quoted script for the inner `sh -c`. Compose keeps
# the newlines inside the quoted string, so every line needs a trailing
# `\` — without it, `redis-server` on the first line runs with no flags
# at all, and the --save/--appendonly/--appendfsync lines are never read.
command: >
sh -c '
redis-server
--save 60 1
--appendonly yes
--appendfsync everysec
redis-server \
--save 60 1 \
--appendonly yes \
--appendfsync everysec \
${REDIS_PASSWORD:+--requirepass "$REDIS_PASSWORD"}'
environment:
- TZ=${TZ:-Asia/Shanghai}
+1
View File
@@ -39,6 +39,7 @@ services:
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
- UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-}
# =======================================================================
# Database Configuration (PostgreSQL) - Required
+10
View File
@@ -49,6 +49,7 @@ services:
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
- UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-}
# =======================================================================
# Database Configuration (PostgreSQL)
@@ -202,6 +203,15 @@ services:
nofile:
soft: 100000
hard: 100000
# Pass the tuning values documented in .env.example to postgres. The
# fallbacks are the postgres:18 stock defaults, so nothing changes when
# the variables are not set.
command: >
postgres
-c max_connections=${POSTGRES_MAX_CONNECTIONS:-100}
-c shared_buffers=${POSTGRES_SHARED_BUFFERS:-128MB}
-c effective_cache_size=${POSTGRES_EFFECTIVE_CACHE_SIZE:-4GB}
-c maintenance_work_mem=${POSTGRES_MAINTENANCE_WORK_MEM:-64MB}
volumes:
- postgres_data:/var/lib/postgresql/data
environment:
+51 -3
View File
@@ -480,10 +480,58 @@ check_dependencies() {
fi
}
# Authenticate only GitHub REST API requests. Release asset downloads must stay anonymous.
github_api_curl() {
local arg
local expect_value=false
local url
if [ "$#" -lt 1 ]; then
echo "github_api_curl requires exactly one GitHub API URL" >&2
return 2
fi
url="${!#}"
# Keep authenticated invocations constrained to the options used below. In
# particular, curl config, --url, and --next could add another destination.
for arg in "${@:1:$#-1}"; do
if [ "$expect_value" = true ]; then
expect_value=false
continue
fi
case "$arg" in
-s|--silent)
;;
--connect-timeout|--max-time|-o|--output|-w|--write-out)
expect_value=true
;;
*)
echo "Unsafe github_api_curl argument: $arg" >&2
return 2
;;
esac
done
if [ "$expect_value" = true ] || [[ "$url" != https://api.github.com/* ]]; then
echo "github_api_curl requires exactly one GitHub API URL" >&2
return 2
fi
if [ -n "${UPDATE_GITHUB_TOKEN:-}" ]; then
if [[ "$UPDATE_GITHUB_TOKEN" == *$'\n'* || "$UPDATE_GITHUB_TOKEN" == *$'\r'* || "$UPDATE_GITHUB_TOKEN" == *'"'* || "$UPDATE_GITHUB_TOKEN" == *'\'* ]]; then
echo "UPDATE_GITHUB_TOKEN contains unsupported characters" >&2
return 2
fi
printf 'header = "Authorization: Bearer %s"\n' "$UPDATE_GITHUB_TOKEN" | UPDATE_GITHUB_TOKEN= GITHUB_TOKEN= GH_TOKEN= curl -q --globoff --config - "$@"
else
UPDATE_GITHUB_TOKEN= GITHUB_TOKEN= GH_TOKEN= curl -q --globoff "$@"
fi
}
# Get latest release version
get_latest_version() {
print_info "$(msg 'fetching_version')"
LATEST_VERSION=$(curl -s --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases/latest" 2>/dev/null | grep '"tag_name"' | sed -E 's/.*"([^"]+)".*/\1/')
LATEST_VERSION=$(github_api_curl -s --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases/latest" 2>/dev/null | grep '"tag_name"' | sed -E 's/.*"([^"]+)".*/\1/')
if [ -z "$LATEST_VERSION" ]; then
print_error "$(msg 'failed_get_version')"
@@ -499,7 +547,7 @@ list_versions() {
print_info "$(msg 'fetching_versions')"
local versions
versions=$(curl -s --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases" 2>/dev/null | grep '"tag_name"' | sed -E 's/.*"([^"]+)".*/\1/' | head -20)
versions=$(github_api_curl -s --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases" 2>/dev/null | grep '"tag_name"' | sed -E 's/.*"([^"]+)".*/\1/' | head -20)
if [ -z "$versions" ]; then
print_error "$(msg 'failed_get_version')"
@@ -536,7 +584,7 @@ validate_version() {
# Check if the release exists
local http_code
http_code=$(curl -s -o /dev/null -w "%{http_code}" --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases/tags/${version}" 2>/dev/null)
http_code=$(github_api_curl -s -o /dev/null -w "%{http_code}" --connect-timeout 10 --max-time 30 "https://api.github.com/repos/${GITHUB_REPO}/releases/tags/${version}" 2>/dev/null)
# Check for network errors (empty or non-numeric response)
if [ -z "$http_code" ] || ! [[ "$http_code" =~ ^[0-9]+$ ]]; then
+103
View File
@@ -0,0 +1,103 @@
#!/bin/bash
set -euo pipefail
ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
TEMP_DIR=$(mktemp -d)
trap 'rm -rf "$TEMP_DIR"' EXIT
cat > "$TEMP_DIR/curl" <<'EOF'
#!/bin/bash
printf '%s\n' "$@" > "$CURL_ARGS_LOG"
env > "${CURL_ARGS_LOG}.env"
cat > "${CURL_ARGS_LOG}.stdin"
EOF
chmod +x "$TEMP_DIR/curl"
mkdir "$TEMP_DIR/home"
cat > "$TEMP_DIR/home/.curlrc" <<'EOF'
url = "https://example.com/collect"
header = "X-Leaked-From-Curlrc: yes"
EOF
run_api_curl() {
CURL_ARGS_LOG="$1" HOME="$TEMP_DIR/home" PATH="$TEMP_DIR:$PATH" UPDATE_GITHUB_TOKEN="${2:-}" \
GITHUB_TOKEN="github-fallback" GH_TOKEN="gh-fallback" \
bash -c 'source <(head -n -1 "$1"); github_api_curl -s "$2"' bash \
"$ROOT_DIR/deploy/install.sh" "https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest"
}
run_api_curl "$TEMP_DIR/authenticated" "update-secret"
test "$(head -n 1 "$TEMP_DIR/authenticated")" = '-q'
grep -Fxq -- '--config' "$TEMP_DIR/authenticated"
grep -Fxq -- '-' "$TEMP_DIR/authenticated"
grep -Fxq -- '--globoff' "$TEMP_DIR/authenticated"
grep -Fxq 'header = "Authorization: Bearer update-secret"' "$TEMP_DIR/authenticated.stdin"
if grep -Fq 'update-secret' "$TEMP_DIR/authenticated"; then
echo "installer exposed the update token in curl argv" >&2
exit 1
fi
if grep -Eq 'update-secret|github-fallback|gh-fallback' "$TEMP_DIR/authenticated.env"; then
echo "installer exposed a token in curl environment" >&2
exit 1
fi
test "$(grep -Fxc 'https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest' "$TEMP_DIR/authenticated")" -eq 1
if grep -Fq 'example.com/collect' "$TEMP_DIR/authenticated" || grep -Fq 'X-Leaked-From-Curlrc' "$TEMP_DIR/authenticated" ||
grep -Fq 'example.com/collect' "$TEMP_DIR/authenticated.stdin" || grep -Fq 'X-Leaked-From-Curlrc' "$TEMP_DIR/authenticated.stdin"; then
echo "installer allowed hostile curl config into authenticated invocation" >&2
exit 1
fi
run_api_curl "$TEMP_DIR/anonymous"
test "$(head -n 1 "$TEMP_DIR/anonymous")" = '-q'
if grep -Eq 'github-fallback|gh-fallback' "$TEMP_DIR/anonymous.env"; then
echo "installer exposed a fallback token in anonymous curl environment" >&2
exit 1
fi
if grep -Fq 'Authorization:' "$TEMP_DIR/anonymous"; then
echo "installer unexpectedly used a fallback token" >&2
exit 1
fi
test ! -s "$TEMP_DIR/anonymous.stdin"
test "$(grep -Fxc 'https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest' "$TEMP_DIR/anonymous")" -eq 1
if grep -Fq 'example.com/collect' "$TEMP_DIR/anonymous" || grep -Fq 'X-Leaked-From-Curlrc' "$TEMP_DIR/anonymous"; then
echo "installer allowed hostile curl config into anonymous invocation" >&2
exit 1
fi
assert_unsafe_invocation_rejected() {
local name=$1
shift
rm -f "$TEMP_DIR/$name" "$TEMP_DIR/$name.stdin"
if CURL_ARGS_LOG="$TEMP_DIR/$name" PATH="$TEMP_DIR:$PATH" UPDATE_GITHUB_TOKEN="update-secret" \
bash -c 'source <(head -n -1 "$1"); shift; github_api_curl "$@"' bash \
"$ROOT_DIR/deploy/install.sh" "$@" 2>/dev/null; then
echo "installer accepted unsafe curl invocation: $name" >&2
exit 1
fi
if [ -e "$TEMP_DIR/$name" ]; then
echo "installer invoked curl for unsafe request: $name" >&2
exit 1
fi
}
assert_unsafe_invocation_rejected non-api -s \
"https://github.com/Wei-Shaw/sub2api/releases/download/v1/asset"
assert_unsafe_invocation_rejected mixed-host -s \
"https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest" \
"https://example.com/collect"
assert_unsafe_invocation_rejected multiple-api -s \
"https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest" \
"https://api.github.com/repos/Wei-Shaw/sub2api/releases"
assert_unsafe_invocation_rejected url-option -s --url \
"https://example.com/collect" \
"https://api.github.com/repos/Wei-Shaw/sub2api/releases/latest"
# Every installer release API request must use the scoped helper.
test "$(grep -c 'github_api_curl .*https://api.github.com/' "$ROOT_DIR/deploy/install.sh")" -eq 3
# Asset and checksum downloads must continue to call curl directly.
grep -Fq 'curl -sL "$download_url"' "$ROOT_DIR/deploy/install.sh"
grep -Fq 'curl -sL "$checksum_url"' "$ROOT_DIR/deploy/install.sh"
echo "install GitHub token checks passed"
+30
View File
@@ -20,6 +20,20 @@ Only OpenAI and Grok groups are supported. Requests use the same JSON or multipa
Asynchronous image tasks are **disabled by default** and gated on object storage. When the switch is off — or the S3 credentials are incomplete — the async endpoints return `404` and never create a task or write to Redis. This is deliberate: without offloading, large `b64_json` results (several MB each, e.g. `gpt-image-1`) would accumulate in Redis and exhaust its memory.
### From the admin UI (recommended)
**Admin → Backup → Async image object storage.** Saving the form takes effect immediately — the object-storage client is rebuilt on the next request, so there is no container restart.
Because the async image storage and the database backup share one S3 client, the form defaults to **reusing the backup S3 configuration**: it borrows the endpoint, region and credentials already configured above and keeps only its own bucket and prefix, so backups stay under `backups/` while images go to `images/`. Leave the bucket empty to use the backup bucket as well. Untick the box to point images at a completely separate account.
Saving requires step-up 2FA when that gate is enabled, for the same reason the backup S3 form does: changing the target redirects generated content to another account.
Turning the switch off stops new submissions but keeps already-accepted tasks pollable, so nothing in flight is stranded.
### From the config file
The admin setting takes precedence. When nothing has ever been saved there, the `image_storage` block in `config.yaml` is used instead, so deployments that enabled the feature before the admin UI existed keep working untouched.
Configure an S3-compatible object store (AWS S3, Cloudflare R2, Aliyun OSS, MinIO, …) in `config.yaml` (all keys also accept the `IMAGE_STORAGE_*` environment overrides):
```yaml
@@ -41,6 +55,22 @@ When a task completes, each generated image is uploaded to the bucket and the re
To support a different vendor beyond the S3-compatible client, implement the `service.ImageStorage` interface (`Save(ctx, key, contentType, data) (url, error)`) and provide it in place of the S3 implementation.
### Troubleshooting: the endpoints return 404 after enabling
`404 async image tasks are not enabled` means `image_storage` did not resolve to a complete configuration, so the feature stayed off. The route exists either way — the 404 comes from the handler, not from an unregistered path, which makes it easy to mistake for a missing build.
Check the startup log for:
```text
WARN image_storage.enabled is true but object storage is not fully configured; async image tasks are disabled missing_keys=[...]
```
`missing_keys` names exactly which credentials were empty when the config was loaded.
Note that releases **before v0.1.161 silently dropped `IMAGE_STORAGE_ENDPOINT`, `_BUCKET`, `_ACCESS_KEY_ID`, `_SECRET_ACCESS_KEY` and `_PUBLIC_BASE_URL`** when they were supplied only through the environment: those keys had no registered default, and viper cannot see an environment variable for a key it does not already know about. Deployments driven purely by `environment:` — which is what `deploy/docker-compose.yml` does by default — therefore reported `enabled: true` with empty credentials and 404'd on every async call. On an affected release the workaround is to also place the `image_storage` block in `/app/data/config.yaml` (copy it from `deploy/config.example.yaml`); once the keys exist in the file, the environment overrides apply normally.
Two further causes of a 404 that are unrelated to storage: the API key's group must be on the **OpenAI or Grok** platform (any other platform, or a key with no group at all, yields `Images API is not supported for this platform`), and a task may only be polled with the **same API key that submitted it** — polling with a different key of the same user returns `image task not found` by design.
## Submit a task
```bash
+49
View File
@@ -60,6 +60,52 @@ export async function testS3Connection(config: BackupS3Config): Promise<TestS3Re
return data
}
// Async image object storage
//
// Shares the S3 client with backups, so `reuse_backup_s3` borrows the endpoint and
// credentials configured above and only keeps its own bucket/prefix.
export interface ImageStorageConfig {
enabled: boolean
reuse_backup_s3: boolean
bucket: string
prefix: string
public_base_url: string
presign_expiry_hours: number
max_download_bytes: number
endpoint: string
region: string
access_key_id: string
secret_access_key?: string
force_path_style: boolean
}
export interface ImageStorageConfigResponse {
config: ImageStorageConfig
secret_configured: boolean
}
export async function getImageStorageConfig(): Promise<ImageStorageConfigResponse> {
const { data } = await apiClient.get<ImageStorageConfigResponse>('/admin/backups/image-storage')
return data
}
export async function updateImageStorageConfig(
config: ImageStorageConfig,
): Promise<ImageStorageConfig> {
const { data } = await apiClient.put<ImageStorageConfig>('/admin/backups/image-storage', config)
return data
}
export async function testImageStorageConnection(
config: ImageStorageConfig,
): Promise<TestS3Response> {
const { data } = await apiClient.post<TestS3Response>(
'/admin/backups/image-storage/test',
config,
)
return data
}
// Schedule
export async function getSchedule(): Promise<BackupScheduleConfig> {
const { data } = await apiClient.get<BackupScheduleConfig>('/admin/backups/schedule')
@@ -106,6 +152,9 @@ export const backupAPI = {
getS3Config,
updateS3Config,
testS3Connection,
getImageStorageConfig,
updateImageStorageConfig,
testImageStorageConnection,
getSchedule,
updateSchedule,
createBackup,
+2
View File
@@ -455,6 +455,7 @@ export interface SystemSettings {
turnstile_site_key: string;
turnstile_secret_key_configured: boolean;
api_key_acl_trust_forwarded_ip: boolean;
forwarded_client_ip_headers: string[];
// LinuxDo Connect OAuth settings
linuxdo_connect_enabled: boolean;
@@ -757,6 +758,7 @@ export interface UpdateSettingsRequest {
turnstile_site_key?: string;
turnstile_secret_key?: string;
api_key_acl_trust_forwarded_ip?: boolean;
forwarded_client_ip_headers?: string[];
linuxdo_connect_enabled?: boolean;
linuxdo_connect_client_id?: string;
linuxdo_connect_client_secret?: string;
Loaded 100 of 117 files, more files were not shown because too many files have changed in this diff. Show more