diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 3b3283e1aa..84c70b45c3 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -1415,6 +1415,9 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) { return } + // Drop SSO/password residue; re-auth must leave only OAuth tokens on disk. + req.Credentials = service.SanitizeStoredCredentials(existing.Platform, req.Credentials) + updatedAccount, err := h.adminService.UpdateAccount(ctx, accountID, &service.UpdateAccountInput{ Type: req.Type, Credentials: req.Credentials, @@ -1442,6 +1445,19 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) { } } + // Successful re-auth clears the soft spending-limit reauth flag for Grok. + if existing.Platform == service.PlatformGrok { + if clearErr := h.adminService.UpdateAccountExtra(ctx, accountID, map[string]any{ + "grok_needs_reauth": false, + "grok_needs_reauth_reason": "", + }); clearErr != nil { + slog.Warn("apply_oauth_credentials.clear_grok_reauth_failed", + "account_id", accountID, + "err", clearErr, + ) + } + } + if cleared, clearErr := h.adminService.ClearAccountError(ctx, accountID); clearErr != nil { slog.Warn("apply_oauth_credentials.clear_error_failed", "account_id", accountID, diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index 7212b62e92..8ed1094029 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -462,8 +462,9 @@ func grokSSOImportCredentials(built map[string]any, reqCredentials map[string]an // Only merge operator config from the request — never free-form secrets // (password / sso_token / cookie / etc.) into stored credentials. allowedReqKeys := map[string]struct{}{ - "base_url": {}, "model_mapping": {}, "header_override": {}, - "header_overrides": {}, "custom_headers": {}, + "base_url": {}, "model_mapping": {}, + "header_override": {}, "header_overrides": {}, "header_override_enabled": {}, + "custom_headers": {}, } ops := map[string]any{} for k, v := range reqCredentials { @@ -489,7 +490,7 @@ func grokSSOImportCredentials(built map[string]any, reqCredentials map[string]an if reqBaseURL, ok := reqCredentials["base_url"].(string); ok && strings.TrimSpace(reqBaseURL) != "" { credentials["base_url"] = strings.TrimSpace(reqBaseURL) } - return credentials + return service.SanitizeStoredCredentials(service.PlatformGrok, credentials) } func grokSSOImportExpiry(requestExpiresAt *int64, requestAutoPause *bool, tokenInfo *service.GrokTokenInfo) (*int64, *bool) { diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 287d59fd9f..1e2d0fbe3d 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -2416,6 +2416,33 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv task(ctx) } +// submitMandatoryUsageRecordTask never silently drops billing work on pool overflow. +func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, task service.UsageRecordTask) { + if task == nil { + return + } + task = wrapUsageRecordTaskContext(parent, task) + if h.usageRecordWorkerPool != nil { + if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() { + return + } + logger.L().With( + zap.String("component", "handler.gateway.usage"), + ).Warn("gateway.usage_record_task_mandatory_sync_fallback") + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + defer func() { + if recovered := recover(); recovered != nil { + logger.L().With( + zap.String("component", "handler.gateway.usage"), + zap.Any("panic", recovered), + ).Error("gateway.usage_record_task_panic_recovered") + } + }() + task(ctx) +} + // getUserMsgQueueMode 获取当前请求的 UMQ 模式 // 返回 "serialize" | "throttle" | "" func (h *GatewayHandler) getUserMsgQueueMode(account *service.Account, parsed *service.ParsedRequest) string { diff --git a/backend/internal/handler/gateway_web_search.go b/backend/internal/handler/gateway_web_search.go index 84b99935b8..8251695e83 100644 --- a/backend/internal/handler/gateway_web_search.go +++ b/backend/internal/handler/gateway_web_search.go @@ -3,19 +3,21 @@ package handler import ( "context" "encoding/json" + "errors" "fmt" "net/http" "net/url" "strconv" "strings" - middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/Wei-Shaw/sub2api/internal/pkg/websearch" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" + "github.com/google/uuid" "github.com/tidwall/gjson" "go.uber.org/zap" ) @@ -79,7 +81,8 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { return } - selected, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, nil, "", 0) + 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", @@ -143,7 +146,11 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { accountReleaseFunc = release } if accountReleaseFunc != nil { - defer accountReleaseFunc() + defer func() { + if accountReleaseFunc != nil { + accountReleaseFunc() + } + }() } // Scheduling is 100% the same as other requests: @@ -156,11 +163,40 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { nativeResp, providerName, err := h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults) if err != nil { - c.JSON(http.StatusBadGateway, gin.H{"error": gin.H{ - "type": "web_search_error", - "message": err.Error(), - }}) - return + 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") @@ -169,15 +205,16 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) requestPayloadHash := service.HashUsageRequestPayload([]byte(req.Query)) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - // Unique per invocation (not hash(query) alone) so identical queries still bill. - searchRequestID := "web_search:" + service.HashUsageRequestPayload([]byte(req.Query+"|"+clientIP+"|"+userAgent)) + // Request IDs are billing idempotency keys, so they must be unique per invocation. + // Query/IP/UA hashes would collapse repeated identical searches into one charge. + searchRequestID := "web_search:" + uuid.NewString() if apiKey.Group != nil && (apiKey.Group.GetSearchPricePer1k() == nil || *apiKey.Group.GetSearchPricePer1k() <= 0) { logger.L().With( zap.String("component", "handler.gateway.web_search"), zap.Int64("group_id", apiKey.Group.ID), ).Warn("gateway.web_search.search_price_per_1k_unset_free") } - h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ Result: &service.ForwardResult{ RequestID: searchRequestID, diff --git a/backend/internal/handler/grok_audio.go b/backend/internal/handler/grok_audio.go index 677bda053e..4a76b0d900 100644 --- a/backend/internal/handler/grok_audio.go +++ b/backend/internal/handler/grok_audio.go @@ -92,11 +92,13 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) { elapsed := time.Since(started) if proxyErr != nil { reqLog.Info("grok_realtime.proxy_failed", zap.Error(proxyErr)) - _ = conn.Close(coderws.StatusInternalError, "upstream realtime websocket failed") - // Do not bill failed realtime sessions (avoids charging idle/error wall-clock). - return + if !isExpectedGrokRealtimeClose(proxyErr) { + _ = conn.Close(coderws.StatusInternalError, "upstream realtime websocket failed") + return + } } - // Bill successful realtime sessions by elapsed minutes when group price is set. + // A relay normally returns a close error when either side closes normally. + // Those sessions still consumed upstream audio time and must be billed. if elapsed > 0 { result := &service.OpenAIForwardResult{ Model: model, @@ -107,6 +109,19 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) { } } +func isExpectedGrokRealtimeClose(err error) bool { + if err == nil { + return true + } + switch coderws.CloseStatus(err) { + case coderws.StatusNormalClosure, coderws.StatusGoingAway, + coderws.StatusNoStatusRcvd, coderws.StatusAbnormalClosure: + return true + default: + return false + } +} + // GrokVoice handles xAI Voice HTTP endpoints. endpoint is "tts", "stt", or "custom-voices". func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) { apiKey, ok := middleware2.GetAPIKeyFromContext(c) diff --git a/backend/internal/handler/grok_audio_billing_test.go b/backend/internal/handler/grok_audio_billing_test.go new file mode 100644 index 0000000000..6b8a6c81be --- /dev/null +++ b/backend/internal/handler/grok_audio_billing_test.go @@ -0,0 +1,25 @@ +//go:build unit + +package handler + +import ( + "testing" + + coderws "github.com/coder/websocket" +) + +func TestIsExpectedGrokRealtimeClose(t *testing.T) { + for _, status := range []coderws.StatusCode{ + coderws.StatusNormalClosure, + coderws.StatusGoingAway, + coderws.StatusNoStatusRcvd, + coderws.StatusAbnormalClosure, + } { + if !isExpectedGrokRealtimeClose(coderws.CloseError{Code: status}) { + t.Fatalf("status %v should be treated as an expected session close", status) + } + } + if isExpectedGrokRealtimeClose(coderws.CloseError{Code: coderws.StatusPolicyViolation}) { + t.Fatal("policy violations must not be treated as billable normal closes") + } +} diff --git a/backend/internal/repository/grok_oauth_client.go b/backend/internal/repository/grok_oauth_client.go index 0b709535b6..e78f26a4f2 100644 --- a/backend/internal/repository/grok_oauth_client.go +++ b/backend/internal/repository/grok_oauth_client.go @@ -35,11 +35,12 @@ const ( ) func NewGrokOAuthClient() service.GrokOAuthClient { - // Prefer validated official/allowlisted token URL so env misconfig cannot - // exfiltrate auth codes / refresh tokens to an arbitrary host. + // Fail closed: never fall back to an unvalidated EffectiveTokenURL (env can + // point at an attacker host and steal code/refresh tokens). tokenURL, err := xai.ValidatedTokenURL() if err != nil || strings.TrimSpace(tokenURL) == "" { - tokenURL = xai.EffectiveTokenURL() + // Official allowlisted endpoint only — never EffectiveTokenURL() (raw env). + tokenURL = xai.DefaultTokenURL } return &grokOAuthClient{tokenURL: tokenURL} } diff --git a/backend/internal/repository/grok_oauth_client_test.go b/backend/internal/repository/grok_oauth_client_test.go index 690a8e2139..44c934c220 100644 --- a/backend/internal/repository/grok_oauth_client_test.go +++ b/backend/internal/repository/grok_oauth_client_test.go @@ -47,6 +47,8 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) { } })) defer server.Close() + // Tests inject a loopback token endpoint; allowlist requires unsafe override. + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") t.Setenv(xai.EnvTokenURL, server.URL) client := NewGrokOAuthClient() @@ -80,6 +82,7 @@ func TestGrokOAuthClientRefreshForbiddenClassifiesOnlyExplicitEntitlement(t *tes _, _ = w.Write([]byte(tt.body)) })) defer server.Close() + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") t.Setenv(xai.EnvTokenURL, server.URL) client := NewGrokOAuthClient() @@ -96,6 +99,7 @@ func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) { _, _ = w.Write([]byte(`{"error":"invalid_grant","access_token":"access-secret","refresh_token":"refresh-secret","code_verifier":"verifier-secret"}`)) })) defer server.Close() + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") t.Setenv(xai.EnvTokenURL, server.URL) client := NewGrokOAuthClient() @@ -119,3 +123,12 @@ func TestGrokOAuthEntitlementDenialRequiresExplicitEvidence(t *testing.T) { require.False(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"forbidden","message":"request forbidden"}`)) require.False(t, grokOAuthHasExplicitEntitlementDenial(`403 Forbidden`)) } + +func TestNewGrokOAuthClient_UnvalidatedTokenURLFallsBackToDefault(t *testing.T) { + // Without unsafe overrides, a random env TokenURL must not be used (fail-closed). + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "") + t.Setenv(xai.EnvTokenURL, "https://evil.example/oauth/token") + + client := NewGrokOAuthClient().(*grokOAuthClient) + require.Equal(t, xai.DefaultTokenURL, client.tokenURL) +} diff --git a/backend/internal/service/account_scheduling_threshold_eval.go b/backend/internal/service/account_scheduling_threshold_eval.go index efa09fd104..aed703e9ca 100644 --- a/backend/internal/service/account_scheduling_threshold_eval.go +++ b/backend/internal/service/account_scheduling_threshold_eval.go @@ -297,36 +297,61 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand // Official billing windows (weekly credits / monthly used%) from probe snapshot. // Complements header-based grok_sched_* so paid SuperGrok accounts can auto-pause // when request/token headers are sparse but billing % is high. + // Do not cross-fallback period ends (weekly until ≠ monthly end) to avoid multi-week parks. + const grokBillingThresholdStaleAfter = 48 * time.Hour if billing, err := grokBillingSnapshotFromExtra(account.Extra); err == nil && billing != nil { - if billing.UsagePercent != nil { - candidates = append(candidates, &accountSchedulingThresholdCandidate{ - window: "seven_day", - scope: "grok_billing", - usedPercent: *billing.UsagePercent, - until: parseSchedulingResetAt(firstNonEmpty(billing.PeriodEnd, billing.BillingPeriodEnd)), - }) - } - monthlyUtil := 0.0 - hasMonthly := false - if billing.UsedPercent != nil { - monthlyUtil = *billing.UsedPercent - hasMonthly = true - } else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil { - monthlyUtil = (*billing.UsedCents / *billing.MonthlyLimitCents) * 100 - hasMonthly = true - } - if hasMonthly { - candidates = append(candidates, &accountSchedulingThresholdCandidate{ - window: "thirty_day", - scope: "grok_billing", - usedPercent: monthlyUtil, - until: parseSchedulingResetAt(firstNonEmpty(billing.BillingPeriodEnd, billing.PeriodEnd)), - }) + fetchedAt := firstNonEmpty(billing.UpdatedAt, billing.FetchedAt, billing.WeeklyUpdatedAt, billing.MonthlyUpdatedAt) + if !grokBillingSnapshotStale(fetchedAt, grokBillingThresholdStaleAfter) { + if billing.UsagePercent != nil { + if until := parseSchedulingResetAt(billing.PeriodEnd); until != nil { + candidates = append(candidates, &accountSchedulingThresholdCandidate{ + window: "seven_day", + scope: "grok_billing", + usedPercent: *billing.UsagePercent, + until: until, + }) + } + } + monthlyUtil := 0.0 + hasMonthly := false + if billing.UsedPercent != nil { + monthlyUtil = *billing.UsedPercent + hasMonthly = true + } else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil { + monthlyUtil = (*billing.UsedCents / *billing.MonthlyLimitCents) * 100 + hasMonthly = true + } + if hasMonthly { + if until := parseSchedulingResetAt(billing.BillingPeriodEnd); until != nil { + candidates = append(candidates, &accountSchedulingThresholdCandidate{ + window: "thirty_day", + scope: "grok_billing", + usedPercent: monthlyUtil, + until: until, + }) + } + } } } return candidates } +func grokBillingSnapshotStale(fetchedAtRFC3339 string, maxAge time.Duration) bool { + if maxAge <= 0 { + return false + } + raw := strings.TrimSpace(fetchedAtRFC3339) + if raw == "" { + // Unknown freshness: still allow (probe may omit timestamps on older snapshots). + return false + } + ts, err := time.Parse(time.RFC3339, raw) + if err != nil { + return false + } + return time.Since(ts) > maxAge +} + func pickLatestResetSchedulingCandidate(candidates []*accountSchedulingThresholdCandidate, threshold int, now time.Time) *accountSchedulingThresholdCandidate { var winner *accountSchedulingThresholdCandidate for _, candidate := range candidates { diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index d6905d7329..65d4d3f268 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -232,7 +232,7 @@ func (s *AccountService) Create(ctx context.Context, req CreateAccountRequest) ( Notes: normalizeAccountNotes(req.Notes), Platform: req.Platform, Type: req.Type, - Credentials: req.Credentials, + Credentials: SanitizeStoredCredentials(req.Platform, req.Credentials), Extra: req.Extra, ProxyID: req.ProxyID, Concurrency: req.Concurrency, @@ -325,7 +325,7 @@ func (s *AccountService) Update(ctx context.Context, id int64, req UpdateAccount } if req.Credentials != nil { - account.Credentials = *req.Credentials + account.Credentials = SanitizeStoredCredentials(account.Platform, *req.Credentials) } if req.Extra != nil { diff --git a/backend/internal/service/admin_account.go b/backend/internal/service/admin_account.go index c32c36af8c..9958846658 100644 --- a/backend/internal/service/admin_account.go +++ b/backend/internal/service/admin_account.go @@ -495,6 +495,8 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { return nil, err } + // Never persist ephemeral SSO/password secrets after OAuth conversion. + input.Credentials = SanitizeStoredCredentials(input.Platform, input.Credentials) account, err := buildAccountForCreate(input, accountExtra) if err != nil { @@ -605,6 +607,8 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U if err := NormalizeHeaderOverrideCredentials(account.Credentials); err != nil { return nil, err } + // Strip SSO/password residue that must never sit next to OAuth tokens. + account.Credentials = SanitizeStoredCredentials(account.Platform, account.Credentials) } // Extra 使用 map:需要区分“未提供(nil)”与“显式清空({})”。 // 关闭配额限制时前端会删除 quota_* 键并提交 extra:{},此时也必须落库。 @@ -1015,6 +1019,11 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil { return nil, err } + // Bulk may mix platforms; always drop ephemeral SSO/password keys (cookie + // only when platform is known Grok — empty platform still strips password/*). + if input.Credentials != nil { + input.Credentials = SanitizeStoredCredentials("", input.Credentials) + } // Prepare bulk updates for columns and JSONB fields. repoUpdates := AccountBulkUpdate{ diff --git a/backend/internal/service/credentials_sanitize.go b/backend/internal/service/credentials_sanitize.go new file mode 100644 index 0000000000..354ff8435d --- /dev/null +++ b/backend/internal/service/credentials_sanitize.go @@ -0,0 +1,24 @@ +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. +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). + for _, key := range []string{ + "password", "sso_token", "sso", "sso-rw", "clearTextPassword", + } { + 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 new file mode 100644 index 0000000000..5467da61e7 --- /dev/null +++ b/backend/internal/service/credentials_sanitize_test.go @@ -0,0 +1,47 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSanitizeStoredCredentials_StripsEphemeralSSOSecrets(t *testing.T) { + creds := map[string]any{ + "access_token": "at", + "refresh_token": "rt", + "password": "secret", + "sso_token": "sso", + "sso": "cookie-sso", + "sso-rw": "rw", + "clearTextPassword": "plain", + "cookie": "jar", + "base_url": "https://api.x.ai", + } + out := SanitizeStoredCredentials(PlatformGrok, creds) + require.Equal(t, "at", out["access_token"]) + require.Equal(t, "rt", out["refresh_token"]) + require.Equal(t, "https://api.x.ai", out["base_url"]) + require.NotContains(t, out, "password") + require.NotContains(t, out, "sso_token") + require.NotContains(t, out, "sso") + require.NotContains(t, out, "sso-rw") + require.NotContains(t, out, "clearTextPassword") + require.NotContains(t, out, "cookie") +} + +func TestSanitizeStoredCredentials_NonGrokKeepsCookieButDropsPassword(t *testing.T) { + creds := map[string]any{ + "cookie": "session", + "password": "x", + "api_key": "k", + } + 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) { + require.Nil(t, SanitizeStoredCredentials(PlatformGrok, nil)) +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index cfc725170c..ba6162c3e3 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1307,7 +1307,7 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account } resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) if err != nil { - return nil, fmt.Errorf("grok native search upstream: %w", err) + return nil, &UpstreamFailoverError{StatusCode: http.StatusBadGateway, Reason: GatewayFailureReason("grok_search_transport")} } defer func() { _ = resp.Body.Close() }() respBytes, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) @@ -1319,6 +1319,9 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account if len(msg) > 200 { msg = msg[:200] } + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusPaymentRequired || resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 { + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBytes} + } return nil, fmt.Errorf("grok upstream %d: %s", resp.StatusCode, msg) } return respBytes, nil diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 1297caed28..fee26dfec9 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -813,19 +813,29 @@ func (s *GatewayService) calculateRecordUsageCost( return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier) } - // Grok native web_search / tool search: group search_price_per_1k when set. - if result.SearchCount > 0 { - price := groupSearchPricePer1kFromAPIKey(apiKey) - return s.billingService.CalculateSearchCost(result.SearchCount, price, multiplier) - } // Voice audio (TTS / STT / realtime) when present on the forward result. if result.AudioUsage != nil { cfg := groupAudioPriceConfigFromAPIKey(apiKey) return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier) } - // Token 计费 - return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) + // Token 计费;SearchCount 为叠加 surcharge(不替代 token)。 + tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) + if result.SearchCount > 0 { + price := groupSearchPricePer1kFromAPIKey(apiKey) + if price == nil || *price <= 0 { + logger.LegacyPrintf("service.gateway", "[Billing] search_price_per_1k unset; search calls free group_model=%s count=%d", billingModel, result.SearchCount) + } + searchCost := s.billingService.CalculateSearchCost(result.SearchCount, price, multiplier) + if searchCost != nil && (searchCost.TotalCost > 0 || searchCost.ActualCost > 0) { + if tokenCost == nil { + return searchCost + } + tokenCost.TotalCost += searchCost.TotalCost + tokenCost.ActualCost += searchCost.ActualCost + } + } + return tokenCost } // compositeBillableModel 决定 composite 分组请求的计费模型:来源覆盖把计费模型 diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index bd3924edf8..f805c8c5f1 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -319,8 +319,9 @@ func isKnownGrokFreeAccount(account *Account) bool { paidSignal = true } } - if billing.UsagePercent != nil || billing.UsedPercent != nil || - (billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) { + // Paid evidence: explicit plan name or monthly dollar cap — NOT mere presence of + // weekly/monthly usage percent (Free accounts also report creditUsagePercent). + if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 { paidSignal = true } // xAI deliberately reports an empty plan for Free accounts; only paid @@ -362,10 +363,23 @@ 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. + 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 && (freeSignal || inferredFreeSignal) + return !paidSignal && inferredFreeSignal } func isGrokFreeSubscriptionTier(tier string) bool { diff --git a/backend/internal/service/openai_gateway_search_surcharge_test.go b/backend/internal/service/openai_gateway_search_surcharge_test.go new file mode 100644 index 0000000000..f77776709d --- /dev/null +++ b/backend/internal/service/openai_gateway_search_surcharge_test.go @@ -0,0 +1,85 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCalculateOpenAIRecordUsageCost_SearchIsAdditiveToTokens(t *testing.T) { + t.Parallel() + + price := 10.0 // $10 / 1k searches → 100 searches = $1.0 + svc := &OpenAIGatewayService{ + billingService: newTestBillingService(), + } + apiKey := &APIKey{ + Group: &Group{ + SearchPricePer1k: &price, + }, + } + + // claude-sonnet-4 fallback: Input $3/MTok, Output $15/MTok + // 1000 in + 500 out → 0.003 + 0.0075 = 0.0105 + // + 100 searches → +1.0 → total 1.0105 + cost, err := svc.calculateOpenAIRecordUsageCost( + context.Background(), + &OpenAIForwardResult{SearchCount: 100}, + apiKey, + []string{"claude-sonnet-4"}, + 1.0, + 1.0, + 1.0, + 1.0, + UsageTokens{InputTokens: 1000, OutputTokens: 500}, + "", + false, + ) + require.NoError(t, err) + require.NotNil(t, cost) + require.InDelta(t, 1.0105, cost.ActualCost, 1e-9) + require.InDelta(t, 1.0105, cost.TotalCost, 1e-9) +} + +func TestCalculateOpenAIRecordUsageCost_SearchOnlyWhenNoTokenPricing(t *testing.T) { + t.Parallel() + + price := 10.0 + svc := &OpenAIGatewayService{ + billingService: newTestBillingService(), + } + apiKey := &APIKey{ + Group: &Group{SearchPricePer1k: &price}, + } + // Empty model list: token path fails; search-only surcharge still bills. + cost, err := svc.calculateOpenAIRecordUsageCost( + context.Background(), + &OpenAIForwardResult{SearchCount: 100}, + apiKey, + nil, + 1.0, + 1.0, + 1.0, + 1.0, + UsageTokens{}, + "", + false, + ) + require.NoError(t, err) + require.NotNil(t, cost) + require.InDelta(t, 1.0, cost.ActualCost, 1e-9) +} + +func TestGroupMediaPricingLooksIncomplete_VideoModelPricesComplete(t *testing.T) { + t.Parallel() + require.True(t, groupMediaPricingLooksIncomplete(nil)) + require.True(t, groupMediaPricingLooksIncomplete(&Group{})) + require.False(t, groupMediaPricingLooksIncomplete(&Group{ + VideoModelPrices: map[string]map[string]float64{ + "grok-imagine-video": {"720p": 0.1}, + }, + })) +} diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 0ad8cfed8b..ec83534242 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -447,11 +447,6 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost( return s.calculateOpenAIVideoCost(ctx, billingModel, apiKey, result, videoMultiplier), nil } } - if result != nil && result.SearchCount > 0 { - // Grok native web_search / tool search: group search_price_per_1k when set. - price := groupSearchPricePer1kFromAPIKey(apiKey) - return s.billingService.CalculateSearchCost(result.SearchCount, price, webSearchMultiplier), nil - } if result != nil && result.AudioUsage != nil { cfg := groupAudioPriceConfigFromAPIKey(apiKey) return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, webSearchMultiplier), nil @@ -463,33 +458,64 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost( return s.calculateOpenAIImageCost(ctx, billingModel, apiKey, result, imageMultiplier), nil } } - if len(billingModels) == 0 || billingModel == "" { - return nil, errors.New("openai usage billing model is empty") - } + + // Token path (optional search surcharge is additive — never replaces token cost). + var tokenCost *CostBreakdown var lastErr error - for _, candidate := range billingModels { - candidate = strings.TrimSpace(candidate) - if candidate == "" { - continue + if len(billingModels) > 0 && billingModel != "" { + for _, candidate := range billingModels { + candidate = strings.TrimSpace(candidate) + if candidate == "" { + continue + } + cost, err := s.calculateOpenAIRecordUsageTokenCost( + ctx, + apiKey, + candidate, + multiplier, + tokens, + serviceTier, + longContextBillingEnabled, + ) + if err == nil { + tokenCost = cost + break + } + lastErr = err } - cost, err := s.calculateOpenAIRecordUsageTokenCost( - ctx, - apiKey, - candidate, - multiplier, - tokens, - serviceTier, - longContextBillingEnabled, - ) - if err == nil { - return cost, nil + } + // Search-only (e.g. no model / zero tokens): still bill search when priced. + searchCost := (*CostBreakdown)(nil) + if result != nil && result.SearchCount > 0 { + price := groupSearchPricePer1kFromAPIKey(apiKey) + if price == nil || *price <= 0 { + logger.L().Warn("openai_usage.grok_search_price_per_1k_unset_free", + zap.Int("search_count", result.SearchCount), + zap.String("model", billingModel), + ) } - lastErr = err + searchCost = s.billingService.CalculateSearchCost(result.SearchCount, price, webSearchMultiplier) } - if lastErr == nil { - lastErr = errors.New("no non-empty billing model candidates") + + if tokenCost == nil && searchCost == nil { + if lastErr == nil { + if len(billingModels) == 0 || billingModel == "" { + return nil, errors.New("openai usage billing model is empty") + } + lastErr = errors.New("no non-empty billing model candidates") + } + return nil, fmt.Errorf("calculate OpenAI usage cost failed for billing models %s: %w", strings.Join(billingModels, ","), lastErr) } - return nil, fmt.Errorf("calculate OpenAI usage cost failed for billing models %s: %w", strings.Join(billingModels, ","), lastErr) + if tokenCost == nil { + return searchCost, nil + } + if searchCost == nil || (searchCost.TotalCost == 0 && searchCost.ActualCost == 0) { + return tokenCost, nil + } + // Additive: tokens + search surcharge. + tokenCost.TotalCost += searchCost.TotalCost + tokenCost.ActualCost += searchCost.ActualCost + return tokenCost, nil } func isGrokVideoBillingModel(model string) bool { @@ -500,6 +526,8 @@ func isGrokVideoUsageResult(result *OpenAIForwardResult, billingModels []string) if result == nil || result.VideoCount <= 0 { return false } + // VideoCount alone is authoritative for async video completion billing. + // Prefer model-family match when present; never drop video mode on rename/mapping. candidates := append([]string{}, billingModels...) candidates = append(candidates, result.BillingModel, result.Model, result.UpstreamModel) for _, candidate := range candidates { @@ -507,7 +535,7 @@ func isGrokVideoUsageResult(result *OpenAIForwardResult, billingModels []string) return true } } - return false + return true } func isUsagePricingUnavailableError(err error) bool { @@ -675,6 +703,11 @@ func groupMediaPricingLooksIncomplete(group *Group) bool { if group.ImageRateMultiplier != 0 || group.VideoRateMultiplier != 0 { return false } + // Per-model video prices are first-class billing config; a projection that + // already carries them is complete enough to skip a DB refresh. + if len(group.VideoModelPrices) > 0 { + return false + } return group.ImagePrice1K == nil && group.ImagePrice2K == nil && group.ImagePrice4K == nil && group.VideoPrice480P == nil && group.VideoPrice720P == nil && group.VideoPrice1080P == nil }