diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 408020f0a6..6b10f7bcbe 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -166,6 +166,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState switchCount := 0 maxAccountSwitches := h.maxAccountSwitches if maxAccountSwitches <= 0 { @@ -300,6 +301,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } switchCount++ + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { + h.handleFailoverExhausted(c, failoverErr, false) + return + } reqLog.Warn("grok_media.upstream_failover_switching", zap.Int64("account_id", account.ID), zap.Int("upstream_status", failoverErr.StatusCode), diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index a808145235..4f5c07d8c5 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -106,6 +106,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError switchCount := 0 + var oauth429FailoverState service.OpenAIOAuth429FailoverState routingStart := time.Now() for { @@ -188,7 +189,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, false) return } diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index fc58b3a25a..004aff2bb7 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -133,6 +133,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { if c.Request.Context().Err() != nil { @@ -275,7 +276,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } diff --git a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go index e7030d9d2f..42a6bf6cd3 100644 --- a/backend/internal/handler/openai_gateway_credential_failover_loop_test.go +++ b/backend/internal/handler/openai_gateway_credential_failover_loop_test.go @@ -31,6 +31,7 @@ type grokCredentialHandlerRepo struct { accounts []service.Account setErrorIDs []int64 setTempIDs []int64 + rateLimitIDs []int64 updateExtraIDs []int64 selectionCalls int setErrorErr error @@ -108,6 +109,34 @@ func (r *grokCredentialHandlerRepo) SetTempUnschedulable(_ context.Context, id i return nil } +func (r *grokCredentialHandlerRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + r.mu.Lock() + defer r.mu.Unlock() + r.rateLimitIDs = append(r.rateLimitIDs, id) + for i := range r.accounts { + if r.accounts[i].ID != id { + continue + } + now := time.Now() + r.accounts[i].RateLimitedAt = &now + value := resetAt + r.accounts[i].RateLimitResetAt = &value + } + return nil +} + +func (r *grokCredentialHandlerRepo) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + r.mu.Lock() + for i := range r.accounts { + if r.accounts[i].ID == id && r.accounts[i].RateLimitResetAt != nil && !resetAt.After(*r.accounts[i].RateLimitResetAt) { + r.mu.Unlock() + return nil + } + } + r.mu.Unlock() + return r.SetRateLimited(ctx, id, resetAt) +} + func (r *grokCredentialHandlerRepo) SetGrokCredentialErrorIfMatch( _ context.Context, id int64, @@ -204,6 +233,12 @@ func (r *grokCredentialHandlerRepo) selectorCalls() int { return r.selectionCalls } +func (r *grokCredentialHandlerRepo) rateLimitedAccountIDs() []int64 { + r.mu.Lock() + defer r.mu.Unlock() + return append([]int64(nil), r.rateLimitIDs...) +} + type grokCredentialHandlerTokenCache struct { service.GrokTokenCache mu sync.Mutex @@ -282,6 +317,8 @@ type grokCredentialHandlerUpstream struct { requestURLs []string authorization []string failAccountID int64 + rateLimitIDs map[int64]bool + failureStatus map[int64]int cancelRequest context.CancelFunc } @@ -295,8 +332,27 @@ func (u *grokCredentialHandlerUpstream) Do(req *http.Request, _ string, accountI u.requestURLs = append(u.requestURLs, req.URL.String()) u.authorization = append(u.authorization, req.Header.Get("Authorization")) failAccountID := u.failAccountID + rateLimited := u.rateLimitIDs[accountID] + failureStatus := u.failureStatus[accountID] cancelRequest := u.cancelRequest u.mu.Unlock() + if rateLimited { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Retry-After": []string{"60"}, + }, + Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"rate limited"}}`)), + }, nil + } + if failureStatus > 0 { + return &http.Response{ + StatusCode: failureStatus, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"upstream unavailable"}}`)), + }, nil + } if accountID == failAccountID { if cancelRequest != nil { cancelRequest() @@ -502,6 +558,115 @@ func TestResponsesCredentialFailoverLoop(t *testing.T) { }) } +func TestResponsesGrok429FailoverIsBounded(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("first rate limited account selects healthy account", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Contains(t, recorder.Body.String(), "resp_healthy") + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.Equal(t, []int64{801}, repo.rateLimitedAccountIDs()) + }) + + t.Run("two rate limited accounts stop without sweeping the pool", func(t *testing.T) { + _, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.Equal(t, []int64{801, 802}, repo.rateLimitedAccountIDs()) + require.NotContains(t, recorder.Body.String(), "expired") + require.NotContains(t, recorder.Body.String(), "healthy-access") + require.NotContains(t, recorder.Body.String(), "rate limited") + }) +} + +func TestResponsesGrok429FailoverHandlesMixedStatuses(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("429 then 500 stops after the bounded followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_429_500") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusBadGateway, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.NotContains(t, recorder.Body.String(), "upstream unavailable") + }) + + t.Run("500 then 429 permits one healthy followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_500_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802, 803}, upstream.accountHits()) + }) + + t.Run("OAuth 429 then API-key failure cannot bypass the bound", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "oauth_429_apikey_500") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)) + req.Header.Set("Content-Type", "application/json") + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusBadGateway, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + }) +} + +func TestGrokMedia429FailoverIsBounded(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("first 429 selects one healthy followup", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil) + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + }) + + t.Run("second 429 stops without sweeping a third account", func(t *testing.T) { + _, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429") + defer cleanup() + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil) + + router.ServeHTTP(recorder, req) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code, recorder.Body.String()) + require.Equal(t, []int64{801, 802}, upstream.accountHits()) + require.NotContains(t, recorder.Body.String(), "rate limited") + }) +} + func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) { gin.SetMode(gin.TestMode) endpoints := []struct { @@ -667,9 +832,23 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway }, }, } - if mode == "postmap_cancel" { + if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" { accounts[0].Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339) } + if mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" { + accounts = append(accounts, service.Account{ + ID: 803, Name: "untried-healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 3, + Credentials: map[string]any{ + "access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh", + "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339), + }, + }) + } + if mode == "oauth_429_apikey_500" { + accounts[1].Type = service.AccountTypeAPIKey + accounts[1].Credentials = map[string]any{"api_key": "third-party-key"} + } if mode == "all_revoked" { accounts[1].Credentials["expires_at"] = time.Now().Add(-time.Minute).UTC().Format(time.RFC3339) } @@ -694,6 +873,21 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway provider.SetRefreshAPI(service.NewOAuthRefreshAPI(repo, tokenCache), refresher) } upstream := &grokCredentialHandlerUpstream{} + switch mode { + case "first_429": + upstream.rateLimitIDs = map[int64]bool{801: true} + case "all_429": + upstream.rateLimitIDs = map[int64]bool{801: true, 802: true} + case "mixed_429_500": + upstream.rateLimitIDs = map[int64]bool{801: true} + upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError} + case "mixed_500_429": + upstream.failureStatus = map[int64]int{801: http.StatusInternalServerError} + upstream.rateLimitIDs = map[int64]bool{802: true} + case "oauth_429_apikey_500": + upstream.rateLimitIDs = map[int64]bool{801: true} + upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError} + } cfg := &config.Config{RunMode: config.RunModeSimple} cfg.Gateway.MaxAccountSwitches = 3 billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 6db72c270b..33bef42cef 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -332,6 +332,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { if c.Request.Context().Err() != nil { @@ -487,7 +488,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } @@ -844,6 +845,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { failedAccountIDs := make(map[int64]struct{}) sameAccountRetryCount := make(map[int64]int) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState effectiveMappedModel := preferredMappedModel for { @@ -991,7 +993,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted) return } @@ -1496,6 +1498,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { switchCount := 0 failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError + var oauth429FailoverState service.OpenAIOAuth429FailoverState handleWSFailover := func(account *service.Account, failoverErr *service.UpstreamFailoverError) bool { if ctx.Err() != nil { return false @@ -1519,7 +1522,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return false } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { closeOpenAIWSFailoverExhausted(wsConn, failoverErr) return false } diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index c5982fb7d1..d18af28966 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -145,6 +145,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { stopJSONKeepalive := func() {} jsonKeepaliveStarted := false defer func() { stopJSONKeepalive() }() + var oauth429FailoverState service.OpenAIOAuth429FailoverState for { reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs))) @@ -299,7 +300,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { return } switchCount++ - if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) { h.handleFailoverExhausted(c, failoverErr, streamStarted) return } diff --git a/backend/internal/pkg/xai/oauth.go b/backend/internal/pkg/xai/oauth.go index a8a549c9dc..ea837fd78a 100644 --- a/backend/internal/pkg/xai/oauth.go +++ b/backend/internal/pkg/xai/oauth.go @@ -5,6 +5,7 @@ import ( "crypto/sha256" "encoding/base64" "encoding/hex" + "errors" "fmt" "net/url" "os" @@ -164,6 +165,24 @@ func ValidatedBaseURL(override string) (string, error) { return ValidateBaseURL(EffectiveBaseURL(override)) } +// BaseURLValidator applies the caller's outbound URL trust policy before xAI +// endpoint paths are appended. The service layer uses this for API-key accounts +// so the global security.url_allowlist policy remains the single source of +// truth; OAuth callers keep using the strict trusted-host validator. +type BaseURLValidator func(string) (string, error) + +func validatedBaseURLWithValidator(override string, validator BaseURLValidator) (string, error) { + if validator == nil { + return ValidatedBaseURL(override) + } + raw := EffectiveBaseURL(override) + validated, err := validator(raw) + if err != nil { + return "", err + } + return normalizeKnownBaseURLPath(validated) +} + type RuntimeSanityCheck struct { Value string `json:"value"` Valid bool `json:"valid"` @@ -282,7 +301,16 @@ func ValidateTrustedBaseURL(raw string) (string, error) { func normalizeKnownBaseURLPath(raw string) (string, error) { parsed, err := url.Parse(raw) if err != nil || parsed.Scheme == "" || parsed.Host == "" { - return "", fmt.Errorf("invalid url: %s", raw) + return "", errors.New("invalid base URL") + } + if parsed.User != nil { + return "", errors.New("base URL must not include userinfo") + } + if parsed.RawQuery != "" { + return "", errors.New("base URL must not include a query") + } + if parsed.Fragment != "" { + return "", errors.New("base URL must not include a fragment") } path := strings.TrimRight(parsed.Path, "/") if path == "" { @@ -435,7 +463,11 @@ func ParseAuthorizationInput(raw string) AuthorizationInput { } func BuildResponsesURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildResponsesURLWithValidator(baseURL, nil) +} + +func BuildResponsesURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -443,7 +475,11 @@ func BuildResponsesURL(baseURL string) (string, error) { } func BuildChatCompletionsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildChatCompletionsURLWithValidator(baseURL, nil) +} + +func BuildChatCompletionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -451,7 +487,11 @@ func BuildChatCompletionsURL(baseURL string) (string, error) { } func BuildImagesGenerationsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildImagesGenerationsURLWithValidator(baseURL, nil) +} + +func BuildImagesGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -459,7 +499,11 @@ func BuildImagesGenerationsURL(baseURL string) (string, error) { } func BuildImagesEditsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildImagesEditsURLWithValidator(baseURL, nil) +} + +func BuildImagesEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -467,7 +511,11 @@ func BuildImagesEditsURL(baseURL string) (string, error) { } func BuildVideosGenerationsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosGenerationsURLWithValidator(baseURL, nil) +} + +func BuildVideosGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -475,7 +523,11 @@ func BuildVideosGenerationsURL(baseURL string) (string, error) { } func BuildVideosEditsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosEditsURLWithValidator(baseURL, nil) +} + +func BuildVideosEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -483,7 +535,11 @@ func BuildVideosEditsURL(baseURL string) (string, error) { } func BuildVideosExtensionsURL(baseURL string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideosExtensionsURLWithValidator(baseURL, nil) +} + +func BuildVideosExtensionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } @@ -491,7 +547,11 @@ func BuildVideosExtensionsURL(baseURL string) (string, error) { } func BuildVideoURL(baseURL, requestID string) (string, error) { - validatedBaseURL, err := ValidatedBaseURL(baseURL) + return BuildVideoURLWithValidator(baseURL, requestID, nil) +} + +func BuildVideoURLWithValidator(baseURL, requestID string, validator BaseURLValidator) (string, error) { + validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator) if err != nil { return "", fmt.Errorf("invalid base url: %w", err) } diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go index 39200ffc39..a82a26e299 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -6,6 +6,7 @@ import ( "net/url" "testing" + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" "github.com/stretchr/testify/require" ) @@ -165,6 +166,44 @@ func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) { require.Error(t, err) } +func TestBuildResponsesURLWithValidatorUsesCallerPolicy(t *testing.T) { + validator := func(raw string) (string, error) { + return urlvalidator.ValidateURLFormat(raw, true) + } + + target, err := BuildResponsesURLWithValidator("http://grok.example.test/v1/", validator) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", target) +} + +func TestBuildResponsesURLPreservesUnsafeOverrideCustomPath(t *testing.T) { + t.Setenv(EnvAllowUnsafeURLOverrides, "true") + + target, err := BuildResponsesURL("http://localhost:8080/custom") + require.NoError(t, err) + require.Equal(t, "http://localhost:8080/custom/responses", target) +} + +func TestBuildResponsesURLWithValidatorRejectsBaseURLComponents(t *testing.T) { + permissive := func(raw string) (string, error) { return raw, nil } + tests := []struct { + name string + raw string + }{ + {name: "userinfo", raw: "https://user:secret@grok.example.test/v1"}, + {name: "query", raw: "https://grok.example.test/v1?token=secret"}, + {name: "fragment", raw: "https://grok.example.test/v1#secret"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := BuildResponsesURLWithValidator(tt.raw, permissive) + require.Error(t, err) + require.NotContains(t, err.Error(), "secret") + }) + } +} + func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) { t.Setenv(EnvAllowUnsafeURLOverrides, "true") diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index a3eba2eaa9..8527b7010d 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1377,6 +1377,35 @@ func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, return nil } +// ClearRateLimitIfObserved clears exactly the Grok rate-limit generation seen +// by a successful request. Matching both timestamps prevents a stale success +// from erasing a later clear/re-arm generation with an equal or shorter reset. +func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) { + updated, err := r.client.Account.Update(). + Where( + dbaccount.IDEQ(id), + dbaccount.PlatformEQ(service.PlatformGrok), + dbaccount.TypeEQ(service.AccountTypeOAuth), + dbaccount.RateLimitedAtEQ(observedLimitedAt), + dbaccount.RateLimitResetAtEQ(observedResetAt), + ). + ClearRateLimitedAt(). + ClearRateLimitResetAt(). + Save(ctx) + if err != nil { + return false, err + } + if updated == 0 { + r.syncSchedulerAccountSnapshot(ctx, id) + return false, nil + } + if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil { + logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue observed rate-limit clear failed: account=%d err=%v", id, err) + } + r.syncSchedulerAccountSnapshot(ctx, id) + return true, nil +} + func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error { if scope == "" { return nil diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index f82ad4ca28..ae519f3df4 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -722,6 +722,56 @@ func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() { s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second) } +func (s *AccountRepoSuite) TestClearRateLimitIfObservedProtectsRearmed429Generation() { + account := mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "acc-rl-conditional-clear", + Platform: service.PlatformGrok, + Type: service.AccountTypeOAuth, + }) + firstReset := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second) + rearmedReset := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second) + + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, firstReset)) + staleGeneration, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(staleGeneration.RateLimitedAt) + s.Require().NotNil(staleGeneration.RateLimitResetAt) + cleared, err := s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt) + s.Require().NoError(err) + s.Require().True(cleared) + + // A newer generation may legitimately re-arm a shorter boundary after the + // first generation was cleared. The stale success must not erase it. + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, rearmedReset)) + cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt) + s.Require().NoError(err) + s.Require().False(cleared) + + got, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(got.RateLimitedAt) + s.Require().NotNil(got.RateLimitResetAt) + s.Require().WithinDuration(rearmedReset, *got.RateLimitResetAt, time.Second) + + // An admin can retype the row while the successful OAuth request is still + // in flight. The stale OAuth recovery must not cross into API-key state even + // when both observed timestamps still match. + _, err = s.client.Account.UpdateOneID(account.ID). + SetType(service.AccountTypeAPIKey). + Save(s.ctx) + s.Require().NoError(err) + cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *got.RateLimitedAt, *got.RateLimitResetAt) + s.Require().NoError(err) + s.Require().False(cleared) + + retyped, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().Equal(service.AccountTypeAPIKey, retyped.Type) + s.Require().NotNil(retyped.RateLimitedAt) + s.Require().NotNil(retyped.RateLimitResetAt) + s.Require().WithinDuration(rearmedReset, *retyped.RateLimitResetAt, time.Second) +} + func (s *AccountRepoSuite) TestClearRateLimit() { account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"}) until := time.Now().Add(1 * time.Hour) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index a9eef998e4..a739bcc120 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -695,7 +695,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Grok account type: %s", account.Type)) } - apiURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) + apiURL, err := buildGrokResponsesURL(account, s.cfg) if err != nil { return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok base URL: %s", err.Error())) } @@ -742,7 +742,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * now := time.Now() snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now) if snapshot != nil && s.accountRepo != nil { - resetAt, limited := grokRateLimitResetAt(snapshot, now) + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) } @@ -751,7 +751,11 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * }) if limited { persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } + } else if s.accountRepo != nil && isSuccessfulGrokRateLimitRecovery(account, &xai.QuotaSnapshot{StatusCode: resp.StatusCode}) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } if resp.StatusCode != http.StatusOK { diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 5cc8e2d05e..f7c6bd8153 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -13,7 +13,6 @@ import ( "strings" "time" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" @@ -268,25 +267,6 @@ func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Cont return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID) } -func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error) { - switch e { - case GrokMediaEndpointImagesGenerations: - return xai.BuildImagesGenerationsURL(baseURL) - case GrokMediaEndpointImagesEdits: - return xai.BuildImagesEditsURL(baseURL) - case GrokMediaEndpointVideosGenerations: - return xai.BuildVideosGenerationsURL(baseURL) - case GrokMediaEndpointVideosEdits: - return xai.BuildVideosEditsURL(baseURL) - case GrokMediaEndpointVideosExtensions: - return xai.BuildVideosExtensionsURL(baseURL) - case GrokMediaEndpointVideoStatus: - return xai.BuildVideoURL(baseURL, requestID) - default: - return "", fmt.Errorf("unsupported grok media endpoint: %s", e) - } -} - func (s *OpenAIGatewayService) ForwardGrokMedia( ctx context.Context, c *gin.Context, @@ -308,7 +288,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( if err != nil { return nil, err } - targetURL, err := endpoint.upstreamURL(account.GetGrokMediaBaseURL(), requestID) + targetURL, err := buildGrokMediaURL(account, s.cfg, endpoint, requestID) if err != nil { return nil, err } @@ -368,7 +348,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index ef8e3e4c03..a95b778707 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -135,7 +135,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err) } - targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) + targetURL, err := buildGrokResponsesURL(account, nil) if err != nil { return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err) } @@ -160,7 +160,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr defer func() { _ = resp.Body.Close() }() snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe") - resetAt, limited := grokRateLimitResetAt(snapshot, time.Now()) + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, time.Now()) if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) } @@ -169,6 +169,8 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr }) if limited { persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } else if isSuccessfulGrokRateLimitRecovery(account, snapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } result := &GrokQuotaProbeResult{ diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 4e06fe156c..ad3c85ba19 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -32,6 +32,10 @@ type grokQuotaAccountRepo struct { lastTempUnschedID int64 lastTempUnschedUntil time.Time lastTempUnschedReason string + recoveryClearCalls int + recoveryObservedAt time.Time + recoveryObservedReset time.Time + recoveryClearResult bool } func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { @@ -54,6 +58,13 @@ func (r *grokQuotaAccountRepo) SetRateLimitedIfLater(ctx context.Context, id int return r.SetRateLimited(ctx, id, resetAt) } +func (r *grokQuotaAccountRepo) ClearRateLimitIfObserved(_ context.Context, _ int64, observedLimitedAt, observedResetAt time.Time) (bool, error) { + r.recoveryClearCalls++ + r.recoveryObservedAt = observedLimitedAt + r.recoveryObservedReset = observedResetAt + return r.recoveryClearResult, nil +} + func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, reason string) error { r.tempUnschedCalls++ r.lastTempUnschedID = id @@ -375,10 +386,15 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { t.Parallel() account := healthyGrokQuotaOAuthAccount(45) + observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + account.RateLimitedAt = &observedLimitedAt + account.RateLimitResetAt = &observedResetAt repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{45: account}, }, + recoveryClearResult: true, } upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -401,6 +417,9 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) { require.True(t, ok) require.False(t, stored.HeadersObserved) require.Equal(t, http.StatusOK, stored.StatusCode) + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, observedLimitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) } func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { diff --git a/backend/internal/service/grok_upstream_url.go b/backend/internal/service/grok_upstream_url.go new file mode 100644 index 0000000000..ca2ba6b64d --- /dev/null +++ b/backend/internal/service/grok_upstream_url.go @@ -0,0 +1,90 @@ +package service + +import ( + "errors" + "fmt" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" +) + +func grokBaseURLValidator(account *Account, cfg *config.Config) (xai.BaseURLValidator, error) { + if account == nil || !account.IsGrok() { + return nil, fmt.Errorf("grok account is required") + } + switch account.Type { + case AccountTypeOAuth: + // Subscription credentials are never governed by the operator's API-key + // URL policy. They stay pinned to the supported CLI gateway. + return redactedGrokBaseURLValidator(xai.ValidateTrustedBaseURL), nil + case AccountTypeAPIKey: + if cfg == nil { + return redactedGrokBaseURLValidator(xai.ValidateBaseURL), nil + } + if !cfg.Security.URLAllowlist.Enabled { + return redactedGrokBaseURLValidator(func(raw string) (string, error) { + return urlvalidator.ValidateURLFormat(raw, cfg.Security.URLAllowlist.AllowInsecureHTTP) + }), nil + } + return redactedGrokBaseURLValidator(func(raw string) (string, error) { + return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{ + AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts, + RequireAllowlist: true, + AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts, + }) + }), nil + default: + return nil, fmt.Errorf("unsupported grok account type: %s", account.Type) + } +} + +func redactedGrokBaseURLValidator(validator xai.BaseURLValidator) xai.BaseURLValidator { + return func(raw string) (string, error) { + validated, err := validator(raw) + if err != nil { + return "", errors.New("base URL rejected by URL security policy") + } + return validated, nil + } +} + +func buildGrokResponsesURL(account *Account, cfg *config.Config) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + return xai.BuildResponsesURLWithValidator(account.GetGrokBaseURL(), validator) +} + +func buildGrokChatCompletionsURL(account *Account, cfg *config.Config) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + return xai.BuildChatCompletionsURLWithValidator(account.GetGrokBaseURL(), validator) +} + +func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaEndpoint, requestID string) (string, error) { + validator, err := grokBaseURLValidator(account, cfg) + if err != nil { + return "", err + } + baseURL := account.GetGrokMediaBaseURL() + switch endpoint { + case GrokMediaEndpointImagesGenerations: + return xai.BuildImagesGenerationsURLWithValidator(baseURL, validator) + case GrokMediaEndpointImagesEdits: + return xai.BuildImagesEditsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosGenerations: + return xai.BuildVideosGenerationsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosEdits: + return xai.BuildVideosEditsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideosExtensions: + return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator) + case GrokMediaEndpointVideoStatus: + return xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + default: + return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint) + } +} diff --git a/backend/internal/service/grok_upstream_url_test.go b/backend/internal/service/grok_upstream_url_test.go new file mode 100644 index 0000000000..7919e47202 --- /dev/null +++ b/backend/internal/service/grok_upstream_url_test.go @@ -0,0 +1,122 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" +) + +func TestGrokAPIKeyURLPolicyFollowsGlobalSecurityConfig(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "http://grok.example.test/v1", + }, + } + + t.Run("insecure HTTP enabled with allowlist disabled", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + responsesURL, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", responsesURL) + + chatURL, err := buildGrokChatCompletionsURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/chat/completions", chatURL) + + mediaURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointImagesGenerations, "") + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/images/generations", mediaURL) + }) + + t.Run("insecure HTTP disabled", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = false + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + }) + + t.Run("enabled allowlist remains HTTPS only", func(t *testing.T) { + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = true + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"} + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + }) +} + +func TestGrokAPIKeyURLPolicyAppliesAllowlistAndPrivateHostControls(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "https://grok.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = true + cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"} + + target, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "https://grok.example.test/v1/responses", target) + + cfg.Security.URLAllowlist.UpstreamHosts = []string{"other.example.test"} + _, err = buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + + account.Credentials["base_url"] = "https://127.0.0.1/v1" + cfg.Security.URLAllowlist.UpstreamHosts = []string{"127.0.0.1"} + _, err = buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + + cfg.Security.URLAllowlist.AllowPrivateHosts = true + target, err = buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, "https://127.0.0.1/v1/responses", target) +} + +func TestGrokAPIKeyURLPolicyRedactsMalformedConfiguredURL(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "base_url": "https://%zz:secret@grok.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + _, err := buildGrokResponsesURL(account, cfg) + require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy") + require.NotContains(t, err.Error(), "secret") +} + +func TestGrokOAuthURLPolicyIgnoresAPIKeyOverrides(t *testing.T) { + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "base_url": "http://attacker.example.test/v1", + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + + target, err := buildGrokResponsesURL(account, cfg) + require.NoError(t, err) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 978b740a56..00d971ba22 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -16,6 +16,13 @@ const ( openAIOAuth429StormMaxAccountSwitches = 1 ) +// OpenAIOAuth429FailoverState tracks the request-local follow-up budget after +// the first Grok OAuth 429. Once that 429 occurs, exactly one different account +// may be attempted; any failure from that follow-up account ends failover. +type OpenAIOAuth429FailoverState struct { + grokOAuth429FollowupPending bool +} + func openAIAccountStateContext(ctx context.Context) (context.Context, context.CancelFunc) { base := context.Background() if ctx != nil { @@ -210,14 +217,28 @@ func (s *OpenAIGatewayService) isOpenAIOAuth429Storm() bool { return s.openaiOAuth429WindowCount.Load() >= openAIOAuth429StormThreshold } -func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int) bool { - if statusCode != http.StatusTooManyRequests || failedSwitches < openAIOAuth429StormMaxAccountSwitches { +func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int, state *OpenAIOAuth429FailoverState) bool { + if failedSwitches < openAIOAuth429StormMaxAccountSwitches { return false } - if isGrokOAuthAccount(account) { + if state != nil && state.grokOAuth429FollowupPending { + // The follow-up budget was armed by a Grok OAuth 429. Consume it on + // any failing follow-up account, even if a mixed pool selected an API-key + // account next. return true } - if !isOpenAIOAuthAccount(account) { + if isGrokOAuthAccount(account) { + if state == nil { + // Preserve the old threshold for callers that have not adopted the + // request-local state contract yet. + return statusCode == http.StatusTooManyRequests && failedSwitches >= 2 + } + if statusCode == http.StatusTooManyRequests { + state.grokOAuth429FollowupPending = true + } + return false + } + if statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) { return false } return s.isOpenAIOAuth429Storm() diff --git a/backend/internal/service/openai_account_runtime_block_fastpath_test.go b/backend/internal/service/openai_account_runtime_block_fastpath_test.go index a51c13c531..7d7cfbda50 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath_test.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath_test.go @@ -135,26 +135,45 @@ func TestShouldStopOpenAIOAuth429Failover_OnlyDuringStorm(t *testing.T) { svc := &OpenAIGatewayService{} account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth} apiKeyAccount := &Account{ID: 43, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + var state OpenAIOAuth429FailoverState - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) for i := 0; i < openAIOAuth429StormThreshold; i++ { svc.recordOpenAIOAuth429() } - require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state)) } -func TestShouldStopOpenAIOAuth429Failover_StopsGrokAfterFirst429Switch(t *testing.T) { +func TestShouldStopOpenAIOAuth429Failover_TracksOneGrokFollowupAttempt(t *testing.T) { svc := &OpenAIGatewayService{} account := &Account{ID: 44, Platform: PlatformGrok, Type: AccountTypeOAuth} apiKeyAccount := &Account{ID: 45, Platform: PlatformGrok, Type: AccountTypeAPIKey} - require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1)) - require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1)) + t.Run("429 then 500 stops after one followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 2, &state)) + }) + + t.Run("500 then 429 still allows one followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 2, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusBadGateway, 3, &state)) + }) + + t.Run("OAuth 429 then API-key failure consumes the same followup", func(t *testing.T) { + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state)) + require.True(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusInternalServerError, 2, &state)) + }) + + var state OpenAIOAuth429FailoverState + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state)) + require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 2, &state)) } diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 095d544f89..800d53be32 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -9,7 +9,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" @@ -202,7 +201,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } // 8. Forward response @@ -222,7 +221,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) { if account.Platform == PlatformGrok { - targetURL, err := xai.BuildChatCompletionsURL(account.GetGrokBaseURL()) + targetURL, err := buildGrokChatCompletionsURL(account, s.cfg) if err != nil { return "", fmt.Errorf("invalid grok base_url: %w", err) } diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 675b4bc9f9..92cf2910c6 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" @@ -25,6 +26,10 @@ const ( grokCLIVersion = "0.2.93" grokDefaultResponsesModel = "grok-4.5" grokRateLimitFallbackCooldown = 2 * time.Minute + grokRateLimitRepeatCooldown = 10 * time.Minute + grokRateLimitSustainedCooldown = 30 * time.Minute + grokRateLimitMaxAdaptiveCooldown = time.Hour + grokRateLimitBackoffQuietPeriod = time.Hour ) func (s *OpenAIGatewayService) forwardGrokResponses( @@ -61,7 +66,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) defer releaseUpstreamCtx() - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) if err != nil { return nil, err } @@ -107,7 +112,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) var usage *OpenAIUsage var firstTokenMs *int @@ -574,7 +579,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) // Image-description probes are auxiliary requests, not conversation turns. // Do not bind them to the caller's Grok prompt-cache identity. - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "") + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "", s.cfg) releaseUpstreamCtx() if err != nil { return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err) @@ -618,7 +623,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) if err != nil { return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) @@ -740,8 +745,8 @@ func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) { dst.ImageOutputTokens += usage.ImageOutputTokens } -func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string) (*http.Request, error) { - targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) +func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string, cfg *config.Config) (*http.Request, error) { + targetURL, err := buildGrokResponsesURL(account, cfg) if err != nil { return nil, err } @@ -780,11 +785,12 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco } accountID := account.ID now := time.Now() - resetAt, hasActiveLimit := grokRateLimitResetAt(snapshot, now) + resetAt, hasActiveLimit := grokRateLimitResetAtForAccount(account, snapshot, now) if hasActiveLimit { normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) } - critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit + recovery := isSuccessfulGrokRateLimitRecovery(account, snapshot) + critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit || recovery if s.codexSnapshotThrottle != nil { allowed := s.codexSnapshotThrottle.Allow(accountID, now) if !critical && !allowed { @@ -810,6 +816,23 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco // on the passive snapshot scheduler check. if hasActiveLimit { s.rateLimitGrok(stateCtx, account, resetAt) + } else if recovery { + clearGrokRateLimitAfterRecovery(stateCtx, s.accountRepo, account) + } +} + +func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) { + snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now()) + if snapshot != nil { + s.updateGrokUsageSnapshot(ctx, account, snapshot) + return + } + // Successful responses are recovery evidence even when the upstream omits + // optional quota headers. Do not replace an informative stored snapshot with + // an empty one; only clear the exact observed cooldown generation. + recoverySnapshot := &xai.QuotaSnapshot{StatusCode: statusCode} + if isSuccessfulGrokRateLimitRecovery(account, recoverySnapshot) { + clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account) } } @@ -900,6 +923,37 @@ func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time return time.Time{}, false } +func grokRateLimitResetAtForAccount(account *Account, snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) { + resetAt, limited := grokRateLimitResetAt(snapshot, now) + if !limited || !isGrokOAuthAccount(account) || snapshot == nil || snapshot.StatusCode != http.StatusTooManyRequests { + return resetAt, limited + } + if account.RateLimitedAt == nil || account.RateLimitResetAt == nil { + return resetAt, true + } + previousResetAt := *account.RateLimitResetAt + if previousResetAt.After(now) || now.Sub(previousResetAt) > grokRateLimitBackoffQuietPeriod { + return resetAt, true + } + previousCooldown := previousResetAt.Sub(*account.RateLimitedAt) + if previousCooldown <= 0 { + return resetAt, true + } + + adaptiveCooldown := grokRateLimitRepeatCooldown + switch { + case previousCooldown >= grokRateLimitSustainedCooldown: + adaptiveCooldown = grokRateLimitMaxAdaptiveCooldown + case previousCooldown >= grokRateLimitRepeatCooldown: + adaptiveCooldown = grokRateLimitSustainedCooldown + } + adaptiveResetAt := now.Add(adaptiveCooldown) + if adaptiveResetAt.After(resetAt) { + resetAt = adaptiveResetAt + } + return resetAt, true +} + func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time { if !resetAt.After(now) { resetAt = now.Add(grokRateLimitFallbackCooldown) @@ -914,6 +968,33 @@ type grokRateLimitExtendingRepository interface { SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error } +type grokRateLimitRecoveryRepository interface { + ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) +} + +func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnapshot) bool { + return isGrokOAuthAccount(account) && + account.RateLimitedAt != nil && + account.RateLimitResetAt != nil && + snapshot != nil && + snapshot.StatusCode >= http.StatusOK && + snapshot.StatusCode < http.StatusMultipleChoices +} + +func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) { + if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil { + return + } + recoveryRepo, ok := repo.(grokRateLimitRecoveryRepository) + if !ok { + return + } + _, err := recoveryRepo.ClearRateLimitIfObserved(ctx, account.ID, *account.RateLimitedAt, *account.RateLimitResetAt) + if err != nil { + slog.Warn("grok_rate_limit_recovery_clear_failed", "account_id", account.ID, "error", err) + } +} + func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) { if repo == nil || account == nil || account.ID <= 0 { return diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index fe9cfb09fd..6fd58b2db1 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -10,7 +10,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" ) @@ -258,7 +257,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return nil, fmt.Errorf("get grok access token: %w", err) } upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity, s.cfg) releaseUpstreamCtx() if err != nil { return nil, fmt.Errorf("build grok responses bridge request: %w", err) @@ -301,7 +300,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel) } - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) var result *OpenAIForwardResult if clientStream { diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 4f7ed05edf..d93415b2e0 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -256,7 +256,7 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil) require.NoError(t, err) require.Equal(t, http.MethodPost, req.Method) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) @@ -280,7 +280,7 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing. }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "", nil) require.NoError(t, err) require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String()) require.Equal(t, "Bearer api-key", req.Header.Get("Authorization")) @@ -299,7 +299,7 @@ func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "", nil) require.NoError(t, err) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String()) } @@ -1077,6 +1077,41 @@ func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) { require.Contains(t, recorder.Body.String(), `"type":"test_complete"`) } +func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 55, + Name: "grok-api-key-http", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "third-party-key", + "base_url": "http://grok.example.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + + "data: {\"type\":\"response.completed\"}\n\n", + )), + }} + svc := &AccountTestService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/55/test", nil) + + err := svc.testGrokAccountConnection(c, account, "grok") + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.Contains(t, recorder.Body.String(), `"type":"test_complete"`) +} + func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) { gin.SetMode(gin.TestMode) @@ -1144,10 +1179,15 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test c.Set("api_key", &APIKey{ID: 5202}) account := healthyGrokOAuthGatewayTestAccount(56, "access-token") + observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + account.RateLimitedAt = &observedLimitedAt + account.RateLimitResetAt = &observedResetAt repo := &grokQuotaAccountRepo{ mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ accountsByID: map[int64]*Account{56: account}, }, + recoveryClearResult: true, } upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -1180,6 +1220,9 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String()) + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, observedLimitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) } func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) { @@ -1565,6 +1608,102 @@ func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokRateLimitResetAtForAccountEscalatesRepeated429s(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + tests := []struct { + name string + previousCooldown time.Duration + wantCooldown time.Duration + }{ + {name: "repeat after short boundary", previousCooldown: 45 * time.Second, wantCooldown: grokRateLimitRepeatCooldown}, + {name: "sustained repeat", previousCooldown: grokRateLimitRepeatCooldown, wantCooldown: grokRateLimitSustainedCooldown}, + {name: "capped repeat", previousCooldown: grokRateLimitSustainedCooldown, wantCooldown: grokRateLimitMaxAdaptiveCooldown}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + previousReset := now.Add(-time.Second) + previousLimited := previousReset.Add(-tt.previousCooldown) + account := &Account{ + ID: 630, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + + require.True(t, limited) + require.WithinDuration(t, now.Add(tt.wantCooldown), resetAt, time.Second) + }) + } +} + +func TestGrokRateLimitResetAtForAccountPreservesAuthoritativeAndQuietRecovery(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + previousReset := now.Add(-grokRateLimitBackoffQuietPeriod - time.Second) + previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown) + account := &Account{ + ID: 631, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second) + + authoritativeReset := now.Add(2 * time.Hour) + remaining := int64(0) + snapshot.Requests = &xai.QuotaWindow{Remaining: &remaining, ResetUnix: grokInt64PtrForTest(authoritativeReset.Unix())} + recentReset := now.Add(-time.Second) + recentLimited := recentReset.Add(-grokRateLimitSustainedCooldown) + account.RateLimitResetAt = &recentReset + account.RateLimitedAt = &recentLimited + + resetAt, limited = grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, authoritativeReset, resetAt, time.Second) +} + +func TestGrokRateLimitResetAtForAccountLeavesAPIKey429PolicyUnchanged(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + retryAfter := 45 + previousReset := now.Add(-time.Second) + previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown) + account := &Account{ + ID: 632, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + RateLimitedAt: &previousLimited, + RateLimitResetAt: &previousReset, + } + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + RetryAfterSeconds: &retryAfter, + UpdatedAt: now.Format(time.RFC3339), + } + + resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) + require.True(t, limited) + require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second) +} + func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) { now := time.Now().UTC().Truncate(time.Second) observedAt := now.Add(-2 * time.Minute) @@ -1666,6 +1805,72 @@ func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing require.Zero(t, repo.rateLimitedCalls) } +func TestUpdateGrokUsageFromResponseHeaderlessSuccessClearsObservedCooldown(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + limitedAt := now.Add(-grokRateLimitRepeatCooldown) + observedResetAt := now.Add(-time.Second) + account := &Account{ + ID: 660, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &limitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{ + accountRepo: repo, + codexSnapshotThrottle: newAccountWriteThrottle(time.Hour), + } + + svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK) + + require.Zero(t, repo.updateCalls, "headerless success must not overwrite an informative quota snapshot") + require.Equal(t, 1, repo.recoveryClearCalls) + require.Equal(t, limitedAt, repo.recoveryObservedAt) + require.Equal(t, observedResetAt, repo.recoveryObservedReset) + require.Same(t, &observedResetAt, account.RateLimitResetAt, "shared account snapshots must not be mutated in place") +} + +func TestUpdateGrokUsageFromResponseRecoveryRespectsCancellationAndAPIKeyBoundary(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + observedResetAt := now.Add(-time.Second) + observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown) + + t.Run("parent cancellation does not mutate account state", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + account := &Account{ + ID: 661, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + RateLimitedAt: &observedLimitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.updateGrokUsageFromResponse(ctx, account, nil, http.StatusOK) + + require.Zero(t, repo.recoveryClearCalls) + }) + + t.Run("API key success does not alter OAuth cooldown state", func(t *testing.T) { + account := &Account{ + ID: 662, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + RateLimitedAt: &observedLimitedAt, + RateLimitResetAt: &observedResetAt, + } + repo := &grokQuotaAccountRepo{recoveryClearResult: true} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK) + + require.Zero(t, repo.recoveryClearCalls) + }) +} + func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) { repo := &grokQuotaAccountRepo{} svc := &OpenAIGatewayService{accountRepo: repo} diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index e2add90ff2..067d458b89 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -14,7 +14,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -273,7 +272,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) var upstreamReq *http.Request if account.Platform == PlatformGrok { - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity) + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity, s.cfg) } else { upstreamReq, err = s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, isStream, promptCacheKey, false) } @@ -344,7 +343,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } if account.Type == AccountTypeOAuth && promptCacheKey != "" { diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index dbc24c850a..77a5e59329 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -12,7 +12,6 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" ) @@ -193,7 +192,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( releaseUpstreamCtx() return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) } - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity) + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity, s.cfg) } else { upstreamReq, err = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) } @@ -236,7 +235,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) } responseID := ""