mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
fix(openai): fail over image server errors
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"}`)
|
||||
|
||||
Reference in New Issue
Block a user