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
2026-11-11 10:56:46 +08:00 · 2026-06-22 08:41:52 +08:00
parent 1f2504ebcc
commit 079ec51f50
3 changed files with 188 additions and 28 deletions
@@ -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)
 	}
@@ -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
+}
@@ -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