mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-26 16:03:35 +00:00
Merge remote-tracking branch 'upstream/main' into fix/api-double-billing
# Conflicts: # frontend/src/components/account/EditAccountModal.vue
This commit is contained in:
@@ -1 +1 @@
|
||||
0.1.152
|
||||
0.1.153
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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 },
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 <> ''")
|
||||
}
|
||||
Reference in New Issue
Block a user