Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/handler/admin/grok_import_probe.go

@@ -0,0 +1,231 @@
+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
+	grokImportProbeQueueLimit  = 64
+)
+
+type grokImportProber interface {
+	QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
+}
+
+type grokImportProbeTask struct {
+	prober    grokImportProber
+	accountID int64
+}
+
+type grokImportProbeScheduler struct {
+	mu          sync.Mutex
+	queue       []grokImportProbeTask
+	pending     map[int64]struct{}
+	inFlight    map[int64]struct{}
+	concurrency int
+	workers     int
+	maxWorkers  int
+	timeout     time.Duration
+}
+
+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,
+		pending:     make(map[int64]struct{}),
+		inFlight:    make(map[int64]struct{}),
+	}
+}
+
+func (s *grokImportProbeScheduler) schedule(prober grokImportProber, 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()
+	if _, exists := s.pending[account.ID]; exists {
+		s.mu.Unlock()
+		return
+	}
+	if _, exists := s.inFlight[account.ID]; exists {
+		s.mu.Unlock()
+		return
+	}
+	if len(s.queue) >= grokImportProbeQueueLimit {
+		s.mu.Unlock()
+		slog.Debug("grok_import_active_probe_dropped", "account_id", account.ID, "reason", "queue_full")
+		return
+	}
+	s.queue = append(s.queue, grokImportProbeTask{prober: prober, accountID: account.ID})
+	s.pending[account.ID] = struct{}{}
+	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)
+		s.finish(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
+	}
+	delete(s.pending, task.accountID)
+	s.inFlight[task.accountID] = struct{}{}
+	return task, true
+}
+
+func (s *grokImportProbeScheduler) finish(accountID int64) {
+	s.mu.Lock()
+	delete(s.inFlight, accountID)
+	s.mu.Unlock()
+}
+
+func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) {
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			slog.Error(
+				"grok_import_active_probe_panic",
+				"account_id", accountID,
+				"recovery_type", panicType(recovered),
+			)
+		}
+	}()
+
+	ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
+	defer cancel()
+	result, err := prober.QueryQuota(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,
+	grokOAuthService service.GrokOAuthTokenService,
+	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,
+		grokOAuthService,
+		rateLimitService,
+		accountUsageService,
+		accountTestService,
+		concurrencyService,
+		crsSyncService,
+		sessionLimitCache,
+		rpmCache,
+		tokenCacheInvalidator,
+	)
+	handler.grokImportProber = grokQuotaService
+	return handler
+}