From 5bfaf8086b268aeed5a2b2e2bc6da1dcd844e484 Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Sun, 12 Apr 2026 13:56:05 +0800
Subject: [PATCH] feat(auth): add configurable worker pool size for
 auto-refresh loop

- Introduced `auth-auto-refresh-workers` config option to override default concurrency.
- Updated `authAutoRefreshLoop` to support customizable worker counts.
- Enhanced token refresh scheduling flexibility by aligning worker pool with runtime configurations.
---
 config.example.yaml                    |  4 ++++
 internal/config/config.go              |  4 ++++
 sdk/cliproxy/auth/auto_refresh_loop.go | 31 +++++++++++++++++---------
 sdk/cliproxy/auth/conductor.go         |  6 ++++-
 4 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/config.example.yaml b/config.example.yaml
index 9d839a87..067910c5 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -90,6 +90,10 @@ max-retry-interval: 30
 # When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states).
 disable-cooling: false
 
+# Core auth auto-refresh worker pool size (OAuth/file-based auth token refresh).
+# When > 0, overrides the default worker count (16).
+# auth-auto-refresh-workers: 16
+
 # Quota exceeded behavior
 quota-exceeded:
   switch-project: true # Whether to automatically switch to another project when a quota is exceeded
diff --git a/internal/config/config.go b/internal/config/config.go
index b1957426..a3dd4c59 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -68,6 +68,10 @@ type Config struct {
 	// DisableCooling disables quota cooldown scheduling when true.
 	DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
 
+	// AuthAutoRefreshWorkers overrides the size of the core auth auto-refresh worker pool.
+	// When <= 0, the default worker count is used.
+	AuthAutoRefreshWorkers int `yaml:"auth-auto-refresh-workers" json:"auth-auto-refresh-workers"`
+
 	// RequestRetry defines the retry times when the request failed.
 	RequestRetry int `yaml:"request-retry" json:"request-retry"`
 	// MaxRetryCredentials defines the maximum number of credentials to try for a failed request.
diff --git a/sdk/cliproxy/auth/auto_refresh_loop.go b/sdk/cliproxy/auth/auto_refresh_loop.go
index 3350b603..9767ee58 100644
--- a/sdk/cliproxy/auth/auto_refresh_loop.go
+++ b/sdk/cliproxy/auth/auto_refresh_loop.go
@@ -11,8 +11,9 @@ import (
 )
 
 type authAutoRefreshLoop struct {
-	manager  *Manager
-	interval time.Duration
+	manager     *Manager
+	interval    time.Duration
+	concurrency int
 
 	mu    sync.Mutex
 	queue refreshMinHeap
@@ -23,21 +24,25 @@ type authAutoRefreshLoop struct {
 	jobs   chan string
 }
 
-func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration) *authAutoRefreshLoop {
+func newAuthAutoRefreshLoop(manager *Manager, interval time.Duration, concurrency int) *authAutoRefreshLoop {
 	if interval <= 0 {
 		interval = refreshCheckInterval
 	}
-	jobBuffer := refreshMaxConcurrency * 4
+	if concurrency <= 0 {
+		concurrency = refreshMaxConcurrency
+	}
+	jobBuffer := concurrency * 4
 	if jobBuffer < 64 {
 		jobBuffer = 64
 	}
 	return &authAutoRefreshLoop{
-		manager:  manager,
-		interval: interval,
-		index:    make(map[string]*refreshHeapItem),
-		dirty:    make(map[string]struct{}),
-		wakeCh:   make(chan struct{}, 1),
-		jobs:     make(chan string, jobBuffer),
+		manager:     manager,
+		interval:    interval,
+		concurrency: concurrency,
+		index:       make(map[string]*refreshHeapItem),
+		dirty:       make(map[string]struct{}),
+		wakeCh:      make(chan struct{}, 1),
+		jobs:        make(chan string, jobBuffer),
 	}
 }
 
@@ -59,7 +64,11 @@ func (l *authAutoRefreshLoop) run(ctx context.Context) {
 		return
 	}
 
-	for i := 0; i < refreshMaxConcurrency; i++ {
+	workers := l.concurrency
+	if workers <= 0 {
+		workers = refreshMaxConcurrency
+	}
+	for i := 0; i < workers; i++ {
 		go l.worker(ctx)
 	}
 
diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go
index 3cf02524..5e7d3161 100644
--- a/sdk/cliproxy/auth/conductor.go
+++ b/sdk/cliproxy/auth/conductor.go
@@ -2912,7 +2912,11 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio
 	}
 
 	ctx, cancelCtx := context.WithCancel(parent)
-	loop := newAuthAutoRefreshLoop(m, interval)
+	workers := refreshMaxConcurrency
+	if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 {
+		workers = cfg.AuthAutoRefreshWorkers
+	}
+	loop := newAuthAutoRefreshLoop(m, interval, workers)
 
 	m.mu.Lock()
 	m.refreshCancel = cancelCtx