diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index fd7eb09c3d..1709e2485f 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -16,6 +16,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -2498,6 +2499,54 @@ func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) { require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) } +func TestOpenAIWSHTTPBridgeSSEErrorSideEffectsRunOncePerPlatform(t *testing.T) { + gin.SetMode(gin.TestMode) + + for _, platform := range []string{PlatformOpenAI, PlatformGrok} { + t.Run(platform, func(t *testing.T) { + repo := &grokQuotaAccountRepo{} + cfg := &config.Config{} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n", + )), + }} + svc := &OpenAIGatewayService{ + cfg: cfg, + accountRepo: repo, + httpUpstream: upstream, + } + if platform == PlatformOpenAI { + svc.rateLimitService = NewRateLimitService(repo, nil, cfg, nil, nil) + } + account := &Account{ID: 70, Platform: platform, Type: AccountTypeOAuth, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + writes := 0 + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", 1, + func([]byte) error { + writes++ + return nil + }, + ) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Zero(t, writes) + require.Equal(t, 1, repo.rateLimitedCalls) + }) + } +} + func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) { repo := &grokQuotaAccountRepo{} resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index cc7b307982..a84968061b 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -230,6 +230,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( turnStart := time.Now() resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) if err != nil { + if turn == 1 { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true) + } safeErr := sanitizeUpstreamErrorMessage(err.Error()) _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(http.StatusBadGateway, "Upstream request failed")) return nil, fmt.Errorf("upstream http bridge request failed: %s", safeErr) @@ -238,16 +241,23 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if resp.StatusCode >= 400 { respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) - if account.Platform == PlatformGrok { - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - } else if shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody) { - canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel) - } upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if upstreamMsg == "" { upstreamMsg = http.StatusText(resp.StatusCode) } + shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) + if account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + if turn == 1 && shouldFailover { + return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, false) + } + } else if turn == 1 && shouldFailover { + return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, respBody) + } + if account.Platform != PlatformGrok && (shouldFailover || shouldCooldownOpenAITransientUpstreamError(resp.StatusCode, respBody)) { + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, canonicalModel) + } _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)) return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } @@ -372,6 +382,31 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } replayCollector.AddEvent(eventType, upstreamMessage) + var upstreamEventErr error + if eventType == "error" { + errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) + errMessage := strings.TrimSpace(errMsgRaw) + if errMessage == "" { + errMessage = "upstream error event" + } + statusCode := openAIWSErrorHTTPStatusFromRaw(errCodeRaw, errTypeRaw) + shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(statusCode, errMessage, upstreamMessage) + if account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, statusCode, resp.Header, upstreamMessage) + } else if shouldFailover { + accountStatus := statusCode + if transientStatus := openAIWSPayloadTransientStatus(upstreamMessage); transientStatus != 0 { + accountStatus = transientStatus + } + canonicalModel := canonicalOpenAIAccountSchedulingModel(account, originalModel) + s.handleOpenAIAccountUpstreamError(ctx, account, accountStatus, resp.Header, upstreamMessage, canonicalModel) + } + if turn == 1 && !wroteDownstream && shouldFailover { + return nil, newOpenAIUpstreamFailoverError(statusCode, resp.Header, upstreamMessage, errMessage, false) + } + upstreamEventErr = errors.New(errMessage) + } + if !clientDisconnected { if err := writeClientMessage(upstreamMessage); err != nil { if isOpenAIWSClientDisconnectError(err) { @@ -396,15 +431,8 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } } - if eventType == "error" { - s.handleOpenAIWSErrorEventTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage) - errCodeRaw, errTypeRaw, errMsgRaw := parseOpenAIWSErrorEventFields(upstreamMessage) - s.persistOpenAIWSRateLimitSignal(ctx, account, resp.Header, upstreamMessage, errCodeRaw, errTypeRaw, errMsgRaw) - errMessage := strings.TrimSpace(errMsgRaw) - if errMessage == "" { - errMessage = "upstream error event" - } - return resultWithUsage(), errors.New(errMessage) + if upstreamEventErr != nil { + return resultWithUsage(), upstreamEventErr } if isOpenAIWSTerminalEvent(eventType) { upstreamTerminalEvent = s.handleOpenAIWSTerminalTransientFailure(ctx, account, canonicalOpenAIAccountSchedulingModel(account, originalModel), resp.Header, upstreamMessage) @@ -432,12 +460,20 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } } if err := scanner.Err(); err != nil { - return resultWithUsage(), fmt.Errorf("read upstream http bridge stream: %w", err) + streamErr := fmt.Errorf("read upstream http bridge stream: %w", err) + if turn == 1 && !wroteDownstream { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, streamErr, true) + } + return resultWithUsage(), streamErr } - if sawDone && eventCount > 0 { - return resultWithUsage(), nil + terminalErr := errors.New("upstream http bridge stream ended before terminal event") + if sawDone { + terminalErr = errors.New("upstream http bridge stream sent [DONE] before terminal event") } - return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event") + if turn == 1 && !wroteDownstream { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, terminalErr, true) + } + return resultWithUsage(), terminalErr } func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte, originalModel string) (string, error) { diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 808a63c6b1..de2e463816 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -62,6 +62,219 @@ func TestOpenAIWSHTTPBridgeDecisionKeepsSmallFramesOnWS(t *testing.T) { require.True(t, svc.shouldBridgeOpenAIWSHTTP(&Account{Platform: PlatformGrok}, 1, "resp_existing")) } +func TestProxyOpenAIWSHTTPBridgeTurnTransportErrorFailoverSafety(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + turn int + wantFailover bool + wantWrites int + }{ + {name: "first_turn_fails_over_before_downstream_event", turn: 1, wantFailover: true}, + {name: "later_turn_does_not_replay_completed_turns", turn: 2, wantWrites: 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstream := &httpUpstreamRecorder{err: io.EOF} + svc := &OpenAIGatewayService{ + cfg: &config.Config{}, + httpUpstream: upstream, + } + account := &Account{ + ID: 8, + Name: "api-key", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", tt.turn, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + return nil + }, + ) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + if tt.wantFailover { + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.JSONEq(t, string(openAITransportFailoverBody), string(failoverErr.ResponseBody)) + } else { + require.Error(t, err) + require.False(t, errors.As(err, &failoverErr)) + } + require.Len(t, writes, tt.wantWrites) + if tt.wantWrites > 0 { + require.Equal(t, "error", gjson.GetBytes(writes[0], "type").String()) + require.Equal(t, int64(http.StatusBadGateway), gjson.GetBytes(writes[0], "status").Int()) + } + }) + } +} + +func TestProxyOpenAIWSHTTPBridgeTurnHTTPStatusFailoverSafety(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + turn int + status int + wantFailover bool + wantWrites int + }{ + {name: "first_turn_401", turn: 1, status: http.StatusUnauthorized, wantFailover: true}, + {name: "first_turn_429", turn: 1, status: http.StatusTooManyRequests, wantFailover: true}, + {name: "first_turn_500", turn: 1, status: http.StatusInternalServerError, wantFailover: true}, + {name: "later_turn_500_does_not_replay", turn: 2, status: http.StatusInternalServerError, wantWrites: 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: tt.status, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"error":{"type":"server_error","message":"temporary upstream failure"}}`)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 9, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", tt.turn, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + return nil + }, + ) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + if tt.wantFailover { + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, tt.status, failoverErr.StatusCode) + } else { + require.Error(t, err) + require.False(t, errors.As(err, &failoverErr)) + } + require.Len(t, writes, tt.wantWrites) + }) + } +} + +func TestProxyOpenAIWSHTTPBridgeTurnSSEErrorFailoverSafety(t *testing.T) { + gin.SetMode(gin.TestMode) + + for _, turn := range []int{1, 2} { + t.Run(fmt.Sprintf("turn_%d", turn), func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n", + )), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", turn, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + return nil + }, + ) + + var failoverErr *UpstreamFailoverError + if turn == 1 { + require.Nil(t, result) + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Empty(t, writes) + } else { + require.NotNil(t, result) + require.Error(t, err) + require.False(t, errors.As(err, &failoverErr)) + require.Len(t, writes, 1) + } + }) + } +} + +func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + body string + wantFailover bool + wantWrites int + }{ + {name: "done_without_events_fails_over", body: "data: [DONE]\n\n", wantFailover: true}, + { + name: "created_then_done_is_truncated_not_success", + body: "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_truncated\"}}\n\n" + + "data: [DONE]\n\n", + wantWrites: 1, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(tt.body)), + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ID: 11, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`) + var writes [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "sk-test", payload, len(payload), + "gpt-5", "", "", "", "", 1, + func(message []byte) error { + writes = append(writes, append([]byte(nil), message...)) + return nil + }, + ) + + var failoverErr *UpstreamFailoverError + if tt.wantFailover { + require.Nil(t, result) + require.ErrorAs(t, err, &failoverErr) + } else { + require.NotNil(t, result) + require.Error(t, err) + require.False(t, errors.As(err, &failoverErr)) + } + require.Len(t, writes, tt.wantWrites) + }) + } +} + func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { gin.SetMode(gin.TestMode)