From 079ec51f50fe3dc7e8e958a03863501764c1b99f Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 22 Jun 2026 00:41:52 +0000 Subject: [PATCH] feat(cliproxy): optimize API key alias rebuild with deferred execution and caching - Added `RefreshAPIKeyModelAlias` for explicit alias table rebuilds. - Introduced deferred rebuild support with `WithDeferredAPIKeyModelAliasRebuild` and context flag validation. - Implemented `openAICompatibilityRegistrationCache` to streamline OpenAI compatibility model registrations. - Updated executor and model registration workflows to utilize cached compatibility data, improving efficiency in batch operations. - Adjusted max worker limits dynamically based on model categories. Closes: #3953 --- sdk/cliproxy/service.go | 180 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------------- sdk/cliproxy/auth/conductor.go | 17 ++++++++++++++--- sdk/cliproxy/auth/persist_policy.go | 19 +++++++++++++++++++ 3 file(s) changed, 188 insertion(s)(+), 28 deletion(s)(-) diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -108,7 +108,10 @@ homeLogForwarder *logging.HomeAppLogForwarder } -const modelRegistrationMaxWorkersPerCategory = 5 +const ( + modelRegistrationMaxWorkersPerCategory = 5 + modelRegistrationMaxWorkersOpenAICompatibility = 20 +) const ( modelRegistrationPhaseConfigAPIKey = iota @@ -118,7 +121,7 @@ type modelRegistrationTask struct { phase int category string - run func() + run func(*openAICompatibilityRegistrationCache) } type executorRegistrationOptions struct { @@ -235,8 +238,8 @@ 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 @@ 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 @@ 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 @@ return default: } - task.run() + task.run(compatCache) } }() } @@ -361,6 +366,14 @@ 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 @@ 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 @@ 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 @@ 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) 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 @@ -799,6 +823,62 @@ return util.OpenAICompatibleProviderKey(providerKey), compatName, true } 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 { @@ -973,6 +1053,13 @@ 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 } @@ -1004,6 +1091,13 @@ !s.hasNativeOpenAICompatExecutorConfig(a, providerKey) { 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 @@ 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 @@ // 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 @@ 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 @@ // 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 @@ 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 @@ // 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 diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1750,6 +1750,11 @@ 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 @@ 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 @@ 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 @@ } 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 --- a/sdk/cliproxy/auth/persist_policy.go +++ b/sdk/cliproxy/auth/persist_policy.go @@ -3,6 +3,7 @@ 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 @@ -19,6 +20,24 @@ return false } v := ctx.Value(skipPersistContextKey{}) + 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 } -- tangled.sh