From de98bdff1e964573c0bf3c32793e2bafb23688d7 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Mon, 21 Sep 2026 10:23:10 +0800 Subject: [PATCH] fix(images): preserve insufficient balance failures --- .../handler/openai_gateway_handler.go | 15 ++- .../handler/openai_images_balance_test.go | 38 +++++++ backend/internal/service/openai_images.go | 16 +++ .../internal/service/openai_images_balance.go | 93 +++++++++++++++ .../service/openai_images_balance_test.go | 107 ++++++++++++++++++ 5 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 backend/internal/handler/openai_images_balance_test.go create mode 100644 backend/internal/service/openai_images_balance.go create mode 100644 backend/internal/service/openai_images_balance_test.go diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 600d2c7372..bf92e3a82e 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -3340,6 +3340,7 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted) return } + copyFailoverRetryAfter(c, failoverErr.ResponseHeaders) if failoverErr.IsOpenAIRequestBodyTooLarge() { service.SetOpsUpstreamError(c, http.StatusRequestEntityTooLarge, service.OpenAIRequestBodyTooLargeClientMessage, "") h.handleStreamingAwareError( @@ -3359,7 +3360,19 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE h.handleStreamingAwareError(c, http.StatusBadRequest, "invalid_request_error", message, streamStarted) return } - copyFailoverRetryAfter(c, failoverErr.ResponseHeaders) + if failoverErr.Reason == service.OpenAIImagesInsufficientBalanceReason { + status := failoverErr.ClientStatusCode + if status <= 0 { + status = http.StatusPaymentRequired + } + message := strings.TrimSpace(failoverErr.ClientMessage) + if message == "" { + message = service.OpenAIImagesInsufficientBalanceMessage + } + service.SetOpsUpstreamError(c, failoverErr.StatusCode, message, "") + h.handleStreamingAwareErrorWithCode(c, status, "upstream_error", service.OpenAIImagesInsufficientBalanceCode, message, streamStarted, false) + return + } if failoverErr.IsCredentialFailure() { status, message := credentialFailoverClientResponse(failoverErr) h.handleStreamingAwareError(c, status, "upstream_error", message, streamStarted) diff --git a/backend/internal/handler/openai_images_balance_test.go b/backend/internal/handler/openai_images_balance_test.go new file mode 100644 index 0000000000..dc6014f469 --- /dev/null +++ b/backend/internal/handler/openai_images_balance_test.go @@ -0,0 +1,38 @@ +//go:build unit + +package handler + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestOpenAIImagesBalanceExhaustionPreservesMachineReadableCause(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, streamStarted := range []bool{false, true} { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", nil) + (&OpenAIGatewayHandler{}).handleFailoverExhausted(c, &service.UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + Reason: service.OpenAIImagesInsufficientBalanceReason, + ClientStatusCode: http.StatusPaymentRequired, + ClientMessage: service.OpenAIImagesInsufficientBalanceMessage, + ResponseHeaders: http.Header{"Retry-After": []string{"23"}, "X-Request-Id": []string{"req-images-balance"}}, + }, streamStarted) + require.Equal(t, "23", recorder.Header().Get("Retry-After")) + + if streamStarted { + require.Contains(t, recorder.Body.String(), `"code":"insufficient_balance"`) + } else { + require.Equal(t, http.StatusPaymentRequired, recorder.Code) + require.Equal(t, service.OpenAIImagesInsufficientBalanceCode, gjson.GetBytes(recorder.Body.Bytes(), "error.code").String()) + } + } +} diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 0fe06a334b..8ef483e537 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -661,6 +661,22 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + if isOpenAIImagesInsufficientBalance(respBody) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + ProxyID: opsUpstreamProxyID(account), + ProxyName: opsUpstreamProxyName(account), + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "failover", + Message: OpenAIImagesInsufficientBalanceMessage, + }) + s.coolOpenAIImagesInsufficientBalance(upstreamCtx, account) + return nil, newOpenAIImagesInsufficientBalanceFailoverError(resp.StatusCode, resp.Header, respBody) + } if s.shouldFailoverOpenAIUpstreamResponse(account, resp.StatusCode, upstreamMsg, respBody) { appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ ProxyID: opsUpstreamProxyID(account), diff --git a/backend/internal/service/openai_images_balance.go b/backend/internal/service/openai_images_balance.go new file mode 100644 index 0000000000..92d4ccf052 --- /dev/null +++ b/backend/internal/service/openai_images_balance.go @@ -0,0 +1,93 @@ +package service + +import ( + "context" + "encoding/json" + "log/slog" + "net/http" + "strings" + "time" +) + +const ( + OpenAIImagesInsufficientBalanceReason GatewayFailureReason = "openai_images_insufficient_balance" + OpenAIImagesInsufficientBalanceCode = "insufficient_balance" + OpenAIImagesInsufficientBalanceMessage = "Upstream image account has insufficient balance" + openAIImagesBalanceCooldown = 5 * time.Minute + openAIImagesBalanceRateLimitReason = "openai_images_insufficient_balance" +) + +// isOpenAIImagesInsufficientBalance only accepts explicit values in structured +// error fields. This avoids treating echoed prompts or other arbitrary text as +// an account balance signal. +func isOpenAIImagesInsufficientBalance(body []byte) bool { + var value any + if len(body) == 0 || json.Unmarshal(body, &value) != nil { + return false + } + return hasOpenAIImagesInsufficientBalance(value, 0) +} + +func hasOpenAIImagesInsufficientBalance(value any, depth int) bool { + if depth > 6 { + return false + } + switch typed := value.(type) { + case map[string]any: + for key, child := range typed { + normalizedKey := strings.ToLower(strings.TrimSpace(key)) + if text, ok := child.(string); ok { + normalizedValue := strings.ToLower(strings.Trim(strings.TrimSpace(text), ".")) + switch normalizedKey { + case "errorkey", "error_key", "code": + if normalizedValue == OpenAIImagesInsufficientBalanceCode { + return true + } + } + } + if isOpenAIImagesErrorContainer(normalizedKey) && hasOpenAIImagesInsufficientBalance(child, depth+1) { + return true + } + } + case []any: + for _, child := range typed { + if hasOpenAIImagesInsufficientBalance(child, depth+1) { + return true + } + } + } + return false +} + +func isOpenAIImagesErrorContainer(key string) bool { + switch key { + case "error", "errors", "cause", "detail", "details", "inner", "innererror", "inner_error", "response": + return true + default: + return false + } +} + +func (s *OpenAIGatewayService) coolOpenAIImagesInsufficientBalance(ctx context.Context, account *Account) { + if s == nil || s.accountRepo == nil || account == nil { + return + } + now := time.Now() + resetAt := now.Add(openAIImagesBalanceCooldown) + setAccountModelRateLimitSnapshot(account, openAIImageGenerationRateLimitKey, resetAt, openAIImagesBalanceRateLimitReason, now) + if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, openAIImageGenerationRateLimitKey, resetAt, openAIImagesBalanceRateLimitReason); err != nil { + slog.Warn("openai_images_balance_cooldown_failed", "account_id", account.ID, "error", err) + } +} + +func newOpenAIImagesInsufficientBalanceFailoverError(statusCode int, headers http.Header, body []byte) *UpstreamFailoverError { + return &UpstreamFailoverError{ + StatusCode: statusCode, + ResponseHeaders: headers.Clone(), + ResponseBody: body, + Reason: OpenAIImagesInsufficientBalanceReason, + ClientStatusCode: 402, + ClientMessage: OpenAIImagesInsufficientBalanceMessage, + NextAccountAction: NextAccountRetry, + } +} diff --git a/backend/internal/service/openai_images_balance_test.go b/backend/internal/service/openai_images_balance_test.go new file mode 100644 index 0000000000..d96fa0ebbf --- /dev/null +++ b/backend/internal/service/openai_images_balance_test.go @@ -0,0 +1,107 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type openAIImagesBalanceRepo struct { + AccountRepository + scope string + reason string + until time.Time +} + +func (r *openAIImagesBalanceRepo) SetModelRateLimit(_ context.Context, _ int64, scope string, resetAt time.Time, reason ...string) error { + r.scope = scope + r.until = resetAt + if len(reason) > 0 { + r.reason = reason[0] + } + return nil +} + +type openAIImagesBalanceUpstream struct{ HTTPUpstream } + +func (openAIImagesBalanceUpstream) Do(*http.Request, string, int64, int) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusBadGateway, + Header: http.Header{ + "X-Request-Id": []string{"req-images-balance"}, + "Retry-After": []string{"23"}, + }, + Body: io.NopCloser(bytes.NewBufferString( + `{"error":{"message":"upstream request failed","cause":{"errorKey":"insufficient_balance","errorMessage":"Insufficient credit balance"}}}`, + )), + }, nil +} + +func TestForwardOpenAIImagesAPIKey_ClassifiesWrappedBalanceBeforeTransientFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + name string + endpoint string + stream bool + }{ + {name: "non-stream generation", endpoint: openAIImagesGenerationsEndpoint}, + {name: "stream edit", endpoint: openAIImagesEditsEndpoint, stream: true}, + } { + t.Run(tc.name, func(t *testing.T) { + repo := &openAIImagesBalanceRepo{} + svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: openAIImagesBalanceUpstream{}, cfg: &config.Config{}} + account := &Account{ + ID: 7, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Status: StatusActive, Schedulable: true, + Credentials: map[string]any{"api_key": "test-key", "base_url": "https://images.example.test"}, + } + body := []byte(`{"model":"gpt-image-1","prompt":"cat"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/"+tc.endpoint, bytes.NewReader(body)) + parsed := &OpenAIImagesRequest{Model: "gpt-image-1", Endpoint: tc.endpoint, ContentType: "application/json", Stream: tc.stream, N: 1} + + _, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "") + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, "req-images-balance", failoverErr.ResponseHeaders.Get("X-Request-Id")) + require.Equal(t, "23", failoverErr.ResponseHeaders.Get("Retry-After")) + require.Equal(t, OpenAIImagesInsufficientBalanceReason, failoverErr.Reason) + require.True(t, failoverErr.ShouldRetryNextAccount()) + require.Equal(t, OpenAIImagesInsufficientBalanceCode, string(failoverErr.Reason)[len("openai_images_"):]) + require.Equal(t, openAIImageGenerationRateLimitKey, repo.scope) + require.Equal(t, openAIImagesBalanceRateLimitReason, repo.reason) + require.WithinDuration(t, time.Now().Add(openAIImagesBalanceCooldown), repo.until, 2*time.Second) + require.True(t, account.isRateLimitActiveForKey(openAIImageGenerationRateLimitKey)) + require.False(t, account.isModelRateLimitedWithContext(context.Background(), "gpt-5.5"), "text requests must remain schedulable") + }) + } +} + +func TestOpenAIImagesInsufficientBalanceDetectionRejectsUnstructuredAndEchoedText(t *testing.T) { + require.True(t, isOpenAIImagesInsufficientBalance([]byte(`{"errorKey":"insufficient_balance"}`))) + require.True(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"cause":{"errorKey":"insufficient_balance"}}}`))) + require.True(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"details":[{"code":"insufficient_balance"}]}}`))) + require.True(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"inner_error":{"error_key":" INSUFFICIENT_BALANCE. "}}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"message":"Insufficient credit balance"}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"detail":"Insufficient credit"}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"message":"Insufficient credit balance"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"details":[{"detail":"Insufficient credit"}]}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"errorMessage":"Insufficient credit balance"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"message":"{\"errorKey\":\"insufficient_balance\"}"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"error":{"detail":"upstream returned 400 {\"errorKey\":\"insufficient_balance\"}"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`gateway says insufficient_balance`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"request":{"prompt":"show the words insufficient_balance"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"request":{"errorKey":"insufficient_balance"}}`))) + require.False(t, isOpenAIImagesInsufficientBalance([]byte(`{"message":"the prompt mentioned insufficient credit balance yesterday"}`))) +}