diff --git a/backend/internal/service/grok_upstream_errors_test.go b/backend/internal/service/grok_upstream_errors_test.go index 74fa10bbd..0dee756b7 100644 --- a/backend/internal/service/grok_upstream_errors_test.go +++ b/backend/internal/service/grok_upstream_errors_test.go @@ -122,6 +122,43 @@ func TestGrokContentPolicy403DoesNotMutateOrFailover(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokNonFailoverDoesNotApplyGenericTempUnschedulablePolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + rateLimitService: NewRateLimitService(repo, nil, nil, nil, nil), + } + account := &Account{ + ID: 5099, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{map[string]any{ + "error_code": float64(http.StatusForbidden), + "keywords": []any{"text is sensitive"}, + "duration_minutes": float64(1), + }}, + }, + } + body := []byte(`{"error":{"code":"new_sensitive","message":"text is sensitive"}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + resp := &http.Response{StatusCode: http.StatusForbidden, Header: http.Header{}} + + got := svc.failoverOpenAIUpstreamHTTPError( + context.Background(), c, account, resp, body, "text is sensitive", "", + ) + + require.Nil(t, got) + require.Zero(t, repo.tempUnschedCalls) + require.Zero(t, repo.rateLimitedCalls) + require.Zero(t, repo.updateCalls) + require.False(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + func TestGrokContentPolicy403SharedErrorFallbackDoesNotMutate(t *testing.T) { gin.SetMode(gin.TestMode) body := []byte(`{"error":{"code":"content_filter","message":"prohibited content"}}`) diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index 635b985ea..8080d49e5 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -89,6 +89,11 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( upstreamModel string, ) *UpstreamFailoverError { shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) + tempUnscheduled := false + if c != nil && account != nil && account.Platform != PlatformGrok && !shouldFailover && !IsResponseCommitted(c) && s.rateLimitService != nil { + tempUnscheduled = s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, upstreamModel) == ErrorPolicyTempUnscheduled + shouldFailover = tempUnscheduled + } if account != nil && account.Platform == PlatformGrok { shouldFailover = s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) } @@ -116,8 +121,8 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( Message: upstreamMsg, Detail: upstreamDetail, }) - shouldDisable := false - if account.Platform != PlatformGrok { + shouldDisable := tempUnscheduled + if account.Platform != PlatformGrok && !tempUnscheduled { shouldDisable = s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) } return newOpenAIUpstreamFailoverError( diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 0e86389d1..d03d48e00 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -31,6 +31,18 @@ type stubOpenAIAccountRepo struct { accounts []Account } +type tempUnschedulableOpenAIAccountRepo struct { + stubOpenAIAccountRepo + modelRateLimitAccountID int64 + modelRateLimitKey string +} + +func (r *tempUnschedulableOpenAIAccountRepo) SetModelRateLimit(_ context.Context, accountID int64, modelKey string, _ time.Time, _ ...string) error { + r.modelRateLimitAccountID = accountID + r.modelRateLimitKey = modelKey + return nil +} + type snapshotUpdateAccountRepo struct { stubOpenAIAccountRepo updateExtraCalls chan map[string]any @@ -80,6 +92,113 @@ func (r stubOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Co return r.ListSchedulableByPlatform(ctx, platform) } +func TestOpenAIGatewayService_ForwardAsAnthropic_TempUnschedulableReturnsFailoverWithoutCommit(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"gpt-5.4","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":false}`) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstreamBody := []byte(`{"error":{"message":"Our servers are currently overloaded. Please try again later."}}`) + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader(upstreamBody)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.output_text.delta","delta":"ok"}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_second","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n"))), + }, + }} + repo := &tempUnschedulableOpenAIAccountRepo{} + rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + svc := &OpenAIGatewayService{ + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{ + Enabled: false, AllowInsecureHTTP: true, + }}}, + httpUpstream: upstream, + rateLimitService: rateLimitService, + } + account := &Account{ + ID: 5099, Name: "temporary-unschedulable", Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "http://upstream.example", + "model_provider": "env-openai", + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{map[string]any{ + "error_code": float64(http.StatusBadRequest), + "keywords": []any{"our servers are currently overloaded", "please try again later"}, + "duration_minutes": float64(1), + }}, + }, + } + + _, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "") + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusBadRequest, failoverErr.StatusCode) + require.True(t, failoverErr.ShouldRetryNextAccount()) + require.Equal(t, account.ID, repo.modelRateLimitAccountID, "temporary unschedulability should exclude this account from reselection") + require.Equal(t, "gpt-5.4", repo.modelRateLimitKey) + require.False(t, IsResponseCommitted(c)) + require.Equal(t, http.StatusOK, rec.Code) + require.Empty(t, rec.Body.String()) + + secondRec := httptest.NewRecorder() + secondContext, _ := gin.CreateTestContext(secondRec) + secondContext.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + secondContext.Request.Header.Set("Content-Type", "application/json") + secondAccount := *account + secondAccount.ID = 5100 + secondAccount.Name = "healthy-failover-account" + result, secondErr := svc.ForwardAsAnthropic(context.Background(), secondContext, &secondAccount, body, "", "") + require.NoError(t, secondErr) + require.NotNil(t, result) + require.Equal(t, "resp_second", result.ResponseID) + require.NotEmpty(t, secondRec.Body.String()) +} + +func TestFailoverOpenAIUpstreamHTTPError_NilContextSkipsTempUnschedulablePolicy(t *testing.T) { + repo := &tempUnschedulableOpenAIAccountRepo{} + svc := &OpenAIGatewayService{ + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + } + account := &Account{ + ID: 5099, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "temp_unschedulable_enabled": true, + "temp_unschedulable_rules": []any{map[string]any{ + "error_code": float64(http.StatusBadRequest), + "keywords": []any{"servers are currently overloaded"}, + "duration_minutes": float64(1), + }}, + }, + } + body := []byte(`{"error":{"message":"Our servers are currently overloaded."}}`) + resp := &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{}} + + got := svc.failoverOpenAIUpstreamHTTPError( + context.Background(), nil, account, resp, body, + "Our servers are currently overloaded.", "gpt-5.4", + ) + + require.Nil(t, got) + require.Zero(t, repo.modelRateLimitAccountID) + require.Empty(t, repo.modelRateLimitKey) +} + type groupAwareStubOpenAIAccountRepo struct { stubOpenAIAccountRepo }