Fixed: #1747
feat(auth): add configurable max-retry-credentials for finer control over cross-credential retries
This commit is contained in:
@@ -75,6 +75,10 @@ passthrough-headers: false
|
|||||||
# Number of times to retry a request. Retries will occur if the HTTP response code is 403, 408, 500, 502, 503, or 504.
|
# Number of times to retry a request. Retries will occur if the HTTP response code is 403, 408, 500, 502, 503, or 504.
|
||||||
request-retry: 3
|
request-retry: 3
|
||||||
|
|
||||||
|
# Maximum number of different credentials to try for one failed request.
|
||||||
|
# Set to 0 to keep legacy behavior (try all available credentials).
|
||||||
|
max-retry-credentials: 0
|
||||||
|
|
||||||
# Maximum wait time in seconds for a cooled-down credential before triggering a retry.
|
# Maximum wait time in seconds for a cooled-down credential before triggering a retry.
|
||||||
max-retry-interval: 30
|
max-retry-interval: 30
|
||||||
|
|
||||||
|
|||||||
@@ -257,7 +257,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
|||||||
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
||||||
s.applyAccessConfig(nil, cfg)
|
s.applyAccessConfig(nil, cfg)
|
||||||
if authManager != nil {
|
if authManager != nil {
|
||||||
authManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second)
|
authManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||||
}
|
}
|
||||||
managementasset.SetCurrentConfig(cfg)
|
managementasset.SetCurrentConfig(cfg)
|
||||||
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||||
@@ -915,7 +915,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if s.handlers != nil && s.handlers.AuthManager != nil {
|
if s.handlers != nil && s.handlers.AuthManager != nil {
|
||||||
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second)
|
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Update log level dynamically when debug flag changes
|
// Update log level dynamically when debug flag changes
|
||||||
|
|||||||
@@ -69,6 +69,9 @@ type Config struct {
|
|||||||
|
|
||||||
// RequestRetry defines the retry times when the request failed.
|
// RequestRetry defines the retry times when the request failed.
|
||||||
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
||||||
|
// MaxRetryCredentials defines the maximum number of credentials to try for a failed request.
|
||||||
|
// Set to 0 or a negative value to keep trying all available credentials (legacy behavior).
|
||||||
|
MaxRetryCredentials int `yaml:"max-retry-credentials" json:"max-retry-credentials"`
|
||||||
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
||||||
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
|
MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"`
|
||||||
|
|
||||||
@@ -609,6 +612,10 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
|||||||
cfg.ErrorLogsMaxFiles = 10
|
cfg.ErrorLogsMaxFiles = 10
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if cfg.MaxRetryCredentials < 0 {
|
||||||
|
cfg.MaxRetryCredentials = 0
|
||||||
|
}
|
||||||
|
|
||||||
// Sanitize Gemini API key configuration and migrate legacy entries.
|
// Sanitize Gemini API key configuration and migrate legacy entries.
|
||||||
cfg.SanitizeGeminiKeys()
|
cfg.SanitizeGeminiKeys()
|
||||||
|
|
||||||
|
|||||||
@@ -127,7 +127,8 @@ func (w *Watcher) reloadConfig() bool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
authDirChanged := oldConfig == nil || oldConfig.AuthDir != newConfig.AuthDir
|
authDirChanged := oldConfig == nil || oldConfig.AuthDir != newConfig.AuthDir
|
||||||
forceAuthRefresh := oldConfig != nil && (oldConfig.ForceModelPrefix != newConfig.ForceModelPrefix || !reflect.DeepEqual(oldConfig.OAuthModelAlias, newConfig.OAuthModelAlias))
|
retryConfigChanged := oldConfig != nil && (oldConfig.RequestRetry != newConfig.RequestRetry || oldConfig.MaxRetryInterval != newConfig.MaxRetryInterval || oldConfig.MaxRetryCredentials != newConfig.MaxRetryCredentials)
|
||||||
|
forceAuthRefresh := oldConfig != nil && (oldConfig.ForceModelPrefix != newConfig.ForceModelPrefix || !reflect.DeepEqual(oldConfig.OAuthModelAlias, newConfig.OAuthModelAlias) || retryConfigChanged)
|
||||||
|
|
||||||
log.Infof("config successfully reloaded, triggering client reload")
|
log.Infof("config successfully reloaded, triggering client reload")
|
||||||
w.reloadClients(authDirChanged, affectedOAuthProviders, forceAuthRefresh)
|
w.reloadClients(authDirChanged, affectedOAuthProviders, forceAuthRefresh)
|
||||||
|
|||||||
@@ -54,6 +54,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
|||||||
if oldCfg.RequestRetry != newCfg.RequestRetry {
|
if oldCfg.RequestRetry != newCfg.RequestRetry {
|
||||||
changes = append(changes, fmt.Sprintf("request-retry: %d -> %d", oldCfg.RequestRetry, newCfg.RequestRetry))
|
changes = append(changes, fmt.Sprintf("request-retry: %d -> %d", oldCfg.RequestRetry, newCfg.RequestRetry))
|
||||||
}
|
}
|
||||||
|
if oldCfg.MaxRetryCredentials != newCfg.MaxRetryCredentials {
|
||||||
|
changes = append(changes, fmt.Sprintf("max-retry-credentials: %d -> %d", oldCfg.MaxRetryCredentials, newCfg.MaxRetryCredentials))
|
||||||
|
}
|
||||||
if oldCfg.MaxRetryInterval != newCfg.MaxRetryInterval {
|
if oldCfg.MaxRetryInterval != newCfg.MaxRetryInterval {
|
||||||
changes = append(changes, fmt.Sprintf("max-retry-interval: %d -> %d", oldCfg.MaxRetryInterval, newCfg.MaxRetryInterval))
|
changes = append(changes, fmt.Sprintf("max-retry-interval: %d -> %d", oldCfg.MaxRetryInterval, newCfg.MaxRetryInterval))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -223,6 +223,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
UsageStatisticsEnabled: false,
|
UsageStatisticsEnabled: false,
|
||||||
DisableCooling: false,
|
DisableCooling: false,
|
||||||
RequestRetry: 1,
|
RequestRetry: 1,
|
||||||
|
MaxRetryCredentials: 1,
|
||||||
MaxRetryInterval: 1,
|
MaxRetryInterval: 1,
|
||||||
WebsocketAuth: false,
|
WebsocketAuth: false,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
||||||
@@ -246,6 +247,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
UsageStatisticsEnabled: true,
|
UsageStatisticsEnabled: true,
|
||||||
DisableCooling: true,
|
DisableCooling: true,
|
||||||
RequestRetry: 2,
|
RequestRetry: 2,
|
||||||
|
MaxRetryCredentials: 3,
|
||||||
MaxRetryInterval: 3,
|
MaxRetryInterval: 3,
|
||||||
WebsocketAuth: true,
|
WebsocketAuth: true,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
||||||
@@ -283,6 +285,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
expectContains(t, details, "disable-cooling: false -> true")
|
expectContains(t, details, "disable-cooling: false -> true")
|
||||||
expectContains(t, details, "request-log: false -> true")
|
expectContains(t, details, "request-log: false -> true")
|
||||||
expectContains(t, details, "request-retry: 1 -> 2")
|
expectContains(t, details, "request-retry: 1 -> 2")
|
||||||
|
expectContains(t, details, "max-retry-credentials: 1 -> 3")
|
||||||
expectContains(t, details, "max-retry-interval: 1 -> 3")
|
expectContains(t, details, "max-retry-interval: 1 -> 3")
|
||||||
expectContains(t, details, "proxy-url: http://old-proxy -> http://new-proxy")
|
expectContains(t, details, "proxy-url: http://old-proxy -> http://new-proxy")
|
||||||
expectContains(t, details, "ws-auth: false -> true")
|
expectContains(t, details, "ws-auth: false -> true")
|
||||||
@@ -309,6 +312,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
UsageStatisticsEnabled: false,
|
UsageStatisticsEnabled: false,
|
||||||
DisableCooling: false,
|
DisableCooling: false,
|
||||||
RequestRetry: 1,
|
RequestRetry: 1,
|
||||||
|
MaxRetryCredentials: 1,
|
||||||
MaxRetryInterval: 1,
|
MaxRetryInterval: 1,
|
||||||
WebsocketAuth: false,
|
WebsocketAuth: false,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
||||||
@@ -361,6 +365,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
UsageStatisticsEnabled: true,
|
UsageStatisticsEnabled: true,
|
||||||
DisableCooling: true,
|
DisableCooling: true,
|
||||||
RequestRetry: 2,
|
RequestRetry: 2,
|
||||||
|
MaxRetryCredentials: 3,
|
||||||
MaxRetryInterval: 3,
|
MaxRetryInterval: 3,
|
||||||
WebsocketAuth: true,
|
WebsocketAuth: true,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
||||||
@@ -419,6 +424,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
expectContains(t, changes, "usage-statistics-enabled: false -> true")
|
expectContains(t, changes, "usage-statistics-enabled: false -> true")
|
||||||
expectContains(t, changes, "disable-cooling: false -> true")
|
expectContains(t, changes, "disable-cooling: false -> true")
|
||||||
expectContains(t, changes, "request-retry: 1 -> 2")
|
expectContains(t, changes, "request-retry: 1 -> 2")
|
||||||
|
expectContains(t, changes, "max-retry-credentials: 1 -> 3")
|
||||||
expectContains(t, changes, "max-retry-interval: 1 -> 3")
|
expectContains(t, changes, "max-retry-interval: 1 -> 3")
|
||||||
expectContains(t, changes, "proxy-url: http://old-proxy -> http://new-proxy")
|
expectContains(t, changes, "proxy-url: http://old-proxy -> http://new-proxy")
|
||||||
expectContains(t, changes, "ws-auth: false -> true")
|
expectContains(t, changes, "ws-auth: false -> true")
|
||||||
|
|||||||
@@ -1239,6 +1239,67 @@ func TestReloadConfigFiltersAffectedOAuthProviders(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestReloadConfigTriggersCallbackForMaxRetryCredentialsChange(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
authDir := filepath.Join(tmpDir, "auth")
|
||||||
|
if err := os.MkdirAll(authDir, 0o755); err != nil {
|
||||||
|
t.Fatalf("failed to create auth dir: %v", err)
|
||||||
|
}
|
||||||
|
configPath := filepath.Join(tmpDir, "config.yaml")
|
||||||
|
|
||||||
|
oldCfg := &config.Config{
|
||||||
|
AuthDir: authDir,
|
||||||
|
MaxRetryCredentials: 0,
|
||||||
|
RequestRetry: 1,
|
||||||
|
MaxRetryInterval: 5,
|
||||||
|
}
|
||||||
|
newCfg := &config.Config{
|
||||||
|
AuthDir: authDir,
|
||||||
|
MaxRetryCredentials: 2,
|
||||||
|
RequestRetry: 1,
|
||||||
|
MaxRetryInterval: 5,
|
||||||
|
}
|
||||||
|
data, errMarshal := yaml.Marshal(newCfg)
|
||||||
|
if errMarshal != nil {
|
||||||
|
t.Fatalf("failed to marshal config: %v", errMarshal)
|
||||||
|
}
|
||||||
|
if errWrite := os.WriteFile(configPath, data, 0o644); errWrite != nil {
|
||||||
|
t.Fatalf("failed to write config: %v", errWrite)
|
||||||
|
}
|
||||||
|
|
||||||
|
callbackCalls := 0
|
||||||
|
callbackMaxRetryCredentials := -1
|
||||||
|
w := &Watcher{
|
||||||
|
configPath: configPath,
|
||||||
|
authDir: authDir,
|
||||||
|
lastAuthHashes: make(map[string]string),
|
||||||
|
reloadCallback: func(cfg *config.Config) {
|
||||||
|
callbackCalls++
|
||||||
|
if cfg != nil {
|
||||||
|
callbackMaxRetryCredentials = cfg.MaxRetryCredentials
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
w.SetConfig(oldCfg)
|
||||||
|
|
||||||
|
if ok := w.reloadConfig(); !ok {
|
||||||
|
t.Fatal("expected reloadConfig to succeed")
|
||||||
|
}
|
||||||
|
|
||||||
|
if callbackCalls != 1 {
|
||||||
|
t.Fatalf("expected reload callback to be called once, got %d", callbackCalls)
|
||||||
|
}
|
||||||
|
if callbackMaxRetryCredentials != 2 {
|
||||||
|
t.Fatalf("expected callback MaxRetryCredentials=2, got %d", callbackMaxRetryCredentials)
|
||||||
|
}
|
||||||
|
|
||||||
|
w.clientsMutex.RLock()
|
||||||
|
defer w.clientsMutex.RUnlock()
|
||||||
|
if w.config == nil || w.config.MaxRetryCredentials != 2 {
|
||||||
|
t.Fatalf("expected watcher config MaxRetryCredentials=2, got %+v", w.config)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestStartFailsWhenAuthDirMissing(t *testing.T) {
|
func TestStartFailsWhenAuthDirMissing(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
configPath := filepath.Join(tmpDir, "config.yaml")
|
configPath := filepath.Join(tmpDir, "config.yaml")
|
||||||
|
|||||||
@@ -138,8 +138,9 @@ type Manager struct {
|
|||||||
providerOffsets map[string]int
|
providerOffsets map[string]int
|
||||||
|
|
||||||
// Retry controls request retry behavior.
|
// Retry controls request retry behavior.
|
||||||
requestRetry atomic.Int32
|
requestRetry atomic.Int32
|
||||||
maxRetryInterval atomic.Int64
|
maxRetryCredentials atomic.Int32
|
||||||
|
maxRetryInterval atomic.Int64
|
||||||
|
|
||||||
// oauthModelAlias stores global OAuth model alias mappings (alias -> upstream name) keyed by channel.
|
// oauthModelAlias stores global OAuth model alias mappings (alias -> upstream name) keyed by channel.
|
||||||
oauthModelAlias atomic.Value
|
oauthModelAlias atomic.Value
|
||||||
@@ -384,18 +385,22 @@ func compileAPIKeyModelAliasForModels[T interface {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetRetryConfig updates retry attempts and cooldown wait interval.
|
// SetRetryConfig updates retry attempts, credential retry limit and cooldown wait interval.
|
||||||
func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration) {
|
func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) {
|
||||||
if m == nil {
|
if m == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if retry < 0 {
|
if retry < 0 {
|
||||||
retry = 0
|
retry = 0
|
||||||
}
|
}
|
||||||
|
if maxRetryCredentials < 0 {
|
||||||
|
maxRetryCredentials = 0
|
||||||
|
}
|
||||||
if maxRetryInterval < 0 {
|
if maxRetryInterval < 0 {
|
||||||
maxRetryInterval = 0
|
maxRetryInterval = 0
|
||||||
}
|
}
|
||||||
m.requestRetry.Store(int32(retry))
|
m.requestRetry.Store(int32(retry))
|
||||||
|
m.maxRetryCredentials.Store(int32(maxRetryCredentials))
|
||||||
m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds())
|
m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -506,11 +511,11 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye
|
|||||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
|
|
||||||
_, maxWait := m.retrySettings()
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
||||||
|
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for attempt := 0; ; attempt++ {
|
for attempt := 0; ; attempt++ {
|
||||||
resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts)
|
resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
||||||
if errExec == nil {
|
if errExec == nil {
|
||||||
return resp, nil
|
return resp, nil
|
||||||
}
|
}
|
||||||
@@ -537,11 +542,11 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip
|
|||||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
|
|
||||||
_, maxWait := m.retrySettings()
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
||||||
|
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for attempt := 0; ; attempt++ {
|
for attempt := 0; ; attempt++ {
|
||||||
resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts)
|
resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
||||||
if errExec == nil {
|
if errExec == nil {
|
||||||
return resp, nil
|
return resp, nil
|
||||||
}
|
}
|
||||||
@@ -568,11 +573,11 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli
|
|||||||
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
|
|
||||||
_, maxWait := m.retrySettings()
|
_, maxRetryCredentials, maxWait := m.retrySettings()
|
||||||
|
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for attempt := 0; ; attempt++ {
|
for attempt := 0; ; attempt++ {
|
||||||
result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts)
|
result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials)
|
||||||
if errStream == nil {
|
if errStream == nil {
|
||||||
return result, nil
|
return result, nil
|
||||||
}
|
}
|
||||||
@@ -591,7 +596,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli
|
|||||||
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) {
|
||||||
if len(providers) == 0 {
|
if len(providers) == 0 {
|
||||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
@@ -600,6 +605,12 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
|
|||||||
tried := make(map[string]struct{})
|
tried := make(map[string]struct{})
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for {
|
for {
|
||||||
|
if maxRetryCredentials > 0 && len(tried) >= maxRetryCredentials {
|
||||||
|
if lastErr != nil {
|
||||||
|
return cliproxyexecutor.Response{}, lastErr
|
||||||
|
}
|
||||||
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||||
|
}
|
||||||
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
||||||
if errPick != nil {
|
if errPick != nil {
|
||||||
if lastErr != nil {
|
if lastErr != nil {
|
||||||
@@ -647,7 +658,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) {
|
||||||
if len(providers) == 0 {
|
if len(providers) == 0 {
|
||||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
@@ -656,6 +667,12 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
|
|||||||
tried := make(map[string]struct{})
|
tried := make(map[string]struct{})
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for {
|
for {
|
||||||
|
if maxRetryCredentials > 0 && len(tried) >= maxRetryCredentials {
|
||||||
|
if lastErr != nil {
|
||||||
|
return cliproxyexecutor.Response{}, lastErr
|
||||||
|
}
|
||||||
|
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||||
|
}
|
||||||
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
||||||
if errPick != nil {
|
if errPick != nil {
|
||||||
if lastErr != nil {
|
if lastErr != nil {
|
||||||
@@ -703,7 +720,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string,
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) {
|
||||||
if len(providers) == 0 {
|
if len(providers) == 0 {
|
||||||
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||||
}
|
}
|
||||||
@@ -712,6 +729,12 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string
|
|||||||
tried := make(map[string]struct{})
|
tried := make(map[string]struct{})
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for {
|
for {
|
||||||
|
if maxRetryCredentials > 0 && len(tried) >= maxRetryCredentials {
|
||||||
|
if lastErr != nil {
|
||||||
|
return nil, lastErr
|
||||||
|
}
|
||||||
|
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||||
|
}
|
||||||
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, opts, tried)
|
||||||
if errPick != nil {
|
if errPick != nil {
|
||||||
if lastErr != nil {
|
if lastErr != nil {
|
||||||
@@ -1108,11 +1131,11 @@ func (m *Manager) normalizeProviders(providers []string) []string {
|
|||||||
return result
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) retrySettings() (int, time.Duration) {
|
func (m *Manager) retrySettings() (int, int, time.Duration) {
|
||||||
if m == nil {
|
if m == nil {
|
||||||
return 0, 0
|
return 0, 0, 0
|
||||||
}
|
}
|
||||||
return int(m.requestRetry.Load()), time.Duration(m.maxRetryInterval.Load())
|
return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load())
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) {
|
func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) {
|
||||||
|
|||||||
@@ -2,13 +2,17 @@ package auth
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testing.T) {
|
func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testing.T) {
|
||||||
m := NewManager(nil, nil, nil)
|
m := NewManager(nil, nil, nil)
|
||||||
m.SetRetryConfig(3, 30*time.Second)
|
m.SetRetryConfig(3, 30*time.Second, 0)
|
||||||
|
|
||||||
model := "test-model"
|
model := "test-model"
|
||||||
next := time.Now().Add(5 * time.Second)
|
next := time.Now().Add(5 * time.Second)
|
||||||
@@ -31,7 +35,7 @@ func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testi
|
|||||||
t.Fatalf("register auth: %v", errRegister)
|
t.Fatalf("register auth: %v", errRegister)
|
||||||
}
|
}
|
||||||
|
|
||||||
_, maxWait := m.retrySettings()
|
_, _, maxWait := m.retrySettings()
|
||||||
wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: 500, Message: "boom"}, 0, []string{"claude"}, model, maxWait)
|
wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: 500, Message: "boom"}, 0, []string{"claude"}, model, maxWait)
|
||||||
if shouldRetry {
|
if shouldRetry {
|
||||||
t.Fatalf("expected shouldRetry=false for request_retry=0, got true (wait=%v)", wait)
|
t.Fatalf("expected shouldRetry=false for request_retry=0, got true (wait=%v)", wait)
|
||||||
@@ -56,6 +60,124 @@ func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testi
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type credentialRetryLimitExecutor struct {
|
||||||
|
id string
|
||||||
|
|
||||||
|
mu sync.Mutex
|
||||||
|
calls int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) Identifier() string {
|
||||||
|
return e.id
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||||
|
e.recordCall()
|
||||||
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: 500, Message: "boom"}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
||||||
|
e.recordCall()
|
||||||
|
return nil, &Error{HTTPStatus: 500, Message: "boom"}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
|
||||||
|
return auth, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||||
|
e.recordCall()
|
||||||
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: 500, Message: "boom"}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) recordCall() {
|
||||||
|
e.mu.Lock()
|
||||||
|
defer e.mu.Unlock()
|
||||||
|
e.calls++
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *credentialRetryLimitExecutor) Calls() int {
|
||||||
|
e.mu.Lock()
|
||||||
|
defer e.mu.Unlock()
|
||||||
|
return e.calls
|
||||||
|
}
|
||||||
|
|
||||||
|
func newCredentialRetryLimitTestManager(t *testing.T, maxRetryCredentials int) (*Manager, *credentialRetryLimitExecutor) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
m := NewManager(nil, nil, nil)
|
||||||
|
m.SetRetryConfig(0, 0, maxRetryCredentials)
|
||||||
|
|
||||||
|
executor := &credentialRetryLimitExecutor{id: "claude"}
|
||||||
|
m.RegisterExecutor(executor)
|
||||||
|
|
||||||
|
auth1 := &Auth{ID: "auth-1", Provider: "claude"}
|
||||||
|
auth2 := &Auth{ID: "auth-2", Provider: "claude"}
|
||||||
|
if _, errRegister := m.Register(context.Background(), auth1); errRegister != nil {
|
||||||
|
t.Fatalf("register auth1: %v", errRegister)
|
||||||
|
}
|
||||||
|
if _, errRegister := m.Register(context.Background(), auth2); errRegister != nil {
|
||||||
|
t.Fatalf("register auth2: %v", errRegister)
|
||||||
|
}
|
||||||
|
|
||||||
|
return m, executor
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManager_MaxRetryCredentials_LimitsCrossCredentialRetries(t *testing.T) {
|
||||||
|
request := cliproxyexecutor.Request{Model: "test-model"}
|
||||||
|
testCases := []struct {
|
||||||
|
name string
|
||||||
|
invoke func(*Manager) error
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "execute",
|
||||||
|
invoke: func(m *Manager) error {
|
||||||
|
_, errExecute := m.Execute(context.Background(), []string{"claude"}, request, cliproxyexecutor.Options{})
|
||||||
|
return errExecute
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "execute_count",
|
||||||
|
invoke: func(m *Manager) error {
|
||||||
|
_, errExecute := m.ExecuteCount(context.Background(), []string{"claude"}, request, cliproxyexecutor.Options{})
|
||||||
|
return errExecute
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "execute_stream",
|
||||||
|
invoke: func(m *Manager) error {
|
||||||
|
_, errExecute := m.ExecuteStream(context.Background(), []string{"claude"}, request, cliproxyexecutor.Options{})
|
||||||
|
return errExecute
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
tc := tc
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
limitedManager, limitedExecutor := newCredentialRetryLimitTestManager(t, 1)
|
||||||
|
if errInvoke := tc.invoke(limitedManager); errInvoke == nil {
|
||||||
|
t.Fatalf("expected error for limited retry execution")
|
||||||
|
}
|
||||||
|
if calls := limitedExecutor.Calls(); calls != 1 {
|
||||||
|
t.Fatalf("expected 1 call with max-retry-credentials=1, got %d", calls)
|
||||||
|
}
|
||||||
|
|
||||||
|
unlimitedManager, unlimitedExecutor := newCredentialRetryLimitTestManager(t, 0)
|
||||||
|
if errInvoke := tc.invoke(unlimitedManager); errInvoke == nil {
|
||||||
|
t.Fatalf("expected error for unlimited retry execution")
|
||||||
|
}
|
||||||
|
if calls := unlimitedExecutor.Calls(); calls != 2 {
|
||||||
|
t.Fatalf("expected 2 calls with max-retry-credentials=0, got %d", calls)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestManager_MarkResult_RespectsAuthDisableCoolingOverride(t *testing.T) {
|
func TestManager_MarkResult_RespectsAuthDisableCoolingOverride(t *testing.T) {
|
||||||
prev := quotaCooldownDisabled.Load()
|
prev := quotaCooldownDisabled.Load()
|
||||||
quotaCooldownDisabled.Store(false)
|
quotaCooldownDisabled.Store(false)
|
||||||
|
|||||||
@@ -336,7 +336,7 @@ func (s *Service) applyRetryConfig(cfg *config.Config) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second
|
maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second
|
||||||
s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval)
|
s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval, cfg.MaxRetryCredentials)
|
||||||
}
|
}
|
||||||
|
|
||||||
func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) {
|
func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) {
|
||||||
|
|||||||
Reference in New Issue
Block a user