Merge pull request #4497 from heathermhuang/agent/fix-grok-media-fallback-4471

fix(grok): fail closed for ineligible OAuth media
This commit is contained in:
Wesley Liddick
2026-07-18 20:41:24 +08:00
committed by GitHub
16 changed files with 317 additions and 47 deletions
@@ -61,7 +61,7 @@ type AccountHandler struct {
sessionLimitCache service.SessionLimitCache
rpmCache service.RPMCache
tokenCacheInvalidator service.TokenCacheInvalidator
grokImportProber grokUsageProber
grokImportProber grokImportProber
upstreamBillingProbe *service.UpstreamBillingProbeService
}
@@ -15,12 +15,12 @@ const (
grokImportProbeTimeout = 25 * time.Second
)
type grokUsageProber interface {
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
type grokImportProber interface {
QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
}
type grokImportProbeTask struct {
prober grokUsageProber
prober grokImportProber
accountID int64
}
@@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm
}
}
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) {
if s == nil || prober == nil || account == nil || account.ID <= 0 {
return
}
@@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
return task, true
}
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error(
@@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64)
// 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)
result, err := prober.QueryQuota(ctx, accountID)
if err != nil {
slog.Warn(
"grok_import_active_probe_failed",
@@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
}
}
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
_, deadlineSeen := ctx.Deadline()
s.mu.Lock()
s.calls[accountID]++
@@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (
return nil, failure
}
return &service.GrokQuotaProbeResult{
Source: "active_probe",
Source: "hybrid_probe",
Model: "grok-4.5",
StatusCode: 200,
ResetSupported: false,
@@ -23,7 +23,7 @@ type GrokOAuthHandler struct {
grokOAuthService *service.GrokOAuthService
adminService service.AdminService
quotaService *service.GrokQuotaService
importProber grokUsageProber
importProber grokImportProber
reconciler service.GrokOAuthReconciler
}