From 3c5ed704752232ba3b2a94ebf5ce7309d30a9e16 Mon Sep 17 00:00:00 2001 From: warelik Date: Sat, 22 Aug 2026 03:13:46 +0300 Subject: [PATCH 1/2] feat(auth): active credential health prober --- internal/config/config.go | 3 + internal/config/prober.go | 49 ++++ sdk/cliproxy/auth/conductor.go | 4 + sdk/cliproxy/auth/conductor_cooldown.go | 1 + sdk/cliproxy/auth/conductor_prober.go | 292 +++++++++++++++++++++ sdk/cliproxy/auth/conductor_prober_test.go | 199 ++++++++++++++ 6 files changed, 548 insertions(+) create mode 100644 internal/config/prober.go create mode 100644 sdk/cliproxy/auth/conductor_prober.go create mode 100644 sdk/cliproxy/auth/conductor_prober_test.go diff --git a/internal/config/config.go b/internal/config/config.go index 0d8fb234f..c72fb3632 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -25,6 +25,9 @@ type Config struct { // CredentialInFlight configures credential observation snapshots. CredentialInFlight CredentialInFlightConfig `yaml:"credential-in-flight" json:"credential-in-flight"` + // CredentialProber configures optional active credential health probing. + CredentialProber CredentialProberConfig `yaml:"credential-prober" json:"credential-prober"` + // RemoteManagement nests management-related options under 'remote-management'. RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"` diff --git a/internal/config/prober.go b/internal/config/prober.go new file mode 100644 index 000000000..200b0960e --- /dev/null +++ b/internal/config/prober.go @@ -0,0 +1,49 @@ +package config + +import "time" + +const ( + defaultCredentialProberInterval = 60 * time.Second + defaultCredentialProberTimeout = 10 * time.Second + defaultCredentialProberMaxConcurrency = 4 + defaultCredentialProberRatePerMinute = 60 + defaultCredentialProberBackoffBase = 5 * time.Second + defaultCredentialProberBackoffMax = 5 * time.Minute + defaultCredentialProberPath = "/v1/models" +) + +// CredentialProberConfig controls optional active health probing for registered credentials. +// When enabled, the conductor periodically issues a lightweight HTTP probe per credential +// and feeds failures into the existing cooldown/suspension machinery. +type CredentialProberConfig struct { + // Enabled turns active credential health probing on. Default false. + Enabled bool `yaml:"enabled" json:"enabled"` + // Interval is the period between probe sweeps. Default 60s. + Interval time.Duration `yaml:"interval" json:"interval"` + // Timeout is the maximum duration a single probe request may take. Default 10s. + Timeout time.Duration `yaml:"timeout" json:"timeout"` + // MaxConcurrency limits the number of in-flight probes. Default 4. + MaxConcurrency int `yaml:"max-concurrency" json:"max-concurrency"` + // RateLimitPerMinute caps the number of probe requests across all credentials per minute. Default 60. + RateLimitPerMinute int `yaml:"rate-limit-per-minute" json:"rate-limit-per-minute"` + // BackoffBase is the initial cooldown applied when a probe fails. Default 5s. + BackoffBase time.Duration `yaml:"backoff-base" json:"backoff-base"` + // BackoffMax is the maximum probe-induced cooldown. Default 5m. + BackoffMax time.Duration `yaml:"backoff-max" json:"backoff-max"` + // DefaultProbePath is the HTTP path appended to the credential base_url for the probe. Default /v1/models. + DefaultProbePath string `yaml:"default-probe-path" json:"default-probe-path"` +} + +// DefaultCredentialProberConfig returns the prober default configuration. +func DefaultCredentialProberConfig() CredentialProberConfig { + return CredentialProberConfig{ + Enabled: false, + Interval: defaultCredentialProberInterval, + Timeout: defaultCredentialProberTimeout, + MaxConcurrency: defaultCredentialProberMaxConcurrency, + RateLimitPerMinute: defaultCredentialProberRatePerMinute, + BackoffBase: defaultCredentialProberBackoffBase, + BackoffMax: defaultCredentialProberBackoffMax, + DefaultProbePath: defaultCredentialProberPath, + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index a4f2f6ac6..0062c0fc8 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -157,6 +157,10 @@ type Manager struct { refreshCancel context.CancelFunc refreshLoop *authAutoRefreshLoop + // Active credential prober state + proberCancel context.CancelFunc + proberLoop *authProberLoop + requestPrepareLocks sync.Map // refreshLocks serializes credential refresh per auth ID so concurrent // 401 recoveries and auto-refresh workers do not race the same refresh_token. diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 931aef868..85bd44670 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -153,6 +153,7 @@ func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { m.clearHomeRuntimeAuths() } m.rebuildAPIKeyModelAliasFromRuntimeConfig() + m.restartProberLocked(cfg) return clearedCooldowns } diff --git a/sdk/cliproxy/auth/conductor_prober.go b/sdk/cliproxy/auth/conductor_prober.go new file mode 100644 index 000000000..7ef0eab35 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_prober.go @@ -0,0 +1,292 @@ +package auth + +import ( + "context" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "sync" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + log "github.com/sirupsen/logrus" +) + +const ( + proberCheckInterval = 60 * time.Second + proberMaxConcurrency = 4 + proberRatePerMinute = 60 + proberTimeout = 10 * time.Second + proberBackoffBase = 5 * time.Second + proberBackoffMax = 5 * time.Minute + proberDefaultPath = "/v1/models" + proberMaxBodyBytes = 1024 +) + +// authProberLoop runs periodic lightweight health probes for registered auths. +// Failures are fed back into the existing MarkResult/cooldown path. +type authProberLoop struct { + manager *Manager + cfg internalconfig.CredentialProberConfig +} + +func newAuthProberLoop(manager *Manager, cfg internalconfig.CredentialProberConfig) *authProberLoop { + if cfg.Interval <= 0 { + cfg.Interval = proberCheckInterval + } + if cfg.Timeout <= 0 { + cfg.Timeout = proberTimeout + } + if cfg.MaxConcurrency <= 0 { + cfg.MaxConcurrency = proberMaxConcurrency + } + if cfg.RateLimitPerMinute <= 0 { + cfg.RateLimitPerMinute = proberRatePerMinute + } + if cfg.BackoffBase <= 0 { + cfg.BackoffBase = proberBackoffBase + } + if cfg.BackoffMax <= 0 { + cfg.BackoffMax = proberBackoffMax + } + if strings.TrimSpace(cfg.DefaultProbePath) == "" { + cfg.DefaultProbePath = proberDefaultPath + } + return &authProberLoop{manager: manager, cfg: cfg} +} + +// StartProber launches a background credential health prober. +// Only one loop is kept alive; starting a new one cancels the previous run. +func (m *Manager) StartProber(parent context.Context, cfg internalconfig.CredentialProberConfig) { + if m == nil { + return + } + + m.StopProber() + + ctx, cancel := context.WithCancel(parent) + loop := newAuthProberLoop(m, cfg) + + m.mu.Lock() + m.proberCancel = cancel + m.proberLoop = loop + m.mu.Unlock() + + go loop.run(ctx) +} + +// StopProber cancels the background prober loop, if running. +func (m *Manager) StopProber() { + if m == nil { + return + } + m.mu.Lock() + cancel := m.proberCancel + m.proberCancel = nil + m.proberLoop = nil + m.mu.Unlock() + if cancel != nil { + cancel() + } +} + +func (m *Manager) restartProberLocked(cfg *internalconfig.Config) { + if m == nil || cfg == nil { + return + } + if cfg.CredentialProber.Enabled { + m.StartProber(context.Background(), cfg.CredentialProber) + } else { + m.StopProber() + } +} + +func (l *authProberLoop) run(ctx context.Context) { + if l == nil || l.manager == nil { + return + } + + timer := time.NewTimer(0) + defer timer.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + l.sweep(ctx) + interval := l.cfg.Interval + if interval <= 0 { + interval = proberCheckInterval + } + timer.Reset(interval) + } + } +} + +func (l *authProberLoop) sweep(ctx context.Context) { + auths := l.snapshotAuths() + if len(auths) == 0 { + return + } + + concurrency := l.cfg.MaxConcurrency + if concurrency <= 0 { + concurrency = proberMaxConcurrency + } + + ratePerMinute := l.cfg.RateLimitPerMinute + if ratePerMinute <= 0 { + ratePerMinute = proberRatePerMinute + } + + var wg sync.WaitGroup + sem := make(chan struct{}, concurrency) + + var ticker *time.Ticker + if ratePerMinute > 0 { + ticker = time.NewTicker(time.Minute / time.Duration(ratePerMinute)) + defer ticker.Stop() + } + + for _, auth := range auths { + if ctx.Err() != nil { + break + } + if ticker != nil { + select { + case <-ctx.Done(): + break + case <-ticker.C: + } + } + + wg.Add(1) + sem <- struct{}{} + go func(a *Auth) { + defer wg.Done() + defer func() { <-sem }() + l.probe(ctx, a) + }(auth) + } + + wg.Wait() +} + +func (l *authProberLoop) snapshotAuths() []*Auth { + l.manager.mu.RLock() + defer l.manager.mu.RUnlock() + + now := time.Now() + out := make([]*Auth, 0, len(l.manager.auths)) + for _, auth := range l.manager.auths { + if auth == nil { + continue + } + if auth.Disabled || auth.Status == StatusDisabled { + continue + } + if auth.Unavailable && auth.NextRetryAfter.After(now) { + continue + } + out = append(out, auth) + } + return out +} + +func (l *authProberLoop) probe(parent context.Context, auth *Auth) { + l.manager.mu.RLock() + exec := l.manager.executors[auth.Provider] + l.manager.mu.RUnlock() + + if exec == nil { + return + } + + baseURL := "" + if auth.Attributes != nil { + baseURL = strings.TrimSpace(auth.Attributes["base_url"]) + } + if baseURL == "" { + return + } + + path := strings.TrimSpace(l.cfg.DefaultProbePath) + if path == "" { + path = proberDefaultPath + } + + probeURL, errParse := url.Parse(strings.TrimRight(baseURL, "/") + "/" + strings.TrimLeft(path, "/")) + if errParse != nil { + return + } + + req, errReq := http.NewRequestWithContext(parent, http.MethodGet, probeURL.String(), nil) + if errReq != nil { + return + } + + timeout := l.cfg.Timeout + if timeout <= 0 { + timeout = proberTimeout + } + ctx, cancel := context.WithTimeout(parent, timeout) + defer cancel() + req = req.WithContext(ctx) + + resp, errExec := exec.HttpRequest(ctx, auth, req) + var bodyBytes int64 + if resp != nil && resp.Body != nil { + bodyBytes, _ = io.CopyN(io.Discard, resp.Body, proberMaxBodyBytes) + _ = resp.Body.Close() + } + + var resultErr *Error + if errExec != nil { + resultErr = &Error{ + Code: ErrorCodeForceCooldown, + Message: "prober: " + errExec.Error(), + HTTPStatus: http.StatusServiceUnavailable, + Retryable: true, + } + } else if resp == nil { + resultErr = &Error{ + Code: ErrorCodeForceCooldown, + Message: "prober: empty upstream response", + HTTPStatus: http.StatusServiceUnavailable, + Retryable: true, + } + } else if resp.StatusCode == http.StatusNoContent || (resp.StatusCode == http.StatusOK && bodyBytes == 0) { + resultErr = &Error{ + Code: ErrorCodeForceCooldown, + Message: "prober: empty 200/204 response", + HTTPStatus: http.StatusServiceUnavailable, + Retryable: true, + } + } else if resp.StatusCode < 200 || resp.StatusCode >= 300 { + resultErr = &Error{ + Code: ErrorCodeForceCooldown, + Message: fmt.Sprintf("prober: upstream returned %d", resp.StatusCode), + HTTPStatus: resp.StatusCode, + Retryable: resp.StatusCode >= 500 || resp.StatusCode == 429, + } + } + + if resultErr == nil { + return + } + + if log.IsLevelEnabled(log.DebugLevel) { + log.Debugf("credential prober failure for %s: %s", auth.ID, resultErr.Message) + } + + l.manager.MarkResult(ctx, Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Success: false, + CredentialScope: true, + Error: resultErr, + }) +} diff --git a/sdk/cliproxy/auth/conductor_prober_test.go b/sdk/cliproxy/auth/conductor_prober_test.go new file mode 100644 index 000000000..650b2fd0e --- /dev/null +++ b/sdk/cliproxy/auth/conductor_prober_test.go @@ -0,0 +1,199 @@ +package auth + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type proberTestExecutor struct { + provider string + statusCode int + body string + err error + calls atomic.Int32 + respondStatus int +} + +func (e *proberTestExecutor) Identifier() string { return e.provider } + +func (e *proberTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *proberTestExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} + +func (e *proberTestExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } + +func (e *proberTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *proberTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + e.calls.Add(1) + if e.err != nil { + return nil, e.err + } + status := e.statusCode + if status <= 0 { + status = e.respondStatus + } + if status <= 0 { + status = http.StatusOK + } + body := e.body + if body == "" { + body = "{}" + } + return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body))}, nil +} + +func TestProberDisabledByDefault(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test"} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: false} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + if exec.calls.Load() != 0 { + t.Fatalf("prober called disabled executor %d times", exec.calls.Load()) + } +} + +func TestProberSkipsDisabledAuth(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test", err: fmt.Errorf("boom")} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusDisabled, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + if exec.calls.Load() != 0 { + t.Fatalf("prober probed disabled auth %d times", exec.calls.Load()) + } +} + +func TestProberMarksAuthUnavailableOnFailure(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test", err: fmt.Errorf("upstream unreachable")} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + // wait for the immediate first sweep + time.Sleep(100 * time.Millisecond) + + m.mu.RLock() + updated := m.auths["a1"] + m.mu.RUnlock() + + if updated == nil { + t.Fatal("auth disappeared") + } + if exec.calls.Load() != 1 { + t.Fatalf("prober calls = %d, want 1", exec.calls.Load()) + } + if !updated.Unavailable { + t.Fatalf("auth.Unavailable = %v, want true", updated.Unavailable) + } + if updated.Status != StatusError { + t.Fatalf("auth.Status = %v, want %v", updated.Status, StatusError) + } + if updated.NextRetryAfter.IsZero() { + t.Fatalf("auth.NextRetryAfter not set after prober failure") + } +} + +func TestProberLeavesAuthActiveOnSuccess(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test", statusCode: http.StatusOK} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + + m.mu.RLock() + updated := m.auths["a1"] + m.mu.RUnlock() + + if updated == nil { + t.Fatal("auth disappeared") + } + if exec.calls.Load() != 1 { + t.Fatalf("prober calls = %d, want 1", exec.calls.Load()) + } + if updated.Unavailable { + t.Fatalf("auth.Unavailable = %v, want false", updated.Unavailable) + } + if updated.Status != StatusActive { + t.Fatalf("auth.Status = %v, want %v", updated.Status, StatusActive) + } +} + +func TestProberMarksAuthUnavailableOnEmptyResponse(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test", statusCode: http.StatusNoContent} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + + m.mu.RLock() + updated := m.auths["a1"] + m.mu.RUnlock() + + if updated == nil { + t.Fatal("auth disappeared") + } + if !updated.Unavailable { + t.Fatalf("auth.Unavailable = %v, want true", updated.Unavailable) + } +} From 25ff7576ed4c4b785fcddd1f9fc41b71c68f9c23 Mon Sep 17 00:00:00 2001 From: warelik Date: Mon, 24 Aug 2026 05:14:30 -0400 Subject: [PATCH 2/2] fix(auth): correct active prober criteria and decouple from cooldown locks --- sdk/cliproxy/auth/conductor_cooldown.go | 22 +++- sdk/cliproxy/auth/conductor_prober.go | 41 ++++-- sdk/cliproxy/auth/conductor_prober_test.go | 141 ++++++++++++++++++++- 3 files changed, 180 insertions(+), 24 deletions(-) diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 85bd44670..23e173ec0 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -110,10 +110,12 @@ func (m *Manager) SetConfig(cfg *internalconfig.Config) { return } m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - if m.setConfigSnapshotLocked(cfg) { + cleared := m.setConfigSnapshotLocked(cfg) + if cleared { m.persistCooldownStatesLocked(context.Background()) } + m.configCooldownMu.Unlock() + m.restartProber(cfg) } // SetConfigSnapshot updates only in-memory configuration state. It reports whether @@ -123,8 +125,10 @@ func (m *Manager) SetConfigSnapshot(cfg *internalconfig.Config) bool { return false } m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - return m.setConfigSnapshotLocked(cfg) + cleared := m.setConfigSnapshotLocked(cfg) + m.configCooldownMu.Unlock() + m.restartProber(cfg) + return cleared } func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { @@ -153,7 +157,6 @@ func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { m.clearHomeRuntimeAuths() } m.rebuildAPIKeyModelAliasFromRuntimeConfig() - m.restartProberLocked(cfg) return clearedCooldowns } @@ -172,26 +175,31 @@ func (m *Manager) ApplyConfigWithCooldownStateStore(ctx context.Context, cfg *in } m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() m.mu.RLock() oldStore := m.cooldownStore m.mu.RUnlock() m.setConfigSnapshotLocked(cfg) if oldStore != nil && !m.persistCooldownStatesToLocked(ctx, oldStore) { + m.configCooldownMu.Unlock() return false } if errContext := ctx.Err(); errContext != nil { + m.configCooldownMu.Unlock() return false } m.mu.Lock() - defer m.mu.Unlock() if m.cooldownStore != oldStore { + m.mu.Unlock() + m.configCooldownMu.Unlock() return false } if m.pendingCooldownStateStore == oldStore { m.pendingCooldownStateStore = nil } m.cooldownStore = store + m.mu.Unlock() + m.configCooldownMu.Unlock() + m.restartProber(cfg) return true } diff --git a/sdk/cliproxy/auth/conductor_prober.go b/sdk/cliproxy/auth/conductor_prober.go index 7ef0eab35..81d81cf12 100644 --- a/sdk/cliproxy/auth/conductor_prober.go +++ b/sdk/cliproxy/auth/conductor_prober.go @@ -92,7 +92,7 @@ func (m *Manager) StopProber() { } } -func (m *Manager) restartProberLocked(cfg *internalconfig.Config) { +func (m *Manager) restartProber(cfg *internalconfig.Config) { if m == nil || cfg == nil { return } @@ -197,10 +197,12 @@ func (l *authProberLoop) snapshotAuths() []*Auth { } func (l *authProberLoop) probe(parent context.Context, auth *Auth) { - l.manager.mu.RLock() - exec := l.manager.executors[auth.Provider] - l.manager.mu.RUnlock() + providerKey := executorKeyFromAuth(auth) + if providerKey == "" { + return + } + exec := l.manager.executorFor(providerKey) if exec == nil { return } @@ -210,7 +212,14 @@ func (l *authProberLoop) probe(parent context.Context, auth *Auth) { baseURL = strings.TrimSpace(auth.Attributes["base_url"]) } if baseURL == "" { - return + switch providerKey { + case "claude": + baseURL = "https://api.anthropic.com" + case "gemini": + baseURL = "https://generativelanguage.googleapis.com" + default: + return + } } path := strings.TrimSpace(l.cfg.DefaultProbePath) @@ -228,6 +237,18 @@ func (l *authProberLoop) probe(parent context.Context, auth *Auth) { return } + if providerKey == "claude" { + if req.Header.Get("Anthropic-Version") == "" { + req.Header.Set("Anthropic-Version", "2023-06-01") + } + isAPIKey := auth.Attributes != nil && strings.TrimSpace(auth.Attributes["api_key"]) != "" + if !isAPIKey { + if req.Header.Get("Anthropic-Beta") == "" { + req.Header.Set("Anthropic-Beta", "oauth-2025-04-20") + } + } + } + timeout := l.cfg.Timeout if timeout <= 0 { timeout = proberTimeout @@ -237,9 +258,8 @@ func (l *authProberLoop) probe(parent context.Context, auth *Auth) { req = req.WithContext(ctx) resp, errExec := exec.HttpRequest(ctx, auth, req) - var bodyBytes int64 if resp != nil && resp.Body != nil { - bodyBytes, _ = io.CopyN(io.Discard, resp.Body, proberMaxBodyBytes) + _, _ = io.CopyN(io.Discard, resp.Body, proberMaxBodyBytes) _ = resp.Body.Close() } @@ -258,13 +278,6 @@ func (l *authProberLoop) probe(parent context.Context, auth *Auth) { HTTPStatus: http.StatusServiceUnavailable, Retryable: true, } - } else if resp.StatusCode == http.StatusNoContent || (resp.StatusCode == http.StatusOK && bodyBytes == 0) { - resultErr = &Error{ - Code: ErrorCodeForceCooldown, - Message: "prober: empty 200/204 response", - HTTPStatus: http.StatusServiceUnavailable, - Retryable: true, - } } else if resp.StatusCode < 200 || resp.StatusCode >= 300 { resultErr = &Error{ Code: ErrorCodeForceCooldown, diff --git a/sdk/cliproxy/auth/conductor_prober_test.go b/sdk/cliproxy/auth/conductor_prober_test.go index 650b2fd0e..84cf85862 100644 --- a/sdk/cliproxy/auth/conductor_prober_test.go +++ b/sdk/cliproxy/auth/conductor_prober_test.go @@ -170,7 +170,7 @@ func TestProberLeavesAuthActiveOnSuccess(t *testing.T) { } } -func TestProberMarksAuthUnavailableOnEmptyResponse(t *testing.T) { +func TestProberLeavesAuthActiveOnEmptyResponse(t *testing.T) { ctx := context.Background() m := NewManager(nil, nil, nil) exec := &proberTestExecutor{provider: "test", statusCode: http.StatusNoContent} @@ -193,7 +193,142 @@ func TestProberMarksAuthUnavailableOnEmptyResponse(t *testing.T) { if updated == nil { t.Fatal("auth disappeared") } - if !updated.Unavailable { - t.Fatalf("auth.Unavailable = %v, want true", updated.Unavailable) + if updated.Unavailable { + t.Fatalf("auth.Unavailable = %v, want false (204 is healthy for probe)", updated.Unavailable) + } + if updated.Status != StatusActive { + t.Fatalf("auth.Status = %v, want %v", updated.Status, StatusActive) + } +} + +func TestProberLeavesAuthActiveOnEmpty200Response(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test", statusCode: http.StatusOK, body: " "} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + + m.mu.RLock() + updated := m.auths["a1"] + m.mu.RUnlock() + + if updated == nil { + t.Fatal("auth disappeared") + } + if updated.Unavailable { + t.Fatalf("auth.Unavailable = %v, want false (empty 200 is healthy for probe)", updated.Unavailable) + } + if updated.Status != StatusActive { + t.Fatalf("auth.Status = %v, want %v", updated.Status, StatusActive) + } +} + +type headerRecordingProberExecutor struct { + proberTestExecutor + lastReq *http.Request +} + +func (e *headerRecordingProberExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { + if req != nil { + e.lastReq = req.Clone(ctx) + } + return e.proberTestExecutor.HttpRequest(ctx, auth, req) +} + +func TestProberSetsClaudeHeaders(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &headerRecordingProberExecutor{proberTestExecutor: proberTestExecutor{provider: "claude"}} + m.RegisterExecutor(exec) + + // OAuth credential without base_url + authOAuth := &Auth{ID: "c-oauth", Provider: "claude", Status: StatusActive} + if _, err := m.Register(ctx, authOAuth); err != nil { + t.Fatalf("Register oauth: %v", err) + } + + cfg := internalconfig.CredentialProberConfig{Enabled: true, Interval: time.Hour, Timeout: time.Second, MaxConcurrency: 1, RateLimitPerMinute: 1000} + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + + time.Sleep(100 * time.Millisecond) + + if exec.lastReq == nil { + t.Fatal("probe request was not executed") + } + if got := exec.lastReq.URL.String(); got != "https://api.anthropic.com/v1/models" { + t.Fatalf("probe URL = %q, want https://api.anthropic.com/v1/models", got) + } + if got := exec.lastReq.Header.Get("Anthropic-Version"); got != "2023-06-01" { + t.Fatalf("Anthropic-Version = %q, want 2023-06-01", got) + } + if got := exec.lastReq.Header.Get("Anthropic-Beta"); got != "oauth-2025-04-20" { + t.Fatalf("Anthropic-Beta = %q, want oauth-2025-04-20", got) + } + + // API key credential should not have Anthropic-Beta set by default + execAPIKey := &headerRecordingProberExecutor{proberTestExecutor: proberTestExecutor{provider: "claude"}} + mAPIKey := NewManager(nil, nil, nil) + mAPIKey.RegisterExecutor(execAPIKey) + + authAPIKey := &Auth{ID: "c-apikey", Provider: "claude", Status: StatusActive, Attributes: map[string]string{"api_key": "sk-ant-xxx"}} + if _, err := mAPIKey.Register(ctx, authAPIKey); err != nil { + t.Fatalf("Register api key: %v", err) + } + mAPIKey.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + time.Sleep(100 * time.Millisecond) + + if execAPIKey.lastReq == nil { + t.Fatal("probe request for api key was not executed") + } + if got := execAPIKey.lastReq.Header.Get("Anthropic-Version"); got != "2023-06-01" { + t.Fatalf("Anthropic-Version = %q, want 2023-06-01", got) + } + if got := execAPIKey.lastReq.Header.Get("Anthropic-Beta"); got != "" { + t.Fatalf("Anthropic-Beta = %q, want empty for API key", got) + } +} + +func TestProberNoDeadlockOnConfigChange(t *testing.T) { + ctx := context.Background() + m := NewManager(nil, nil, nil) + exec := &proberTestExecutor{provider: "test"} + m.RegisterExecutor(exec) + + auth := &Auth{ID: "a1", Provider: "test", Status: StatusActive, Attributes: map[string]string{"base_url": "https://example.com"}} + if _, err := m.Register(ctx, auth); err != nil { + t.Fatalf("Register: %v", err) + } + + done := make(chan struct{}) + go func() { + for i := 0; i < 50; i++ { + enabled := (i % 2) == 0 + cfg := internalconfig.CredentialProberConfig{ + Enabled: enabled, + Interval: 10 * time.Millisecond, + Timeout: time.Second, + MaxConcurrency: 2, + RateLimitPerMinute: 600, + } + m.SetConfig(&internalconfig.Config{CredentialProber: cfg}) + m.SetConfigSnapshot(&internalconfig.Config{CredentialProber: cfg}) + } + close(done) + }() + + select { + case <-done: + // success + case <-time.After(5 * time.Second): + t.Fatal("deadlock detected during SetConfig / SetConfigSnapshot with prober enabled/disabled") } }