From a1b5c75ca334c972a6bc62ef99baf35ab1eee716 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Tue, 14 Jul 2026 12:23:43 +0800 Subject: [PATCH] feat(grok): probe newly imported OAuth accounts --- backend/cmd/server/wire_gen.go | 2 +- .../internal/handler/admin/account_data.go | 1 + .../internal/handler/admin/account_handler.go | 3 + .../handler/admin/grok_import_probe.go | 222 ++++++++++++++++++ .../admin/grok_import_probe_handler_test.go | 119 ++++++++++ .../handler/admin/grok_import_probe_test.go | 213 +++++++++++++++++ .../handler/admin/grok_oauth_handler.go | 4 + backend/internal/handler/wire.go | 2 +- .../internal/service/grok_quota_service.go | 20 +- .../service/grok_quota_service_test.go | 50 ++++ 10 files changed, 630 insertions(+), 6 deletions(-) create mode 100644 backend/internal/handler/admin/grok_import_probe.go create mode 100644 backend/internal/handler/admin/grok_import_probe_handler_test.go create mode 100644 backend/internal/handler/admin/grok_import_probe_test.go diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index ba5b78a72f..75786b9c70 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -196,7 +196,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService) accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService) crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig) - accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator) + accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService) adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService) dataManagementService := service.NewDataManagementService() dataManagementHandler := admin.NewDataManagementHandler(dataManagementService) diff --git a/backend/internal/handler/admin/account_data.go b/backend/internal/handler/admin/account_data.go index bf872c4826..e44d726fd6 100644 --- a/backend/internal/handler/admin/account_data.go +++ b/backend/internal/handler/admin/account_data.go @@ -460,6 +460,7 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest) if created.Platform == service.PlatformAntigravity && created.Type == service.AccountTypeOAuth { privacyAccounts = append(privacyAccounts, created) } + h.scheduleGrokImportProbe(created) result.AccountCreated++ } diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index b886728159..e4ed5b46b0 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -60,6 +60,7 @@ type AccountHandler struct { sessionLimitCache service.SessionLimitCache rpmCache service.RPMCache tokenCacheInvalidator service.TokenCacheInvalidator + grokImportProber grokUsageProber } // NewAccountHandler creates a new admin account handler @@ -855,6 +856,7 @@ func (h *AccountHandler) Create(c *gin.Context) { // OpenAI APIKey 账号创建后异步探测上游 /v1/responses 能力。 // 探测失败不影响账号创建响应。 h.scheduleOpenAIResponsesProbe(createdAccount) + h.scheduleGrokImportProbe(createdAccount) response.Success(c, result.Data) } @@ -1667,6 +1669,7 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) { } // OpenAI APIKey 账号异步探测 /v1/responses 能力。 h.scheduleOpenAIResponsesProbe(account) + h.scheduleGrokImportProbe(account) success++ results = append(results, gin.H{ "name": item.Name, diff --git a/backend/internal/handler/admin/grok_import_probe.go b/backend/internal/handler/admin/grok_import_probe.go new file mode 100644 index 0000000000..d301776419 --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe.go @@ -0,0 +1,222 @@ +package admin + +import ( + "context" + "log/slog" + "sync" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +const ( + grokImportProbeConcurrency = 3 + grokImportProbeTimeout = 25 * time.Second +) + +type grokUsageProber interface { + ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) +} + +type grokImportProbeTask struct { + prober grokUsageProber + accountID int64 +} + +type grokImportProbeScheduler struct { + mu sync.Mutex + queue []grokImportProbeTask + concurrency int + workers int + maxWorkers int + timeout time.Duration +} + +type grokImportProbeSchedulerSnapshot struct { + Queued int + Workers int + MaxWorkers int +} + +var defaultGrokImportProbeScheduler = newGrokImportProbeScheduler( + grokImportProbeConcurrency, + grokImportProbeTimeout, +) + +func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokImportProbeScheduler { + if concurrency <= 0 { + concurrency = 1 + } + if timeout <= 0 { + timeout = grokImportProbeTimeout + } + return &grokImportProbeScheduler{ + concurrency: concurrency, + timeout: timeout, + } +} + +func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) { + if s == nil || prober == nil || account == nil || account.ID <= 0 { + return + } + if account.Platform != service.PlatformGrok || account.Type != service.AccountTypeOAuth { + return + } + + s.mu.Lock() + s.queue = append(s.queue, grokImportProbeTask{prober: prober, accountID: account.ID}) + if s.workers < s.concurrency { + s.workers++ + if s.workers > s.maxWorkers { + s.maxWorkers = s.workers + } + go s.worker() + } + s.mu.Unlock() +} + +func (s *grokImportProbeScheduler) worker() { + for { + task, ok := s.nextTask() + if !ok { + return + } + s.run(task.prober, task.accountID) + } +} + +func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.queue) == 0 { + s.workers-- + return grokImportProbeTask{}, false + } + task := s.queue[0] + s.queue[0] = grokImportProbeTask{} + s.queue = s.queue[1:] + if len(s.queue) == 0 { + s.queue = nil + } + return task, true +} + +func (s *grokImportProbeScheduler) snapshot() grokImportProbeSchedulerSnapshot { + if s == nil { + return grokImportProbeSchedulerSnapshot{} + } + s.mu.Lock() + defer s.mu.Unlock() + return grokImportProbeSchedulerSnapshot{ + Queued: len(s.queue), + Workers: s.workers, + MaxWorkers: s.maxWorkers, + } +} + +func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) { + defer func() { + if recovered := recover(); recovered != nil { + slog.Error( + "grok_import_active_probe_panic", + "account_id", accountID, + "recovery_type", panicType(recovered), + ) + } + }() + + // Queue time is intentionally excluded: every imported account is probed, + // while this timeout only bounds the actual upstream probe execution. + ctx, cancel := context.WithTimeout(context.Background(), s.timeout) + defer cancel() + result, err := prober.ProbeUsage(ctx, accountID) + if err != nil { + slog.Warn( + "grok_import_active_probe_failed", + "account_id", accountID, + "status", infraerrors.Code(err), + "reason", infraerrors.Reason(err), + ) + return + } + if result == nil { + slog.Warn( + "grok_import_active_probe_failed", + "account_id", accountID, + "reason", "empty_result", + ) + return + } + + slog.Info( + "grok_import_active_probe_completed", + "account_id", accountID, + "model", result.Model, + "status", result.StatusCode, + "headers_observed", result.HeadersObserved, + ) +} + +func panicType(value any) string { + switch value.(type) { + case string: + return "string" + case error: + return "error" + default: + return "unknown" + } +} + +func (h *AccountHandler) scheduleGrokImportProbe(account *service.Account) { + if h == nil { + return + } + defaultGrokImportProbeScheduler.schedule(h.grokImportProber, account) +} + +func (h *GrokOAuthHandler) scheduleGrokImportProbe(account *service.Account) { + if h == nil { + return + } + defaultGrokImportProbeScheduler.schedule(h.importProber, account) +} + +// ProvideAccountHandler injects the Grok active prober for production while +// keeping NewAccountHandler convenient for focused unit tests. +func ProvideAccountHandler( + adminService service.AdminService, + oauthService *service.OAuthService, + openaiOAuthService *service.OpenAIOAuthService, + geminiOAuthService *service.GeminiOAuthService, + antigravityOAuthService *service.AntigravityOAuthService, + rateLimitService *service.RateLimitService, + accountUsageService *service.AccountUsageService, + accountTestService *service.AccountTestService, + concurrencyService *service.ConcurrencyService, + crsSyncService *service.CRSSyncService, + sessionLimitCache service.SessionLimitCache, + rpmCache service.RPMCache, + tokenCacheInvalidator service.TokenCacheInvalidator, + grokQuotaService *service.GrokQuotaService, +) *AccountHandler { + handler := NewAccountHandler( + adminService, + oauthService, + openaiOAuthService, + geminiOAuthService, + antigravityOAuthService, + rateLimitService, + accountUsageService, + accountTestService, + concurrencyService, + crsSyncService, + sessionLimitCache, + rpmCache, + tokenCacheInvalidator, + ) + handler.grokImportProber = grokQuotaService + return handler +} diff --git a/backend/internal/handler/admin/grok_import_probe_handler_test.go b/backend/internal/handler/admin/grok_import_probe_handler_test.go new file mode 100644 index 0000000000..489a13ae6d --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe_handler_test.go @@ -0,0 +1,119 @@ +//go:build unit + +package admin + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type grokImportAdminService struct { + *stubAdminService + mu sync.Mutex + nextID int64 +} + +func newGrokImportAdminService() *grokImportAdminService { + return &grokImportAdminService{ + stubAdminService: newStubAdminService(), + nextID: 500, + } +} + +func (s *grokImportAdminService) CreateAccount(_ context.Context, input *service.CreateAccountInput) (*service.Account, error) { + s.mu.Lock() + s.nextID++ + id := s.nextID + s.mu.Unlock() + return &service.Account{ + ID: id, + Name: input.Name, + Platform: input.Platform, + Type: input.Type, + Credentials: input.Credentials, + Extra: input.Extra, + ProxyID: input.ProxyID, + Concurrency: input.Concurrency, + Status: service.StatusActive, + Schedulable: true, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }, nil +} + +type grokImportOAuthClientStub struct{} + +func (grokImportOAuthClientStub) ExchangeCode(context.Context, string, string, string, string, string) (*xai.TokenResponse, error) { + return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil +} + +func (grokImportOAuthClientStub) RefreshToken(context.Context, string, string, string) (*xai.TokenResponse, error) { + return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil +} + +func (grokImportOAuthClientStub) ConvertSSOToBuild(context.Context, string, string) (*xai.TokenResponse, error) { + return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil +} + +func TestGrokSSOBatchImportKeepsCreatedAccountsWhenOneAutomaticProbeFails(t *testing.T) { + gin.SetMode(gin.TestMode) + adminService := newGrokImportAdminService() + oauthService := service.NewGrokOAuthService(nil, grokImportOAuthClientStub{}) + defer oauthService.Stop() + prober := newGrokImportProbeStub(3) + prober.failures[502] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body") + handler := NewGrokOAuthHandler(oauthService, adminService, nil) + handler.importProber = prober + + router := gin.New() + router.POST("/api/v1/admin/grok/sso-to-oauth", handler.CreateAccountsFromSSO) + recorder := httptest.NewRecorder() + request := httptest.NewRequest( + http.MethodPost, + "/api/v1/admin/grok/sso-to-oauth", + strings.NewReader(`{"sso_tokens":["sso-one","sso-two","sso-three"]}`), + ) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Body.String(), `"created"`) + require.NotContains(t, recorder.Body.String(), `GROK_TEST_PROBE_FAILED`) + for i := 0; i < 3; i++ { + awaitGrokProbeSignal(t, prober.done) + } + calls, _, _ := prober.snapshot() + require.Equal(t, map[int64]int{501: 1, 502: 1, 503: 1}, calls) +} + +func TestAccountCreateWithoutAutomaticGrokProbeServiceStillSucceeds(t *testing.T) { + gin.SetMode(gin.TestMode) + handler := NewAccountHandler( + newGrokImportAdminService(), + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + + router := gin.New() + router.POST("/api/v1/admin/accounts", handler.Create) + recorder := httptest.NewRecorder() + request := httptest.NewRequest( + http.MethodPost, + "/api/v1/admin/accounts", + strings.NewReader(`{"name":"grok-rt","platform":"grok","type":"oauth","credentials":{"refresh_token":"secret"}}`), + ) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusOK, recorder.Code) +} diff --git a/backend/internal/handler/admin/grok_import_probe_test.go b/backend/internal/handler/admin/grok_import_probe_test.go new file mode 100644 index 0000000000..27d0b744e2 --- /dev/null +++ b/backend/internal/handler/admin/grok_import_probe_test.go @@ -0,0 +1,213 @@ +//go:build unit + +package admin + +import ( + "bytes" + "context" + "log/slog" + "sync" + "testing" + "time" + + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +type grokImportProbeStub struct { + mu sync.Mutex + calls map[int64]int + failures map[int64]error + active int + maxActive int + deadlineSeen bool + block <-chan struct{} + started chan int64 + done chan int64 +} + +func newGrokImportProbeStub(buffer int) *grokImportProbeStub { + return &grokImportProbeStub{ + calls: make(map[int64]int), + failures: make(map[int64]error), + started: make(chan int64, buffer), + done: make(chan int64, buffer), + } +} + +func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) { + _, deadlineSeen := ctx.Deadline() + s.mu.Lock() + s.calls[accountID]++ + s.active++ + if s.active > s.maxActive { + s.maxActive = s.active + } + s.deadlineSeen = s.deadlineSeen || deadlineSeen + s.mu.Unlock() + + s.started <- accountID + var ctxErr error + if s.block != nil { + select { + case <-s.block: + case <-ctx.Done(): + ctxErr = ctx.Err() + } + } + + s.mu.Lock() + s.active-- + failure := s.failures[accountID] + s.mu.Unlock() + s.done <- accountID + if ctxErr != nil { + return nil, ctxErr + } + if failure != nil { + return nil, failure + } + return &service.GrokQuotaProbeResult{ + Source: "active_probe", + Model: "grok-4.5", + StatusCode: 200, + ResetSupported: false, + }, nil +} + +func (s *grokImportProbeStub) snapshot() (map[int64]int, int, bool) { + s.mu.Lock() + defer s.mu.Unlock() + calls := make(map[int64]int, len(s.calls)) + for id, count := range s.calls { + calls[id] = count + } + return calls, s.maxActive, s.deadlineSeen +} + +func newGrokOAuthImportAccount(id int64) *service.Account { + return &service.Account{ + ID: id, + Platform: service.PlatformGrok, + Type: service.AccountTypeOAuth, + } +} + +func awaitGrokProbeSignal(t *testing.T, signals <-chan int64) int64 { + t.Helper() + select { + case id := <-signals: + return id + case <-time.After(time.Second): + t.Fatal("timed out waiting for Grok import probe") + return 0 + } +} + +func TestGrokImportProbeSchedulerProbesSingleAccountOnce(t *testing.T) { + scheduler := newGrokImportProbeScheduler(1, time.Second) + prober := newGrokImportProbeStub(1) + + scheduler.schedule(prober, newGrokOAuthImportAccount(101)) + require.Equal(t, int64(101), awaitGrokProbeSignal(t, prober.done)) + + calls, maxActive, deadlineSeen := prober.snapshot() + require.Equal(t, map[int64]int{101: 1}, calls) + require.Equal(t, 1, maxActive) + require.True(t, deadlineSeen) + require.Eventually(t, func() bool { + snapshot := scheduler.snapshot() + return snapshot.Queued == 0 && snapshot.Workers == 0 + }, time.Second, 10*time.Millisecond) +} + +func TestGrokImportProbeSchedulerQueuesBatchWithoutPerTaskGoroutines(t *testing.T) { + const taskCount = 100 + release := make(chan struct{}) + scheduler := newGrokImportProbeScheduler(3, time.Second) + prober := newGrokImportProbeStub(taskCount) + prober.block = release + prober.failures[150] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body") + + for id := int64(101); id < 101+taskCount; id++ { + scheduler.schedule(prober, newGrokOAuthImportAccount(id)) + } + for i := 0; i < 3; i++ { + awaitGrokProbeSignal(t, prober.started) + } + snapshot := scheduler.snapshot() + require.Equal(t, 97, snapshot.Queued) + require.Equal(t, 3, snapshot.Workers) + require.Equal(t, 3, snapshot.MaxWorkers) + select { + case id := <-prober.started: + t.Fatalf("probe %d started before a concurrency slot was released", id) + case <-time.After(75 * time.Millisecond): + } + close(release) + for i := 0; i < taskCount; i++ { + awaitGrokProbeSignal(t, prober.done) + } + + calls, maxActive, _ := prober.snapshot() + require.Len(t, calls, taskCount) + for id := int64(101); id < 101+taskCount; id++ { + require.Equal(t, 1, calls[id]) + } + require.Equal(t, 3, maxActive) + require.Eventually(t, func() bool { + snapshot = scheduler.snapshot() + return snapshot.Queued == 0 && snapshot.Workers == 0 + }, time.Second, 10*time.Millisecond) + require.Equal(t, 3, snapshot.MaxWorkers) +} + +func TestGrokImportProbeSchedulerTimeoutCancelsProbe(t *testing.T) { + neverRelease := make(chan struct{}) + scheduler := newGrokImportProbeScheduler(1, 20*time.Millisecond) + prober := newGrokImportProbeStub(1) + prober.block = neverRelease + + scheduler.schedule(prober, newGrokOAuthImportAccount(201)) + require.Equal(t, int64(201), awaitGrokProbeSignal(t, prober.done)) + + calls, _, _ := prober.snapshot() + require.Equal(t, 1, calls[201]) +} + +func TestGrokImportProbeSchedulerSkipsMissingServiceAndNonGrokAccounts(t *testing.T) { + scheduler := newGrokImportProbeScheduler(1, time.Second) + prober := newGrokImportProbeStub(1) + + scheduler.schedule(nil, newGrokOAuthImportAccount(301)) + scheduler.schedule(prober, &service.Account{ID: 302, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}) + scheduler.schedule(prober, &service.Account{ID: 303, Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey}) + + select { + case id := <-prober.started: + t.Fatalf("unexpected probe for account %d", id) + case <-time.After(50 * time.Millisecond): + } + calls, _, _ := prober.snapshot() + require.Empty(t, calls) +} + +func TestGrokImportProbeFailureLogDoesNotIncludeErrorMessage(t *testing.T) { + var logs bytes.Buffer + previousLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil))) + defer slog.SetDefault(previousLogger) + + scheduler := newGrokImportProbeScheduler(1, time.Second) + prober := newGrokImportProbeStub(1) + prober.failures[401] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "refresh-token-secret") + scheduler.schedule(prober, newGrokOAuthImportAccount(401)) + awaitGrokProbeSignal(t, prober.done) + + require.Eventually(t, func() bool { + return bytes.Contains(logs.Bytes(), []byte("grok_import_active_probe_failed")) + }, time.Second, 10*time.Millisecond) + require.Contains(t, logs.String(), "GROK_TEST_PROBE_FAILED") + require.NotContains(t, logs.String(), "refresh-token-secret") +} diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index dfe5632e30..1a309b7c9a 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -22,6 +22,7 @@ type GrokOAuthHandler struct { grokOAuthService *service.GrokOAuthService adminService service.AdminService quotaService *service.GrokQuotaService + importProber grokUsageProber } func NewGrokOAuthHandler( @@ -33,6 +34,7 @@ func NewGrokOAuthHandler( grokOAuthService: grokOAuthService, adminService: adminService, quotaService: quotaService, + importProber: quotaService, } } @@ -209,6 +211,7 @@ func (h *GrokOAuthHandler) CreateAccountFromOAuth(c *gin.Context) { response.ErrorFrom(c, err) return } + h.scheduleGrokImportProbe(account) response.Success(c, dto.AccountFromService(account)) } @@ -345,6 +348,7 @@ func (h *GrokOAuthHandler) createAccountFromSSOToken(ctx context.Context, req Gr if err != nil { return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Name: name, Email: tokenInfo.Email, Error: grokSSOImportErrorMessage(err)}} } + h.scheduleGrokImportProbe(account) return grokSSOImportWorkerResult{ created: true, item: GrokSSOToOAuthItemResult{ diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index cfbb72554c..67380fed33 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet( admin.NewDashboardHandler, admin.NewUserHandler, admin.NewGroupHandler, - admin.NewAccountHandler, + admin.ProvideAccountHandler, admin.NewAnnouncementHandler, admin.NewDataManagementHandler, admin.NewBackupHandler, diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 7817a00ec4..2219a75c15 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -183,10 +183,22 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr return result, nil } if resp.StatusCode >= 400 { - bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 240)) - bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240) - slog.Warn("grok_quota_probe_failed", "account_id", account.ID, "model", probeModel, "status", resp.StatusCode, "body", bodyText) - return nil, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "upstream returned %d for probe model %q: %s", resp.StatusCode, probeModel, bodyText) + const reason = "GROK_QUOTA_PROBE_UPSTREAM_ERROR" + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10)) + slog.Warn( + "grok_quota_probe_failed", + "account_id", account.ID, + "model", probeModel, + "status", resp.StatusCode, + "reason", reason, + ) + return nil, infraerrors.Newf( + mapUpstreamStatus(resp.StatusCode), + reason, + "upstream returned %d for probe model %q", + resp.StatusCode, + probeModel, + ) } return result, nil } diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 7a72f80e9f..ba9b6cbcb8 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -3,8 +3,10 @@ package service import ( + "bytes" "context" "io" + "log/slog" "net/http" "strconv" "strings" @@ -298,6 +300,54 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T require.Contains(t, infraerrors.Message(err), `probe model "grok-4.5"`) } +func TestGrokQuotaServiceProbeUsageRedactsUpstreamErrorBodyFromErrorAndLogs(t *testing.T) { + const upstreamSecret = "upstream-secret-refresh-token" + account := &Account{ + ID: 49, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{49: account}, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader( + `{"error":"` + upstreamSecret + `","detail":"credential rejected"}`, + )), + }} + svc := NewGrokQuotaService( + repo, + nil, + NewGrokTokenProvider(repo, nil), + upstream, + ) + + var logs bytes.Buffer + previousLogger := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil))) + defer slog.SetDefault(previousLogger) + + _, err := svc.ProbeUsage(context.Background(), account.ID) + require.Error(t, err) + require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) + require.Contains(t, infraerrors.Message(err), `probe model "grok-4.5"`) + require.NotContains(t, err.Error(), upstreamSecret) + require.NotContains(t, infraerrors.Message(err), upstreamSecret) + require.Contains(t, logs.String(), "GROK_QUOTA_PROBE_UPSTREAM_ERROR") + require.NotContains(t, logs.String(), upstreamSecret) + require.NotContains(t, logs.String(), "credential rejected") + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) +} + func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T) { t.Parallel()