diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index d440c2776..e94c10c2e 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -892,6 +892,15 @@ func openAIStreamFailedEventShouldFailover(payload []byte, message string) bool return true } +func openAIStreamFailedEventRetryableOnSameAccount(account *Account, payload []byte, message string) bool { + if account == nil || !account.IsPoolMode() { + return false + } + semanticStatus := openAIStreamFailedEventSemanticStatus(payload, message) + return account.IsPoolModeRetryableStatus(semanticStatus) || + isOpenAITransientProcessingError(http.StatusBadRequest, message, payload) +} + func (s *OpenAIGatewayService) recordOpenAIStreamUpstreamError( c *gin.Context, account *Account, @@ -954,8 +963,9 @@ func (s *OpenAIGatewayService) newOpenAIStreamFailoverError( }, }) return &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: body, + StatusCode: http.StatusBadGateway, + ResponseBody: body, + RetryableOnSameAccount: openAIStreamFailedEventRetryableOnSameAccount(account, payload, message), } } diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index d8cfba73b..05979533f 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -1474,6 +1474,7 @@ func TestOpenAIStreamingResponseFailedBeforeOutputReturnsFailover(t *testing.T) var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.False(t, failoverErr.RetryableOnSameAccount) require.Contains(t, string(failoverErr.ResponseBody), "An error occurred while processing your request") require.False(t, c.Writer.Written()) require.Empty(t, rec.Body.String()) @@ -1510,11 +1511,21 @@ func TestOpenAIStreamingResponseFailedBeforeOutputCapacityErrorReturnsFailover(t Header: http.Header{"X-Request-Id": []string{"rid-capacity-failed"}}, } - _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Name: "acc"}, time.Now(), "model", "model") + account := &Account{ + ID: 1, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Name: "pool-account", + Credentials: map[string]any{ + "pool_mode": true, + }, + } + _, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model") require.Error(t, err) var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode) + require.True(t, failoverErr.RetryableOnSameAccount) require.Contains(t, string(failoverErr.ResponseBody), "Selected model is at capacity") require.False(t, c.Writer.Written()) require.Empty(t, rec.Body.String())