From da30c59923489adcaacc3b079547974792002365 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Tue, 9 Jun 2026 13:57:45 +0800 Subject: [PATCH] fix(openai): fail over image server errors --- backend/internal/handler/openai_images.go | 17 +- .../handler/openai_images_failover_test.go | 186 ++++++++++++++++++ .../service/openai_images_responses.go | 142 ++++++++++++- .../internal/service/openai_images_test.go | 149 ++++++++++++++ 4 files changed, 484 insertions(+), 10 deletions(-) create mode 100644 backend/internal/handler/openai_images_failover_test.go diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index bfa93de977..8b430cc271 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -226,8 +226,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } else { var imageUpstreamErr *service.OpenAIImagesUpstreamError if errors.As(err, &imageUpstreamErr) { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) - reqLog.Warn("openai.images.upstream_user_error", + retryableServerError := service.IsOpenAIImagesRetryableUpstreamError(imageUpstreamErr) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, !retryableServerError, nil) + logEvent := "openai.images.upstream_user_error" + if retryableServerError { + logEvent = "openai.images.upstream_server_error_after_flush" + } + reqLog.Warn(logEvent, zap.Int64("account_id", account.ID), zap.Int("status_code", imageUpstreamErr.StatusCode), zap.String("error_type", imageUpstreamErr.ErrorType), @@ -239,6 +244,14 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if c.Writer.Size() != writerSizeBeforeForward { + reqLog.Warn("openai.images.upstream_failover_skipped_after_flush", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", failoverErr.StatusCode), + ) + h.handleFailoverExhausted(c, failoverErr, true) + return + } if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() if sameAccountRetryCount[account.ID] < retryLimit { diff --git a/backend/internal/handler/openai_images_failover_test.go b/backend/internal/handler/openai_images_failover_test.go new file mode 100644 index 0000000000..79392d31d1 --- /dev/null +++ b/backend/internal/handler/openai_images_failover_test.go @@ -0,0 +1,186 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + middleware2 "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" + "github.com/tidwall/gjson" +) + +type openAIImagesFailoverAccountRepo struct { + service.AccountRepository + accounts []service.Account +} + +func (r openAIImagesFailoverAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) { + for i := range r.accounts { + if r.accounts[i].ID == id { + account := r.accounts[i] + return &account, nil + } + } + return nil, service.ErrNoAvailableAccounts +} + +func (r openAIImagesFailoverAccountRepo) ListSchedulableByGroupIDAndPlatform(_ context.Context, _ int64, platform string) ([]service.Account, error) { + return r.accountsForPlatform(platform), nil +} + +func (r openAIImagesFailoverAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) { + return r.accountsForPlatform(platform), nil +} + +func (r openAIImagesFailoverAccountRepo) ListSchedulableUngroupedByPlatform(_ context.Context, platform string) ([]service.Account, error) { + return r.accountsForPlatform(platform), nil +} + +func (r openAIImagesFailoverAccountRepo) accountsForPlatform(platform string) []service.Account { + out := make([]service.Account, 0, len(r.accounts)) + for _, account := range r.accounts { + if account.Platform == platform { + out = append(out, account) + } + } + return out +} + +type openAIImagesFailoverHTTPUpstream struct { + service.HTTPUpstream + mu sync.Mutex + accountIDs []int64 +} + +func (u *openAIImagesFailoverHTTPUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) { + u.mu.Lock() + u.accountIDs = append(u.accountIDs, accountID) + u.mu.Unlock() + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"req_img_failover"}, + }, + Body: io.NopCloser(bytes.NewBufferString( + "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"image backend unavailable\"}}\n\n", + )), + }, nil +} + +func (u *openAIImagesFailoverHTTPUpstream) calls() []int64 { + u.mu.Lock() + defer u.mu.Unlock() + return append([]int64(nil), u.accountIDs...) +} + +func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhenExhausted(t *testing.T) { + gin.SetMode(gin.TestMode) + groupID := int64(3130) + accounts := []service.Account{ + { + ID: 1, + Name: "image-account-1", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 0, + Priority: 0, + Credentials: map[string]any{"access_token": "token-1"}, + }, + { + ID: 2, + Name: "image-account-2", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 0, + Priority: 1, + Credentials: map[string]any{"access_token": "token-2"}, + }, + } + accountRepo := openAIImagesFailoverAccountRepo{accounts: accounts} + upstream := &openAIImagesFailoverHTTPUpstream{} + cfg := &config.Config{RunMode: config.RunModeSimple} + gatewayService := service.NewOpenAIGatewayService( + accountRepo, + nil, + nil, + nil, + nil, + nil, + nil, + cfg, + nil, + nil, + nil, + nil, + nil, + upstream, + nil, + nil, + nil, + nil, + nil, + nil, + nil, + ) + billingService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingService.Stop) + concurrencyService := service.NewConcurrencyService(nil) + handler := NewOpenAIGatewayHandler( + gatewayService, + concurrencyService, + billingService, + service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg), + nil, + nil, + nil, + cfg, + ) + handler.maxAccountSwitches = 10 + + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`) + 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 + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ + ID: 99, + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + AllowImageGeneration: true, + }, + User: &service.User{ID: 100}, + }) + c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 100, Concurrency: 0}) + + handler.Images(c) + + require.Equal(t, []int64{1, 2}, upstream.calls()) + require.Equal(t, http.StatusBadGateway, rec.Code) + require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String()) + require.Equal(t, "Upstream service temporarily unavailable", gjson.GetBytes(rec.Body.Bytes(), "error.message").String()) + + rawEvents, ok := c.Get(service.OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*service.OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 2) + require.Equal(t, "failover", events[0].Kind) + require.Equal(t, "failover", events[1].Kind) +} diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 56f98a999d..9782cc070b 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -5,6 +5,7 @@ import ( "bytes" "context" "encoding/base64" + "errors" "fmt" "io" "net/http" @@ -92,6 +93,53 @@ func (e *OpenAIImagesUpstreamError) clientMessage() string { return "Upstream request failed" } +// IsOpenAIImagesRetryableUpstreamError reports whether an Images error is an +// upstream server failure that may be retried on another account. +func IsOpenAIImagesRetryableUpstreamError(err *OpenAIImagesUpstreamError) bool { + return err != nil && err.StatusCode >= http.StatusInternalServerError +} + +func openAIImagesSSEErrorStatus(errType, code string) int { + errType = strings.ToLower(strings.TrimSpace(errType)) + code = strings.ToLower(strings.TrimSpace(code)) + + switch { + case strings.Contains(errType, "rate_limit"), strings.Contains(code, "rate_limit"): + return http.StatusTooManyRequests + case strings.Contains(errType, "authentication"), strings.Contains(code, "invalid_api_key"), code == "unauthorized": + return http.StatusUnauthorized + case strings.Contains(errType, "permission"), code == "forbidden": + return http.StatusForbidden + case strings.Contains(errType, "not_found"), strings.Contains(code, "not_found"): + return http.StatusNotFound + case strings.Contains(errType, "invalid_request"), + errType == "image_generation_user_error", + code == "moderation_blocked", + strings.Contains(code, "content_policy"), + strings.Contains(code, "policy_violation"), + strings.Contains(code, "safety_violation"): + return http.StatusBadRequest + default: + return http.StatusBadGateway + } +} + +func openAIImagesUpstreamErrorResponseBody(err *OpenAIImagesUpstreamError) []byte { + if err == nil { + return nil + } + body := []byte(`{"error":{"type":"","message":""}}`) + body, _ = sjson.SetBytes(body, "error.type", err.clientErrorType()) + body, _ = sjson.SetBytes(body, "error.message", err.clientMessage()) + if code := strings.TrimSpace(err.Code); code != "" { + body, _ = sjson.SetBytes(body, "error.code", code) + } + if param := strings.TrimSpace(err.Param); param != "" { + body, _ = sjson.SetBytes(body, "error.param", param) + } + return body +} + func openAIResponsesImageResultKey(itemID string, result openAIResponsesImageResult) string { if strings.TrimSpace(result.Result) != "" { return strings.TrimSpace(result.OutputFormat) + "|" + strings.TrimSpace(result.Result) @@ -562,10 +610,7 @@ func openAIImagesUpstreamErrorFromGJSON(errorObj gjson.Result, upstreamRequestID errType := strings.TrimSpace(errorObj.Get("type").String()) message := strings.TrimSpace(errorObj.Get("message").String()) param := strings.TrimSpace(errorObj.Get("param").String()) - statusCode := http.StatusBadGateway - if strings.EqualFold(code, "moderation_blocked") || strings.EqualFold(errType, "image_generation_user_error") { - statusCode = http.StatusBadRequest - } + statusCode := openAIImagesSSEErrorStatus(errType, code) if message == "" { message = "Upstream request failed" } @@ -909,7 +954,9 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthNonStreamingResponse( if len(results) == 0 { if upstreamErr := extractOpenAIImagesUpstreamError(body); upstreamErr != nil { setOpsUpstreamError(c, upstreamErr.clientStatusCode(), upstreamErr.clientMessage(), "") - writeOpenAIImagesUpstreamErrorResponse(c, upstreamErr) + if !IsOpenAIImagesRetryableUpstreamError(upstreamErr) { + writeOpenAIImagesUpstreamErrorResponse(c, upstreamErr) + } return OpenAIUsage{}, 0, nil, upstreamErr } return OpenAIUsage{}, 0, nil, fmt.Errorf("upstream did not return image output") @@ -965,6 +1012,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( var sseData openAISSEDataAccumulator var processDataErr error processDataDone := false + writerSizeBeforeResponse := c.Writer.Size() processData := func(dataBytes []byte) { if processDataDone || processDataErr != nil { @@ -1068,7 +1116,8 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse( processDataDone = true case "error", "response.failed": if upstreamErr := openAIImagesUpstreamErrorFromSSEPayload(dataBytes); upstreamErr != nil { - if !clientDisconnected { + retryable := IsOpenAIImagesRetryableUpstreamError(upstreamErr) + if !clientDisconnected && (!retryable || c.Writer.Size() != writerSizeBeforeResponse) { s.tryWriteOpenAIImagesStreamEvent(c, flusher, &clientDisconnected, &lastDownstreamWriteAt, "error", buildOpenAIImagesStreamErrorBodyFromUpstream(upstreamErr)) } setOpsUpstreamError(c, upstreamErr.clientStatusCode(), upstreamErr.clientMessage(), "") @@ -1375,6 +1424,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( imageOutputSizes []string firstTokenMs *int ) + writerSizeBeforeResponse := c.Writer.Size() if parsed.Stream { usage, imageCount, imageOutputSizes, firstTokenMs, err = s.handleOpenAIImagesOAuthStreamingResponse(resp, c, startTime, parsed.ResponseFormat, openAIImagesStreamPrefix(parsed), requestModel) if err != nil { @@ -1394,12 +1444,30 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( ImageOutputSizes: imageOutputSizes, }, err } - return nil, err + return nil, s.handleOpenAIImagesOAuthResponseError( + upstreamCtx, + c, + account, + requestModel, + safeUpstreamURL(upstreamReq.URL.String()), + resp, + writerSizeBeforeResponse, + err, + ) } } else { usage, imageCount, imageOutputSizes, err = s.handleOpenAIImagesOAuthNonStreamingResponse(resp, c, parsed.ResponseFormat, requestModel) if err != nil { - return nil, err + return nil, s.handleOpenAIImagesOAuthResponseError( + upstreamCtx, + c, + account, + requestModel, + safeUpstreamURL(upstreamReq.URL.String()), + resp, + writerSizeBeforeResponse, + err, + ) } } if imageCount <= 0 { @@ -1420,3 +1488,61 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( ImageOutputSizes: imageOutputSizes, }, nil } + +func (s *OpenAIGatewayService) handleOpenAIImagesOAuthResponseError( + ctx context.Context, + c *gin.Context, + account *Account, + requestedModel string, + upstreamURL string, + resp *http.Response, + writerSizeBeforeResponse int, + err error, +) error { + var upstreamErr *OpenAIImagesUpstreamError + if !errors.As(err, &upstreamErr) { + return err + } + + retryable := IsOpenAIImagesRetryableUpstreamError(upstreamErr) + responseWritten := c != nil && c.Writer != nil && c.Writer.Size() != writerSizeBeforeResponse + kind := "http_error" + if retryable { + kind = "failover" + if responseWritten { + kind = "retry_exhausted_failover" + } + } + + requestID := strings.TrimSpace(upstreamErr.UpstreamRequestID) + headers := http.Header(nil) + if resp != nil { + headers = resp.Header.Clone() + if requestID == "" { + requestID = strings.TrimSpace(resp.Header.Get("x-request-id")) + } + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: upstreamErr.StatusCode, + UpstreamRequestID: requestID, + UpstreamURL: upstreamURL, + Kind: kind, + Message: upstreamErr.clientMessage(), + }) + + if !retryable || responseWritten { + return err + } + + responseBody := openAIImagesUpstreamErrorResponseBody(upstreamErr) + s.handleOpenAIAccountUpstreamError(ctx, account, upstreamErr.StatusCode, headers, responseBody, requestedModel) + return &UpstreamFailoverError{ + StatusCode: upstreamErr.StatusCode, + ResponseBody: responseBody, + ResponseHeaders: headers, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(upstreamErr.StatusCode), + } +} diff --git a/backend/internal/service/openai_images_test.go b/backend/internal/service/openai_images_test.go index c6083bd0b9..74846e0a23 100644 --- a/backend/internal/service/openai_images_test.go +++ b/backend/internal/service/openai_images_test.go @@ -767,6 +767,155 @@ func TestOpenAIGatewayServiceForwardImages_OAuthNonStreamModerationBlockedReturn require.Contains(t, gjson.Get(rec.Body.String(), "error.message").String(), "safety system") } +func TestOpenAIGatewayServiceForwardImages_OAuthNonStreamServerErrorReturnsFailoverBeforeFlush(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`) + + 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{ + httpUpstream: &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"req_img_server_error"}, + }, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000021}}\n\n" + + "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service is temporarily unavailable.\"}}\n\n", + )), + }, + }, + } + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 21, + Name: "openai-oauth-server-error", + 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.StatusBadGateway, failoverErr.StatusCode) + require.Contains(t, string(failoverErr.ResponseBody), "temporarily unavailable") + require.False(t, c.Writer.Written()) + require.Empty(t, rec.Body.String()) + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 1) + require.Equal(t, "failover", events[0].Kind) + require.Equal(t, account.ID, events[0].AccountID) + require.Equal(t, http.StatusBadGateway, events[0].UpstreamStatusCode) +} + +func TestOpenAIGatewayServiceForwardImages_OAuthStreamServerErrorAfterFlushDoesNotFailover(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","stream":true,"response_format":"b64_json"}`) + + 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{ + httpUpstream: &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"req_img_server_error_after_partial"}, + }, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_b64\":\"cGFydGlhbA==\",\"partial_image_index\":0,\"output_format\":\"png\"}\n\n" + + "data: {\"type\":\"error\",\"error\":{\"type\":\"server_error\",\"code\":\"server_error\",\"message\":\"The image service failed after partial output.\"}}\n\n", + )), + }, + }, + } + parsed, err := svc.ParseOpenAIImagesRequest(c, body) + require.NoError(t, err) + account := &Account{ + ID: 22, + Name: "openai-oauth-partial-server-error", + 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.False(t, errors.As(err, &failoverErr)) + var upstreamErr *OpenAIImagesUpstreamError + require.ErrorAs(t, err, &upstreamErr) + require.True(t, IsOpenAIImagesRetryableUpstreamError(upstreamErr)) + require.True(t, c.Writer.Written()) + require.Contains(t, rec.Body.String(), "event: image_generation.partial_image") + require.Contains(t, rec.Body.String(), "event: error") + require.Contains(t, rec.Body.String(), "failed after partial output") + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.Len(t, events, 1) + require.Equal(t, "retry_exhausted_failover", events[0].Kind) + require.Equal(t, account.ID, events[0].AccountID) +} + +func TestOpenAIImagesSSEClientErrorsAreNotRetryable(t *testing.T) { + tests := []struct { + name string + payload string + wantStatus int + }{ + { + name: "invalid request", + payload: `{"type":"error","error":{"type":"invalid_request_error","code":"invalid_value","message":"bad size"}}`, + wantStatus: http.StatusBadRequest, + }, + { + name: "content policy", + payload: `{"type":"error","error":{"type":"image_generation_user_error","code":"content_policy_violation","message":"blocked"}}`, + wantStatus: http.StatusBadRequest, + }, + { + name: "rate limit remains distinct from server error", + payload: `{"type":"error","error":{"type":"rate_limit_exceeded","code":"rate_limit_exceeded","message":"try again"}}`, + wantStatus: http.StatusTooManyRequests, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstreamErr := openAIImagesUpstreamErrorFromSSEPayload([]byte(tt.payload)) + require.NotNil(t, upstreamErr) + require.Equal(t, tt.wantStatus, upstreamErr.StatusCode) + require.False(t, IsOpenAIImagesRetryableUpstreamError(upstreamErr)) + }) + } +} + func TestOpenAIGatewayServiceForwardImages_APIKeyGenerationUsesConfiguredV1BaseURL(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","response_format":"b64_json"}`)