Merge pull request #7442 from wucm667/fix/issue-7436-image-balance-error

fix(images): preserve insufficient balance failures

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)

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())
+		}
+	}
+}

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),

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,
+	}
+}

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"}`)))
+}