From 0f55e550cf18988f45f6a997701e2a476980598a Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 26 Sep 2025 19:38:44 +0800 Subject: [PATCH] refactor(registry): Preserve duplicate models in client registration The `RegisterClient` function previously deduplicated the list of models provided by a client. This could lead to an inaccurate representation of the client's state if it intentionally registered the same model ID multiple times. This change refactors the registration logic to store the raw, unfiltered list of models, preserving their original order and count. A new `rawModelIDs` slice tracks the complete list for storage in `clientModels`, while the logic for processing changes continues to use a unique set of model IDs for efficiency. This ensures the registry's state accurately reflects what the client provides. --- internal/registry/model_registry.go | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index b52e48b2..44b6991c 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -102,24 +102,24 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ defer r.mutex.Unlock() provider := strings.ToLower(clientProvider) - seen := make(map[string]struct{}) - modelIDs := make([]string, 0, len(models)) + uniqueModelIDs := make([]string, 0, len(models)) + rawModelIDs := make([]string, 0, len(models)) newModels := make(map[string]*ModelInfo, len(models)) newCounts := make(map[string]int, len(models)) for _, model := range models { if model == nil || model.ID == "" { continue } + rawModelIDs = append(rawModelIDs, model.ID) newCounts[model.ID]++ - if _, exists := seen[model.ID]; exists { + if _, exists := newModels[model.ID]; exists { continue } - seen[model.ID] = struct{}{} - modelIDs = append(modelIDs, model.ID) newModels[model.ID] = model + uniqueModelIDs = append(uniqueModelIDs, model.ID) } - if len(modelIDs) == 0 { + if len(uniqueModelIDs) == 0 { // No models supplied; unregister existing client state if present. r.unregisterClientInternal(clientID) delete(r.clientModels, clientID) @@ -135,17 +135,17 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ providerChanged := oldProvider != provider if !hadExisting { // Pure addition path. - for _, modelID := range modelIDs { + for _, modelID := range rawModelIDs { model := newModels[modelID] r.addModelRegistration(modelID, provider, model, now) } - r.clientModels[clientID] = modelIDs + r.clientModels[clientID] = append([]string(nil), rawModelIDs...) if provider != "" { r.clientProviders[clientID] = provider } else { delete(r.clientProviders, clientID) } - log.Debugf("Registered client %s from provider %s with %d models", clientID, clientProvider, len(modelIDs)) + log.Debugf("Registered client %s from provider %s with %d models", clientID, clientProvider, len(rawModelIDs)) misc.LogCredentialSeparator() return } @@ -156,7 +156,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ } added := make([]string, 0) - for _, id := range modelIDs { + for _, id := range uniqueModelIDs { if oldCounts[id] == 0 { added = append(added, id) } @@ -232,7 +232,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ for _, id := range added { addedSet[id] = struct{}{} } - for _, id := range modelIDs { + for _, id := range uniqueModelIDs { model := newModels[id] if reg, ok := r.models[id]; ok { reg.Info = cloneModelInfo(model) @@ -263,8 +263,8 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ } // Update client bookkeeping. - if len(modelIDs) > 0 { - r.clientModels[clientID] = modelIDs + if len(rawModelIDs) > 0 { + r.clientModels[clientID] = append([]string(nil), rawModelIDs...) } if provider != "" { r.clientProviders[clientID] = provider