diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 84c70b45c3..debee444ec 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -1450,6 +1450,7 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) { if clearErr := h.adminService.UpdateAccountExtra(ctx, accountID, map[string]any{ "grok_needs_reauth": false, "grok_needs_reauth_reason": "", + "grok_needs_reauth_at": "", }); clearErr != nil { slog.Warn("apply_oauth_credentials.clear_grok_reauth_failed", "account_id", accountID, diff --git a/backend/internal/handler/gateway_web_search.go b/backend/internal/handler/gateway_web_search.go index 8251695e83..e7e02d2822 100644 --- a/backend/internal/handler/gateway_web_search.go +++ b/backend/internal/handler/gateway_web_search.go @@ -82,122 +82,88 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { } failedAccounts := make(map[int64]struct{}) - selected, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0) - if err != nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{ - "type": "scheduling_error", - "message": err.Error(), - }}) + var account *service.Account + var accountReleaseFunc func() + var nativeResp *websearch.SearchResponse + var providerName string + var err error + + // Acquire + release holder for the whole handler (including failover retries). + defer func() { + if accountReleaseFunc != nil { + accountReleaseFunc() + } + }() + + // First attempt + up to 3 failover accounts (max 4 total). + for attempt := 0; attempt < 4; attempt++ { + selected, selectErr := h.gatewayService.SelectAccountWithLoadAwareness( + c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0, + ) + if selectErr != nil { + if attempt == 0 { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{ + "type": "scheduling_error", + "message": selectErr.Error(), + }}) + return + } + break + } + if selected == nil || selected.Account == nil { + if attempt == 0 { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{ + "type": "scheduling_error", + "message": "No available accounts", + }}) + return + } + break + } + + release, acquireOK, acquireErr := h.acquireWebSearchAccountSlot(c, selected) + if !acquireOK { + // First hop: surface concurrency errors; later hops try another account. + if attempt == 0 && acquireErr != nil { + h.handleConcurrencyError(c, acquireErr, "account", false) + return + } + failedAccounts[selected.Account.ID] = struct{}{} + continue + } + account = selected.Account + accountReleaseFunc = release + + nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults) + if err == nil { + break + } + var failoverErr *service.UpstreamFailoverError + if !errors.As(err, &failoverErr) || !failoverErr.ShouldRetryNextAccount() { + break + } + failedAccounts[account.ID] = struct{}{} + if accountReleaseFunc != nil { + accountReleaseFunc() + accountReleaseFunc = nil + } + account = nil + } + if err != nil || nativeResp == nil { + msg := "web search failed" + if err != nil { + msg = err.Error() + } + c.JSON(http.StatusBadGateway, gin.H{"error": gin.H{"type": "web_search_error", "message": msg}}) return } - if selected == nil || selected.Account == nil { + if account == nil { c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{ "type": "scheduling_error", "message": "No available accounts", }}) return } - account := selected.Account - accountReleaseFunc := selected.ReleaseFunc - if !selected.Acquired { - if selected.WaitPlan == nil || h.concurrencyHelper == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{ - "type": "scheduling_error", - "message": "No available accounts", - }}) - return - } - accountWaitCounted := false - canWait, waitErr := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selected.WaitPlan.MaxWaiting) - if waitErr != nil { - logger.L().Warn("gateway.web_search.account_wait_counter_increment_failed", - zap.Int64("account_id", account.ID), - zap.Error(waitErr), - ) - } else if !canWait { - c.JSON(http.StatusTooManyRequests, gin.H{"error": gin.H{ - "type": "rate_limit_error", - "message": "Too many pending requests, please retry later", - }}) - return - } else { - accountWaitCounted = true - } - releaseWait := func() { - if accountWaitCounted { - h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID) - accountWaitCounted = false - } - } - streamStarted := false - release, acquireErr := h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout( - c, - account.ID, - selected.WaitPlan.MaxConcurrency, - selected.WaitPlan.Timeout, - false, - &streamStarted, - ) - releaseWait() - if acquireErr != nil { - h.handleConcurrencyError(c, acquireErr, "account", streamStarted) - return - } - accountReleaseFunc = release - } - if accountReleaseFunc != nil { - defer func() { - if accountReleaseFunc != nil { - accountReleaseFunc() - } - }() - } - - // Scheduling is 100% the same as other requests: - // SelectAccountWithLoadAwareness handles load balancing, rate limits, failover, sticky sessions, concurrency, proxies etc. - // Downstream rate limiting, billing etc. can be wired the same way. - - // Use Grok *native* web search via the selected Grok account + responses API + web_search tool. - // This ensures results come from Grok's own search (not third-party emulation like Tavily/Brave). - // Output is normalized to the same unified format for clients/agents/MCP. - - nativeResp, providerName, err := h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults) - if err != nil { - var failoverErr *service.UpstreamFailoverError - if errors.As(err, &failoverErr) && failoverErr.ShouldRetryNextAccount() { - failedAccounts[account.ID] = struct{}{} - if accountReleaseFunc != nil { - accountReleaseFunc() - accountReleaseFunc = nil - } - for attempt := 1; attempt < 4; attempt++ { - selected, selectErr := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0) - if selectErr != nil || selected == nil || selected.Account == nil { - break - } - account = selected.Account - if selected.Acquired { - accountReleaseFunc = selected.ReleaseFunc - } else { - break - } - nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults) - if err == nil { - break - } - if !errors.As(err, &failoverErr) || !failoverErr.ShouldRetryNextAccount() { - break - } - failedAccounts[account.ID] = struct{}{} - accountReleaseFunc() - accountReleaseFunc = nil - } - } - if err != nil { - c.JSON(http.StatusBadGateway, gin.H{"error": gin.H{"type": "web_search_error", "message": err.Error()}}) - return - } - } userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) @@ -251,6 +217,59 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { }) } +// acquireWebSearchAccountSlot resolves an immediate slot or WaitPlan wait. +// On failure returns (nil, false, err); err is non-nil for concurrency acquire +// failures so the first hop can map them to HTTP. Wait-queue full returns +// (nil, false, nil) so failover can try another account. +func (h *GatewayHandler) acquireWebSearchAccountSlot( + c *gin.Context, + selected *service.AccountSelectionResult, +) (release func(), ok bool, acquireErr error) { + if selected == nil || selected.Account == nil { + return nil, false, nil + } + if selected.Acquired { + return selected.ReleaseFunc, true, nil + } + if selected.WaitPlan == nil || h.concurrencyHelper == nil { + return nil, false, nil + } + account := selected.Account + accountWaitCounted := false + canWait, waitErr := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selected.WaitPlan.MaxWaiting) + if waitErr != nil { + logger.L().Warn("gateway.web_search.account_wait_counter_increment_failed", + zap.Int64("account_id", account.ID), + zap.Error(waitErr), + ) + // Best-effort wait without counter (same as first-hop legacy path). + } else if !canWait { + return nil, false, nil + } else { + accountWaitCounted = true + } + releaseWait := func() { + if accountWaitCounted { + h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID) + accountWaitCounted = false + } + } + streamStarted := false + slotRelease, err := h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout( + c, + account.ID, + selected.WaitPlan.MaxConcurrency, + selected.WaitPlan.Timeout, + false, + &streamStarted, + ) + releaseWait() + if err != nil { + return nil, false, err + } + return slotRelease, true, nil +} + // doGrokNativeWebSearch executes web search using the Grok account's native capability // by calling the responses endpoint with web_search tool, then normalizes sources to unified format. func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Context, account *service.Account, query string, maxResults int) (*websearch.SearchResponse, string, error) { diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 9e0f253059..723bf901ce 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2415,8 +2415,9 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas } func (h *OpenAIGatewayHandler) submitOpenAIUsageRecordTask(parent context.Context, result *service.OpenAIForwardResult, task service.UsageRecordTask) { - // Image and async video bills are money-critical: never drop on pool overflow. - if result != nil && (result.ImageCount > 0 || result.VideoCount > 0) { + // Money-critical bills never drop on pool overflow: media, search surcharge, voice. + if result != nil && (result.ImageCount > 0 || result.VideoCount > 0 || + result.SearchCount > 0 || result.WebSearchCalls > 0 || result.AudioUsage != nil) { h.submitMandatoryUsageRecordTask(parent, task) return } diff --git a/backend/internal/handler/usage_record_submit_task_test.go b/backend/internal/handler/usage_record_submit_task_test.go index ebe5c3df82..66b3aca0d3 100644 --- a/backend/internal/handler/usage_record_submit_task_test.go +++ b/backend/internal/handler/usage_record_submit_task_test.go @@ -189,3 +189,33 @@ func TestOpenAIGatewayHandlerSubmitOpenAIUsageRecordTask_ImageResultUsesMandator require.True(t, called.Load(), "image usage task must be mandatory when async submit is dropped") } + +func TestOpenAIGatewayHandlerSubmitOpenAIUsageRecordTask_SearchCountUsesMandatoryFallback(t *testing.T) { + pool := service.NewUsageRecordWorkerPoolWithOptions(service.UsageRecordWorkerPoolOptions{ + WorkerCount: 1, + QueueSize: 1, + TaskTimeout: time.Second, + OverflowPolicy: "drop", + OverflowSamplePercent: 0, + AutoScaleEnabled: false, + }) + t.Cleanup(pool.Stop) + h := &OpenAIGatewayHandler{usageRecordWorkerPool: pool} + + block := make(chan struct{}) + release := make(chan struct{}) + pool.Submit(func(ctx context.Context) { + close(block) + <-release + }) + <-block + pool.Submit(func(ctx context.Context) {}) + + var called atomic.Bool + h.submitOpenAIUsageRecordTask(context.Background(), &service.OpenAIForwardResult{SearchCount: 3}, func(ctx context.Context) { + called.Store(true) + }) + close(release) + + require.True(t, called.Load(), "search surcharge usage task must be mandatory when async submit is dropped") +} diff --git a/backend/internal/pkg/xai/billing.go b/backend/internal/pkg/xai/billing.go index a942de84d2..44e4772d44 100644 --- a/backend/internal/pkg/xai/billing.go +++ b/backend/internal/pkg/xai/billing.go @@ -173,18 +173,15 @@ func BuildBillingSummary(config *BillingConfig) *BillingSummary { periodType := resolvePeriodType(period) creditUsage := cloneFloat(config.CreditUsagePercent) + // Weekly period bounds must not fall back to monthly billing period ends — + // that would park accounts on a multi-week horizon when weekly UsagePercent + // is high (scheduler seven_day uses PeriodEnd). periodStart := "" periodEnd := "" if period != nil { periodStart = strings.TrimSpace(period.Start) periodEnd = strings.TrimSpace(period.End) } - if periodStart == "" { - periodStart = strings.TrimSpace(config.BillingPeriodStart) - } - if periodEnd == "" { - periodEnd = strings.TrimSpace(config.BillingPeriodEnd) - } products := make([]BillingProductSummary, 0, len(config.ProductUsage)) for _, item := range config.ProductUsage { diff --git a/backend/internal/pkg/xai/billing_test.go b/backend/internal/pkg/xai/billing_test.go index ffa282907d..3dbfc89bf2 100644 --- a/backend/internal/pkg/xai/billing_test.go +++ b/backend/internal/pkg/xai/billing_test.go @@ -119,6 +119,27 @@ func TestBuildBillingSummaryMonthlyOnlyKeepsWeeklyUsageEmpty(t *testing.T) { require.InDelta(t, 50, *summary.UsedPercent, 1e-9) } +func TestBuildBillingSummaryWeeklyDoesNotInheritMonthlyPeriodEnd(t *testing.T) { + t.Parallel() + // Weekly usage without currentPeriod.end must not copy billingPeriodEnd (monthly). + payload, err := ParseBillingPayload([]byte(`{ + "config": { + "creditUsagePercent": 95.0, + "productUsage": [{"product":"Api","usagePercent":95.0}], + "billingPeriodStart": "2026-07-01T00:00:00Z", + "billingPeriodEnd": "2026-08-01T00:00:00Z", + "monthlyLimit": {"val": 15000}, + "used": {"val": 1000} + } + }`)) + require.NoError(t, err) + summary := BuildBillingSummary(payload.Config) + require.NotNil(t, summary) + require.Equal(t, "weekly", summary.PeriodType) + require.Equal(t, "", summary.PeriodEnd) + require.Equal(t, "2026-08-01T00:00:00Z", summary.BillingPeriodEnd) +} + func TestMergeBillingProbeResultRetainsFailedWindow(t *testing.T) { t.Parallel() previous := &BillingSummary{ diff --git a/backend/internal/service/account_scheduling_threshold_eval_test.go b/backend/internal/service/account_scheduling_threshold_eval_test.go index 284ca6945f..95c8c41aac 100644 --- a/backend/internal/service/account_scheduling_threshold_eval_test.go +++ b/backend/internal/service/account_scheduling_threshold_eval_test.go @@ -323,9 +323,9 @@ func TestEvaluateAccountSchedulingThreshold_GrokBillingWindows(t *testing.T) { "grok_sched_reset_at": now.Add(time.Hour).Format(time.RFC3339), // Official weekly billing above threshold (extra key used by probe persist). grokBillingExtraKey: &xai.BillingSummary{ - UsagePercent: &weeklyPct, - PeriodEnd: weeklyEnd.Format(time.RFC3339), - UsedPercent: &monthlyPct, + UsagePercent: &weeklyPct, + PeriodEnd: weeklyEnd.Format(time.RFC3339), + UsedPercent: &monthlyPct, BillingPeriodEnd: now.Add(10 * 24 * time.Hour).Format(time.RFC3339), }, }, @@ -342,3 +342,36 @@ func TestEvaluateAccountSchedulingThreshold_GrokBillingWindows(t *testing.T) { require.NotNil(t, decision.Until) require.True(t, weeklyEnd.Equal(*decision.Until)) } + +func TestEvaluateAccountSchedulingThreshold_GrokBillingStaleIgnored(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC) + weeklyEnd := now.Add(3 * time.Hour) + weeklyPct := 99.0 + // Snapshot older than 48h → billing windows skipped; header util below threshold. + account := &Account{ + Platform: PlatformGrok, + Extra: map[string]any{ + "grok_sched_utilization": 50.0, + "grok_sched_reset_at": now.Add(time.Hour).Format(time.RFC3339), + grokBillingExtraKey: &xai.BillingSummary{ + UsagePercent: &weeklyPct, + PeriodEnd: weeklyEnd.Format(time.RFC3339), + UpdatedAt: now.Add(-49 * time.Hour).Format(time.RFC3339), + }, + }, + } + + decision := EvaluateAccountSchedulingThreshold(account, map[string]int{ + PlatformGrok: 90, + }, now) + require.False(t, decision.ShouldPause) +} + +func TestGrokBillingSnapshotStale(t *testing.T) { + t.Parallel() + require.False(t, grokBillingSnapshotStale("", 48*time.Hour)) + require.False(t, grokBillingSnapshotStale(time.Now().UTC().Format(time.RFC3339), 48*time.Hour)) + require.True(t, grokBillingSnapshotStale(time.Now().UTC().Add(-49*time.Hour).Format(time.RFC3339), 48*time.Hour)) +} diff --git a/backend/internal/service/credentials_sanitize.go b/backend/internal/service/credentials_sanitize.go index 354ff8435d..fb1b500e9d 100644 --- a/backend/internal/service/credentials_sanitize.go +++ b/backend/internal/service/credentials_sanitize.go @@ -1,24 +1,21 @@ package service -import "strings" - // SanitizeStoredCredentials strips secrets that must never be persisted on the // account credentials map after conversion to OAuth tokens (Grok Web SSO / password). // Call from admin create/update/import/apply-oauth paths. +// +// Cookie is always stripped: bulk paths may pass an empty platform label, and +// session-jar residue must never sit next to OAuth tokens on any platform. +// The platform argument is retained for call-site clarity / future scrubbing. func SanitizeStoredCredentials(platform string, creds map[string]any) map[string]any { if creds == nil { return nil } - // Always strip Grok ephemeral secrets regardless of platform label (import - // may set platform later; better never store them). + _ = platform for _, key := range []string{ - "password", "sso_token", "sso", "sso-rw", "clearTextPassword", + "password", "sso_token", "sso", "sso-rw", "clearTextPassword", "cookie", } { delete(creds, key) } - // Platform-specific: also strip cookie jar residue on Grok OAuth stores. - if strings.EqualFold(strings.TrimSpace(platform), PlatformGrok) { - delete(creds, "cookie") - } return creds } diff --git a/backend/internal/service/credentials_sanitize_test.go b/backend/internal/service/credentials_sanitize_test.go index 5467da61e7..c704102d5f 100644 --- a/backend/internal/service/credentials_sanitize_test.go +++ b/backend/internal/service/credentials_sanitize_test.go @@ -30,16 +30,19 @@ func TestSanitizeStoredCredentials_StripsEphemeralSSOSecrets(t *testing.T) { require.NotContains(t, out, "cookie") } -func TestSanitizeStoredCredentials_NonGrokKeepsCookieButDropsPassword(t *testing.T) { - creds := map[string]any{ - "cookie": "session", - "password": "x", - "api_key": "k", +func TestSanitizeStoredCredentials_AlwaysStripsCookie(t *testing.T) { + // Bulk paths may pass empty platform; cookie must never persist next to tokens. + for _, platform := range []string{PlatformOpenAI, PlatformGrok, ""} { + creds := map[string]any{ + "cookie": "session", + "password": "x", + "api_key": "k", + } + out := SanitizeStoredCredentials(platform, creds) + require.Equal(t, "k", out["api_key"], platform) + require.NotContains(t, out, "password", platform) + require.NotContains(t, out, "cookie", platform) } - out := SanitizeStoredCredentials(PlatformOpenAI, creds) - require.Equal(t, "session", out["cookie"]) - require.Equal(t, "k", out["api_key"]) - require.NotContains(t, out, "password") } func TestSanitizeStoredCredentials_NilSafe(t *testing.T) { diff --git a/backend/internal/service/gateway_scheduling.go b/backend/internal/service/gateway_scheduling.go index 3c5f1cb96e..aee0ad7f75 100644 --- a/backend/internal/service/gateway_scheduling.go +++ b/backend/internal/service/gateway_scheduling.go @@ -962,6 +962,9 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i accounts, useMixed, err := s.schedulerSnapshot.ListSchedulableAccounts(ctx, groupID, platform, hasForcePlatform) if err == nil { accounts = s.filterAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok || strings.EqualFold(platform, PlatformGrok) { + accounts = s.filterGrokFreeQuotaAccountsForGateway(ctx, accounts) + } slog.Debug("account_scheduling_list_snapshot", "group_id", derefGroupID(groupID), "platform", platform, @@ -1058,7 +1061,11 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i "tls_fingerprint", acc.IsTLSFingerprintEnabled()) } } - return s.filterAccountsBySchedulingThreshold(ctx, accounts), useMixed, nil + accounts = s.filterAccountsBySchedulingThreshold(ctx, accounts) + if platform == PlatformGrok || strings.EqualFold(platform, PlatformGrok) { + accounts = s.filterGrokFreeQuotaAccountsForGateway(ctx, accounts) + } + return accounts, useMixed, nil } // IsSingleAntigravityAccountGroup 检查指定分组是否只有一个 antigravity 平台的可调度账号。 diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index ba6162c3e3..d6f71e3e64 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1278,7 +1278,11 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account } token, _, err := s.GetAccessToken(ctx, account) if err != nil { - return nil, fmt.Errorf("get grok token: %w", err) + // Credential/token failures should try the next Grok account in the pool. + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusUnauthorized, + Reason: GatewayFailureReason("grok_search_token"), + } } targetURL, err := buildGrokResponsesURL(account, nil) if err != nil { @@ -1312,7 +1316,10 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account defer func() { _ = resp.Body.Close() }() respBytes, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) if readErr != nil { - return nil, fmt.Errorf("read grok native search response: %w", readErr) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + Reason: GatewayFailureReason("grok_search_read"), + } } if resp.StatusCode >= 400 { msg := string(respBytes) diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index fee26dfec9..f3f0a16311 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -204,6 +204,13 @@ func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *bill } func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) string { + // Forced durable money-event IDs must win over client/local context IDs so + // standalone web_search / async video cannot collapse under a reused client id. + if requestID := strings.TrimSpace(upstreamRequestID); requestID != "" { + if isForcedUsageBillingRequestID(requestID) { + return requestID + } + } if ctx != nil { if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { return "client:" + strings.TrimSpace(clientRequestID) @@ -218,6 +225,14 @@ func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) return "generated:" + generateRequestID() } +func isForcedUsageBillingRequestID(requestID string) bool { + id := strings.TrimSpace(requestID) + return strings.HasPrefix(id, "web_search:") || + strings.HasPrefix(id, "grok-video:") || + strings.HasPrefix(id, "grok_audio:") || + strings.HasPrefix(id, "grok_realtime:") +} + func resolveUsageBillingPayloadFingerprint(ctx context.Context, requestPayloadHash string) string { if payloadHash := strings.TrimSpace(requestPayloadHash); payloadHash != "" { return payloadHash diff --git a/backend/internal/service/gateway_usage_billing_request_id_test.go b/backend/internal/service/gateway_usage_billing_request_id_test.go new file mode 100644 index 0000000000..b1b97b941e --- /dev/null +++ b/backend/internal/service/gateway_usage_billing_request_id_test.go @@ -0,0 +1,32 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func TestResolveUsageBillingRequestID_ForcedWebSearchBeatsClientID(t *testing.T) { + t.Parallel() + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + got := resolveUsageBillingRequestID(ctx, "web_search:uuid-1") + require.Equal(t, "web_search:uuid-1", got) +} + +func TestResolveUsageBillingRequestID_ClientWinsOverPlainUpstream(t *testing.T) { + t.Parallel() + ctx := context.WithValue(context.Background(), ctxkey.ClientRequestID, "client-shared-id") + got := resolveUsageBillingRequestID(ctx, "resp_abc") + require.Equal(t, "client:client-shared-id", got) +} + +func TestIsForcedUsageBillingRequestID(t *testing.T) { + t.Parallel() + require.True(t, isForcedUsageBillingRequestID("web_search:x")) + require.True(t, isForcedUsageBillingRequestID("grok-video:task-1")) + require.False(t, isForcedUsageBillingRequestID("resp_abc")) +} diff --git a/backend/internal/service/grok_free_quota_gate.go b/backend/internal/service/grok_free_quota_gate.go index affd3e5fcb..424f12e53d 100644 --- a/backend/internal/service/grok_free_quota_gate.go +++ b/backend/internal/service/grok_free_quota_gate.go @@ -3,6 +3,7 @@ package service import ( "context" "log/slog" + "sync" "sync/atomic" "time" @@ -92,15 +93,41 @@ func isExplicitGrokFreeOAuthAccount(account *Account) bool { } // filterGrokFreeQuotaAccounts applies a local, rolling soft gate only to -// explicitly FREE Grok OAuth accounts on the scheduling hot path. +// FREE Grok OAuth accounts on the OpenAI scheduling hot path. // Missing or failed statistics always fail open; upstream quota/rate-limit // handling remains authoritative. Admin quota/import probes never call this. func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context.Context, accounts []Account) []Account { if s == nil || s.service == nil { return accounts } - settings, enabled := resolveGrokFreeQuotaGateSettings(s.service.cfg) - if !enabled || len(accounts) == 0 || s.service.usageLogRepo == nil { + return filterGrokFreeQuotaAccountsCore(ctx, s.service.cfg, s.service.usageLogRepo, &s.grokFreeQuotaGateCache, accounts) +} + +// filterGrokFreeQuotaAccountsForGateway applies the same soft gate on Gateway +// scheduling (e.g. /v1/web_search) so free accounts near local 95%/1M are not +// still selected for native search while Responses soft-gates them out. +func (s *GatewayService) filterGrokFreeQuotaAccountsForGateway(ctx context.Context, accounts []Account) []Account { + if s == nil { + return accounts + } + return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &gatewayGrokFreeQuotaGateCache, accounts) +} + +// Shared cache for Gateway path (OpenAI scheduler keeps its own map). +var gatewayGrokFreeQuotaGateCache sync.Map + +func filterGrokFreeQuotaAccountsCore( + ctx context.Context, + cfg *config.Config, + usageLogRepo UsageLogRepository, + cache *sync.Map, + accounts []Account, +) []Account { + if cache == nil { + return accounts + } + settings, enabled := resolveGrokFreeQuotaGateSettings(cfg) + if !enabled || len(accounts) == 0 || usageLogRepo == nil { return accounts } now := time.Now().UTC() @@ -112,7 +139,7 @@ func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context. if !isExplicitGrokFreeOAuthAccount(account) || account.ID <= 0 { continue } - if cached, ok := s.grokFreeQuotaGateCache.Load(account.ID); ok { + if cached, ok := cache.Load(account.ID); ok { entry, valid := cached.(grokFreeQuotaGateCacheEntry) age := now.Sub(entry.checkedAt) if valid && settings.cacheTTL > 0 && age >= 0 && age < settings.cacheTTL { @@ -129,12 +156,12 @@ func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context. } if len(missingIDs) > 0 { - statsByID, err := s.queryGrokFreeQuotaWindowStats(ctx, missingIDs, now.Add(-settings.window)) + statsByID, err := queryGrokFreeQuotaWindowStats(ctx, usageLogRepo, missingIDs, now.Add(-settings.window)) if err != nil { grokFreeQuotaGateQueryFailureTotal.Add(1) if settings.cacheTTL > 0 { for _, accountID := range missingIDs { - s.grokFreeQuotaGateCache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) + cache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) } } slog.Warn("grok_free_quota_soft_gate_stats_failed", @@ -148,7 +175,7 @@ func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context. tokens = stats.Tokens } tokensByID[accountID] = tokens - s.grokFreeQuotaGateCache.Store(accountID, grokFreeQuotaGateCacheEntry{tokens: tokens, checkedAt: now, known: true}) + cache.Store(accountID, grokFreeQuotaGateCacheEntry{tokens: tokens, checkedAt: now, known: true}) if tokens >= settings.gateTokens { grokFreeQuotaGateBlockedTotal.Add(1) slog.Info("grok_free_quota_soft_gate_blocked", @@ -175,13 +202,16 @@ func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context. return filtered } -func (s *defaultOpenAIAccountScheduler) queryGrokFreeQuotaWindowStats(ctx context.Context, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { - if batch, ok := s.service.usageLogRepo.(accountWindowStatsBatchReader); ok { +func queryGrokFreeQuotaWindowStats(ctx context.Context, usageLogRepo UsageLogRepository, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { + if usageLogRepo == nil { + return nil, nil + } + if batch, ok := usageLogRepo.(accountWindowStatsBatchReader); ok { return batch.GetAccountWindowStatsBatch(ctx, accountIDs, start) } statsByID := make(map[int64]*usagestats.AccountStats, len(accountIDs)) for _, accountID := range accountIDs { - stats, err := s.service.usageLogRepo.GetAccountWindowStats(ctx, accountID, start) + stats, err := usageLogRepo.GetAccountWindowStats(ctx, accountID, start) if err != nil { return nil, err } diff --git a/backend/internal/service/grok_search_count.go b/backend/internal/service/grok_search_count.go index ad4e8b581b..80dc372168 100644 --- a/backend/internal/service/grok_search_count.go +++ b/backend/internal/service/grok_search_count.go @@ -24,34 +24,56 @@ func countGrokNativeSearchCallsFromSSEBody(body string) int { if strings.TrimSpace(body) == "" { return 0 } - total := 0 seen := make(map[string]struct{}) + total := 0 forEachOpenAISSEDataPayload(body, func(data []byte) { - n, keys := countGrokNativeSearchCallsInSSEDataWithKeys(data) - if n <= 0 { - return - } - // Dedup stream item.added / item.done for the same call id. - if len(keys) == 0 { - total += n - return - } - for _, k := range keys { - if _, ok := seen[k]; ok { - continue - } - seen[k] = struct{}{} - total++ - } + total += countGrokNativeSearchCallsInSSEDataDedup(data, seen) }) return total } +// countGrokNativeSearchCallsInSSEData counts search tool calls in one SSE +// payload without cross-event dedup. Prefer countGrokNativeSearchCallsInSSEDataDedup +// for live streams so item.done + response.completed do not double-bill. func countGrokNativeSearchCallsInSSEData(data []byte) int { n, _ := countGrokNativeSearchCallsInSSEDataWithKeys(data) return n } +// countGrokNativeSearchCallsInSSEDataDedup increments only unseen call ids. +// Callers must reuse the same seen map for the full stream lifetime. +func countGrokNativeSearchCallsInSSEDataDedup(data []byte, seen map[string]struct{}) int { + if seen == nil { + return countGrokNativeSearchCallsInSSEData(data) + } + n, keys := countGrokNativeSearchCallsInSSEDataWithKeys(data) + if n <= 0 { + return 0 + } + // No stable id: fall back to raw count once (cannot dedup across events). + if len(keys) == 0 { + return n + } + // Intra-event + cross-event dedup by call_id/id. + added := 0 + local := make(map[string]struct{}, len(keys)) + for _, k := range keys { + if k == "" { + continue + } + if _, ok := local[k]; ok { + continue + } + local[k] = struct{}{} + if _, ok := seen[k]; ok { + continue + } + seen[k] = struct{}{} + added++ + } + return added +} + func countGrokNativeSearchCallsInSSEDataWithKeys(data []byte) (int, []string) { if len(data) == 0 || !gjson.ValidBytes(data) { return 0, nil diff --git a/backend/internal/service/grok_search_count_test.go b/backend/internal/service/grok_search_count_test.go index 457d8917e7..b626e29c6d 100644 --- a/backend/internal/service/grok_search_count_test.go +++ b/backend/internal/service/grok_search_count_test.go @@ -29,6 +29,20 @@ func TestCountGrokNativeSearchCallsFromSSEBodyDedups(t *testing.T) { require.Equal(t, 2, countGrokNativeSearchCallsFromSSEBody(sse)) } +func TestCountGrokNativeSearchCallsInSSEDataDedup_LiveStreamPath(t *testing.T) { + t.Parallel() + // Mirrors the live streaming accumulator: item.done then response.completed + // for the same call_id must bill once (regression for ~2× surcharge). + seen := make(map[string]struct{}) + done := []byte(`{"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`) + completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"web_search_call","id":"ws1","call_id":"c1"},{"type":"x_search_call","id":"xs1","call_id":"c2"}]}}`) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(done, seen)) + require.Equal(t, 1, countGrokNativeSearchCallsInSSEDataDedup(completed, seen)) + // Raw (no-dedup) path still double-counts the same envelope pair. + require.Equal(t, 1, countGrokNativeSearchCallsInSSEData(done)) + require.Equal(t, 2, countGrokNativeSearchCallsInSSEData(completed)) +} + func stringsJoin(lines ...string) string { out := "" for _, l := range lines { diff --git a/backend/internal/service/grok_spending_reauth.go b/backend/internal/service/grok_spending_reauth.go index d58bb2f67f..14adf62cde 100644 --- a/backend/internal/service/grok_spending_reauth.go +++ b/backend/internal/service/grok_spending_reauth.go @@ -48,6 +48,7 @@ func clearGrokNeedsReauthExtra(ctx context.Context, repo AccountRepository, acco _ = repo.UpdateExtra(stateCtx, accountID, map[string]any{ "grok_needs_reauth": false, "grok_needs_reauth_reason": "", + "grok_needs_reauth_at": "", }) } diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index f805c8c5f1..f2b4186d84 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -363,23 +363,18 @@ func isKnownGrokFreeAccount(account *Account) bool { } } } - // Explicit free credentials (subscription_tier / plan_type free|basic) win over - // weak paid inference so weekly credit % cannot disable soft-gate / media block - // for real Free OAuth accounts. Strong paid evidence (named paid plan, monthly - // limit cents) still wins over *inferred* free only. + // Strong paid evidence always wins — including over stale free credentials + // left after SuperGrok upgrade (subscription_tier still "free" until re-import). + if paidSignal { + return false + } + // Explicit free/basic tier (credentials or free plan string). if freeSignal { - // Explicit free/basic tier from credentials or free plan string. - // Only override if we saw a clearly paid plan name (not mere usage %). - if paidSignal { - // paidSignal from MonthlyLimitCents or non-free Plan — treat as paid. - return false - } return true } - // Explicit paid evidence always wins over an inferred Free signal. This - // protects upgraded/stale accounts whose previous quota snapshot still - // carries the historical 2M Free token limit. - return !paidSignal && inferredFreeSignal + // Inferred free (empty plan + successful monthly probe, or free token limit) + // only when there is no paid evidence. + return inferredFreeSignal } func isGrokFreeSubscriptionTier(tier string) bool { diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 3251bb7485..4cc3521c38 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -888,6 +888,17 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing. return a }(), }, + { + name: "paid billing overrides stale free credentials", + account: func() *Account { + a := healthyGrokOAuthGatewayTestAccount(9123, "access-token") + a.Credentials["subscription_tier"] = "free" + a.Extra = map[string]any{ + grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": http.StatusOK}, + } + return a + }(), + }, { name: "partial billing without monthly evidence remains unknown", account: func() *Account { diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index 2be20306c1..ac3025d1d2 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -301,6 +301,9 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. streamImageOutputs := make([]json.RawMessage, 0, 1) streamSeenImages := make(map[string]struct{}) searchCounter := 0 + // Dedup search tool calls across SSE events (item.done + response.completed + // both list the same call_id — counting both would ~2× the surcharge). + streamSearchSeen := make(map[string]struct{}) resultWithUsage := func() *openaiStreamingResult { return &openaiStreamingResult{ usage: usage, @@ -476,7 +479,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. line = "data: " + data } imageCounter.AddSSEData(dataBytes) - searchCounter += countGrokNativeSearchCallsInSSEData(dataBytes) + searchCounter += countGrokNativeSearchCallsInSSEDataDedup(dataBytes, streamSearchSeen) // Correct Codex tool calls if needed (apply_patch -> edit, etc.) if correctedData, corrected := s.toolCorrector.CorrectToolCallsInSSEBytes(dataBytes); corrected {