From ccf7ba3ba768fadfd57a41d6ee37f23289777288 Mon Sep 17 00:00:00 2001 From: IanShaw027 Date: Fri, 7 Aug 2026 17:21:01 +0800 Subject: [PATCH] =?UTF-8?q?fix(grok):=20P2=20=E5=88=86=E5=B1=82=E6=B8=85?= =?UTF-8?q?=E7=90=86=E3=80=81=E4=B8=BB=E5=AF=B9=E8=AF=9D=20SearchCount=20?= =?UTF-8?q?=E4=B8=8E=E8=B0=83=E5=BA=A6=207d/30d?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - DoGrokNativeResponsesJSON 去 gin.Context，UA 固定 CLI 身份 - GrokOAuthService 去掉 redis 直依赖，session store 由 wire 注入 - 主对话/流式响应统计 web_search_call/x_search/tool_search 写入 SearchCount - Grok 调度阈值候选并入官方 weekly/monthly billing % --- .../internal/handler/gateway_web_search.go | 2 +- .../account_scheduling_threshold_eval.go | 33 ++++- .../account_scheduling_threshold_eval_test.go | 36 +++++ backend/internal/service/gateway_service.go | 6 +- .../internal/service/grok_oauth_service.go | 11 +- backend/internal/service/grok_search_count.go | 132 ++++++++++++++++++ .../service/grok_search_count_test.go | 38 +++++ .../service/openai_gateway_forward.go | 9 ++ .../openai_gateway_response_handling.go | 7 + backend/internal/service/wire.go | 8 +- 10 files changed, 271 insertions(+), 11 deletions(-) create mode 100644 backend/internal/service/grok_search_count.go create mode 100644 backend/internal/service/grok_search_count_test.go diff --git a/backend/internal/handler/gateway_web_search.go b/backend/internal/handler/gateway_web_search.go index 039a9c2487..84b99935b8 100644 --- a/backend/internal/handler/gateway_web_search.go +++ b/backend/internal/handler/gateway_web_search.go @@ -231,7 +231,7 @@ func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Conte } bodyBytes, _ := json.Marshal(searchBody) - respBytes, err := h.gatewayService.DoGrokNativeResponsesJSON(ctx, c, account, bodyBytes) + respBytes, err := h.gatewayService.DoGrokNativeResponsesJSON(ctx, account, bodyBytes) if err != nil { return nil, "", err } diff --git a/backend/internal/service/account_scheduling_threshold_eval.go b/backend/internal/service/account_scheduling_threshold_eval.go index 49ca80db35..efa09fd104 100644 --- a/backend/internal/service/account_scheduling_threshold_eval.go +++ b/backend/internal/service/account_scheduling_threshold_eval.go @@ -286,7 +286,7 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand if account == nil { return nil } - return []*accountSchedulingThresholdCandidate{ + candidates := []*accountSchedulingThresholdCandidate{ { window: "quota", scope: "grok", @@ -294,6 +294,37 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand until: parseSchedulingResetAt(account.Extra["grok_sched_reset_at"]), }, } + // 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. + 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)), + }) + } + } + return candidates } func pickLatestResetSchedulingCandidate(candidates []*accountSchedulingThresholdCandidate, threshold int, now time.Time) *accountSchedulingThresholdCandidate { diff --git a/backend/internal/service/account_scheduling_threshold_eval_test.go b/backend/internal/service/account_scheduling_threshold_eval_test.go index 43546be923..284ca6945f 100644 --- a/backend/internal/service/account_scheduling_threshold_eval_test.go +++ b/backend/internal/service/account_scheduling_threshold_eval_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" ) @@ -306,3 +307,38 @@ func TestEvaluateAccountSchedulingThreshold_GrokUsesConfiguredThresholds(t *test require.NotNil(t, decision.Until) require.True(t, wantUntil.Equal(*decision.Until)) } + +func TestEvaluateAccountSchedulingThreshold_GrokBillingWindows(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 6, 3, 12, 0, 0, 0, time.UTC) + weeklyEnd := now.Add(3 * time.Hour) + weeklyPct := 95.0 + monthlyPct := 40.0 + account := &Account{ + Platform: PlatformGrok, + Extra: map[string]any{ + // Header-based quota below threshold. + "grok_sched_utilization": 50.0, + "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, + BillingPeriodEnd: now.Add(10 * 24 * time.Hour).Format(time.RFC3339), + }, + }, + } + + decision := EvaluateAccountSchedulingThreshold(account, map[string]int{ + PlatformGrok: 90, + }, now) + + require.True(t, decision.ShouldPause) + require.Equal(t, "grok_billing", decision.Scope) + require.Equal(t, "seven_day", decision.Window) + require.InDelta(t, 95.0, decision.UsedPercent, 1e-9) + require.NotNil(t, decision.Until) + require.True(t, weeklyEnd.Equal(*decision.Until)) +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 0b48c00bb0..cfc725170c 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -24,7 +24,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" "github.com/cespare/xxhash/v2" - "github.com/gin-gonic/gin" gocache "github.com/patrickmn/go-cache" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -1266,7 +1265,8 @@ func (s *GatewayService) getOAuthToken(ctx context.Context, account *Account) (s // DoGrokNativeResponsesJSON POSTs a non-streaming Responses body to the account's // Grok upstream and returns the raw JSON body. Used by /v1/web_search. -func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, c *gin.Context, account *Account, body []byte) ([]byte, error) { +// Gin-free: UA is always the pinned Grok CLI identity (resolveGrokUpstreamUserAgent ignores inbound). +func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account *Account, body []byte) ([]byte, error) { if s == nil || s.httpUpstream == nil { return nil, errors.New("http upstream not configured") } @@ -1298,7 +1298,7 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, c *gin.C upstreamReq.Header.Set("Authorization", "Bearer "+token) upstreamReq.Header.Set("Content-Type", "application/json") upstreamReq.Header.Set("Accept", "application/json") - upstreamReq.Header.Set("User-Agent", resolveGrokUpstreamUserAgent(c)) + upstreamReq.Header.Set("User-Agent", defaultGrokUpstreamUserAgent()) applyGrokCLIHeaders(upstreamReq.Header) account.ApplyHeaderOverrides(upstreamReq.Header) proxyURL := "" diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index 480124fe5a..24f1799093 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -11,7 +11,6 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" - "github.com/redis/go-redis/v9" ) const grokDefaultAccessTokenTTL = 6 * time.Hour @@ -35,13 +34,15 @@ func NewGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient, return service } -// WithRedisSessionStore enables cross-instance, single-use OAuth callbacks. -func (s *GrokOAuthService) WithRedisSessionStore(rdb *redis.Client) *GrokOAuthService { - if s != nil && rdb != nil { +// WithSessionStore replaces the in-memory OAuth session store (e.g. Redis-backed +// for cross-instance single-use callbacks). Redis wiring stays in Wire providers +// so this service package does not import go-redis (depguard). +func (s *GrokOAuthService) WithSessionStore(store *xai.SessionStore) *GrokOAuthService { + if s != nil && store != nil { if s.sessionStore != nil { s.sessionStore.Stop() } - s.sessionStore = xai.NewRedisSessionStore(rdb) + s.sessionStore = store } return s } diff --git a/backend/internal/service/grok_search_count.go b/backend/internal/service/grok_search_count.go new file mode 100644 index 0000000000..ad4e8b581b --- /dev/null +++ b/backend/internal/service/grok_search_count.go @@ -0,0 +1,132 @@ +package service + +import ( + "strings" + + "github.com/tidwall/gjson" +) + +// countGrokNativeSearchCallsFromJSONBytes counts completed native search tool +// calls in a Responses-style JSON body (output array or nested response.output). +// Counts: web_search_call, x_search_call, tool_search_call, and function_call +// named tool_search / web_search / x_search. +func countGrokNativeSearchCallsFromJSONBytes(body []byte) int { + if len(body) == 0 || !gjson.ValidBytes(body) { + return 0 + } + count := 0 + count += countGrokNativeSearchCallsInOutputArray(gjson.GetBytes(body, "output")) + count += countGrokNativeSearchCallsInOutputArray(gjson.GetBytes(body, "response.output")) + return count +} + +func countGrokNativeSearchCallsFromSSEBody(body string) int { + if strings.TrimSpace(body) == "" { + return 0 + } + total := 0 + seen := make(map[string]struct{}) + 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++ + } + }) + return total +} + +func countGrokNativeSearchCallsInSSEData(data []byte) int { + n, _ := countGrokNativeSearchCallsInSSEDataWithKeys(data) + return n +} + +func countGrokNativeSearchCallsInSSEDataWithKeys(data []byte) (int, []string) { + if len(data) == 0 || !gjson.ValidBytes(data) { + return 0, nil + } + eventType := strings.TrimSpace(gjson.GetBytes(data, "type").String()) + // Count once on item completion / completed response, not on every delta. + switch eventType { + case "response.output_item.done", "response.completed", "response.done": + default: + // Also allow bare item objects without envelope. + if eventType != "" { + return 0, nil + } + } + var keys []string + n := 0 + consider := func(item gjson.Result) { + if !isGrokNativeSearchOutputItem(item) { + return + } + n++ + key := firstNonEmpty( + strings.TrimSpace(item.Get("call_id").String()), + strings.TrimSpace(item.Get("id").String()), + strings.TrimSpace(item.Get("item.call_id").String()), + strings.TrimSpace(item.Get("item.id").String()), + ) + if key != "" { + keys = append(keys, key) + } + } + if item := gjson.GetBytes(data, "item"); item.Exists() { + consider(item) + } + gjson.GetBytes(data, "response.output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + gjson.GetBytes(data, "output").ForEach(func(_, item gjson.Result) bool { + consider(item) + return true + }) + // Bare output item event without nested item key. + if n == 0 && isGrokNativeSearchOutputItem(gjson.ParseBytes(data)) { + consider(gjson.ParseBytes(data)) + } + return n, keys +} + +func countGrokNativeSearchCallsInOutputArray(output gjson.Result) int { + if !output.IsArray() { + return 0 + } + count := 0 + output.ForEach(func(_, item gjson.Result) bool { + if isGrokNativeSearchOutputItem(item) { + count++ + } + return true + }) + return count +} + +func isGrokNativeSearchOutputItem(item gjson.Result) bool { + if !item.Exists() { + return false + } + itemType := strings.ToLower(strings.TrimSpace(item.Get("type").String())) + switch itemType { + case "web_search_call", "x_search_call", "tool_search_call": + return true + case "function_call", "custom_tool_call": + name := strings.ToLower(strings.TrimSpace(item.Get("name").String())) + return name == "web_search" || name == "x_search" || name == "tool_search" + default: + return false + } +} diff --git a/backend/internal/service/grok_search_count_test.go b/backend/internal/service/grok_search_count_test.go new file mode 100644 index 0000000000..457d8917e7 --- /dev/null +++ b/backend/internal/service/grok_search_count_test.go @@ -0,0 +1,38 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCountGrokNativeSearchCallsFromJSONBytes(t *testing.T) { + t.Parallel() + require.Equal(t, 0, countGrokNativeSearchCallsFromJSONBytes(nil)) + require.Equal(t, 0, countGrokNativeSearchCallsFromJSONBytes([]byte(`{"output":[]}`))) + body := []byte(`{"output":[ + {"type":"web_search_call","id":"ws1","status":"completed"}, + {"type":"x_search_call","id":"xs1"}, + {"type":"function_call","name":"tool_search","call_id":"ts1"}, + {"type":"function_call","name":"lookup","call_id":"other"} + ]}`) + require.Equal(t, 3, countGrokNativeSearchCallsFromJSONBytes(body)) +} + +func TestCountGrokNativeSearchCallsFromSSEBodyDedups(t *testing.T) { + t.Parallel() + sse := stringsJoin( + `data: {"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`, + `data: {"type":"response.output_item.done","item":{"type":"web_search_call","id":"ws1","call_id":"c1"}}`, + `data: {"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, 2, countGrokNativeSearchCallsFromSSEBody(sse)) +} + +func stringsJoin(lines ...string) string { + out := "" + for _, l := range lines { + out += l + "\n\n" + } + return out +} diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 0484ea7d36..f7a4e608ad 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -928,6 +928,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco var firstTokenMs *int responseID := "" imageCount := 0 + searchCount := 0 var imageOutputSizes []string if reqStream { streamResult, err := s.handleStreamingResponseWithReasoning(ctx, resp, c, account, startTime, originalModel, upstreamModel, reasoningEffortValue) @@ -939,6 +940,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco responseID = strings.TrimSpace(streamResult.responseID) imageCount = streamResult.imageCount imageOutputSizes = streamResult.imageOutputSizes + searchCount = streamResult.searchCount } else { nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel) if err != nil { @@ -948,6 +950,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco responseID = strings.TrimSpace(nonStreamResult.responseID) imageCount = nonStreamResult.imageCount imageOutputSizes = nonStreamResult.imageOutputSizes + searchCount = nonStreamResult.searchCount } s.bindHTTPResponseAccount(ctx, c, account, responseID) @@ -984,6 +987,12 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco forwardResult.ImageOutputSizes = imageOutputSizes forwardResult.BillingModel = imageBillingModel } + // Grok-native web_search / x_search / tool_search tool invocations (per-1k pricing). + // Token cost still applies separately when usage is present; search is additive only + // when search_price_per_1k is configured (nil price → $0 from CalculateSearchCost). + if searchCount > 0 && account != nil && account.IsGrok() { + forwardResult.SearchCount = searchCount + } return forwardResult, nil } } diff --git a/backend/internal/service/openai_gateway_response_handling.go b/backend/internal/service/openai_gateway_response_handling.go index dd83e3357d..2be20306c1 100644 --- a/backend/internal/service/openai_gateway_response_handling.go +++ b/backend/internal/service/openai_gateway_response_handling.go @@ -28,6 +28,7 @@ type openaiStreamingResult struct { responseID string imageCount int imageOutputSizes []string + searchCount int } type openaiNonStreamingResult struct { @@ -36,6 +37,7 @@ type openaiNonStreamingResult struct { responseID string imageCount int imageOutputSizes []string + searchCount int } func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string) (*openaiStreamingResult, error) { @@ -298,6 +300,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. streamOutputAccumulator := apicompat.NewBufferedResponseAccumulator() streamImageOutputs := make([]json.RawMessage, 0, 1) streamSeenImages := make(map[string]struct{}) + searchCounter := 0 resultWithUsage := func() *openaiStreamingResult { return &openaiStreamingResult{ usage: usage, @@ -305,6 +308,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. responseID: responseID, imageCount: imageCounter.Count(), imageOutputSizes: imageCounter.Sizes(), + searchCount: searchCounter, } } flushPending := func(disconnectMessage string) { @@ -472,6 +476,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context. line = "data: " + data } imageCounter.AddSSEData(dataBytes) + searchCounter += countGrokNativeSearchCallsInSSEData(dataBytes) // Correct Codex tool calls if needed (apply_patch -> edit, etc.) if correctedData, corrected := s.toolCorrector.CorrectToolCallsInSSEBytes(dataBytes); corrected { @@ -1200,6 +1205,7 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r responseID: extractOpenAIResponseIDFromJSONBytes(body), imageCount: countOpenAIResponseImageOutputsFromJSONBytes(body), imageOutputSizes: collectOpenAIResponseImageOutputSizesFromJSONBytes(body), + searchCount: countGrokNativeSearchCallsFromJSONBytes(body), }, nil } @@ -1294,6 +1300,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte responseID: extractOpenAIResponseIDFromJSONBytes(body), imageCount: countOpenAIImageOutputsFromSSEBody(bodyText), imageOutputSizes: collectOpenAIImageOutputSizesFromSSEBody(bodyText), + searchCount: countGrokNativeSearchCallsFromSSEBody(bodyText), }, nil } diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 4f60425899..0a285d99b7 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -10,13 +10,19 @@ import ( "github.com/Wei-Shaw/sub2api/internal/payment" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/google/wire" "github.com/redis/go-redis/v9" "go.uber.org/zap" ) func ProvideGrokOAuthService(proxyRepo ProxyRepository, oauthClient GrokOAuthClient, cfg *config.Config, redisClient *redis.Client) *GrokOAuthService { - return NewGrokOAuthService(proxyRepo, oauthClient, cfg).WithRedisSessionStore(redisClient) + svc := NewGrokOAuthService(proxyRepo, oauthClient, cfg) + // wire.go is depguard-exempt for redis; construct the Redis session store here. + if redisClient != nil { + svc = svc.WithSessionStore(xai.NewRedisSessionStore(redisClient)) + } + return svc } // BuildInfo contains build information