chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-03 04:50:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(registry): Preserve duplicate models in client registration

Browse Source 
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.
This commit is contained in:
hkfires
2025-09-26 19:38:44 +08:00
parent e1de04230f
commit 0f55e550cf
1 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
internal/registry/model_registry.go
View File

@@ -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