package auth import ( "context" "fmt" "strings" "time" "github.com/google/uuid" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" ) // SetRetryConfig updates additional credential retry rounds, the per-round credential limit, and the cooldown wait interval. func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) { if m == nil { return } if retry < 0 { retry = 0 } if maxRetryCredentials < 0 { maxRetryCredentials = 0 } if maxRetryInterval < 0 { maxRetryInterval = 0 } m.requestRetry.Store(int32(retry)) m.maxRetryCredentials.Store(int32(maxRetryCredentials)) m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds()) } // RegisterExecutor registers a provider executor with the manager. func (m *Manager) RegisterExecutor(executor ProviderExecutor) { if executor == nil { return } provider := strings.TrimSpace(executor.Identifier()) if provider == "" { return } var replaced ProviderExecutor m.mu.Lock() replaced = m.executors[provider] m.executors[provider] = executor m.mu.Unlock() if replaced == nil || replaced == executor { return } if closer, ok := replaced.(ExecutionSessionCloser); ok && closer != nil { closer.CloseExecutionSession(CloseAllExecutionSessionsID) } } // UnregisterExecutor removes the executor associated with the provider key. func (m *Manager) UnregisterExecutor(provider string) { provider = strings.ToLower(strings.TrimSpace(provider)) if provider == "" { return } m.mu.Lock() delete(m.executors, provider) m.mu.Unlock() } // Register inserts a new auth entry into the manager. func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { if auth == nil { return nil, nil } NormalizeCredentialMetadata(auth.Metadata) if errWeight := ValidateAuthWeight(auth); errWeight != nil { return nil, fmt.Errorf("register auth: %w", errWeight) } if auth.ID == "" { auth.ID = uuid.NewString() } now := time.Now() cooldownStateChanged := normalizeModelStates(auth) if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged } auth.EnsureIndex() authClone := auth.Clone() m.mu.Lock() m.auths[auth.ID] = authClone m.mu.Unlock() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } if m.scheduler != nil { m.scheduler.upsertAuth(authClone) } m.queueRefreshReschedule(auth.ID) _ = m.persist(ctx, auth) m.hook.OnAuthRegistered(ctx, auth.Clone()) if cooldownStateChanged { m.persistCooldownStates(ctx) } return auth.Clone(), nil } // Update replaces an existing auth entry and notifies hooks. func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { if auth == nil || auth.ID == "" { return nil, nil } NormalizeCredentialMetadata(auth.Metadata) if errWeight := ValidateAuthWeight(auth); errWeight != nil { return nil, fmt.Errorf("update auth: %w", errWeight) } m.mu.Lock() existing, ok := m.auths[auth.ID] if !ok || existing == nil { m.mu.Unlock() return nil, nil } if !auth.indexAssigned && auth.Index == "" { auth.Index = existing.Index auth.indexAssigned = existing.indexAssigned } auth.Success = existing.Success auth.Failed = existing.Failed auth.recentRequests = existing.recentRequests if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { auth.ModelStates = existing.ModelStates } if existing.Quota.Exceeded && existing.Quota.Reason == "credential_quota" && existing.Quota.NextRecoverAt.After(time.Now()) { auth.Unavailable = existing.Unavailable auth.NextRetryAfter = existing.NextRetryAfter auth.Quota = existing.Quota if auth.Status == StatusActive { auth.Status = existing.Status } } } now := time.Now() cooldownStateChanged := normalizeModelStates(auth) if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged } auth.EnsureIndex() authClone := auth.Clone() m.auths[auth.ID] = authClone m.mu.Unlock() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } if m.scheduler != nil { m.scheduler.upsertAuth(authClone) } m.queueRefreshReschedule(auth.ID) _ = m.persist(ctx, auth) m.hook.OnAuthUpdated(ctx, auth.Clone()) if cooldownStateChanged { m.persistCooldownStates(ctx) } return auth.Clone(), nil } // Remove deletes an auth from runtime state without persisting. // Disk and token-store deletion must be handled by the caller. func (m *Manager) Remove(ctx context.Context, id string) { if m == nil { return } id = strings.TrimSpace(id) if id == "" { return } _ = ctx m.mu.Lock() existing := m.auths[id] if existing == nil { m.mu.Unlock() return } provider := strings.TrimSpace(existing.Provider) delete(m.auths, id) if m.modelPoolOffsets != nil { delete(m.modelPoolOffsets, id) } for sessionID, sessionAuths := range m.homeRuntimeAuths { if sessionAuths == nil { continue } delete(sessionAuths, id) if len(sessionAuths) == 0 { delete(m.homeRuntimeAuths, sessionID) } } m.mu.Unlock() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } if m.scheduler != nil { m.scheduler.removeAuth(id) } m.queueRefreshUnschedule(id) m.invalidateSessionAffinity(id) if provider != "" { if exec, ok := m.Executor(provider); ok && exec != nil { if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { closer.CloseExecutionSession(CloseAllExecutionSessionsID) } } } m.persistCooldownStates(ctx) } func (m *Manager) invalidateSessionAffinity(authID string) { if m == nil || authID == "" { return } if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil { invalidator.InvalidateAuth(authID) } } // Load resets manager state from the backing store. func (m *Manager) Load(ctx context.Context) error { m.mu.Lock() if m.store == nil { m.mu.Unlock() return nil } items, err := m.store.List(ctx) if err != nil { m.mu.Unlock() return err } m.auths = make(map[string]*Auth, len(items)) for _, auth := range items { if auth == nil || auth.ID == "" { continue } NormalizeCredentialMetadata(auth.Metadata) if errWeight := ValidateAuthWeight(auth); errWeight != nil { continue } auth.EnsureIndex() m.auths[auth.ID] = auth.Clone() } cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) if cfg == nil { cfg = &internalconfig.Config{} } m.rebuildAPIKeyModelAliasLocked(cfg) m.mu.Unlock() m.syncScheduler() return nil } func (m *Manager) persist(ctx context.Context, auth *Auth) error { if m.store == nil || auth == nil { return nil } if errWeight := ValidateAuthWeight(auth); errWeight != nil { return fmt.Errorf("persist auth: %w", errWeight) } if shouldSkipPersist(ctx) { return nil } if IsConfigAPIKeyAuth(auth) { return nil } if auth.Attributes != nil { if v := strings.ToLower(strings.TrimSpace(auth.Attributes["runtime_only"])); v == "true" { return nil } } if IsPluginVirtualAuth(auth) { return nil } // Skip persistence when metadata is absent (e.g., runtime-only auths). if auth.Metadata == nil { return nil } _, err := m.store.Save(ctx, auth) return err }