diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 49f80d19c1..b75515e962 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -80,9 +80,13 @@ func (h *GatewayHandler) Responses(c *gin.Context) { setOpsRequestContext(c, reqModel, reqStream) setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false))) + requestCtx := c.Request.Context() + if service.IsImageGenerationIntent("/v1/responses", reqModel, body) { + requestCtx = service.WithOpenAIImageGenerationIntent(requestCtx) + } // 解析渠道级模型映射 - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel) + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel) // Claude Code only restriction: // /v1/responses is never a Claude Code endpoint. @@ -145,7 +149,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { } // 2. Re-check billing - if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { + if err := h.billingCacheService.CheckBillingEligibility(requestCtx, apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(requestCtx, apiKey)); err != nil { reqLog.Info("gateway.responses.billing_check_failed", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) if retryAfter > 0 { @@ -172,14 +176,14 @@ func (h *GatewayHandler) Responses(c *gin.Context) { fs := NewFailoverState(h.maxAccountSwitches, false) for { - selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0)) + selection, err := h.gatewayService.SelectAccountWithLoadAwareness(requestCtx, apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0)) if err != nil { if len(fs.FailedAccountIDs) == 0 { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) h.responsesErrorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error()) return } - action := fs.HandleSelectionExhausted(c.Request.Context()) + action := fs.HandleSelectionExhausted(requestCtx) switch action { case FailoverContinue: continue @@ -227,7 +231,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { if channelMapping.Mapped { forwardBody = h.gatewayService.ReplaceModelInBody(body, channelMapping.MappedModel) } - result, err := h.gatewayService.ForwardAsResponses(c.Request.Context(), c, account, forwardBody, parsedReq) + result, err := h.gatewayService.ForwardAsResponses(requestCtx, c, account, forwardBody, parsedReq) if accountReleaseFunc != nil { accountReleaseFunc() @@ -241,7 +245,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { h.handleResponsesFailoverExhausted(c, failoverErr, true) return } - action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr) + action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, failoverErr) switch action { case FailoverContinue: continue diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 1e3b530601..580ecab992 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -135,6 +135,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body) + requestCtx := service.WithOpenAIImageGenerationIntent(c.Request.Context()) maxAccountSwitches := h.maxAccountSwitches switchCount := 0 @@ -145,7 +146,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { for { reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForImages( - c.Request.Context(), + requestCtx, apiKey.GroupID, sessionHash, requestModel, @@ -202,7 +203,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { accountReleaseFunc() } }() - return h.gatewayService.ForwardImages(c.Request.Context(), c, account, body, parsed, channelMapping.MappedModel) + return h.gatewayService.ForwardImages(requestCtx, c, account, body, parsed, channelMapping.MappedModel) }() forwardDurationMs := time.Since(forwardStart).Milliseconds() upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey) @@ -248,7 +249,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { zap.Int("retry_count", sameAccountRetryCount[account.ID]), ) select { - case <-c.Request.Context().Done(): + case <-requestCtx.Done(): return case <-time.After(sameAccountRetryDelay): } diff --git a/backend/internal/pkg/ctxkey/ctxkey.go b/backend/internal/pkg/ctxkey/ctxkey.go index 25782c5517..dacd1bd1cf 100644 --- a/backend/internal/pkg/ctxkey/ctxkey.go +++ b/backend/internal/pkg/ctxkey/ctxkey.go @@ -34,6 +34,10 @@ const ( // ThinkingEnabled 标识当前请求是否开启 thinking（用于 Antigravity 最终模型名推导与模型维度限流） ThinkingEnabled Key = "ctx_thinking_enabled" + + // OpenAIImageGenerationIntent 标识 OpenAI 请求会触发生图能力（用于图片能力维度限流） + OpenAIImageGenerationIntent Key = "ctx_openai_image_generation_intent" + // Group 认证后的分组信息，由 API Key 认证中间件设置 Group Key = "ctx_group" diff --git a/backend/internal/service/model_rate_limit.go b/backend/internal/service/model_rate_limit.go index 420f14467c..1568a00f78 100644 --- a/backend/internal/service/model_rate_limit.go +++ b/backend/internal/service/model_rate_limit.go @@ -4,11 +4,14 @@ import ( "context" "strings" "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" ) const ( modelRateLimitsKey = "model_rate_limits" antigravityGeminiModelRateLimitKey = "antigravity:gemini" + openAIImageGenerationRateLimitKey = "openai:image_generation" ) // isRateLimitActiveForKey 检查指定 key 的限流是否生效 @@ -31,22 +34,12 @@ func (a *Account) getRateLimitRemainingForKey(key string) time.Duration { } func (a *Account) isModelRateLimitedWithContext(ctx context.Context, requestedModel string) bool { - if a == nil { - return false - } - - modelKey := a.GetMappedModel(requestedModel) - if a.Platform == PlatformAntigravity { - modelKey = resolveFinalAntigravityModelKey(ctx, a, requestedModel) - if isAntigravityGeminiModel(modelKey) && a.isRateLimitActiveForKey(antigravityGeminiModelRateLimitKey) { + for _, key := range a.modelRateLimitKeysForRequest(ctx, requestedModel) { + if a.isRateLimitActiveForKey(key) { return true } } - modelKey = strings.TrimSpace(modelKey) - if modelKey == "" { - return false - } - return a.isRateLimitActiveForKey(modelKey) + return false } // GetModelRateLimitRemainingTime 获取模型限流剩余时间 @@ -56,8 +49,18 @@ func (a *Account) GetModelRateLimitRemainingTime(requestedModel string) time.Dur } func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration { + remaining := time.Duration(0) + for _, key := range a.modelRateLimitKeysForRequest(ctx, requestedModel) { + if keyRemaining := a.getRateLimitRemainingForKey(key); keyRemaining > remaining { + remaining = keyRemaining + } + } + return remaining +} + +func (a *Account) modelRateLimitKeysForRequest(ctx context.Context, requestedModel string) []string { if a == nil { - return 0 + return nil } modelKey := a.GetMappedModel(requestedModel) @@ -66,15 +69,43 @@ func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context, } modelKey = strings.TrimSpace(modelKey) if modelKey == "" { - return 0 + return nil } - remaining := a.getRateLimitRemainingForKey(modelKey) - if a.Platform == PlatformAntigravity && isAntigravityGeminiModel(modelKey) { - if familyRemaining := a.getRateLimitRemainingForKey(antigravityGeminiModelRateLimitKey); familyRemaining > remaining { - return familyRemaining + + keys := []string{modelKey} + switch a.Platform { + case PlatformAntigravity: + if isAntigravityGeminiModel(modelKey) && modelKey != antigravityGeminiModelRateLimitKey { + keys = append(keys, antigravityGeminiModelRateLimitKey) + } + case PlatformOpenAI: + if openAIImageGenerationRateLimitApplies(ctx, requestedModel, modelKey) && modelKey != openAIImageGenerationRateLimitKey { + keys = append(keys, openAIImageGenerationRateLimitKey) } } - return remaining + return keys +} + +func openAIImageGenerationRateLimitApplies(ctx context.Context, requestedModel, modelKey string) bool { + if isOpenAIImageGenerationModel(requestedModel) || isOpenAIImageGenerationModel(modelKey) { + return true + } + return OpenAIImageGenerationIntentFromContext(ctx) +} + +func WithOpenAIImageGenerationIntent(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, ctxkey.OpenAIImageGenerationIntent, true) +} + +func OpenAIImageGenerationIntentFromContext(ctx context.Context) bool { + if ctx == nil { + return false + } + enabled, ok := ctx.Value(ctxkey.OpenAIImageGenerationIntent).(bool) + return ok && enabled } func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requestedModel string) string { diff --git a/backend/internal/service/model_rate_limit_test.go b/backend/internal/service/model_rate_limit_test.go index 3cce645981..c62e4c4b9b 100644 --- a/backend/internal/service/model_rate_limit_test.go +++ b/backend/internal/service/model_rate_limit_test.go @@ -6,6 +6,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" ) func TestIsModelRateLimited(t *testing.T) { @@ -195,6 +196,36 @@ func TestIsModelRateLimited(t *testing.T) { requestedModel: "claude-3-5-sonnet-20241022", expected: false, }, + { + name: "openai image generation family key blocks image model", + account: &Account{ + Platform: PlatformOpenAI, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + openAIImageGenerationRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "gpt-image-2", + expected: true, + }, + { + name: "openai image generation family key does not block text model", + account: &Account{ + Platform: PlatformOpenAI, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + openAIImageGenerationRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + }, + requestedModel: "gpt-5.4", + expected: false, + }, } for _, tt := range tests { @@ -207,6 +238,23 @@ func TestIsModelRateLimited(t *testing.T) { } } +func TestIsModelRateLimited_OpenAIImageGenerationIntentBlocksTextModelImageTool(t *testing.T) { + future := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + account := &Account{ + Platform: PlatformOpenAI, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + openAIImageGenerationRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + } + + require.False(t, account.isModelRateLimitedWithContext(context.Background(), "gpt-5.4")) + require.True(t, account.isModelRateLimitedWithContext(WithOpenAIImageGenerationIntent(context.Background()), "gpt-5.4")) +} + func TestIsModelRateLimited_Antigravity_ThinkingAffectsModelKey(t *testing.T) { now := time.Now() future := now.Add(10 * time.Minute).Format(time.RFC3339) diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index a4e905e893..bdd97f479a 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -35,6 +35,13 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont stateCtx, cancel := openAIAccountStateContext(ctx) defer cancel() + if isOpenAIImageRateLimitError(statusCode, responseBody) { + if s != nil && s.rateLimitService != nil { + _ = s.rateLimitService.HandleOpenAIImageRateLimit(stateCtx, account, statusCode, headers, responseBody) + } + return false + } + if statusCode == http.StatusTooManyRequests { s.markOpenAIOAuth429RateLimited(stateCtx, account, headers, responseBody) } diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index e57777d270..440f378055 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -424,6 +424,57 @@ func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) } } +func TestOpenAISelectAccountWithLoadAwareness_ImageRateLimitSkipsOnlyImageRequests(t *testing.T) { + future := time.Now().Add(10 * time.Minute).Format(time.RFC3339) + groupID := int64(1) + + imageLimited := Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + openAIImageGenerationRateLimitKey: map[string]any{ + "rate_limit_reset_at": future, + }, + }, + }, + } + available := Account{ + ID: 2, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + } + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{imageLimited, available}}, + concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), + } + + imageSelection, err := svc.SelectAccountWithLoadAwareness(WithOpenAIImageGenerationIntent(context.Background()), &groupID, "", "gpt-5.4", nil) + require.NoError(t, err) + require.NotNil(t, imageSelection) + require.Equal(t, available.ID, imageSelection.Account.ID) + if imageSelection.ReleaseFunc != nil { + imageSelection.ReleaseFunc() + } + + textSelection, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gpt-5.4", nil) + require.NoError(t, err) + require.NotNil(t, textSelection) + require.Equal(t, imageLimited.ID, textSelection.Account.ID) + if textSelection.ReleaseFunc != nil { + textSelection.ReleaseFunc() + } +} + func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurrencyService(t *testing.T) { now := time.Now() resetAt := now.Add(10 * time.Minute) diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index ecbd86d1ae..e2704f3c20 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -7,6 +7,7 @@ import ( "fmt" "log/slog" "net/http" + "regexp" "strconv" "strings" "sync" @@ -66,6 +67,13 @@ const ( maxRateLimit429CooldownSeconds = 7200 ) +const ( + openAIImageRateLimitDefaultCooldown = time.Minute + openAIImageRateLimitReason = "openai_image_rate_limited" +) + +var openAIImageTryAgainPattern = regexp.MustCompile(`(?i)try again in\s+([0-9]+(?:\.[0-9]+)?)\s*(ms|s|sec|secs|second|seconds|m|min|mins|minute|minutes)`) + const ( openAI403CooldownMinutesDefault = 10 openAI403DisableThreshold = 3 @@ -1618,6 +1626,109 @@ func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account return s.tryTempUnschedulable(ctx, account, statusCode, responseBody) } +func (s *RateLimitService) HandleOpenAIImageRateLimit(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) bool { + if s == nil || account == nil || s.accountRepo == nil { + return false + } + if account.Platform != PlatformOpenAI { + return false + } + if !account.ShouldHandleErrorCode(statusCode) { + slog.Info("openai_image_rate_limit_skipped_by_error_code_policy", "account_id", account.ID, "status_code", statusCode) + return false + } + if !isOpenAIImageRateLimitError(statusCode, responseBody) { + return false + } + + resetAt := openAIImageRateLimitResetAt(headers, responseBody) + if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, openAIImageGenerationRateLimitKey, resetAt, openAIImageRateLimitReason); err != nil { + slog.Warn("openai_image_rate_limit_set_model_rate_limit_failed", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "error", err) + return true + } + slog.Info("openai_image_rate_limited", "account_id", account.ID, "scope", openAIImageGenerationRateLimitKey, "reset_at", resetAt, "reset_in", time.Until(resetAt).Truncate(time.Second)) + return true +} + +func isOpenAIImageRateLimitError(statusCode int, body []byte) bool { + if statusCode != http.StatusTooManyRequests || len(body) == 0 { + return false + } + lower := strings.ToLower(string(body)) + for _, marker := range []string{ + "for limit gpt-image", + "input-images per min", + "gpt-image-2-codex", + "gpt-image", + } { + if strings.Contains(lower, marker) { + return true + } + } + return false +} + +func openAIImageRateLimitResetAt(headers http.Header, body []byte) time.Time { + now := time.Now() + if resetAt := parseRetryAfterResetTime(headers, now); resetAt != nil && resetAt.After(now) { + return *resetAt + } + if resetAt := calculateOpenAI429ResetTime(headers); resetAt != nil && resetAt.After(now) { + return *resetAt + } + if resetUnix := parseOpenAIRateLimitResetTime(body); resetUnix != nil { + if resetAt := time.Unix(*resetUnix, 0); resetAt.After(now) { + return resetAt + } + } + if cooldown := parseOpenAIImageTryAgainCooldown(body); cooldown > 0 { + return now.Add(cooldown) + } + return now.Add(openAIImageRateLimitDefaultCooldown) +} + +func parseRetryAfterResetTime(headers http.Header, now time.Time) *time.Time { + if headers == nil { + return nil + } + raw := strings.TrimSpace(headers.Get("Retry-After")) + if raw == "" { + return nil + } + if seconds, err := strconv.ParseFloat(raw, 64); err == nil { + resetAt := now.Add(time.Duration(seconds * float64(time.Second))) + return &resetAt + } + if parsed, err := http.ParseTime(raw); err == nil { + return &parsed + } + return nil +} + +func parseOpenAIImageTryAgainCooldown(body []byte) time.Duration { + if len(body) == 0 { + return 0 + } + match := openAIImageTryAgainPattern.FindSubmatch(body) + if len(match) != 3 { + return 0 + } + value, err := strconv.ParseFloat(string(match[1]), 64) + if err != nil || value <= 0 { + return 0 + } + switch strings.ToLower(string(match[2])) { + case "ms": + return time.Duration(value * float64(time.Millisecond)) + case "s", "sec", "secs", "second", "seconds": + return time.Duration(value * float64(time.Second)) + case "m", "min", "mins", "minute", "minutes": + return time.Duration(value * float64(time.Minute)) + default: + return 0 + } +} + const upstreamModelNotFoundCooldown = 30 * time.Minute const upstreamModelNotFoundReason = "upstream_404_model_not_found" const tempUnschedBodyMaxBytes = 64 << 10 diff --git a/backend/internal/service/ratelimit_service_openai_image_test.go b/backend/internal/service/ratelimit_service_openai_image_test.go new file mode 100644 index 0000000000..76d3bb764b --- /dev/null +++ b/backend/internal/service/ratelimit_service_openai_image_test.go @@ -0,0 +1,121 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestIsOpenAIImageRateLimitError(t *testing.T) { + imageBody := []byte(`{"error":{"message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 467ms."}}`) + textBody := []byte(`{"error":{"message":"Rate limit reached for gpt-5.4 in organization org on tokens per min: Limit 30000, Used 30000. Please try again in 1s."}}`) + + require.True(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, imageBody)) + require.False(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, textBody)) + require.False(t, isOpenAIImageRateLimitError(http.StatusBadRequest, imageBody)) +} + +func TestRateLimitService_HandleOpenAIImageRateLimit_ParsesTryAgainCooldown(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := &Account{ID: 201, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 2s."}}`) + + before := time.Now() + handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body) + + require.True(t, handled) + require.Len(t, repo.modelRateLimitCalls, 1) + call := repo.modelRateLimitCalls[0] + require.Equal(t, account.ID, call.accountID) + require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) + require.Equal(t, openAIImageRateLimitReason, call.reason) + require.WithinDuration(t, before.Add(2*time.Second), call.resetAt, time.Second) +} + +func TestRateLimitService_HandleOpenAIImageRateLimit_DefaultsToOneMinute(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &RateLimitService{accountRepo: repo} + account := &Account{ID: 202, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min."}}`) + + before := time.Now() + handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body) + + require.True(t, handled) + require.Len(t, repo.modelRateLimitCalls, 1) + call := repo.modelRateLimitCalls[0] + require.Equal(t, openAIImageGenerationRateLimitKey, call.scope) + require.Equal(t, openAIImageRateLimitReason, call.reason) + require.WithinDuration(t, before.Add(time.Minute), call.resetAt, time.Second) +} + +func TestOpenAIGatewayService_HandleOpenAIAccountUpstreamError_ImageRateLimitDoesNotBlockWholeAccount(t *testing.T) { + repo := &modelNotFoundAccountRepoStub{} + svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}} + account := &Account{ID: 203, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 1s."}}`) + + disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body, "gpt-image-2") + + require.False(t, disabled) + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope) + _, wholeAccountBlocked := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) + require.False(t, wholeAccountBlocked) +} + +func TestOpenAIGatewayServiceForwardImages_ImageRateLimitReturnsFailoverAndCoolsCapability(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := &modelNotFoundAccountRepoStub{} + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) + errorBody := `{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 1s."}}` + + req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = req + + svc := &OpenAIGatewayService{ + rateLimitService: &RateLimitService{accountRepo: repo}, + httpUpstream: &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"X-Request-Id": []string{"req_img_rate_limited"}}, + Body: io.NopCloser(strings.NewReader(errorBody)), + }, + }, + } + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 204, + Name: "openai-oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "token-123", + }, + } + + result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Contains(t, string(failoverErr.ResponseBody), "input-images per min") + require.Len(t, repo.modelRateLimitCalls, 1) + require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope) +}