diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 54a52559..8a890912 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1750,6 +1750,11 @@ func (m *Manager) rebuildAPIKeyModelAliasFromRuntimeConfig() { m.rebuildAPIKeyModelAliasLocked(cfg) } +// RefreshAPIKeyModelAlias rebuilds the API-key model alias table from the current runtime config. +func (m *Manager) RefreshAPIKeyModelAlias() { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() +} + func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) { if m == nil { return @@ -1931,7 +1936,9 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { m.mu.Lock() m.auths[auth.ID] = authClone m.mu.Unlock() - m.rebuildAPIKeyModelAliasFromRuntimeConfig() + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } if m.scheduler != nil { m.scheduler.upsertAuth(authClone) } @@ -1976,7 +1983,9 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { authClone := auth.Clone() m.auths[auth.ID] = authClone m.mu.Unlock() - m.rebuildAPIKeyModelAliasFromRuntimeConfig() + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } if m.scheduler != nil { m.scheduler.upsertAuth(authClone) } @@ -2023,7 +2032,9 @@ func (m *Manager) Remove(ctx context.Context, id string) { } m.mu.Unlock() - m.rebuildAPIKeyModelAliasFromRuntimeConfig() + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } if m.scheduler != nil { m.scheduler.removeAuth(id) } diff --git a/sdk/cliproxy/auth/persist_policy.go b/sdk/cliproxy/auth/persist_policy.go index 35423c30..3c9e612c 100644 --- a/sdk/cliproxy/auth/persist_policy.go +++ b/sdk/cliproxy/auth/persist_policy.go @@ -3,6 +3,7 @@ package auth import "context" type skipPersistContextKey struct{} +type deferAPIKeyModelAliasRebuildContextKey struct{} // WithSkipPersist returns a derived context that disables persistence for Manager Update/Register calls. // It is intended for code paths that are reacting to file watcher events, where the file on disk is @@ -22,3 +23,21 @@ func shouldSkipPersist(ctx context.Context) bool { enabled, ok := v.(bool) return ok && enabled } + +// WithDeferredAPIKeyModelAliasRebuild returns a derived context that defers API-key model alias table rebuilds. +// Callers that use this for a batch of Register/Update/Remove operations must call RefreshAPIKeyModelAlias once. +func WithDeferredAPIKeyModelAliasRebuild(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, deferAPIKeyModelAliasRebuildContextKey{}, true) +} + +func shouldDeferAPIKeyModelAliasRebuild(ctx context.Context) bool { + if ctx == nil { + return false + } + v := ctx.Value(deferAPIKeyModelAliasRebuildContextKey{}) + enabled, ok := v.(bool) + return ok && enabled +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 6f2d9967..55a0365f 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -108,7 +108,10 @@ type Service struct { homeLogForwarder *logging.HomeAppLogForwarder } -const modelRegistrationMaxWorkersPerCategory = 5 +const ( + modelRegistrationMaxWorkersPerCategory = 5 + modelRegistrationMaxWorkersOpenAICompatibility = 20 +) const ( modelRegistrationPhaseConfigAPIKey = iota @@ -118,7 +121,7 @@ const ( type modelRegistrationTask struct { phase int category string - run func() + run func(*openAICompatibilityRegistrationCache) } type executorRegistrationOptions struct { @@ -235,8 +238,8 @@ func (s *Service) registerModelsForAuthBatch(ctx context.Context, auths []*corea tasks = append(tasks, modelRegistrationTask{ phase: modelRegistrationPhase(authForRegistration), category: modelRegistrationCategory(authForRegistration), - run: func() { - s.completeModelRegistrationForAuth(ctx, authForRegistration) + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(ctx, authForRegistration, compatCache) }, }) } @@ -261,11 +264,12 @@ func (s *Service) runModelRegistrationTasks(ctx context.Context, tasks []modelRe otherTasks = append(otherTasks, task) } - s.runModelRegistrationTaskPhase(ctx, configAPIKeyTasks) - s.runModelRegistrationTaskPhase(ctx, otherTasks) + compatCache := s.newOpenAICompatibilityRegistrationCache() + s.runModelRegistrationTaskPhase(ctx, configAPIKeyTasks, compatCache) + s.runModelRegistrationTaskPhase(ctx, otherTasks, compatCache) } -func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []modelRegistrationTask) { +func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []modelRegistrationTask, compatCache *openAICompatibilityRegistrationCache) { if len(tasks) == 0 { return } @@ -290,8 +294,9 @@ func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []mod for _, category := range order { group := grouped[category] workers := len(group) - if workers > modelRegistrationMaxWorkersPerCategory { - workers = modelRegistrationMaxWorkersPerCategory + maxWorkers := modelRegistrationMaxWorkersForCategory(category) + if workers > maxWorkers { + workers = maxWorkers } if workers <= 0 { continue @@ -308,7 +313,7 @@ func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []mod return default: } - task.run() + task.run(compatCache) } }() } @@ -361,6 +366,14 @@ func modelRegistrationCategory(auth *coreauth.Auth) string { return provider + ":" + authKind } +func modelRegistrationMaxWorkersForCategory(category string) int { + category = strings.ToLower(strings.TrimSpace(category)) + if strings.HasPrefix(category, "openai-compatible-") || strings.HasPrefix(category, "openai-compatibility") { + return modelRegistrationMaxWorkersOpenAICompatibility + } + return modelRegistrationMaxWorkersPerCategory +} + func (s *Service) registerModelRefreshCallback() { // Register callback for startup and periodic model catalog refresh. // When remote model definitions change, re-register models for affected providers. @@ -396,8 +409,8 @@ func (s *Service) registerModelRefreshCallback() { tasks = append(tasks, modelRegistrationTask{ phase: modelRegistrationPhase(authForRefresh), category: modelRegistrationCategory(authForRefresh), - run: func() { - if s.refreshModelRegistrationForAuth(authForRefresh) { + run: func(compatCache *openAICompatibilityRegistrationCache) { + if s.refreshModelRegistrationForAuthWithCache(authForRefresh, compatCache) { refreshedMu.Lock() refreshed++ refreshedMu.Unlock() @@ -500,24 +513,27 @@ func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthU return } + registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) tasks := make([]modelRegistrationTask, 0, len(updates)) needsPluginSync := false + needsAliasRebuild := false for _, update := range updates { switch update.Action { case watcher.AuthUpdateActionAdd, watcher.AuthUpdateActionModify: if update.Auth == nil || update.Auth.ID == "" { continue } - auth := s.prepareCoreAuthForModelRegistration(ctx, update.Auth) + auth := s.prepareCoreAuthForModelRegistration(registrationCtx, update.Auth) if auth == nil { continue } + needsAliasRebuild = true authForRegistration := auth tasks = append(tasks, modelRegistrationTask{ phase: modelRegistrationPhase(authForRegistration), category: modelRegistrationCategory(authForRegistration), - run: func() { - s.completeModelRegistrationForAuth(ctx, authForRegistration) + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) }, }) needsPluginSync = true @@ -529,15 +545,19 @@ func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthU if id == "" { continue } - s.applyCoreAuthRemoval(ctx, id) + s.applyCoreAuthRemoval(registrationCtx, id) + needsAliasRebuild = true default: log.Debugf("received unknown auth update action: %v", update.Action) } } - s.runModelRegistrationTasks(ctx, tasks) + if needsAliasRebuild { + s.coreManager.RefreshAPIKeyModelAlias() + } + s.runModelRegistrationTasks(registrationCtx, tasks) if needsPluginSync { - s.syncPluginRuntime(ctx) + s.syncPluginRuntime(registrationCtx) } } @@ -699,10 +719,14 @@ func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth } func (s *Service) completeModelRegistrationForAuth(ctx context.Context, auth *coreauth.Auth) { + s.completeModelRegistrationForAuthWithCache(ctx, auth, nil) +} + +func (s *Service) completeModelRegistrationForAuthWithCache(ctx context.Context, auth *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { return } - s.registerModelsForAuth(ctx, auth) + s.registerModelsForAuthWithCache(ctx, auth, compatCache) s.coreManager.ReconcileRegistryModelStates(ctx, auth.ID) // Refresh the scheduler entry so that the auth's supportedModelSet is rebuilt @@ -801,6 +825,62 @@ func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName return "", "", false } +type openAICompatibilityRegistrationCache struct { + byName map[string]*openAICompatibilityRegistrationEntry +} + +type openAICompatibilityRegistrationEntry struct { + providerKey string + models []*ModelInfo +} + +func (s *Service) newOpenAICompatibilityRegistrationCache() *openAICompatibilityRegistrationCache { + if s == nil { + return nil + } + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + if cfg == nil || len(cfg.OpenAICompatibility) == 0 { + return nil + } + + cache := &openAICompatibilityRegistrationCache{ + byName: make(map[string]*openAICompatibilityRegistrationEntry, len(cfg.OpenAICompatibility)), + } + for i := range cfg.OpenAICompatibility { + compat := &cfg.OpenAICompatibility[i] + if compat.Disabled { + continue + } + compatName := strings.TrimSpace(compat.Name) + key := strings.ToLower(compatName) + if _, exists := cache.byName[key]; exists { + continue + } + providerName := strings.ToLower(compatName) + if providerName == "" { + providerName = "openai-compatibility" + } + cache.byName[key] = &openAICompatibilityRegistrationEntry{ + providerKey: util.OpenAICompatibleProviderKey(providerName), + models: buildOpenAICompatibilityConfigModels(compat), + } + } + if len(cache.byName) == 0 { + return nil + } + return cache +} + +func (c *openAICompatibilityRegistrationCache) lookup(compatName string) (*openAICompatibilityRegistrationEntry, bool) { + if c == nil || len(c.byName) == 0 { + return nil, false + } + entry, ok := c.byName[strings.ToLower(strings.TrimSpace(compatName))] + return entry, ok +} + func (s *Service) hasNativeOpenAICompatExecutorConfig(a *coreauth.Auth, providerKey string) bool { if a == nil { return false @@ -973,6 +1053,13 @@ func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { if compatProviderKey == "" { compatProviderKey = "openai-compatibility" } + if !forceReplace { + if existingExecutor, hasExecutor := s.coreManager.Executor(compatProviderKey); hasExecutor { + if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { + return + } + } + } s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, s.cfg)) return } @@ -1005,6 +1092,13 @@ func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { s.unregisterOpenAICompatExecutor(providerKey) return } + if !forceReplace { + if existingExecutor, hasExecutor := s.coreManager.Executor(providerKey); hasExecutor { + if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { + return + } + } + } s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, s.cfg)) } } @@ -1289,25 +1383,31 @@ func (s *Service) registerConfigAPIKeyAuths(ctx context.Context, cfg *config.Con return } + registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) tasks := make([]modelRegistrationTask, 0, len(auths)) + needsAliasRebuild := false for _, auth := range auths { if !coreauth.IsConfigAPIKeyAuth(auth) { continue } - prepared := s.prepareCoreAuthForModelRegistration(ctx, auth) + prepared := s.prepareCoreAuthForModelRegistration(registrationCtx, auth) if prepared == nil { continue } + needsAliasRebuild = true authForRegistration := prepared tasks = append(tasks, modelRegistrationTask{ phase: modelRegistrationPhaseConfigAPIKey, category: modelRegistrationCategory(authForRegistration), - run: func() { - s.completeModelRegistrationForAuth(ctx, authForRegistration) + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) }, }) } - s.runModelRegistrationTasks(ctx, tasks) + if needsAliasRebuild { + s.coreManager.RefreshAPIKeyModelAlias() + } + s.runModelRegistrationTasks(registrationCtx, tasks) } func forceHomeRuntimeConfig(cfg *config.Config) { @@ -1780,6 +1880,10 @@ func (s *Service) ensureAuthDir() error { // registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier. func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { + s.registerModelsForAuthWithCache(ctx, a, nil) +} + +func (s *Service) registerModelsForAuthWithCache(ctx context.Context, a *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { if a == nil || a.ID == "" { return } @@ -1934,6 +2038,28 @@ func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { isCompatAuth = true } } + if cached, ok := compatCache.lookup(compatName); ok { + isCompatAuth = true + if providerKey == "" { + providerKey = cached.providerKey + } + if providerKey == "" { + providerKey = "openai-compatibility" + } + ms := cached.models + if len(ms) > 0 { + ms = s.appendPluginModels(providerKey, ms) + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + ms = s.appendPluginModels(providerKey, nil) + if len(ms) > 0 { + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + GlobalModelRegistry().UnregisterClient(a.ID) + } + } + return + } for i := range s.cfg.OpenAICompatibility { compat := &s.cfg.OpenAICompatibility[i] if compat.Disabled { @@ -1995,6 +2121,10 @@ func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { // as part of the previous registration snapshot and is cleared when the auth is // rebound to the refreshed model catalog. func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { + return s.refreshModelRegistrationForAuthWithCache(current, nil) +} + +func (s *Service) refreshModelRegistrationForAuthWithCache(current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { if s == nil || s.coreManager == nil || current == nil || current.ID == "" { return false } @@ -2003,7 +2133,7 @@ func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { if !current.Disabled { s.ensureExecutorsForAuth(current) } - s.registerModelsForAuth(ctx, current) + s.registerModelsForAuthWithCache(ctx, current, compatCache) s.coreManager.ReconcileRegistryModelStates(ctx, current.ID) latest, ok := s.latestAuthForModelRegistration(current.ID) @@ -2017,7 +2147,7 @@ func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { // stale model registrations behind. This may duplicate registration work when // no auth fields changed, but keeps the refresh path simple and correct. s.ensureExecutorsForAuth(latest) - s.registerModelsForAuth(ctx, latest) + s.registerModelsForAuthWithCache(ctx, latest, compatCache) s.coreManager.ReconcileRegistryModelStates(ctx, latest.ID) s.coreManager.RefreshSchedulerEntry(current.ID) return true