fix(grok): 二轮评审残留 — Search 叠加计费与 fail-closed 安全

- SearchCost 叠加 token(openai/gateway),未定价 warn
- Token URL 校验失败回落 DefaultTokenURL,禁止 Effective 旁路
- free 判定收窄 paidSignal(仅 plan/月额度),usage% 不否决 free
- web_search: mandatory 计费、uuid request_id、上游 failover 重选账号
- 调度阈值:7d/30d 不跨期 until + 48h stale 可选跳过
- SanitizeStoredCredentials 接入 create/update/bulk/SSO/ApplyOAuth
- ApplyOAuth 成功清除 grok_needs_reauth;SSO 允许 header_override_enabled
- VideoModelPrices 视为媒体定价完整;realtime 正常关闭仍计费
This commit is contained in:
IanShaw027
2026-08-07 17:40:21 +08:00
parent ccf7ba3ba7
commit 245d069602
18 changed files with 471 additions and 86 deletions
@@ -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,
@@ -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) {
@@ -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 {
+48 -11
View File
@@ -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,
+19 -4
View File
@@ -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)
@@ -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")
}
}
@@ -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}
}
@@ -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(`<html>403 Forbidden</html>`))
}
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)
}
@@ -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 {
+2 -2
View File
@@ -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 {
@@ -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{
@@ -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
}
@@ -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))
}
+4 -1
View File
@@ -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
@@ -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 分组请求的计费模型:来源覆盖把计费模型
@@ -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 {
@@ -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},
},
}))
}
@@ -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
}