Skip to content

Commit 1cb1c99

Browse files
author
song duan
committed
feat(auth): add OAuth refresh configuration for background credential management
Introduced new settings for OAuth credential refresh behavior, allowing control over refresh on startup and batch size for credential updates. Updated relevant structures and methods to support this functionality.
1 parent ea09694 commit 1cb1c99

4 files changed

Lines changed: 266 additions & 12 deletions

File tree

‎config.example.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,13 @@ remote-management:
3535
# Authentication directory (supports ~ for home directory)
3636
auth-dir: "~/.cli-proxy-api"
3737

38+
# Controls background refresh behavior for OAuth / auth-file credentials.
39+
# Set on-startup: false to avoid an immediate refresh pass on every process start.
40+
# Set batch-size to limit how many due credentials are refreshed in one scheduler pass.
41+
# oauth-refresh:
42+
# on-startup: false
43+
# batch-size: 2
44+
3845
# API keys for authentication
3946
api-keys:
4047
- "your-api-key-1"

‎internal/config/sdk_config.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@ type SDKConfig struct {
2828
// Default is false (disabled).
2929
PassthroughHeaders bool `yaml:"passthrough-headers" json:"passthrough-headers"`
3030

31+
// OAuthRefresh controls background refresh behavior for OAuth / auth-file credentials.
32+
OAuthRefresh OAuthRefreshConfig `yaml:"oauth-refresh" json:"oauth-refresh"`
33+
3134
// Streaming configures server-side streaming behavior (keep-alives and safe bootstrap retries).
3235
Streaming StreamingConfig `yaml:"streaming" json:"streaming"`
3336

@@ -36,6 +39,17 @@ type SDKConfig struct {
3639
NonStreamKeepAliveInterval int `yaml:"nonstream-keepalive-interval,omitempty" json:"nonstream-keepalive-interval,omitempty"`
3740
}
3841

42+
// OAuthRefreshConfig controls background refresh scheduling for OAuth/file-backed auths.
43+
type OAuthRefreshConfig struct {
44+
// OnStartup controls whether the auto-refresh loop immediately checks credentials when the service starts.
45+
// Nil preserves the legacy default (enabled).
46+
OnStartup *bool `yaml:"on-startup,omitempty" json:"on-startup,omitempty"`
47+
48+
// BatchSize limits how many due credentials may be refreshed in one scheduler pass.
49+
// <= 0 preserves the legacy behavior (no per-pass cap).
50+
BatchSize int `yaml:"batch-size,omitempty" json:"batch-size,omitempty"`
51+
}
52+
3953
// StreamingConfig holds server streaming behavior configuration.
4054
type StreamingConfig struct {
4155
// KeepAliveSeconds controls how often the server emits SSE heartbeats (": keep-alive\n\n").

‎sdk/cliproxy/auth/conductor.go‎

Lines changed: 77 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,9 @@ type Manager struct {
162162
rtProvider RoundTripperProvider
163163

164164
// Auto refresh state
165-
refreshCancel context.CancelFunc
166-
refreshSemaphore chan struct{}
165+
refreshCancel context.CancelFunc
166+
refreshSemaphore chan struct{}
167+
refreshBatchCursor int
167168
}
168169

169170
// NewManager constructs a manager with optional custom selector and hook.
@@ -355,6 +356,33 @@ func (m *Manager) SetConfig(cfg *internalconfig.Config) {
355356
m.rebuildAPIKeyModelAliasFromRuntimeConfig()
356357
}
357358

359+
func (m *Manager) oauthRefreshConfig() internalconfig.OAuthRefreshConfig {
360+
if m == nil {
361+
return internalconfig.OAuthRefreshConfig{}
362+
}
363+
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
364+
if cfg == nil {
365+
return internalconfig.OAuthRefreshConfig{}
366+
}
367+
return cfg.OAuthRefresh
368+
}
369+
370+
func (m *Manager) refreshOnStartupEnabled() bool {
371+
cfg := m.oauthRefreshConfig()
372+
if cfg.OnStartup == nil {
373+
return true
374+
}
375+
return *cfg.OnStartup
376+
}
377+
378+
func (m *Manager) refreshBatchSize() int {
379+
cfg := m.oauthRefreshConfig()
380+
if cfg.BatchSize <= 0 {
381+
return 0
382+
}
383+
return cfg.BatchSize
384+
}
385+
358386
func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string {
359387
if m == nil {
360388
return ""
@@ -2879,7 +2907,9 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio
28792907
go func() {
28802908
ticker := time.NewTicker(interval)
28812909
defer ticker.Stop()
2882-
m.checkRefreshes(ctx)
2910+
if m.refreshOnStartupEnabled() {
2911+
m.checkRefreshes(ctx)
2912+
}
28832913
for {
28842914
select {
28852915
case <-ctx.Done():
@@ -2903,22 +2933,57 @@ func (m *Manager) checkRefreshes(ctx context.Context) {
29032933
// log.Debugf("checking refreshes")
29042934
now := time.Now()
29052935
snapshot := m.snapshotAuths()
2936+
dueIDs := make([]string, 0, len(snapshot))
29062937
for _, a := range snapshot {
29072938
typ, _ := a.AccountInfo()
29082939
if typ != "api_key" {
29092940
if !m.shouldRefresh(a, now) {
29102941
continue
29112942
}
2912-
log.Debugf("checking refresh for %s, %s, %s", a.Provider, a.ID, typ)
2913-
2914-
if exec := m.executorFor(a.Provider); exec == nil {
2915-
continue
2916-
}
2917-
if !m.markRefreshPending(a.ID, now) {
2918-
continue
2919-
}
2920-
go m.refreshAuthWithLimit(ctx, a.ID)
2943+
dueIDs = append(dueIDs, a.ID)
2944+
}
2945+
}
2946+
if len(dueIDs) == 0 {
2947+
m.mu.Lock()
2948+
m.refreshBatchCursor = 0
2949+
m.mu.Unlock()
2950+
return
2951+
}
2952+
sort.Strings(dueIDs)
2953+
limit := m.refreshBatchSize()
2954+
selectedIDs := dueIDs
2955+
if limit > 0 && len(dueIDs) > limit {
2956+
m.mu.Lock()
2957+
start := 0
2958+
if len(dueIDs) > 0 {
2959+
start = m.refreshBatchCursor % len(dueIDs)
2960+
}
2961+
selectedIDs = make([]string, 0, limit)
2962+
for i := 0; i < limit; i++ {
2963+
idx := (start + i) % len(dueIDs)
2964+
selectedIDs = append(selectedIDs, dueIDs[idx])
2965+
}
2966+
m.refreshBatchCursor = (start + limit) % len(dueIDs)
2967+
m.mu.Unlock()
2968+
} else {
2969+
m.mu.Lock()
2970+
m.refreshBatchCursor = 0
2971+
m.mu.Unlock()
2972+
}
2973+
for _, id := range selectedIDs {
2974+
auth, ok := m.GetByID(id)
2975+
if !ok || auth == nil {
2976+
continue
2977+
}
2978+
typ, _ := auth.AccountInfo()
2979+
log.Debugf("checking refresh for %s, %s, %s", auth.Provider, auth.ID, typ)
2980+
if exec := m.executorFor(auth.Provider); exec == nil {
2981+
continue
2982+
}
2983+
if !m.markRefreshPending(auth.ID, now) {
2984+
continue
29212985
}
2986+
go m.refreshAuthWithLimit(ctx, auth.ID)
29222987
}
29232988
}
29242989

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
package auth
2+
3+
import (
4+
"context"
5+
"net/http"
6+
"strconv"
7+
"sync"
8+
"testing"
9+
"time"
10+
11+
internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
12+
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
13+
)
14+
15+
type autoRefreshTestExecutor struct {
16+
provider string
17+
mu sync.Mutex
18+
ids []string
19+
}
20+
21+
func (e *autoRefreshTestExecutor) Identifier() string { return e.provider }
22+
23+
func (e *autoRefreshTestExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
24+
return cliproxyexecutor.Response{}, nil
25+
}
26+
27+
func (e *autoRefreshTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
28+
return nil, nil
29+
}
30+
31+
func (e *autoRefreshTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) {
32+
e.mu.Lock()
33+
e.ids = append(e.ids, auth.ID)
34+
e.mu.Unlock()
35+
return auth, nil
36+
}
37+
38+
func (e *autoRefreshTestExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
39+
return cliproxyexecutor.Response{}, nil
40+
}
41+
42+
func (e *autoRefreshTestExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) {
43+
return nil, nil
44+
}
45+
46+
func (e *autoRefreshTestExecutor) refreshedIDs() []string {
47+
e.mu.Lock()
48+
defer e.mu.Unlock()
49+
out := make([]string, len(e.ids))
50+
copy(out, e.ids)
51+
return out
52+
}
53+
54+
func TestManager_StartAutoRefresh_ImmediateByDefault(t *testing.T) {
55+
manager := NewManager(nil, &RoundRobinSelector{}, nil)
56+
executor := &autoRefreshTestExecutor{provider: "test"}
57+
manager.RegisterExecutor(executor)
58+
manager.SetConfig(&internalconfig.Config{})
59+
60+
auth := &Auth{
61+
ID: "startup-default",
62+
Provider: "test",
63+
Metadata: map[string]any{"refresh_interval_seconds": 3600},
64+
}
65+
if _, err := manager.Register(context.Background(), auth); err != nil {
66+
t.Fatalf("register auth: %v", err)
67+
}
68+
69+
ctx, cancel := context.WithCancel(context.Background())
70+
defer cancel()
71+
manager.StartAutoRefresh(ctx, time.Hour)
72+
t.Cleanup(manager.StopAutoRefresh)
73+
74+
waitForRefreshCalls(t, executor, 1, time.Second)
75+
}
76+
77+
func TestManager_StartAutoRefresh_CanSkipStartupCheck(t *testing.T) {
78+
manager := NewManager(nil, &RoundRobinSelector{}, nil)
79+
executor := &autoRefreshTestExecutor{provider: "test"}
80+
manager.RegisterExecutor(executor)
81+
82+
disabled := false
83+
manager.SetConfig(&internalconfig.Config{
84+
SDKConfig: internalconfig.SDKConfig{
85+
OAuthRefresh: internalconfig.OAuthRefreshConfig{
86+
OnStartup: &disabled,
87+
},
88+
},
89+
})
90+
91+
auth := &Auth{
92+
ID: "startup-disabled",
93+
Provider: "test",
94+
Metadata: map[string]any{"refresh_interval_seconds": 3600},
95+
}
96+
if _, err := manager.Register(context.Background(), auth); err != nil {
97+
t.Fatalf("register auth: %v", err)
98+
}
99+
100+
ctx, cancel := context.WithCancel(context.Background())
101+
defer cancel()
102+
manager.StartAutoRefresh(ctx, time.Hour)
103+
t.Cleanup(manager.StopAutoRefresh)
104+
105+
time.Sleep(100 * time.Millisecond)
106+
if got := len(executor.refreshedIDs()); got != 0 {
107+
t.Fatalf("startup refresh calls = %d, want 0", got)
108+
}
109+
}
110+
111+
func TestManager_CheckRefreshes_HonorsBatchSize(t *testing.T) {
112+
manager := NewManager(nil, &RoundRobinSelector{}, nil)
113+
executor := &autoRefreshTestExecutor{provider: "test"}
114+
manager.RegisterExecutor(executor)
115+
manager.SetConfig(&internalconfig.Config{
116+
SDKConfig: internalconfig.SDKConfig{
117+
OAuthRefresh: internalconfig.OAuthRefreshConfig{
118+
BatchSize: 2,
119+
},
120+
},
121+
})
122+
123+
ctx := context.Background()
124+
for i := 0; i < 5; i++ {
125+
auth := &Auth{
126+
ID: "batch-auth-" + strconv.Itoa(i),
127+
Provider: "test",
128+
Metadata: map[string]any{"refresh_interval_seconds": 3600},
129+
}
130+
if _, err := manager.Register(ctx, auth); err != nil {
131+
t.Fatalf("register auth %d: %v", i, err)
132+
}
133+
}
134+
135+
manager.checkRefreshes(ctx)
136+
waitForRefreshCalls(t, executor, 2, time.Second)
137+
138+
manager.checkRefreshes(ctx)
139+
waitForRefreshCalls(t, executor, 4, time.Second)
140+
141+
manager.checkRefreshes(ctx)
142+
waitForRefreshCalls(t, executor, 5, time.Second)
143+
144+
seen := make(map[string]int)
145+
for _, id := range executor.refreshedIDs() {
146+
seen[id]++
147+
}
148+
if len(seen) != 5 {
149+
t.Fatalf("refreshed unique auths = %d, want 5; got=%v", len(seen), seen)
150+
}
151+
for id, count := range seen {
152+
if count != 1 {
153+
t.Fatalf("auth %s refreshed %d times, want 1", id, count)
154+
}
155+
}
156+
}
157+
158+
func waitForRefreshCalls(t *testing.T, executor *autoRefreshTestExecutor, want int, timeout time.Duration) {
159+
t.Helper()
160+
deadline := time.Now().Add(timeout)
161+
for time.Now().Before(deadline) {
162+
if got := len(executor.refreshedIDs()); got >= want {
163+
return
164+
}
165+
time.Sleep(10 * time.Millisecond)
166+
}
167+
t.Fatalf("timed out waiting for %d refresh calls; got %d", want, len(executor.refreshedIDs()))
168+
}

0 commit comments

Comments
 (0)