Merge remote-tracking branch 'upstream/main' into fix/api-double-billing

# Conflicts:
#	frontend/src/components/account/EditAccountModal.vue
This commit is contained in:
benjamin
2026-07-13 18:06:44 +08:00
101 changed files with 4480 additions and 221 deletions
+1 -1
View File
@@ -1 +1 @@
0.1.152
0.1.153
+1 -1
View File
@@ -264,7 +264,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig)
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
totpHandler := handler.NewTotpHandler(totpService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
+14
View File
@@ -923,6 +923,12 @@ type GatewayOpenAIWSConfig struct {
ModeRouterV2Enabled bool `mapstructure:"mode_router_v2_enabled"`
// IngressModeDefault: ingress 默认模式(off/ctx_pool/passthrough/http_bridge)
IngressModeDefault string `mapstructure:"ingress_mode_default"`
// IngressInterTurnIdleTimeoutSeconds bounds the time a client may remain idle
// between completed ingress turns. Zero disables this protection.
IngressInterTurnIdleTimeoutSeconds int `mapstructure:"ingress_inter_turn_idle_timeout_seconds"`
// MaxIngressConnectionsPerAPIKey bounds live client WebSocket ingress sessions
// per API key across all instances. Zero disables this protection.
MaxIngressConnectionsPerAPIKey int `mapstructure:"max_ingress_connections_per_api_key"`
// Enabled: 全局总开关(默认 true)
Enabled bool `mapstructure:"enabled"`
// OAuthEnabled: 是否允许 OpenAI OAuth 账号使用 WS
@@ -1945,6 +1951,8 @@ func setDefaults() {
viper.SetDefault("gateway.openai_ws.enabled", true)
viper.SetDefault("gateway.openai_ws.mode_router_v2_enabled", false)
viper.SetDefault("gateway.openai_ws.ingress_mode_default", "ctx_pool")
viper.SetDefault("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds", 300)
viper.SetDefault("gateway.openai_ws.max_ingress_connections_per_api_key", 64)
viper.SetDefault("gateway.openai_ws.oauth_enabled", true)
viper.SetDefault("gateway.openai_ws.apikey_enabled", true)
viper.SetDefault("gateway.openai_ws.force_http", false)
@@ -2717,6 +2725,12 @@ func (c *Config) Validate() error {
if c.Gateway.OpenAIWS.MaxConnsPerAccount <= 0 {
return fmt.Errorf("gateway.openai_ws.max_conns_per_account must be positive")
}
if c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds < 0 {
return fmt.Errorf("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds must be non-negative")
}
if c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey < 0 {
return fmt.Errorf("gateway.openai_ws.max_ingress_connections_per_api_key must be non-negative")
}
if c.Gateway.OpenAIWS.MinIdlePerAccount < 0 {
return fmt.Errorf("gateway.openai_ws.min_idle_per_account must be non-negative")
}
+16
View File
@@ -182,6 +182,12 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) {
if cfg.Gateway.OpenAIWS.IngressModeDefault != "ctx_pool" {
t.Fatalf("Gateway.OpenAIWS.IngressModeDefault = %q, want %q", cfg.Gateway.OpenAIWS.IngressModeDefault, "ctx_pool")
}
if cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds != 300 {
t.Fatalf("Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = %d, want 300", cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds)
}
if cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey != 64 {
t.Fatalf("Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = %d, want 64", cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey)
}
}
func TestLoadDefaultOpenAICompactModel(t *testing.T) {
@@ -1640,6 +1646,16 @@ func TestValidateConfig_OpenAIWSRules(t *testing.T) {
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxConnsPerAccount = 0 },
wantErr: "gateway.openai_ws.max_conns_per_account",
},
{
name: "ingress_inter_turn_idle_timeout_seconds 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = -1 },
wantErr: "gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds",
},
{
name: "max_ingress_connections_per_api_key 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = -1 },
wantErr: "gateway.openai_ws.max_ingress_connections_per_api_key",
},
{
name: "min_idle_per_account 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MinIdlePerAccount = -1 },
+7 -1
View File
@@ -24,6 +24,8 @@ const (
EndpointImagesGenerations = "/v1/images/generations"
EndpointImagesEdits = "/v1/images/edits"
EndpointVideosGenerations = "/v1/videos/generations"
EndpointVideosEdits = "/v1/videos/edits"
EndpointVideosExtensions = "/v1/videos/extensions"
EndpointVideos = "/v1/videos"
EndpointGeminiModels = "/v1beta/models"
)
@@ -88,6 +90,10 @@ func NormalizeInboundEndpoint(path string) string {
return EndpointImagesEdits
case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"):
return EndpointVideosGenerations
case strings.Contains(path, EndpointVideosEdits) || strings.Contains(path, "/videos/edits"):
return EndpointVideosEdits
case strings.Contains(path, EndpointVideosExtensions) || strings.Contains(path, "/videos/extensions"):
return EndpointVideosExtensions
case strings.Contains(path, EndpointVideos) || strings.Contains(path, "/videos/"):
return EndpointVideos
case strings.Contains(path, EndpointResponsesCompact) || isResponsesCompactAliasPath(path):
@@ -173,7 +179,7 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string {
switch platform {
case service.PlatformOpenAI, service.PlatformGrok:
if inbound == EndpointEmbeddings || inbound == EndpointAlphaSearch || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideos {
if inbound == EndpointEmbeddings || inbound == EndpointAlphaSearch || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideosEdits || inbound == EndpointVideosExtensions || inbound == EndpointVideos {
return inbound
}
// OpenAI forwards everything to the Responses API.
+7 -4
View File
@@ -29,7 +29,8 @@ const (
)
const (
// maxSameAccountRetries 同账号重试次数上限(针对 RetryableOnSameAccount 错误)
// maxSameAccountRetries 同账号重试次数默认上限(针对 RetryableOnSameAccount 错误)。
// 生产调用方通常传入账号级配置 account.GetPoolModeRetryCount(),该常量仅作兜底/测试默认值。
maxSameAccountRetries = 3
// sameAccountRetryDelay 同账号重试间隔
sameAccountRetryDelay = 500 * time.Millisecond
@@ -67,6 +68,7 @@ func (s *FailoverState) HandleFailoverError(
gatewayService TempUnscheduler,
accountID int64,
platform string,
retryLimit int,
failoverErr *service.UpstreamFailoverError,
) FailoverAction {
s.LastFailoverErr = failoverErr
@@ -76,14 +78,15 @@ func (s *FailoverState) HandleFailoverError(
s.ForceCacheBilling = true
}
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < maxSameAccountRetries {
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试。
// 重试次数上限 retryLimit 由调用方传入(账号级 pool_mode_retry_count 配置)。
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit {
s.SameAccountRetryCount[accountID]++
logger.FromContext(ctx).Warn("gateway.failover_same_account_retry",
zap.Int64("account_id", accountID),
zap.Int("upstream_status", failoverErr.StatusCode),
zap.Int("same_account_retry_count", s.SameAccountRetryCount[accountID]),
zap.Int("same_account_retry_max", maxSameAccountRetries),
zap.Int("same_account_retry_max", retryLimit),
)
if !sleepWithContext(ctx, sameAccountRetryDelay) {
return FailoverCanceled
+87 -50
View File
@@ -133,7 +133,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
@@ -150,7 +150,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -166,7 +166,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -181,19 +181,19 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
// 第一次切换:0→1
err1 := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
// 第二次切换:1→2
err2 := newTestFailoverErr(502, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 第三次已耗尽:SwitchCount(2) >= MaxSwitches(2)
err3 := newTestFailoverErr(503, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 2, fs.SwitchCount, "耗尽时不应继续递增")
@@ -212,7 +212,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(0, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 0, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -229,7 +229,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, true) // hasBoundSession=true
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -238,7 +238,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, true) // ForceCacheBilling=true
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -247,7 +247,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.False(t, fs.ForceCacheBilling)
})
@@ -257,12 +257,12 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
// 第一次:ForceCacheBilling=true → 设置
err1 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.True(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=false → 仍然保持 true
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "ForceCacheBilling 一旦设置不应被重置")
})
}
@@ -278,7 +278,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -297,7 +297,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 1; i <= maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, i, fs.SameAccountRetryCount[100])
}
@@ -311,12 +311,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
require.Equal(t, maxSameAccountRetries, fs.SameAccountRetryCount[100])
// 第 maxSameAccountRetries+1 次:重试耗尽,应切换账号
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -333,12 +333,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
// 账号 100 第一次重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
// 账号 200 第一次重试(独立计数)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[200])
require.Equal(t, 1, fs.SameAccountRetryCount[100], "账号 100 的计数不应受影响")
@@ -351,17 +351,54 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
// 耗尽账号 100 的重试
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
// 第 maxSameAccountRetries+1 次: 重试耗尽 → 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
// 再次遇到账号 100,计数仍为 maxSameAccountRetries,条件不满足 → 直接切换
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Len(t, mock.calls, 2, "第二次耗尽也应调用 TempUnschedule")
})
t.Run("尊重账号级retryLimit_配置1次只重试1次", func(t *testing.T) {
// 回归测试:Anthropic 等路径此前硬编码同账号重试 3 次,忽略账号
// pool_mode_retry_count 配置。此处验证传入 retryLimit=1 时只重试 1 次即切换。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
const retryLimit = 1
// 第 1 次:同账号重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
require.Equal(t, 0, fs.SwitchCount, "首次重试不应切换账号")
require.Empty(t, mock.calls, "未耗尽前不应 TempUnschedule")
// 第 2 次:已达上限 1 → 不再同账号重试,直接切换 + TempUnschedule
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100], "重试计数不应超过 retryLimit")
require.Equal(t, 1, fs.SwitchCount, "重试耗尽应切换账号")
require.Contains(t, fs.FailedAccountIDs, int64(100))
require.Len(t, mock.calls, 1, "重试耗尽应触发 TempUnschedule")
})
t.Run("retryLimit为0时立即切换不重试", func(t *testing.T) {
// pool_mode_retry_count=0 表示关闭同账号重试(如 GPT Image 账号)。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", 0, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 0, fs.SameAccountRetryCount[100], "retryLimit=0 不应发生同账号重试")
require.Equal(t, 1, fs.SwitchCount, "应立即切换账号")
require.Len(t, mock.calls, 1, "应立即 TempUnschedule")
})
}
// ---------------------------------------------------------------------------
@@ -374,7 +411,7 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false) // RetryableOnSameAccount=false
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Empty(t, mock.calls)
})
@@ -384,10 +421,10 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
err := newTestFailoverErr(502, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
}
// 再次触发时才会执行 TempUnschedule + 切换
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
require.Len(t, mock.calls, 1)
require.Equal(t, int64(42), mock.calls[0].accountID)
@@ -410,7 +447,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, "openai", err)
action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -429,7 +466,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -446,10 +483,10 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(100))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", newTestFailoverErr(502, false, false))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, newTestFailoverErr(502, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(200))
require.Len(t, fs.FailedAccountIDs, 2)
})
@@ -458,7 +495,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(0, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Equal(t, FailoverExhausted, action)
require.Contains(t, fs.FailedAccountIDs, int64(100))
})
@@ -467,7 +504,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(400, true, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(400, true, false))
require.Equal(t, FailoverContinue, action)
require.NotContains(t, fs.FailedAccountIDs, int64(100))
})
@@ -476,8 +513,8 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Len(t, fs.FailedAccountIDs, 1, "map 天然去重")
})
}
@@ -492,11 +529,11 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, err1, fs.LastFailoverErr)
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, err2, fs.LastFailoverErr)
})
@@ -505,7 +542,7 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(400, true, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, err, fs.LastFailoverErr)
})
}
@@ -522,30 +559,30 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 1. 账号 100 遇到可重试错误,同账号重试 maxSameAccountRetries 次
retryErr := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
}
require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling")
// 2. 账号 100 超过重试上限 → TempUnschedule + 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Len(t, mock.calls, 1)
// 3. 账号 200 遇到不可重试错误 → 直接切换
switchErr := newTestFailoverErr(500, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 4. 账号 300 遇到不可重试错误 → 再切换
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 3, fs.SwitchCount)
// 5. 账号 400 → 已耗尽 (SwitchCount=3 >= MaxSwitches=3)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverExhausted, action)
// 最终状态验证
@@ -563,21 +600,21 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次切换:delay = 0s
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
require.Less(t, elapsed, 200*time.Millisecond, "第一次切换延迟为 0")
// 第二次切换:delay = 1s
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverContinue, action)
require.GreaterOrEqual(t, elapsed, 800*time.Millisecond, "第二次切换延迟约 1s")
// 第三次:耗尽(无延迟,因为在检查延迟之前就返回了)
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverExhausted, action)
require.Less(t, elapsed, 200*time.Millisecond, "耗尽时不应有延迟")
@@ -589,17 +626,17 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次:ForceCacheBilling=false
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.False(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=true(Antigravity 粘性会话切换)
err2 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "错误标志应触发 ForceCacheBilling")
// 第三次:ForceCacheBilling=false,但状态仍保持 true
err3 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.True(t, fs.ForceCacheBilling, "不应重置")
})
}
@@ -614,7 +651,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(0, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
})
@@ -623,7 +660,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[0])
})
@@ -633,7 +670,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[-1])
})
@@ -645,7 +682,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
+10 -9
View File
@@ -448,7 +448,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, service.PlatformGemini, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -868,7 +868,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, account.Platform, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -1454,13 +1454,14 @@ func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context,
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
resp["remaining"] = remaining
resp["subscription"] = gin.H{
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"expires_at": subscription.ExpiresAt,
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"weekly_window_start": subscription.WeeklyWindowStart,
"expires_at": subscription.ExpiresAt,
}
}
@@ -254,7 +254,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
h.handleCCFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -233,7 +233,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
h.handleResponsesFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -0,0 +1,51 @@
package handler
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"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"
)
func TestUsageUnrestrictedIncludesWeeklyWindowStart(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/usage", nil)
weeklyWindowStart := time.Date(2026, time.July, 13, 0, 30, 0, 0, time.FixedZone("UTC+8", 8*60*60))
c.Set(string(middleware.ContextKeySubscription), &service.UserSubscription{
WeeklyWindowStart: &weeklyWindowStart,
})
handler := &GatewayHandler{}
handler.usageUnrestricted(
c,
context.Background(),
&service.APIKey{Group: &service.Group{
Name: "Weekly plan",
SubscriptionType: service.SubscriptionTypeSubscription,
}},
middleware.AuthSubject{},
nil,
nil,
nil,
)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Subscription struct {
WeeklyWindowStart *time.Time `json:"weekly_window_start"`
} `json:"subscription"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
require.NotNil(t, response.Subscription.WeeklyWindowStart)
require.True(t, weeklyWindowStart.Equal(*response.Subscription.WeeklyWindowStart))
}
@@ -220,6 +220,15 @@ func (h *ConcurrencyHelper) TryAcquireUserSlotForAPIKey(ctx context.Context, use
return h.withAPIKeySlot(ctx, apiKeyID, releaseFunc), true, nil
}
// AcquireOpenAIWSIngressLease bounds the whole client WebSocket lifecycle,
// independently from per-turn user and account slots.
func (h *ConcurrencyHelper) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*service.OpenAIWSIngressLease, bool, error) {
if h == nil || h.concurrencyService == nil {
return nil, false, fmt.Errorf("concurrency service is unavailable")
}
return h.concurrencyService.AcquireOpenAIWSIngressLease(ctx, apiKeyID, maxConnections)
}
// TryAcquireAccountSlot 尝试立即获取账号并发槽位。
// 返回值: (releaseFunc, acquired, error)
func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (func(), bool, error) {
@@ -11,10 +11,13 @@ import (
)
type concurrencyCacheMock struct {
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
releaseUserCalled int32
releaseAccountCalled int32
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
acquireIngressLeaseFn func(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error)
releaseIngressLeaseFn func(ctx context.Context, apiKeyID int64, leaseID string) error
releaseUserCalled int32
releaseAccountCalled int32
releaseIngressCalled int32
}
func (m *concurrencyCacheMock) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -97,6 +100,25 @@ func (m *concurrencyCacheMock) CleanupStaleProcessSlots(ctx context.Context, act
return nil
}
func (m *concurrencyCacheMock) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
if m.acquireIngressLeaseFn != nil {
return m.acquireIngressLeaseFn(ctx, apiKeyID, maxConnections, leaseID)
}
return false, nil
}
func (m *concurrencyCacheMock) RefreshOpenAIWSIngressLease(context.Context, int64, string) (bool, error) {
return true, nil
}
func (m *concurrencyCacheMock) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
atomic.AddInt32(&m.releaseIngressCalled, 1)
if m.releaseIngressLeaseFn != nil {
return m.releaseIngressLeaseFn(ctx, apiKeyID, leaseID)
}
return nil
}
func TestConcurrencyHelper_TryAcquireUserSlot(t *testing.T) {
cache := &concurrencyCacheMock{
acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -482,7 +482,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
if err != nil {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch failoverAction {
case FailoverContinue:
continue
+11 -1
View File
@@ -31,6 +31,16 @@ func (h *OpenAIGatewayHandler) GrokVideoGeneration(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosGenerations, "")
}
// GrokVideoEdit handles asynchronous xAI video edits through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoEdit(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosEdits, "")
}
// GrokVideoExtension handles asynchronous xAI video extensions through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoExtension(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosExtensions, "")
}
// GrokVideoStatus handles xAI video status retrieval through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
@@ -298,7 +308,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
if endpoint == service.GrokMediaEndpointVideosGenerations && strings.TrimSpace(result.ResponseID) != "" {
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
reqLog.Warn("grok_media.bind_video_request_account_failed",
zap.Int64("account_id", account.ID),
@@ -1299,10 +1299,36 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg))
ctx := c.Request.Context()
maxIngressConnections := 0
if h.cfg != nil {
maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey
}
ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections)
if ingressLeaseErr != nil {
reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr))
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to reserve websocket ingress capacity")
return
}
if !ingressLeaseAcquired {
reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "too many open websocket connections, please retry later")
return
}
if ingressLease != nil {
defer ingressLease.Release()
ctx = ingressLease.Context()
c.Request = c.Request.WithContext(ctx)
}
readCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
msgType, firstMessage, err := wsConn.Read(readCtx)
cancel()
if err != nil {
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost_before_first_message", zap.Error(err))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_read_first_message_failed",
zap.Error(err),
@@ -1692,6 +1718,25 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
continue
}
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost",
zap.Int64("account_id", account.ID),
zap.Error(err),
)
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) && closeErr.StatusCode() == coderws.StatusNormalClosure {
reqLog.Info("openai.websocket_ingress_closed_normally",
zap.Int64("account_id", account.ID),
zap.String("reason", closeErr.Reason()),
)
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_proxy_failed",
@@ -1700,7 +1745,6 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
zap.String("close_status", closeStatus),
zap.String("close_reason", closeReason),
)
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) {
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
@@ -8,6 +8,7 @@ import (
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
@@ -711,6 +712,68 @@ func TestOpenAIResponsesWebSocket_InvalidUpgradeDoesNotSetTransport(t *testing.T
require.Equal(t, service.OpenAIClientTransportUnknown, service.GetOpenAIClientTransport(c))
}
func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return false, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code)
}
func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return true, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageBinary, []byte("not a response.create frame"))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
require.Eventually(t, func() bool {
return atomic.LoadInt32(&cache.releaseIngressCalled) == 1
}, time.Second, 10*time.Millisecond)
}
func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
gin.SetMode(gin.TestMode)
+1 -15
View File
@@ -9,7 +9,6 @@ import (
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/payment"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -19,15 +18,13 @@ import (
// PaymentHandler handles user-facing payment requests.
type PaymentHandler struct {
channelService *service.ChannelService
paymentService *service.PaymentService
configService *service.PaymentConfigService
}
// NewPaymentHandler creates a new PaymentHandler.
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService, channelService *service.ChannelService) *PaymentHandler {
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService) *PaymentHandler {
return &PaymentHandler{
channelService: channelService,
paymentService: paymentService,
configService: configService,
}
@@ -91,17 +88,6 @@ func (h *PaymentHandler) GetPlans(c *gin.Context) {
response.Success(c, result)
}
// GetChannels returns enabled payment channels.
// GET /api/v1/payment/channels
func (h *PaymentHandler) GetChannels(c *gin.Context) {
channels, _, err := h.channelService.List(c.Request.Context(), pagination.PaginationParams{Page: 1, PageSize: 1000}, "active", "")
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, channels)
}
// GetCheckoutInfo returns all data the payment page needs in a single call:
// payment methods with limits, subscription plans, and configuration.
// GET /api/v1/payment/checkout-info
@@ -119,7 +119,7 @@ func TestVerifyOrderPublicReturnsLegacyOrderState(t *testing.T) {
require.NoError(t, err)
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -219,7 +219,7 @@ func TestResolveOrderPublicByResumeTokenReturnsFrontendContractFields(t *testing
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -307,7 +307,7 @@ func TestResolveOrderPublicByResumeTokenReturnsBadRequestForMismatchedToken(t *t
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -347,7 +347,7 @@ func TestVerifyOrderPublicRejectsBlankOutTradeNo(t *testing.T) {
t.Cleanup(func() { _ = client.Close() })
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -718,7 +718,7 @@ func TestStreamingToolCallDoneWithoutDeltaEmitsArguments(t *testing.T) {
assert.Equal(t, "content_block_stop", events[1].Type)
}
func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
func TestStreamingReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
@@ -739,18 +739,17 @@ func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
OutputIndex: 0,
Delta: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
assert.Len(t, events, 0)
require.Len(t, events, 1, "Read tool deltas must be streamed like any other tool")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
Type: "response.function_call_arguments.done",
OutputIndex: 0,
Arguments: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
require.Len(t, events, 2)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.JSONEq(t, `{"file_path":"/tmp/demo.py","limit":2000,"offset":0}`, events[0].Delta.PartialJSON)
assert.Equal(t, "content_block_stop", events[1].Type)
require.Len(t, events, 1, "after streaming deltas, .done should just close the block")
assert.Equal(t, "content_block_stop", events[0].Type)
}
func TestStreamingReasoning(t *testing.T) {
@@ -164,6 +164,8 @@ type AnthropicEventToResponsesState struct {
OutputTokens int
CacheReadInputTokens int
CacheCreationInputTokens int
StopReason string
}
// NewAnthropicEventToResponsesState returns an initialised stream state.
@@ -405,7 +407,6 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
}
func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
// Update usage
if evt.Usage != nil {
state.OutputTokens = evt.Usage.OutputTokens
if evt.Usage.InputTokens > 0 {
@@ -418,6 +419,9 @@ func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEven
state.CacheCreationInputTokens = evt.Usage.CacheCreationInputTokens
}
}
if evt.Delta != nil && evt.Delta.StopReason != "" {
state.StopReason = evt.Delta.StopReason
}
return nil
}
@@ -428,15 +432,15 @@ func anthToResHandleMessageStop(state *AnthropicEventToResponsesState) []Respons
}
var events []ResponsesStreamEvent
// Close any open item
events = append(events, closeCurrentResponsesItem(state)...)
// Determine status
status := "completed"
var incompleteDetails *ResponsesIncompleteDetails
if state.StopReason == "max_tokens" {
status = "incomplete"
incompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
}
// Emit response.completed
events = append(events, makeResponsesCompletedEvent(state, status, incompleteDetails))
state.CompletedSent = true
return events
@@ -509,15 +513,20 @@ func makeResponsesCompletedEvent(
}
}
eventType := "response.completed"
if status == "incomplete" {
eventType = "response.incomplete"
}
return ResponsesStreamEvent{
Type: "response.completed",
Type: eventType,
SequenceNumber: seq,
Response: &ResponsesResponse{
ID: state.ResponseID,
Object: "response",
Model: state.Model,
Status: status,
Output: []ResponsesOutput{}, // Simplified; full output tracking would add complexity
Output: []ResponsesOutput{},
Usage: usage,
IncompleteDetails: incompleteDetails,
},
@@ -35,8 +35,12 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
if req.Reasoning != nil {
out.ReasoningEffort = req.Reasoning.Effort
}
if len(req.Tools) > 0 {
tools, err := responsesToolsToChatTools(req.Tools)
effectiveTools, err := EffectiveResponsesTools(req)
if err != nil {
return nil, err
}
if len(effectiveTools) > 0 {
tools, err := responsesToolsToChatTools(effectiveTools)
if err != nil {
return nil, err
}
@@ -63,6 +67,44 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
return out, nil
}
// EffectiveResponsesTools returns every client-executable tool declared by a
// Responses request. Newer Codex clients place their runtime tools in an
// input item shaped as {"type":"additional_tools","tools":[...]} instead of
// the top-level tools field. Chat-only upstreams must receive both forms.
func EffectiveResponsesTools(req *ResponsesRequest) ([]ResponsesTool, error) {
if req == nil {
return nil, nil
}
tools := append([]ResponsesTool(nil), req.Tools...)
inputRaw := bytesTrimSpace(req.Input)
if len(inputRaw) == 0 || string(inputRaw) == "null" || inputRaw[0] != '[' {
return tools, nil
}
var items []json.RawMessage
if err := json.Unmarshal(inputRaw, &items); err != nil {
return nil, fmt.Errorf("parse responses input for additional tools: %w", err)
}
for _, raw := range items {
raw = bytesTrimSpace(raw)
if len(raw) == 0 || raw[0] != '{' {
continue
}
var item struct {
Type string `json:"type"`
Tools []ResponsesTool `json:"tools"`
}
if err := json.Unmarshal(raw, &item); err != nil {
return nil, fmt.Errorf("parse responses additional tools item: %w", err)
}
if item.Type == "additional_tools" {
tools = append(tools, item.Tools...)
}
}
return tools, nil
}
// CustomToolNames 收集 Responses 请求中 custom/freeform 工具的名字。chat 桥回程时
// 需要据此把模型对这些工具的调用还原为 custom_tool_call 项(codex 只按该类型路由)。
func CustomToolNames(tools []ResponsesTool) map[string]bool {
@@ -34,6 +34,51 @@ func TestResponsesToChatCompletionsRequest_CustomToolBecomesFunctionTool(t *test
assert.Equal(t, "wait", out.Tools[1].Function.Name)
}
func TestResponsesToChatCompletionsRequest_AdditionalToolsItem(t *testing.T) {
req := &ResponsesRequest{
Model: "gpt-test",
Input: json.RawMessage(`[
{"type":"additional_tools","role":"developer","tools":[
{"type":"custom","name":"exec","description":"Run PowerShell","format":{"type":"text"}},
{"type":"function","name":"wait","parameters":{"type":"object","properties":{}}},
{"type":"namespace","name":"collaboration","tools":[
{"type":"function","name":"send_message","parameters":{"type":"object","properties":{}}}
]}
]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"run Get-Location"}]}
]`),
ToolChoice: json.RawMessage(`"auto"`),
}
effective, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, effective, 3)
assert.True(t, CustomToolNames(effective)["exec"])
assert.Equal(t, NamespacedToolName{Namespace: "collaboration", Name: "send_message"}, NamespaceToolNames(effective)["collaboration__send_message"])
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 3)
assert.Equal(t, "exec", out.Tools[0].Function.Name)
assert.Equal(t, "wait", out.Tools[1].Function.Name)
assert.Equal(t, "collaboration__send_message", out.Tools[2].Function.Name)
assert.JSONEq(t, `"auto"`, string(out.ToolChoice))
require.Len(t, out.Messages, 1, "additional_tools must not become a chat message")
assert.Equal(t, "user", out.Messages[0].Role)
}
func TestEffectiveResponsesTools_SkipsStringInputItems(t *testing.T) {
req := &ResponsesRequest{
Input: json.RawMessage(`["plain input",{"type":"additional_tools","tools":[{"type":"custom","name":"exec"}]}]`),
}
tools, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, tools, 1)
assert.Equal(t, "exec", tools[0].Name)
}
func TestResponsesToChatCompletionsRequest_DropsToolChoiceWhenNoConvertibleTools(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
@@ -413,10 +413,6 @@ func resToAnthHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve
return nil
}
if state.CurrentBlockType == "tool_use" && state.CurrentToolName == "Read" {
state.CurrentToolArgs += evt.Delta
return nil
}
if state.CurrentBlockType == "tool_use" {
state.CurrentToolHadDelta = true
}
@@ -0,0 +1,84 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResToAnthFuncArgsDelta_ReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1, "Read tool delta must produce content_block_delta")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.Equal(t, `{"file_path":"/tmp/test.go"}`, events[0].Delta.PartialJSON)
assert.True(t, state.CurrentToolHadDelta, "Read deltas should set CurrentToolHadDelta")
}
func TestResToAnthFuncArgsDelta_ReadToolWithoutDone(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.ContentBlockIndex = 0
state.ContentBlockOpen = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
delta := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(delta, state)
require.Len(t, events, 1, "delta should be streamed")
completed := &ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
Status: "completed",
},
}
events = ResponsesEventToAnthropicEvents(completed, state)
hasStop := false
for _, e := range events {
if e.Type == "content_block_stop" {
hasStop = true
}
}
assert.True(t, hasStop, "block should be closed even without .done event")
}
func TestResToAnthFuncArgsDelta_NonReadToolUnchanged(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Write"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/out.txt","content":"hello"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.True(t, state.CurrentToolHadDelta)
}
@@ -89,8 +89,13 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
func responsesStatusToChatFinishReason(status string, details *ResponsesIncompleteDetails, toolCalls []ChatToolCall) string {
switch status {
case "incomplete":
if details != nil && details.Reason == "max_output_tokens" {
return "length"
if details != nil {
switch details.Reason {
case "max_output_tokens":
return "length"
case "content_filter":
return "content_filter"
}
}
return "stop"
case "completed":
@@ -299,8 +304,13 @@ func resToChatHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
switch evt.Response.Status {
case "incomplete":
if evt.Response.IncompleteDetails != nil && evt.Response.IncompleteDetails.Reason == "max_output_tokens" {
finishReason = "length"
if evt.Response.IncompleteDetails != nil {
switch evt.Response.IncompleteDetails.Reason {
case "max_output_tokens":
finishReason = "length"
case "content_filter":
finishReason = "content_filter"
}
}
case "completed":
if state.SawToolCall {
@@ -0,0 +1,122 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicStreamingMaxTokens_MapsToIncomplete(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{
StopReason: "max_tokens",
},
Usage: &AnthropicUsage{OutputTokens: 4096},
}, state)
require.Equal(t, "max_tokens", state.StopReason)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" || events[i].Type == "response.incomplete" {
completed = &events[i]
break
}
}
require.NotNil(t, completed, "should have terminal event")
assert.Equal(t, "response.incomplete", completed.Type)
require.NotNil(t, completed.Response)
assert.Equal(t, "incomplete", completed.Response.Status)
require.NotNil(t, completed.Response.IncompleteDetails)
assert.Equal(t, "max_output_tokens", completed.Response.IncompleteDetails.Reason)
}
func TestAnthropicStreamingEndTurn_MapsToCompleted(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{StopReason: "end_turn"},
Usage: &AnthropicUsage{OutputTokens: 100},
}, state)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" {
completed = &events[i]
break
}
}
require.NotNil(t, completed)
assert.Equal(t, "completed", completed.Response.Status)
assert.Nil(t, completed.Response.IncompleteDetails)
}
func TestResponsesToChatCompletions_ContentFilter(t *testing.T) {
resp := &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
Output: []ResponsesOutput{{
Type: "message",
Content: []ResponsesContentPart{{Type: "output_text", Text: "partial"}},
}},
Usage: &ResponsesUsage{InputTokens: 10, OutputTokens: 5},
}
cc := ResponsesToChatCompletions(resp, "gpt-5.5")
require.Len(t, cc.Choices, 1)
assert.Equal(t, "content_filter", cc.Choices[0].FinishReason)
}
func TestResponsesToChatCompletionsStreaming_ContentFilter(t *testing.T) {
state := NewResponsesEventToChatState()
state.ID = "resp_cf"
state.Model = "gpt-5.5"
state.SentRole = true
events := ResponsesEventToChatChunks(&ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
},
}, state)
hasContentFilter := false
for _, chunk := range events {
for _, choice := range chunk.Choices {
if choice.FinishReason != nil && *choice.FinishReason == "content_filter" {
hasContentFilter = true
}
}
}
assert.True(t, hasContentFilter, "streaming content_filter should map to finish_reason content_filter")
}
@@ -38,7 +38,7 @@ func (p PaginationParams) Offset() int {
if p.Page < 1 {
p.Page = 1
}
return (p.Page - 1) * p.PageSize
return (p.Page - 1) * p.Limit()
}
// Limit 获取限制数
@@ -69,3 +69,30 @@ func TestPaginationParamsLimit(t *testing.T) {
})
}
}
func TestPaginationParamsOffsetUsesNormalizedLimit(t *testing.T) {
t.Parallel()
tests := []struct {
name string
page int
pageSize int
want int
}{
{name: "invalid page uses first page", page: 0, pageSize: 50, want: 0},
{name: "zero page size uses default", page: 2, pageSize: 0, want: 20},
{name: "negative page size uses default", page: 2, pageSize: -1, want: 20},
{name: "normal values", page: 3, pageSize: 50, want: 100},
{name: "page size beyond max is clamped", page: 2, pageSize: 1500, want: 1000},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
params := PaginationParams{Page: tt.page, PageSize: tt.pageSize}
if got := params.Offset(); got != tt.want {
t.Fatalf("Offset() for Page=%d, PageSize=%d = %d, want %d", tt.page, tt.pageSize, got, tt.want)
}
})
}
}
+30 -1
View File
@@ -191,7 +191,7 @@ func RuntimeSanity() RuntimeSanityReport {
UnsafeURLOverrides: AllowUnsafeURLOverrides(),
UnsafeHighConcurrency: AllowUnsafeHighConcurrency(),
PublicGatewayScope: "responses_only",
ProxyPolicy: "account_proxy_optional; upstream URL allowlists enforced unless unsafe overrides are enabled",
ProxyPolicy: "account_proxy_optional; OAuth URLs use trusted-host allowlists; API-key base URLs require public HTTPS unless unsafe overrides are enabled",
}
}
@@ -252,6 +252,19 @@ func ValidateOAuthEndpointURL(raw string) (string, error) {
}
func ValidateBaseURL(raw string) (string, error) {
if AllowUnsafeURLOverrides() {
return urlvalidator.ValidateURLFormat(raw, true)
}
normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
AllowPrivate: false,
})
if err != nil {
return "", err
}
return normalizeKnownBaseURLPath(normalized)
}
func ValidateTrustedBaseURL(raw string) (string, error) {
if AllowUnsafeURLOverrides() {
return urlvalidator.ValidateURLFormat(raw, true)
}
@@ -461,6 +474,22 @@ func BuildVideosGenerationsURL(baseURL string) (string, error) {
return validatedBaseURL + "/videos/generations", nil
}
func BuildVideosEditsURL(baseURL string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
return "", fmt.Errorf("invalid base url: %w", err)
}
return validatedBaseURL + "/videos/edits", nil
}
func BuildVideosExtensionsURL(baseURL string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
return "", fmt.Errorf("invalid base url: %w", err)
}
return validatedBaseURL + "/videos/extensions", nil
}
func BuildVideoURL(baseURL, requestID string) (string, error) {
validatedBaseURL, err := ValidatedBaseURL(baseURL)
if err != nil {
+19 -4
View File
@@ -129,6 +129,14 @@ func TestBuildGrokMediaURLs(t *testing.T) {
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/generations", videosURL)
videoEditsURL, err := BuildVideosEditsURL(DefaultBaseURL)
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/edits", videoEditsURL)
videoExtensionsURL, err := BuildVideosExtensionsURL(DefaultBaseURL)
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/extensions", videoExtensionsURL)
videoURL, err := BuildVideoURL(DefaultBaseURL, "req 123")
require.NoError(t, err)
require.Equal(t, DefaultBaseURL+"/videos/req%20123", videoURL)
@@ -137,13 +145,10 @@ func TestBuildGrokMediaURLs(t *testing.T) {
require.Error(t, err)
}
func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) {
func TestValidateXAIURLsRejectUntrustedOAuthAndUnsafeBaseURLsByDefault(t *testing.T) {
_, err := ValidateOAuthEndpointURL("https://auth.example.test/oauth2/token")
require.Error(t, err)
_, err = ValidateBaseURL("https://xai.test/v1")
require.Error(t, err)
_, err = ValidateBaseURL("http://127.0.0.1:8080/v1")
require.Error(t, err)
@@ -151,6 +156,15 @@ func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) {
require.Error(t, err)
}
func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) {
baseURL, err := ValidateBaseURL("https://grok.example.test/v1/")
require.NoError(t, err)
require.Equal(t, "https://grok.example.test/v1", baseURL)
_, err = ValidateTrustedBaseURL("https://grok.example.test/v1")
require.Error(t, err)
}
func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) {
t.Setenv(EnvAllowUnsafeURLOverrides, "true")
@@ -182,6 +196,7 @@ func TestRuntimeSanityReportsSafeDefaults(t *testing.T) {
require.False(t, report.UnsafeHighConcurrency)
require.Equal(t, "responses_only", report.PublicGatewayScope)
require.Contains(t, report.ProxyPolicy, "account_proxy_optional")
require.Contains(t, report.ProxyPolicy, "API-key base URLs require public HTTPS")
}
func TestRuntimeSanityReportsInvalidOverridesWithoutSecrets(t *testing.T) {
+12 -10
View File
@@ -525,17 +525,19 @@ func (r *apiKeyRepository) latestUsageLogIPs(ctx context.Context, apiKeyIDs []in
func latestUsageLogIPsQuery(apiKeyIDs []int64, dialectName string) (string, []any) {
if dialectName == dialect.Postgres {
// Keep each key lookup bounded to one ordered index probe instead of ranking its full history.
return `
SELECT api_key_id, ip_address
FROM (
SELECT api_key_id, ip_address,
ROW_NUMBER() OVER (PARTITION BY api_key_id ORDER BY created_at DESC, id DESC) AS rn
FROM usage_logs
WHERE api_key_id = ANY($1::bigint[])
AND ip_address IS NOT NULL
AND ip_address <> ''
) ranked
WHERE rn = 1`, []any{pq.Array(apiKeyIDs)}
SELECT requested.api_key_id, latest.ip_address
FROM unnest($1::bigint[]) AS requested(api_key_id)
CROSS JOIN LATERAL (
SELECT ul.ip_address
FROM usage_logs AS ul
WHERE ul.api_key_id = requested.api_key_id
AND ul.ip_address IS NOT NULL
AND ul.ip_address <> ''
ORDER BY ul.created_at DESC, ul.id DESC
LIMIT 1
) AS latest`, []any{pq.Array(apiKeyIDs)}
}
placeholders := make([]string, len(apiKeyIDs))
@@ -3,6 +3,7 @@ package repository
import (
"context"
"database/sql"
"strings"
"testing"
"time"
@@ -125,6 +126,20 @@ func TestAPIKeyRepositoryListByUserIDAttachesLastUsedIP(t *testing.T) {
require.Nil(t, byID[noLogs.ID].LastUsedIP)
}
func TestLatestUsageLogIPsQueryPostgresUsesPerKeyLateralLookup(t *testing.T) {
query, args := latestUsageLogIPsQuery([]int64{11, 22}, dialect.Postgres)
normalizedQuery := strings.Join(strings.Fields(query), " ")
require.Contains(t, normalizedQuery, "FROM unnest($1::bigint[]) AS requested(api_key_id)")
require.Contains(t, normalizedQuery, "CROSS JOIN LATERAL")
require.Contains(t, normalizedQuery, "WHERE ul.api_key_id = requested.api_key_id")
require.Contains(t, normalizedQuery, "AND ul.ip_address IS NOT NULL")
require.Contains(t, normalizedQuery, "AND ul.ip_address <> ''")
require.Contains(t, normalizedQuery, "ORDER BY ul.created_at DESC, ul.id DESC LIMIT 1")
require.NotContains(t, normalizedQuery, "ROW_NUMBER")
require.Len(t, args, 1)
}
func TestAPIKeyRepository_CreateWithLastUsedAt(t *testing.T) {
repo, client := newAPIKeyRepoSQLite(t)
ctx := context.Background()
@@ -30,6 +30,10 @@ const (
userSlotKeyPrefix = "concurrency:user:"
// 格式: concurrency:api_key:{apiKeyID}
apiKeySlotKeyPrefix = "concurrency:api_key:"
// API-key-scoped client WebSocket ingress leases use a shorter TTL than
// ordinary request slots, because idle ingress sessions do not hold a turn slot.
openAIWSIngressLeaseKeyPrefix = "concurrency:openai_ws_ingress:api_key:"
openAIWSIngressLeaseTTLSeconds = 60
// 等待队列计数器格式: concurrency:wait:{userID}
waitQueueKeyPrefix = "concurrency:wait:"
// 账号级等待队列计数器格式: wait:account:{accountID}
@@ -138,6 +142,49 @@ var (
return 1
`)
// acquireOpenAIWSIngressLeaseScript atomically reaps crashed members and
// acquires or refreshes one API-key-scoped ingress lease using Redis TIME.
acquireOpenAIWSIngressLeaseScript = redis.NewScript(`
redis.replicate_commands()
local key = KEYS[1]
local maxConnections = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local leaseID = ARGV[3]
local now = tonumber(redis.call('TIME')[1])
local expireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZSCORE', key, leaseID) ~= false then
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
end
if redis.call('ZCARD', key) < maxConnections then
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
end
return 0
`)
// refreshOpenAIWSIngressLeaseScript does not recreate a missing member: a
// process that lost its lease must terminate its local WebSocket instead of
// silently continuing beyond the distributed cap.
refreshOpenAIWSIngressLeaseScript = redis.NewScript(`
redis.replicate_commands()
local key = KEYS[1]
local ttl = tonumber(ARGV[1])
local leaseID = ARGV[2]
local now = tonumber(redis.call('TIME')[1])
local expireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZSCORE', key, leaseID) == false then
return 0
end
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
return 1
`)
// incrementWaitScript - refreshes TTL on each increment to keep queue depth accurate
// KEYS[1] = wait queue key
// ARGV[1] = maxWait
@@ -283,6 +330,10 @@ func apiKeySlotKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
}
func openAIWSIngressLeaseKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", openAIWSIngressLeaseKeyPrefix, apiKeyID)
}
func waitQueueKey(userID int64) string {
return fmt.Sprintf("%s%d", waitQueueKeyPrefix, userID)
}
@@ -623,6 +674,48 @@ func (c *concurrencyCache) ReleaseAPIKeySlot(ctx context.Context, apiKeyID int64
return c.rdb.ZRem(ctx, key, requestID).Err()
}
func (c *concurrencyCache) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
if c == nil || c.rdb == nil || apiKeyID <= 0 || maxConnections <= 0 || leaseID == "" {
return false, nil
}
result, err := acquireOpenAIWSIngressLeaseScript.Run(
ctx,
c.rdb,
[]string{openAIWSIngressLeaseKey(apiKeyID)},
maxConnections,
openAIWSIngressLeaseTTLSeconds,
leaseID,
).Int()
if err != nil {
return false, err
}
return result == 1, nil
}
func (c *concurrencyCache) RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) {
if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
return false, nil
}
result, err := refreshOpenAIWSIngressLeaseScript.Run(
ctx,
c.rdb,
[]string{openAIWSIngressLeaseKey(apiKeyID)},
openAIWSIngressLeaseTTLSeconds,
leaseID,
).Int()
if err != nil {
return false, err
}
return result == 1, nil
}
func (c *concurrencyCache) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
if c == nil || c.rdb == nil || apiKeyID <= 0 || leaseID == "" {
return nil
}
return c.rdb.ZRem(ctx, openAIWSIngressLeaseKey(apiKeyID), leaseID).Err()
}
func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
if len(apiKeyIDs) == 0 {
return map[int64]int{}, nil
@@ -50,6 +50,53 @@ func (s *ConcurrencyCacheSuite) apiKeyConcurrencyCache() apiKeyConcurrencyCacheF
return cache
}
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_HardLimitRefreshAndRelease() {
apiKeyID := int64(9011)
firstLeaseID := "ingress-first"
secondLeaseID := "ingress-second"
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, firstLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok)
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
require.NoError(s.T(), err)
require.False(s.T(), ok, "a second live session must not exceed the API key limit")
ok, err = s.rawCache.RefreshOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok, "the current owner must be able to refresh its lease")
require.NoError(s.T(), s.rawCache.ReleaseOpenAIWSIngressLease(s.ctx, apiKeyID, firstLeaseID))
ok, err = s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 1, secondLeaseID)
require.NoError(s.T(), err)
require.True(s.T(), ok, "released capacity must become available immediately")
}
func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_ReapsCrashedLeaseWithoutDeletingLiveOtherInstance() {
apiKeyID := int64(9012)
key := openAIWSIngressLeaseKey(apiKeyID)
now, err := s.rawCache.redisUnixSeconds(s.ctx)
require.NoError(s.T(), err)
require.NoError(s.T(), s.rdb.ZAdd(s.ctx, key,
redis.Z{Score: float64(now - openAIWSIngressLeaseTTLSeconds - 1), Member: "crashed-instance"},
redis.Z{Score: float64(now), Member: "live-other-instance"},
).Err())
require.NoError(s.T(), s.rdb.Expire(s.ctx, key, time.Duration(openAIWSIngressLeaseTTLSeconds)*time.Second).Err())
ok, err := s.rawCache.AcquireOpenAIWSIngressLease(s.ctx, apiKeyID, 2, "new-instance")
require.NoError(s.T(), err)
require.True(s.T(), ok, "the crashed member should be reaped before enforcing the limit")
_, err = s.rdb.ZScore(s.ctx, key, "crashed-instance").Result()
require.ErrorIs(s.T(), err, redis.Nil)
_, err = s.rdb.ZScore(s.ctx, key, "live-other-instance").Result()
require.NoError(s.T(), err, "a live lease owned by another instance must be preserved")
count, err := s.rdb.ZCard(s.ctx, key).Result()
require.NoError(s.T(), err)
require.Equal(s.T(), int64(2), count)
}
func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
accountID := int64(10)
reqID1, reqID2, reqID3 := "req1", "req2", "req3"
@@ -55,6 +55,8 @@ const paymentOrdersOutTradeNoUniqueMigration = "120_enforce_payment_orders_out_t
const paymentOrdersOutTradeNoUniqueIndex = "paymentorder_out_trade_no_unique"
const schedulerOutboxPendingDedupKeyMigration = "153_scheduler_outbox_pending_dedup_key_index_notx.sql"
const schedulerOutboxPendingDedupKeyIndex = "idx_scheduler_outbox_pending_dedup_key"
const latestAPIKeyIPIndexMigration = "174_add_usage_logs_api_key_latest_ip_index_notx.sql"
const latestAPIKeyIPIndex = "idx_usage_logs_api_key_latest_ip"
type migrationChecksumCompatibilityRule struct {
fileChecksum string
@@ -264,6 +266,8 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri
return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db)
case schedulerOutboxPendingDedupKeyMigration:
return dropInvalidIndexIfPresent(ctx, db, schedulerOutboxPendingDedupKeyIndex)
case latestAPIKeyIPIndexMigration:
return dropInvalidIndexIfPresent(ctx, db, latestAPIKeyIPIndex)
default:
return nil
}
@@ -116,6 +116,45 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_t_b ON t(b);
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyMigrationsFS_NonTransactionalMigration_LatestAPIKeyIPIndexDropsInvalidIndexBeforeRetry(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
defer func() { _ = db.Close() }()
prepareMigrationsBootstrapExpectations(mock)
mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1").
WithArgs(latestAPIKeyIPIndexMigration).
WillReturnError(sql.ErrNoRows)
mock.ExpectQuery("SELECT EXISTS \\(").
WithArgs(latestAPIKeyIPIndex).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectExec("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_logs_api_key_latest_ip").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec("CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip").
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectExec("INSERT INTO schema_migrations \\(filename, checksum\\) VALUES \\(\\$1, \\$2\\)").
WithArgs(latestAPIKeyIPIndexMigration, sqlmock.AnyArg()).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)").
WithArgs(migrationsAdvisoryLockID).
WillReturnResult(sqlmock.NewResult(0, 1))
fsys := fstest.MapFS{
latestAPIKeyIPIndexMigration: &fstest.MapFile{
Data: []byte(`
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip
ON usage_logs (api_key_id, created_at DESC, id DESC)
INCLUDE (ip_address)
WHERE ip_address IS NOT NULL AND ip_address <> '';
`),
},
}
err = applyMigrationsFS(context.Background(), db, fsys)
require.NoError(t, err)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyMigrationsFS_PaymentOrdersOutTradeNoUniqueMigration_FailsFastOnDuplicatePrecheck(t *testing.T) {
db, mock, err := sqlmock.New()
require.NoError(t, err)
+48 -21
View File
@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"log/slog"
"strconv"
"time"
@@ -163,14 +164,15 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul
versionStr := strconv.FormatInt(version, 10)
snapshotKey := schedulerSnapshotKey(bucket, versionStr)
if err := c.writeAccounts(ctx, accounts); err != nil {
cacheableAccounts, err := c.writeAccounts(ctx, accounts)
if err != nil {
return err
}
if len(accounts) > 0 {
if len(cacheableAccounts) > 0 {
// 使用序号作为 score,保持数据库返回的排序语义。
members := make([]redis.Z, 0, len(accounts))
for idx, account := range accounts {
members := make([]redis.Z, 0, len(cacheableAccounts))
for idx, account := range cacheableAccounts {
members = append(members, redis.Z{
Score: float64(idx),
Member: strconv.FormatInt(account.ID, 10),
@@ -224,7 +226,14 @@ func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Accoun
if account == nil || account.ID <= 0 {
return nil
}
return c.writeAccounts(ctx, []service.Account{*account})
cacheableAccounts, err := c.writeAccounts(ctx, []service.Account{*account})
if err != nil {
return err
}
if len(cacheableAccounts) == 0 {
return c.DeleteAccount(ctx, account.ID)
}
return nil
}
func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) error {
@@ -262,13 +271,14 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t
return err
}
account.LastUsedAt = ptrTime(updates[ids[i]])
updated, err := json.Marshal(account)
updated, metaPayload, err := marshalSchedulerCacheAccount(*account)
if err != nil {
return err
}
metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(*account))
if err != nil {
return err
slog.Warn("scheduler cache removes account with unencodable payload",
"account_id", ids[i],
"error", err,
)
pipe.Del(ctx, keys[i], schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10)))
continue
}
pipe.Set(ctx, keys[i], updated, 0)
pipe.Set(ctx, schedulerAccountMetaKey(strconv.FormatInt(ids[i], 10)), metaPayload, 0)
@@ -359,12 +369,13 @@ func decodeCachedAccount(val any) (*service.Account, error) {
return &account, nil
}
func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) error {
func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) {
if len(accounts) == 0 {
return nil
return nil, nil
}
pipe := c.rdb.Pipeline()
cacheableAccounts := make([]service.Account, 0, len(accounts))
pending := 0
flush := func() error {
if pending == 0 {
@@ -379,27 +390,43 @@ func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.A
}
for _, account := range accounts {
fullPayload, err := json.Marshal(account)
fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account)
if err != nil {
return err
}
metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account))
if err != nil {
return err
slog.Warn("scheduler cache skips account with unencodable payload",
"account_id", account.ID,
"error", err,
)
continue
}
id := strconv.FormatInt(account.ID, 10)
pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0)
pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0)
cacheableAccounts = append(cacheableAccounts, account)
pending++
if pending >= c.writeChunkSize {
if err := flush(); err != nil {
return err
return nil, err
}
}
}
return flush()
if err := flush(); err != nil {
return nil, err
}
return cacheableAccounts, nil
}
func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) {
fullPayload, err := json.Marshal(account)
if err != nil {
return nil, nil, fmt.Errorf("marshal account: %w", err)
}
metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account))
if err != nil {
return nil, nil, fmt.Errorf("marshal account metadata: %w", err)
}
return fullPayload, metaPayload, nil
}
func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, error) {
@@ -3,12 +3,78 @@
package repository
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func newSchedulerCacheUnit(t *testing.T) *schedulerCache {
t.Helper()
mr := miniredis.RunT(t)
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = rdb.Close() })
cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache)
require.True(t, ok)
return cache
}
func TestSchedulerCacheWriteAccountsSkipsUnencodableTimes(t *testing.T) {
ctx := context.Background()
cache := newSchedulerCacheUnit(t)
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
cacheable, err := cache.writeAccounts(ctx, []service.Account{
{ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey},
{ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime},
})
require.NoError(t, err)
require.Len(t, cacheable, 1)
require.Equal(t, int64(111), cacheable[0].ID)
cached, err := cache.GetAccount(ctx, 111)
require.NoError(t, err)
require.NotNil(t, cached)
invalid, err := cache.GetAccount(ctx, 112)
require.NoError(t, err)
require.Nil(t, invalid)
}
func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) {
ctx := context.Background()
cache := newSchedulerCacheUnit(t)
account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
require.NoError(t, cache.SetAccount(ctx, &account))
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
account.ExpiresAt = &invalidTime
require.NoError(t, cache.SetAccount(ctx, &account))
cached, err := cache.GetAccount(ctx, account.ID)
require.NoError(t, err)
require.Nil(t, cached)
}
func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) {
ctx := context.Background()
cache := newSchedulerCacheUnit(t)
account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
require.NoError(t, cache.SetAccount(ctx, &account))
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime}))
cached, err := cache.GetAccount(ctx, account.ID)
require.NoError(t, err)
require.Nil(t, cached)
}
func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
account := service.Account{
ID: 42,
+20
View File
@@ -84,6 +84,22 @@ func RegisterGatewayRoutes(
},
})
}
videoEditHandler := func(c *gin.Context) {
if getGroupPlatform(c) == service.PlatformGrok {
h.OpenAIGateway.GrokVideoEdit(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}})
}
videoExtensionHandler := func(c *gin.Context) {
if getGroupPlatform(c) == service.PlatformGrok {
h.OpenAIGateway.GrokVideoExtension(c)
return
}
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Videos API is not supported for this platform"}})
}
// API网关(Claude API兼容)
gateway := r.Group("/v1")
gateway.Use(bodyLimit)
@@ -185,6 +201,8 @@ func RegisterGatewayRoutes(
gateway.DELETE("/images/batches/:id", h.BatchImage.DeleteRecord)
gateway.DELETE("/images/batches/:id/outputs", h.BatchImage.DeleteOutputs)
gateway.POST("/videos/generations", videoGenerationHandler)
gateway.POST("/videos/edits", videoEditHandler)
gateway.POST("/videos/extensions", videoExtensionHandler)
gateway.GET("/videos/:request_id", videoStatusHandler)
}
@@ -252,6 +270,8 @@ func RegisterGatewayRoutes(
r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler)
r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler)
r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler)
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler)
r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler)
// Antigravity 模型列表
@@ -123,6 +123,10 @@ func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) {
"/images/edits",
"/v1/videos/generations",
"/videos/generations",
"/v1/videos/edits",
"/videos/edits",
"/v1/videos/extensions",
"/videos/extensions",
} {
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok-imagine","prompt":"draw a cat"}`))
req.Header.Set("Content-Type", "application/json")
@@ -156,6 +160,10 @@ func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) {
}{
{http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`},
{http.MethodPost, "/v1/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/videos/edits", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/v1/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodPost, "/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
{http.MethodGet, "/v1/videos/request-123", ""},
{http.MethodGet, "/videos/request-123", ""},
} {
@@ -28,7 +28,6 @@ func RegisterPaymentRoutes(
authenticated.GET("/config", paymentHandler.GetPaymentConfig)
authenticated.GET("/checkout-info", paymentHandler.GetCheckoutInfo)
authenticated.GET("/plans", paymentHandler.GetPlans)
authenticated.GET("/channels", paymentHandler.GetChannels)
authenticated.GET("/limits", paymentHandler.GetLimits)
orders := authenticated.Group("/orders")
+41 -1
View File
@@ -1265,6 +1265,9 @@ func (a *Account) GetOpenAIRefreshToken() string {
return a.GetCredential("refresh_token")
}
// GetGrokBaseURL selects the upstream used by Grok text and Responses traffic.
// Grok media traffic has a different transport contract and must use
// GetGrokMediaBaseURL instead.
func (a *Account) GetGrokBaseURL() string {
if !a.IsGrok() {
return ""
@@ -1274,6 +1277,10 @@ func (a *Account) GetGrokBaseURL() string {
if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) {
return xai.DefaultCLIBaseURL
}
if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil {
return baseURL
}
return xai.DefaultCLIBaseURL
}
if baseURL != "" {
return baseURL
@@ -1281,12 +1288,45 @@ func (a *Account) GetGrokBaseURL() string {
return xai.DefaultBaseURL
}
// GetGrokMediaBaseURL selects the upstream used by Grok Imagine APIs.
//
// OAuth text requests need the CLI subscription proxy, but that proxy has a
// smaller request-body limit than the official Imagine API. Media requests can
// contain large base64 inputs, so default OAuth accounts must use api.x.ai.
// API-key accounts and explicit unsafe development overrides retain their
// configured base URL.
func (a *Account) GetGrokMediaBaseURL() string {
if !a.IsGrok() {
return ""
}
if !a.IsGrokOAuth() {
return a.GetGrokBaseURL()
}
baseURL := a.GetCredential("base_url")
if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) || isOfficialGrokCLIBaseURL(baseURL) {
return xai.DefaultBaseURL
}
if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil {
return baseURL
}
return xai.DefaultBaseURL
}
func isOfficialGrokAPIBaseURL(raw string) bool {
return isOfficialGrokBaseURL(raw, xai.DefaultBaseURL)
}
func isOfficialGrokCLIBaseURL(raw string) bool {
return isOfficialGrokBaseURL(raw, xai.DefaultCLIBaseURL)
}
func isOfficialGrokBaseURL(raw, expected string) bool {
parsed, err := url.Parse(strings.TrimSpace(raw))
if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
return false
}
defaultURL, err := url.Parse(xai.DefaultBaseURL)
defaultURL, err := url.Parse(expected)
if err != nil {
return false
}
@@ -266,7 +266,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
expected: "https://api.x.ai:8443/v1",
},
{
name: "oauth explicit custom base_url remains supported",
name: "oauth explicit custom base_url stays pinned to CLI proxy by default",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
@@ -274,7 +274,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
"base_url": "https://custom.example.com/v1",
},
},
expected: "https://custom.example.com/v1",
expected: xai.DefaultCLIBaseURL,
},
{
name: "API key without base_url uses official credit-backed API",
@@ -293,3 +293,117 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
})
}
}
func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
account := Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
}
require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL())
}
func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
tests := []struct {
name string
account Account
expected string
}{
{
name: "oauth without base_url uses official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored CLI proxy uses official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultCLIBaseURL,
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth stored CLI proxy variant uses official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/",
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth legacy official API remains on official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": xai.DefaultBaseURL,
},
},
expected: xai.DefaultBaseURL,
},
{
name: "oauth untrusted custom base_url is pinned to official media API",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
},
expected: xai.DefaultBaseURL,
},
{
name: "API key retains its configured media API",
account: Account{
Type: AccountTypeAPIKey,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://grok.example.com/v1",
},
},
expected: "https://grok.example.com/v1",
},
{
name: "non-Grok account has no Grok media base URL",
account: Account{
Type: AccountTypeOAuth,
Platform: PlatformOpenAI,
Credentials: map[string]any{},
},
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.expected, tt.account.GetGrokMediaBaseURL())
})
}
}
func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
account := Account{
Type: AccountTypeOAuth,
Platform: PlatformGrok,
Credentials: map[string]any{
"base_url": "https://custom.example.com/v1",
},
}
require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL())
}
@@ -582,8 +582,13 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
c.Writer.Header().Set("X-Accel-Buffering", "no")
c.Writer.Flush()
// Create OpenAI Responses API payload
payload := createOpenAITestPayload(testModelID, isOAuth)
// Create OpenAI Responses API payload. OAuth accounts use ChatGPT Codex
// upstream and must apply the same model normalization as real forwarding.
upstreamTestModelID := testModelID
if isOAuth {
upstreamTestModelID = normalizeOpenAIModelForUpstream(credentialAccount, testModelID)
}
payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
payloadBytes, _ := json.Marshal(payload)
// Send test_start event
@@ -137,6 +137,34 @@ func TestAccountTestService_OpenAISuccessPersistsSnapshotFromHeaders(t *testing.
require.Contains(t, recorder.Body.String(), "test_complete")
}
func TestAccountTestService_OpenAIOAuthTestNormalizesGPT56Alias(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := newTestContext()
resp := newJSONResponse(http.StatusOK, "")
resp.Body = io.NopCloser(strings.NewReader(`data: {"type":"response.completed"}
`))
upstream := &queuedHTTPUpstream{responses: []*http.Response{resp}}
svc := &AccountTestService{httpUpstream: upstream}
account := &Account{
ID: 90,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "test-token"},
}
err := svc.testOpenAIAccountConnection(ctx, account, "gpt-5.6", "", "")
require.NoError(t, err)
require.Len(t, upstream.requests, 1)
body, err := io.ReadAll(upstream.requests[0].Body)
require.NoError(t, err)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(body, "model").String())
}
func TestAccountTestService_OpenAIShadowUsesParentCredentialsAndShadowModel(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, recorder := newTestContext()
@@ -6,6 +6,7 @@ import (
"crypto/sha256"
"encoding/binary"
"encoding/hex"
"errors"
"os"
"strconv"
"sync"
@@ -13,6 +14,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"go.uber.org/zap"
"golang.org/x/sync/singleflight"
)
@@ -59,6 +61,131 @@ type APIKeyConcurrencyCache interface {
GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error)
}
// OpenAIWSIngressLeaseCache owns the short-lived distributed lease used to
// bound live client WebSocket sessions. It is deliberately independent of the
// request-slot namespace: idle ingress connections do not occupy turn slots.
type OpenAIWSIngressLeaseCache interface {
AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error)
RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error)
ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error
}
const (
openAIWSIngressLeaseTTL = 60 * time.Second
openAIWSIngressLeaseRefreshInterval = 20 * time.Second
openAIWSIngressLeaseOperationTO = 2 * time.Second
)
var ErrOpenAIWSIngressLeaseLost = errors.New("openai websocket ingress lease lost")
// OpenAIWSIngressLease keeps a Redis-backed ingress lease alive and cancels
// its context if Redis cannot confirm ownership for a full lease lifetime.
// Call Release on every handler exit to reclaim capacity immediately.
type OpenAIWSIngressLease struct {
ctx context.Context
cancel context.CancelCauseFunc
cache OpenAIWSIngressLeaseCache
apiKeyID int64
leaseID string
stopOnce sync.Once
stopCh chan struct{}
refreshDone chan struct{}
}
func (l *OpenAIWSIngressLease) Context() context.Context {
if l == nil || l.ctx == nil {
return context.Background()
}
return l.ctx
}
func (l *OpenAIWSIngressLease) Release() {
if l == nil {
return
}
l.stopOnce.Do(func() {
if l.stopCh != nil {
close(l.stopCh)
}
if l.cancel != nil {
l.cancel(nil)
}
if l.refreshDone != nil {
<-l.refreshDone
}
if l.cache == nil || l.apiKeyID <= 0 || l.leaseID == "" {
return
}
releaseCtx, releaseCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO)
defer releaseCancel()
if err := l.cache.ReleaseOpenAIWSIngressLease(releaseCtx, l.apiKeyID, l.leaseID); err != nil {
logger.L().Warn("openai_ws_ingress_lease_release_failed",
zap.Int64("api_key_id", l.apiKeyID),
zap.Error(err),
)
}
})
}
func (l *OpenAIWSIngressLease) refreshLoop() {
defer func() {
if l != nil && l.refreshDone != nil {
close(l.refreshDone)
}
}()
if l == nil || l.cache == nil {
return
}
ticker := time.NewTicker(openAIWSIngressLeaseRefreshInterval)
defer ticker.Stop()
lastConfirmedAt := time.Now()
for {
select {
case <-l.ctx.Done():
return
case <-l.stopCh:
return
case <-ticker.C:
var lost bool
lastConfirmedAt, lost = l.refresh(lastConfirmedAt)
if lost {
l.cancel(ErrOpenAIWSIngressLeaseLost)
return
}
}
}
}
// refresh confirms the lease is still owned. A missing member is an immediate
// lease loss; transient Redis errors are tolerated only for one full lease TTL.
func (l *OpenAIWSIngressLease) refresh(lastConfirmedAt time.Time) (time.Time, bool) {
refreshCtx, refreshCancel := context.WithTimeout(context.Background(), openAIWSIngressLeaseOperationTO)
owned, err := l.cache.RefreshOpenAIWSIngressLease(refreshCtx, l.apiKeyID, l.leaseID)
refreshCancel()
if err == nil && owned {
return time.Now(), false
}
if err == nil {
err = ErrOpenAIWSIngressLeaseLost
}
elapsed := time.Since(lastConfirmedAt)
logger.L().Warn("openai_ws_ingress_lease_refresh_failed",
zap.Int64("api_key_id", l.apiKeyID),
zap.Duration("unconfirmed_for", elapsed),
zap.Error(err),
)
if errors.Is(err, ErrOpenAIWSIngressLeaseLost) || elapsed >= openAIWSIngressLeaseTTL {
logger.L().Error("openai_ws_ingress_lease_lost",
zap.Int64("api_key_id", l.apiKeyID),
zap.Duration("unconfirmed_for", elapsed),
zap.Error(err),
)
return lastConfirmedAt, true
}
return lastConfirmedAt, false
}
var (
requestIDPrefix = initRequestIDPrefix()
requestIDCounter atomic.Uint64
@@ -125,6 +252,47 @@ func NewConcurrencyService(cache ConcurrencyCache) *ConcurrencyService {
return svc
}
// AcquireOpenAIWSIngressLease atomically reserves one live ingress connection
// for an API key. A non-positive limit explicitly disables this protection.
func (s *ConcurrencyService) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*OpenAIWSIngressLease, bool, error) {
if maxConnections <= 0 {
return nil, true, nil
}
if s == nil || s.cache == nil || apiKeyID <= 0 {
return nil, false, errors.New("openai websocket ingress lease cache is unavailable")
}
cache, ok := s.cache.(OpenAIWSIngressLeaseCache)
if !ok {
return nil, false, errors.New("openai websocket ingress lease cache is unsupported")
}
leaseID := generateRequestID()
baseCtx := context.Background()
if ctx != nil {
baseCtx = context.WithoutCancel(ctx)
}
acquireCtx, acquireCancel := context.WithTimeout(baseCtx, openAIWSIngressLeaseOperationTO)
acquired, err := cache.AcquireOpenAIWSIngressLease(acquireCtx, apiKeyID, maxConnections, leaseID)
acquireCancel()
if err != nil || !acquired {
return nil, acquired, err
}
if ctx == nil {
ctx = context.Background()
}
leaseCtx, leaseCancel := context.WithCancelCause(ctx)
lease := &OpenAIWSIngressLease{
ctx: leaseCtx,
cancel: leaseCancel,
cache: cache,
apiKeyID: apiKeyID,
leaseID: leaseID,
stopCh: make(chan struct{}),
refreshDone: make(chan struct{}),
}
go lease.refreshLoop()
return lease, true, nil
}
// SetAccountLoadBatchCacheTTL 设置账号负载批量读取的极短 TTL 缓存;非正数表示禁用缓存。
func (s *ConcurrencyService) SetAccountLoadBatchCacheTTL(ttl time.Duration) {
if s == nil {
@@ -45,7 +45,47 @@ type stubConcurrencyCacheForTest struct {
releasedAPIKeyRequestIDs []string
}
type ingressLeaseCacheForTest struct {
stubConcurrencyCacheForTest
acquireIngressResult bool
acquireIngressErr error
acquireIngressFn func(context.Context, int64, int, string) (bool, error)
refreshIngressResult bool
refreshIngressErr error
refreshIngressFn func(context.Context, int64, string) (bool, error)
releaseIngressErr error
releaseIngressFn func(context.Context, int64, string) error
acquireIngressCalls int
refreshIngressCalls int
releaseIngressCalls int
}
func (c *ingressLeaseCacheForTest) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
c.acquireIngressCalls++
if c.acquireIngressFn != nil {
return c.acquireIngressFn(ctx, apiKeyID, maxConnections, leaseID)
}
return c.acquireIngressResult, c.acquireIngressErr
}
func (c *ingressLeaseCacheForTest) RefreshOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) (bool, error) {
c.refreshIngressCalls++
if c.refreshIngressFn != nil {
return c.refreshIngressFn(ctx, apiKeyID, leaseID)
}
return c.refreshIngressResult, c.refreshIngressErr
}
func (c *ingressLeaseCacheForTest) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
c.releaseIngressCalls++
if c.releaseIngressFn != nil {
return c.releaseIngressFn(ctx, apiKeyID, leaseID)
}
return c.releaseIngressErr
}
var _ ConcurrencyCache = (*stubConcurrencyCacheForTest)(nil)
var _ OpenAIWSIngressLeaseCache = (*ingressLeaseCacheForTest)(nil)
func (c *stubConcurrencyCacheForTest) AcquireAccountSlot(_ context.Context, _ int64, _ int, _ string) (bool, error) {
return c.acquireResult, c.acquireErr
@@ -285,6 +325,114 @@ func TestGetAPIKeyConcurrencyBatch_Fallbacks(t *testing.T) {
})
}
func TestAcquireOpenAIWSIngressLease(t *testing.T) {
t.Run("zero value release is safe", func(t *testing.T) {
var lease OpenAIWSIngressLease
require.NotPanics(t, lease.Release)
})
t.Run("disabled", func(t *testing.T) {
cache := &ingressLeaseCacheForTest{}
lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(nil, 1, 0)
require.NoError(t, err)
require.True(t, acquired)
require.Nil(t, lease)
require.Zero(t, cache.acquireIngressCalls)
})
t.Run("unsupported cache fails closed", func(t *testing.T) {
lease, acquired, err := NewConcurrencyService(&stubConcurrencyCacheForTest{}).AcquireOpenAIWSIngressLease(context.Background(), 1, 1)
require.Error(t, err)
require.False(t, acquired)
require.Nil(t, lease)
})
t.Run("capacity rejected", func(t *testing.T) {
cache := &ingressLeaseCacheForTest{acquireIngressResult: false}
lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(context.Background(), 1, 1)
require.NoError(t, err)
require.False(t, acquired)
require.Nil(t, lease)
})
t.Run("release returns capacity", func(t *testing.T) {
cache := &ingressLeaseCacheForTest{acquireIngressResult: true, refreshIngressResult: true}
lease, acquired, err := NewConcurrencyService(cache).AcquireOpenAIWSIngressLease(nil, 1, 1)
require.NoError(t, err)
require.True(t, acquired)
require.NotNil(t, lease)
lease.Release()
lease.Release()
require.Equal(t, 1, cache.releaseIngressCalls)
})
}
func TestOpenAIWSIngressLeaseRefreshLoss(t *testing.T) {
t.Run("missing lease is lost immediately", func(t *testing.T) {
cache := &ingressLeaseCacheForTest{refreshIngressResult: false}
lease := &OpenAIWSIngressLease{cache: cache, apiKeyID: 1, leaseID: "missing"}
_, lost := lease.refresh(time.Now())
require.True(t, lost)
require.Equal(t, 1, cache.refreshIngressCalls)
})
t.Run("persistent redis errors lose lease after ttl", func(t *testing.T) {
cache := &ingressLeaseCacheForTest{refreshIngressErr: errors.New("redis unavailable")}
lease := &OpenAIWSIngressLease{cache: cache, apiKeyID: 1, leaseID: "unconfirmed"}
_, lost := lease.refresh(time.Now().Add(-openAIWSIngressLeaseTTL))
require.True(t, lost)
require.Equal(t, 1, cache.refreshIngressCalls)
})
}
func TestOpenAIWSIngressLeaseReleaseWaitsForInFlightRefresh(t *testing.T) {
refreshStarted := make(chan struct{})
allowRefresh := make(chan struct{})
cache := &ingressLeaseCacheForTest{
refreshIngressFn: func(context.Context, int64, string) (bool, error) {
close(refreshStarted)
<-allowRefresh
return true, nil
},
}
ctx, cancel := context.WithCancelCause(context.Background())
lease := &OpenAIWSIngressLease{
ctx: ctx,
cancel: cancel,
cache: cache,
apiKeyID: 1,
leaseID: "in-flight-refresh",
stopCh: make(chan struct{}),
refreshDone: make(chan struct{}),
}
go func() {
defer close(lease.refreshDone)
_, _ = lease.refresh(time.Now())
}()
<-refreshStarted
released := make(chan struct{})
go func() {
lease.Release()
close(released)
}()
select {
case <-released:
t.Fatal("release returned before the in-flight refresh completed")
case <-time.After(20 * time.Millisecond):
}
require.Zero(t, cache.releaseIngressCalls)
close(allowRefresh)
select {
case <-released:
case <-time.After(time.Second):
t.Fatal("release did not complete after the refresh returned")
}
require.Equal(t, 1, cache.releaseIngressCalls)
}
func TestGenerateRequestID_UsesStablePrefixAndMonotonicCounter(t *testing.T) {
id1 := generateRequestID()
id2 := generateRequestID()
+9 -3
View File
@@ -26,6 +26,8 @@ const (
GrokMediaEndpointImagesGenerations GrokMediaEndpoint = "images_generations"
GrokMediaEndpointImagesEdits GrokMediaEndpoint = "images_edits"
GrokMediaEndpointVideosGenerations GrokMediaEndpoint = "videos_generations"
GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits"
GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions"
GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status"
)
@@ -35,7 +37,7 @@ func (e GrokMediaEndpoint) RequiresRequestBody() bool {
func (e GrokMediaEndpoint) IsGenerationRequest() bool {
switch e {
case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations:
case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions:
return true
default:
return false
@@ -274,6 +276,10 @@ func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error
return xai.BuildImagesEditsURL(baseURL)
case GrokMediaEndpointVideosGenerations:
return xai.BuildVideosGenerationsURL(baseURL)
case GrokMediaEndpointVideosEdits:
return xai.BuildVideosEditsURL(baseURL)
case GrokMediaEndpointVideosExtensions:
return xai.BuildVideosExtensionsURL(baseURL)
case GrokMediaEndpointVideoStatus:
return xai.BuildVideoURL(baseURL, requestID)
default:
@@ -302,7 +308,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
if err != nil {
return nil, err
}
targetURL, err := endpoint.upstreamURL(account.GetGrokBaseURL(), requestID)
targetURL, err := endpoint.upstreamURL(account.GetGrokMediaBaseURL(), requestID)
if err != nil {
return nil, err
}
@@ -531,7 +537,7 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi
meta.ImageSize = requestInfo.SizeTier
meta.ImageInputSize = requestInfo.Size
meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody)
case GrokMediaEndpointVideosGenerations:
case GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions:
meta.ResponseID = extractGrokMediaVideoRequestID(responseBody)
meta.VideoCount = 1
meta.VideoResolution = requestInfo.Resolution
@@ -22,6 +22,30 @@ func isModelNotFoundError(statusCode int, body []byte) bool {
return isUpstreamModelNotFoundError(statusCode, body) || statusCode == http.StatusNotFound
}
// openAICodexPlanGatedModelPhrase matches the deterministic Codex 400 returned
// when a ChatGPT OAuth account's plan cannot serve the requested model, e.g.
// {"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}
// The phrase is compared against the normalized body (lowercased, "_"/"-"
// folded to spaces), so it also matches the same message embedded in
// error.message-style payloads.
const openAICodexPlanGatedModelPhrase = "model is not supported when using codex"
// isOpenAICodexPlanGatedModelError reports whether the upstream response is the
// deterministic Codex rejection of a plan-gated model on a ChatGPT account.
// Unlike transient failures, retrying the same account cannot succeed until the
// account's plan changes, so callers should treat it like model-not-found and
// cool the (account, model) pair down instead of re-selecting the account.
func isOpenAICodexPlanGatedModelError(statusCode int, body []byte) bool {
if statusCode != http.StatusBadRequest {
return false
}
normalized := normalizeModelNotFoundBody(body)
if normalized == "" {
return false
}
return strings.Contains(normalized, openAICodexPlanGatedModelPhrase)
}
func containsModelNotFoundKeyword(normalizedBody string) bool {
if normalizedBody == "" {
return false
@@ -64,3 +64,51 @@ func TestAntigravityModelNotFoundKeepsBare404Fallback(t *testing.T) {
t.Fatal("antigravity model-not-found helper should keep bare 404 fallback")
}
}
func TestIsOpenAICodexPlanGatedModelError(t *testing.T) {
tests := []struct {
name string
statusCode int
body []byte
want bool
}{
{
name: "400 codex plan gated detail payload",
statusCode: http.StatusBadRequest,
body: []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
want: true,
},
{
name: "400 codex plan gated error message payload",
statusCode: http.StatusBadRequest,
body: []byte(`{"error":{"message":"The 'gpt-5.4' model is not supported when using Codex with a ChatGPT account."}}`),
want: true,
},
{
name: "400 unrelated invalid request does not match",
statusCode: http.StatusBadRequest,
body: []byte(`{"error":{"message":"Invalid schema for response_format 'agentic_plan'"}}`),
want: false,
},
{
name: "404 with plan gated message does not match",
statusCode: http.StatusNotFound,
body: []byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
want: false,
},
{
name: "400 empty body does not match",
statusCode: http.StatusBadRequest,
body: nil,
want: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := isOpenAICodexPlanGatedModelError(tt.statusCode, tt.body); got != tt.want {
t.Fatalf("isOpenAICodexPlanGatedModelError() = %v, want %v", got, tt.want)
}
})
}
}
@@ -271,7 +271,22 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T)
require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data)))
}
func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) {
func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.T) {
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"base_url": "https://grok.example.test/v1/",
},
}
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "")
require.NoError(t, err)
require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String())
require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
}
func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) {
t.Parallel()
account := &Account{
@@ -282,9 +297,9 @@ func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) {
},
}
_, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "")
require.Error(t, err)
require.Contains(t, err.Error(), "invalid base url")
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "")
require.NoError(t, err)
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String())
}
func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) {
@@ -596,6 +611,42 @@ func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T)
require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds)
}
func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
imageData := strings.Repeat("A", 2*1024*1024)
body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,` + imageData + `"}}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 66,
Name: "grok-oauth",
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-access-token",
"base_url": xai.DefaultCLIBaseURL,
},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"Content-Type": []string{"application/json"},
},
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
_, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String())
require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.image_url").String())
}
func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
@@ -637,6 +688,49 @@ func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) {
require.Equal(t, "xai-video-req", result.RequestID)
}
func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
tests := []struct {
name string
endpoint GrokMediaEndpoint
path string
}{
{name: "edit", endpoint: GrokMediaEndpointVideosEdits, path: "/videos/edits"},
{name: "extension", endpoint: GrokMediaEndpointVideosExtensions, path: "/videos/extensions"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok-imagine-video","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1"+tt.path, bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
account := &Account{
ID: 71, Name: "grok", Platform: PlatformGrok, Type: AccountTypeAPIKey, Concurrency: 1,
Credentials: map[string]any{"api_key": "api-key", "base_url": "https://xai.test/v1"},
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-mutation-123"}`)),
}}
svc := &OpenAIGatewayService{httpUpstream: upstream}
result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", body, "application/json")
require.NoError(t, err)
require.Equal(t, "https://xai.test/v1"+tt.path, upstream.lastReq.URL.String())
require.Equal(t, http.MethodPost, upstream.lastReq.Method)
require.JSONEq(t, string(body), string(upstream.lastBody))
require.Equal(t, "video-mutation-123", result.ResponseID)
require.Equal(t, 1, result.VideoCount)
require.Equal(t, 6, result.VideoDurationSeconds)
})
}
}
func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) {
ctx := context.Background()
groupID := int64(7)
@@ -43,9 +43,14 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions(
// custom_tool_call 项,先记下名字集合;tool_search 工具同理,回程还原为
// tool_search_call 项;namespace 子工具(如 MCP 工具)摊平转发,回程按映射还原
// 为带 namespace 字段的 function_call 项。
customTools := apicompat.CustomToolNames(responsesReq.Tools)
toolSearch := apicompat.HasToolSearchTool(responsesReq.Tools)
namespaceTools := apicompat.NamespaceToolNames(responsesReq.Tools)
effectiveTools, err := apicompat.EffectiveResponsesTools(&responsesReq)
if err != nil {
writeOpenAIResponsesFallbackError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
return nil, fmt.Errorf("resolve responses tools: %w", err)
}
customTools := apicompat.CustomToolNames(effectiveTools)
toolSearch := apicompat.HasToolSearchTool(effectiveTools)
namespaceTools := apicompat.NamespaceToolNames(effectiveTools)
chatReq, err := apicompat.ResponsesToChatCompletionsRequest(&responsesReq)
if err != nil {
@@ -39,6 +39,13 @@ type openAIWSClientConn interface {
Close() error
}
// openAIWSIdlePingCapable is intentionally separate from openAIWSClientConn.
// A pool probe happens while no goroutine is reading an idle connection, which
// is not safe for every WebSocket implementation.
type openAIWSIdlePingCapable interface {
SupportsIdlePingWithoutReader() bool
}
// openAIWSClientDialer 抽象 WS 建连器。
type openAIWSClientDialer interface {
Dial(ctx context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error)
@@ -301,6 +308,14 @@ func (c *coderOpenAIWSClientConn) Ping(ctx context.Context) error {
return c.conn.Ping(ctx)
}
// SupportsIdlePingWithoutReader reports the actual coder/websocket contract.
// Conn.Ping waits for a pong, while control frames are only consumed by Read.
// The pool deliberately has no reader on an idle connection, so using Ping as
// a health probe would deterministically time out a healthy socket.
func (*coderOpenAIWSClientConn) SupportsIdlePingWithoutReader() bool {
return false
}
func (c *coderOpenAIWSClientConn) Close() error {
if c == nil || c.conn == nil {
return nil
@@ -110,3 +110,7 @@ func TestCoderOpenAIWSClientDialer_ProxyTransportTLSHandshakeTimeout(t *testing.
require.NotNil(t, transport)
require.Equal(t, 10*time.Second, transport.TLSHandshakeTimeout)
}
func TestCoderOpenAIWSClientConn_DoesNotSupportIdlePingWithoutReader(t *testing.T) {
require.False(t, (&coderOpenAIWSClientConn{}).SupportsIdlePingWithoutReader())
}
@@ -18,6 +18,13 @@ import (
"github.com/tidwall/sjson"
)
func (s *OpenAIGatewayService) openAIWSIngressInterTurnIdleTimeout() time.Duration {
if s == nil || s.cfg == nil || s.cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds <= 0 {
return 0
}
return time.Duration(s.cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds) * time.Second
}
func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
ctx context.Context,
c *gin.Context,
@@ -360,8 +367,23 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
}
readClientMessage := func() ([]byte, error) {
msgType, payload, readErr := clientConn.Read(ctx)
readCtx := ctx
idleTimeout := s.openAIWSIngressInterTurnIdleTimeout()
cancelRead := func() {}
if idleTimeout > 0 {
readCtx, cancelRead = context.WithTimeout(ctx, idleTimeout)
}
msgType, payload, readErr := clientConn.Read(readCtx)
cancelRead()
if readErr != nil {
if idleTimeout > 0 && errors.Is(readErr, context.DeadlineExceeded) && ctx.Err() == nil {
logOpenAIWSModeInfo("ingress_ws_inter_turn_idle_timeout account_id=%d timeout_seconds=%d", account.ID, int(idleTimeout.Seconds()))
return nil, NewOpenAIWSClientCloseError(
coderws.StatusNormalClosure,
"websocket idle timeout",
readErr,
)
}
return nil, readErr
}
if msgType != coderws.MessageText && msgType != coderws.MessageBinary {
@@ -1318,7 +1340,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
unpinSessionConn(sessionConnID)
}
}
shouldPreflightPing := turn > 1 && sessionLease != nil && turnRetry == 0
shouldPreflightPing := turn > 1 && sessionLease != nil && sessionLease.SupportsIdlePingWithoutReader() && turnRetry == 0
if shouldPreflightPing && openAIWSIngressPreflightPingIdle > 0 && !lastTurnFinishedAt.IsZero() {
if time.Since(lastTurnFinishedAt) < openAIWSIngressPreflightPingIdle {
shouldPreflightPing = false
@@ -164,6 +164,111 @@ func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_KeepLeaseAcrossT
require.Len(t, captureConn.writes, 2, "应向同一上游连接发送两轮 response.create")
}
func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_IdleTimeoutReleasesStoreDisabledSession(t *testing.T) {
gin.SetMode(gin.TestMode)
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
cfg.Gateway.OpenAIWS.Enabled = true
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
captureConn := &openAIWSCaptureConn{events: [][]byte{
[]byte(`{"type":"response.completed","response":{"id":"resp_idle_timeout","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`),
}}
captureDialer := &openAIWSCaptureDialer{conn: captureConn}
pool := newOpenAIWSConnPool(cfg)
pool.setClientDialerForTest(captureDialer)
defer pool.Close()
svc := &OpenAIGatewayService{
cfg: cfg,
httpUpstream: &httpUpstreamRecorder{},
cache: &stubGatewayCache{},
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
toolCorrector: NewCodexToolCorrector(),
openaiWSPool: pool,
}
account := &Account{
ID: 116,
Name: "openai-ingress-idle-timeout",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{"api_key": "sk-test"},
Extra: map[string]any{"responses_websockets_v2_enabled": true},
}
serverErrCh := make(chan error, 1)
wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
if err != nil {
serverErrCh <- err
return
}
defer func() { _ = conn.CloseNow() }()
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
_, firstMessage, err := conn.Read(readCtx)
cancelRead()
if err != nil {
serverErrCh <- err
return
}
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
ginCtx.Request = r.Clone(r.Context())
serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
}))
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"store":false}`))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, event, err := clientConn.Read(readCtx)
cancelRead()
require.NoError(t, err)
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
select {
case proxyErr := <-serverErrCh:
var closeErr *OpenAIWSClientCloseError
require.ErrorAs(t, proxyErr, &closeErr)
require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode())
require.Equal(t, "websocket idle timeout", closeErr.Reason())
case <-time.After(4 * time.Second):
t.Fatal("timed out waiting for idle ingress session to close")
}
ap, ok := pool.getAccountPool(account.ID)
require.True(t, ok)
ap.mu.Lock()
require.Empty(t, ap.pinnedConns, "idle close must unpin a store=false session")
for _, conn := range ap.conns {
require.False(t, conn.isLeased(), "idle close must release the upstream lease")
}
ap.mu.Unlock()
}
func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_FollowupCreateCanOmitModel(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -632,3 +632,100 @@ func TestOpenAIWSHTTPBridgeKeepsContinuationFramesOnHTTPWithoutPreviousResponseI
require.Equal(t, 0, captureDialer.DialCount())
require.Empty(t, captureConn.writes)
}
func TestOpenAIWSHTTPBridge_IdleTimeoutClosesClientSession(t *testing.T) {
gin.SetMode(gin.TestMode)
sseBody := strings.Join([]string{
`data: {"type":"response.completed","response":{"id":"resp_bridge_idle","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
"",
}, "\n")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(sseBody)),
}}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
cfg.Gateway.OpenAIWS.Enabled = true
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
cfg.Gateway.OpenAIWS.HTTPBridgeEnabled = true
cfg.Gateway.OpenAIWS.HTTPBridgeThresholdBytes = 1
cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 1
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8
svc := &OpenAIGatewayService{
cfg: cfg,
httpUpstream: upstream,
cache: &stubGatewayCache{},
openaiWSResolver: NewOpenAIWSProtocolResolver(cfg),
toolCorrector: NewCodexToolCorrector(),
}
account := &Account{
ID: 20,
Name: "api-key-bridge-idle-timeout",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "sk-upstream"},
Extra: map[string]any{"responses_websockets_v2_enabled": true},
Concurrency: 1,
Status: StatusActive,
Schedulable: true,
}
errCh := make(chan error, 1)
wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
if err != nil {
errCh <- err
return
}
defer func() { _ = conn.CloseNow() }()
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
_, firstMessage, err := conn.Read(readCtx)
cancelRead()
if err != nil {
errCh <- err
return
}
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
ginCtx.Request = r.Clone(r.Context())
errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil)
}))
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":"hello"}`))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, event, err := clientConn.Read(readCtx)
cancelRead()
require.NoError(t, err)
require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
select {
case proxyErr := <-errCh:
var closeErr *OpenAIWSClientCloseError
require.ErrorAs(t, proxyErr, &closeErr)
require.Equal(t, coderws.StatusNormalClosure, closeErr.StatusCode())
require.Equal(t, "websocket idle timeout", closeErr.Reason())
case <-time.After(4 * time.Second):
t.Fatal("timed out waiting for idle HTTP bridge session to close")
}
require.Len(t, upstream.bodies, 1, "an idle client must not leave a continuation request running")
}
+21 -3
View File
@@ -203,6 +203,14 @@ func (l *openAIWSConnLease) PingWithTimeout(timeout time.Duration) error {
return conn.pingWithTimeout(timeout)
}
func (l *openAIWSConnLease) SupportsIdlePingWithoutReader() bool {
conn, err := l.activeConn()
if err != nil {
return false
}
return conn.supportsIdlePingWithoutReader()
}
func (l *openAIWSConnLease) MarkBroken() {
if l == nil || l.pool == nil || l.conn == nil || l.released.Load() {
return
@@ -437,6 +445,16 @@ func (c *openAIWSConn) pingWithTimeout(timeout time.Duration) error {
return nil
}
func (c *openAIWSConn) supportsIdlePingWithoutReader() bool {
if c == nil || c.ws == nil {
return false
}
capable, ok := c.ws.(openAIWSIdlePingCapable)
// Test and alternate implementations keep the historical probe behavior
// unless they explicitly declare it unsafe.
return !ok || capable.SupportsIdlePingWithoutReader()
}
func (c *openAIWSConn) touch() {
if c == nil {
return
@@ -707,7 +725,7 @@ func (p *openAIWSConnPool) runBackgroundPingSweep() {
g.SetLimit(10)
for _, item := range candidates {
item := item
if item.conn == nil || item.conn.isLeased() || item.conn.waiters.Load() > 0 {
if item.conn == nil || item.conn.isLeased() || item.conn.waiters.Load() > 0 || !item.conn.supportsIdlePingWithoutReader() {
continue
}
g.Go(func() error {
@@ -1613,7 +1631,7 @@ func (p *openAIWSConnPool) nextConnID(accountID int64) string {
}
func (p *openAIWSConnPool) shouldHealthCheckConn(conn *openAIWSConn) bool {
if conn == nil {
if conn == nil || !conn.supportsIdlePingWithoutReader() {
return false
}
return conn.idleDuration(time.Now()) >= openAIWSConnHealthCheckIdle
@@ -1669,7 +1687,7 @@ func (p *openAIWSConnPool) effectiveMaxConnsByAccount(account *Account) int {
if account.Concurrency <= 0 {
return 0
}
return account.Concurrency
return min(account.Concurrency, hardCap)
}
if account == nil || !p.dynamicMaxConnsEnabled() {
return hardCap
@@ -739,7 +739,7 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T)
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为")
}
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2UsesAccountConcurrency(t *testing.T) {
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
@@ -750,7 +750,7 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2UsesAccountConc
pool := newOpenAIWSConnPool(cfg)
high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
require.Equal(t, 20, pool.effectiveMaxConnsByAccount(high), "v2 路径应直接使用账号并发数作为池上限")
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
@@ -1319,6 +1319,9 @@ func TestOpenAIWSConnPool_UtilityBranches(t *testing.T) {
conn := newOpenAIWSConn("health", 1, &openAIWSFakeConn{}, nil)
conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
require.True(t, pool.shouldHealthCheckConn(conn))
unsafeConn := newOpenAIWSConn("unsafe_health", 1, &openAIWSIdlePingUnsupportedConn{}, nil)
unsafeConn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
require.False(t, pool.shouldHealthCheckConn(unsafeConn))
}
func TestOpenAIWSConn_LeaseAndTimeHelpers_NilAndClosedBranches(t *testing.T) {
@@ -1610,6 +1613,14 @@ type openAIWSPingBlockingConn struct {
release <-chan struct{}
}
type openAIWSIdlePingUnsupportedConn struct {
openAIWSFakeConn
}
func (c *openAIWSIdlePingUnsupportedConn) SupportsIdlePingWithoutReader() bool {
return false
}
func (c *openAIWSPingBlockingConn) WriteJSON(context.Context, any) error {
return nil
}
+22 -5
View File
@@ -2006,9 +2006,19 @@ func parseOpenAIImageTryAgainCooldown(body []byte) time.Duration {
const upstreamModelNotFoundCooldown = 30 * time.Minute
const upstreamModelNotFoundReason = "upstream_404_model_not_found"
const upstreamCodexPlanGatedModelCooldown = 30 * time.Minute
const upstreamCodexPlanGatedModelReason = "upstream_400_codex_plan_gated_model"
const tempUnschedBodyMaxBytes = 64 << 10
const tempUnschedMessageMaxBytes = 2048
// HandleUpstreamModelNotFound marks the requested model as temporarily
// unavailable on the account when the upstream deterministically reports it
// cannot serve that model: a 404 model-not-found, or the Codex 400 rejecting a
// plan-gated model on a ChatGPT OAuth account. Returning true tells the caller
// to fail the current attempt over to another account; the scheduler skips the
// (account, model) pair via IsSchedulableForModelWithContext until the
// cooldown expires, instead of re-selecting an account that can never serve
// the model.
func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, account *Account, requestedModel string, statusCode int, responseBody []byte) bool {
if s == nil || account == nil || s.accountRepo == nil {
return false
@@ -2016,19 +2026,26 @@ func (s *RateLimitService) HandleUpstreamModelNotFound(ctx context.Context, acco
if !account.ShouldHandleErrorCode(statusCode) {
return false
}
if !isUpstreamModelNotFoundError(statusCode, responseBody) {
var cooldown time.Duration
var reason string
switch {
case isUpstreamModelNotFoundError(statusCode, responseBody):
cooldown, reason = upstreamModelNotFoundCooldown, upstreamModelNotFoundReason
case isOpenAIOAuthAccount(account) && isOpenAICodexPlanGatedModelError(statusCode, responseBody):
cooldown, reason = upstreamCodexPlanGatedModelCooldown, upstreamCodexPlanGatedModelReason
default:
return false
}
modelKey := modelRateLimitKeyForUpstreamModelNotFound(ctx, account, requestedModel)
if modelKey == "" {
return false
}
resetAt := time.Now().Add(upstreamModelNotFoundCooldown)
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, upstreamModelNotFoundReason); err != nil {
slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "error", err)
resetAt := time.Now().Add(cooldown)
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, resetAt, reason); err != nil {
slog.Warn("upstream_model_not_found_set_model_rate_limit_failed", "account_id", account.ID, "model", modelKey, "reason", reason, "error", err)
return true
}
slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reset_at", resetAt)
slog.Info("upstream_model_not_found_model_rate_limited", "account_id", account.ID, "model", modelKey, "reason", reason, "reset_at", resetAt)
return true
}
@@ -125,3 +125,77 @@ func openAIModelNotFoundTempAccount() *Account {
},
}
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelUsesModelRateLimit(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
call := repo.modelRateLimitCalls[0]
require.Equal(t, account.ID, call.accountID)
require.Equal(t, "gpt-5.6-sol", call.scope)
require.Equal(t, upstreamCodexPlanGatedModelReason, call.reason)
require.WithinDuration(t, time.Now().Add(upstreamCodexPlanGatedModelCooldown), call.resetAt, 5*time.Second)
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelRespectsModelMapping(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
account.Credentials["model_mapping"] = map[string]any{"gpt-5.6-sol": "gpt-5.6-sol-upstream"}
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol-upstream' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.6-sol-upstream", repo.modelRateLimitCalls[0].scope)
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelIgnoresAPIKeyAccount(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
account.Type = AccountTypeAPIKey
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.False(t, handled)
require.Empty(t, repo.modelRateLimitCalls)
}
func openAICodexPlanGatedOAuthAccount() *Account {
return &Account{
ID: 202,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{},
}
}
@@ -131,6 +131,8 @@ func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, acc
switch {
case account.Platform == PlatformAntigravity:
return s.buildAntigravityAPIKeyModelsRequest(ctx, account)
case account.IsGrok():
return s.buildGrokUpstreamModelsRequest(ctx, account)
case account.IsOpenAI():
return s.buildOpenAIUpstreamModelsRequest(ctx, account)
case account.IsGemini():
@@ -144,6 +146,36 @@ func (s *AccountTestService) buildUpstreamModelsRequest(ctx context.Context, acc
}
}
func (s *AccountTestService) buildGrokUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
if account.Type != AccountTypeAPIKey {
return nil, newUpstreamModelSyncUnsupportedError(
fmt.Sprintf("Unsupported Grok account type for upstream model sync: %s", account.Type), nil,
)
}
apiKey := strings.TrimSpace(account.GetCredential("api_key"))
if apiKey == "" {
return nil, newUpstreamModelSyncConfigError("No Grok API key is available", nil)
}
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
if baseURL == "" {
baseURL = "https://api.x.ai"
}
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
if err != nil {
return nil, newUpstreamModelSyncConfigError("Invalid Grok base URL", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, buildOpenAIModelsURL(normalizedBaseURL), nil)
if err != nil {
return nil, newUpstreamModelSyncConfigError("Invalid Grok model list URL", err)
}
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
account.ApplyHeaderOverrides(req.Header)
return req, nil
}
func (s *AccountTestService) buildAnthropicUpstreamModelsRequest(ctx context.Context, account *Account) (*http.Request, error) {
if account.IsBedrock() || account.Type == AccountTypeServiceAccount {
return nil, newUpstreamModelSyncUnsupportedError(
@@ -177,6 +177,18 @@ func TestBuildUpstreamModelsRequestsForAPIKeyAccounts(t *testing.T) {
require.Equal(t, "https://openai.example.com/v1/models", openAIReq.URL.String())
require.Equal(t, "Bearer openai-key", openAIReq.Header.Get("Authorization"))
grokReq, err := svc.buildUpstreamModelsRequest(ctx, &Account{
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"api_key": "xai-key",
"base_url": "https://xai.example.com/v1",
},
})
require.NoError(t, err)
require.Equal(t, "https://xai.example.com/v1/models", grokReq.URL.String())
require.Equal(t, "Bearer xai-key", grokReq.Header.Get("Authorization"))
geminiReq, err := svc.buildGeminiUpstreamModelsRequest(ctx, &Account{
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
@@ -202,6 +214,22 @@ func TestBuildUpstreamModelsRequestsForAPIKeyAccounts(t *testing.T) {
require.Equal(t, "antigravity-key", antigravityReq.Header.Get("x-api-key"))
}
func TestBuildUpstreamModelsRequestRejectsGrokOAuth(t *testing.T) {
t.Parallel()
svc := &AccountTestService{cfg: upstreamModelSyncTestConfig()}
_, err := svc.buildUpstreamModelsRequest(context.Background(), &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
})
require.Error(t, err)
var syncErr *UpstreamModelSyncError
require.True(t, errors.As(err, &syncErr))
require.Equal(t, UpstreamModelSyncErrorUnsupported, syncErr.Kind)
require.Contains(t, syncErr.SafeMessage(), "Unsupported Grok account type")
}
func TestBuildAntigravityAPIKeyModelsRequestRejectsOfficialCloudCodeBase(t *testing.T) {
t.Parallel()
@@ -265,6 +293,34 @@ func TestFetchUpstreamSupportedModelsParsesOpenAIResponse(t *testing.T) {
require.Equal(t, "Bearer openai-key", upstream.lastReq.Header.Get("Authorization"))
}
func TestFetchUpstreamSupportedModelsParsesGrokAPIKeyResponse(t *testing.T) {
t.Parallel()
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"grok-4.5"},{"id":"grok-4.5"},{"id":"grok-imagine"}]}`)),
}}
svc := &AccountTestService{
httpUpstream: upstream,
cfg: upstreamModelSyncTestConfig(),
}
models, err := svc.FetchUpstreamSupportedModels(context.Background(), &Account{
ID: 9,
Platform: PlatformGrok,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"api_key": "xai-key",
"base_url": "https://xai.example.com/v1",
},
})
require.NoError(t, err)
require.Equal(t, []string{"grok-4.5", "grok-imagine"}, models)
require.Equal(t, "https://xai.example.com/v1/models", upstream.lastReq.URL.String())
require.Equal(t, "Bearer xai-key", upstream.lastReq.Header.Get("Authorization"))
}
func TestFetchUpstreamSupportedModelsDoesNotExposeUpstreamBody(t *testing.T) {
t.Parallel()
+8 -2
View File
@@ -109,7 +109,8 @@ func (s *FrontendServer) Middleware() gin.HandlerFunc {
return
}
// Serve static files normally
// Serve static files normally (hashed assets get long-lived cache headers)
applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath)
s.fileServer.ServeHTTP(c.Writer, c.Request)
c.Abort()
}
@@ -135,6 +136,7 @@ func (s *FrontendServer) tryServeOverride(c *gin.Context, cleanPath string) bool
if err != nil || info.IsDir() {
return false
}
applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath)
c.File(filePath)
c.Abort()
return true
@@ -273,6 +275,7 @@ func ServeEmbeddedFrontend() gin.HandlerFunc {
if tryServeOverrideFile(c, overrideDir, cleanPath) {
return
}
applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath)
fileServer.ServeHTTP(c.Writer, c.Request)
c.Abort()
return
@@ -292,6 +295,7 @@ func tryServeOverrideFile(c *gin.Context, overrideDir, cleanPath string) bool {
if err != nil || info.IsDir() {
return false
}
applyStaticAssetCacheHeaders(c.Writer.Header(), cleanPath)
c.File(filePath)
c.Abort()
return true
@@ -308,7 +312,9 @@ func shouldBypassEmbeddedFrontend(path string) bool {
trimmed == "/health" ||
trimmed == "/responses" ||
strings.HasPrefix(trimmed, "/responses/") ||
strings.HasPrefix(trimmed, "/images/")
trimmed == "/alpha/search" ||
strings.HasPrefix(trimmed, "/images/") ||
strings.HasPrefix(trimmed, "/videos/")
}
func serveIndexHTML(c *gin.Context, fsys fs.FS) {
+37
View File
@@ -507,6 +507,32 @@ func TestFrontendServer_Middleware(t *testing.T) {
assert.JSONEq(t, `{"ok":true}`, w.Body.String())
})
t.Run("skips_alpha_search_post_route", func(t *testing.T) {
provider := &mockSettingsProvider{
settings: map[string]string{"test": "value"},
}
server, err := NewFrontendServer(provider)
require.NoError(t, err)
router := gin.New()
router.Use(server.Middleware())
nextCalled := false
router.POST("/alpha/search", func(c *gin.Context) {
nextCalled = true
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/alpha/search", strings.NewReader(`{"model":"gpt-5.6-sol"}`))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(w, req)
assert.True(t, nextCalled, "next handler should be called for alpha search API route")
assert.Equal(t, http.StatusOK, w.Code)
assert.JSONEq(t, `{"ok":true}`, w.Body.String())
})
t.Run("serves_index_for_spa_routes", func(t *testing.T) {
provider := &mockSettingsProvider{
settings: map[string]string{"test": "value"},
@@ -562,6 +588,17 @@ func TestFrontendServer_Middleware(t *testing.T) {
})
}
func TestEmbeddedFrontendBypassesBareVideoAPIRoutes(t *testing.T) {
for _, path := range []string{
"/videos/generations",
"/videos/edits",
"/videos/extensions",
"/videos/request-123",
} {
require.True(t, shouldBypassEmbeddedFrontend(path), "path=%s", path)
}
}
func TestNewFrontendServer(t *testing.T) {
t.Run("creates_server_successfully", func(t *testing.T) {
provider := &mockSettingsProvider{
+31
View File
@@ -0,0 +1,31 @@
//go:build embed || unit
package web
import (
"net/http"
"strings"
)
// staticAssetsCacheControl matches deploy/Caddyfile for hashed frontend assets.
// Vite emits content-hashed filenames under assets/, so long-lived immutable
// caching is safe without relying on a reverse proxy.
const staticAssetsCacheControl = "public, max-age=31536000, immutable"
// isLongCacheStaticPath reports whether a cleaned URL path (no leading slash)
// should receive long-lived Cache-Control headers. Aligned with deploy/Caddyfile.
func isLongCacheStaticPath(cleanPath string) bool {
cleanPath = strings.TrimPrefix(cleanPath, "/")
return strings.HasPrefix(cleanPath, "assets/") ||
cleanPath == "logo.png" ||
cleanPath == "favicon.ico"
}
// applyStaticAssetCacheHeaders sets Cache-Control for long-cacheable static paths.
// index.html / SPA routes must keep no-cache and are not handled here.
func applyStaticAssetCacheHeaders(header http.Header, cleanPath string) {
if header == nil || !isLongCacheStaticPath(cleanPath) {
return
}
header.Set("Cache-Control", staticAssetsCacheControl)
}
+71
View File
@@ -0,0 +1,71 @@
//go:build unit
package web
import (
"net/http"
"testing"
"github.com/stretchr/testify/assert"
)
func TestIsLongCacheStaticPath(t *testing.T) {
t.Parallel()
cases := []struct {
name string
path string
want bool
}{
{name: "hashed_js", path: "assets/index-abc123.js", want: true},
{name: "hashed_css", path: "assets/app-def456.css", want: true},
{name: "nested_asset", path: "assets/vendor/chunk.js", want: true},
{name: "leading_slash_asset", path: "/assets/index.js", want: true},
{name: "logo", path: "logo.png", want: true},
{name: "favicon", path: "favicon.ico", want: true},
{name: "index_html", path: "index.html", want: false},
{name: "spa_route", path: "dashboard", want: false},
{name: "assets_prefix_only", path: "assets", want: false},
{name: "similar_name", path: "assets-backup/x.js", want: false},
{name: "empty", path: "", want: false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
assert.Equal(t, tc.want, isLongCacheStaticPath(tc.path))
})
}
}
func TestApplyStaticAssetCacheHeaders(t *testing.T) {
t.Parallel()
t.Run("sets_immutable_cache_for_assets", func(t *testing.T) {
t.Parallel()
header := make(http.Header)
applyStaticAssetCacheHeaders(header, "assets/index-abc.js")
assert.Equal(t, staticAssetsCacheControl, header.Get("Cache-Control"))
})
t.Run("sets_immutable_cache_for_logo", func(t *testing.T) {
t.Parallel()
header := make(http.Header)
applyStaticAssetCacheHeaders(header, "logo.png")
assert.Equal(t, staticAssetsCacheControl, header.Get("Cache-Control"))
})
t.Run("skips_index_html", func(t *testing.T) {
t.Parallel()
header := make(http.Header)
applyStaticAssetCacheHeaders(header, "index.html")
assert.Empty(t, header.Get("Cache-Control"))
})
t.Run("nil_header_is_noop", func(t *testing.T) {
t.Parallel()
assert.NotPanics(t, func() {
applyStaticAssetCacheHeaders(nil, "assets/x.js")
})
})
}
@@ -0,0 +1,5 @@
-- Support the per-key latest non-empty source IP lookup without scanning full key history.
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip
ON usage_logs (api_key_id, created_at DESC, id DESC)
INCLUDE (ip_address)
WHERE ip_address IS NOT NULL AND ip_address <> '';
@@ -0,0 +1,19 @@
package migrations
import (
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func TestLatestAPIKeyIPIndexMigration(t *testing.T) {
content, err := FS.ReadFile("174_add_usage_logs_api_key_latest_ip_index_notx.sql")
require.NoError(t, err)
sql := strings.Join(strings.Fields(string(content)), " ")
require.Contains(t, sql, "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip")
require.Contains(t, sql, "ON usage_logs (api_key_id, created_at DESC, id DESC)")
require.Contains(t, sql, "INCLUDE (ip_address)")
require.Contains(t, sql, "WHERE ip_address IS NOT NULL AND ip_address <> ''")
}