Unified · Split
| @@ -90,6 +90,9 @@ type ModelRegistry struct { · 89 unmodified lines | ||
| 90 | 90 | models map[string]*ModelRegistration |
| 91 | 91 | // clientModels maps client ID to the models it provides |
| 92 | 92 | clientModels map[string][]string |
| 93 | + // clientModelInfos maps client ID to a map of model ID -> ModelInfo | |
| 94 | + // This preserves the original model info provided by each client | |
| 95 | + clientModelInfos map[string]map[string]*ModelInfo | |
| 93 | 96 | // clientProviders maps client ID to its provider identifier |
| 94 | 97 | clientProviders map[string]string |
| 95 | 98 | // mutex ensures thread-safe access to the registry |
| @@ -104,10 +107,11 @@ var registryOnce sync.Once · 8 unmodified lines | ||
| 104 | 107 | func GetGlobalRegistry() *ModelRegistry { |
| 105 | 108 | registryOnce.Do(func() { |
| 106 | 109 | globalRegistry = &ModelRegistry{ |
| 107 | − models: make(map[string]*ModelRegistration), | |
| 108 | − clientModels: make(map[string][]string), | |
| 109 | − clientProviders: make(map[string]string), | |
| 110 | − mutex: &sync.RWMutex{}, | |
| 110 | + models: make(map[string]*ModelRegistration), | |
| 111 | + clientModels: make(map[string][]string), | |
| 112 | + clientModelInfos: make(map[string]map[string]*ModelInfo), | |
| 113 | + clientProviders: make(map[string]string), | |
| 114 | + mutex: &sync.RWMutex{}, | |
| 111 | 115 | } |
| 112 | 116 | }) |
| 113 | 117 | return globalRegistry |
| @@ -144,6 +148,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ · 30 unmodified lines | ||
| 144 | 148 | // No models supplied; unregister existing client state if present. |
| 145 | 149 | r.unregisterClientInternal(clientID) |
| 146 | 150 | delete(r.clientModels, clientID) |
| 151 | + delete(r.clientModelInfos, clientID) | |
| 147 | 152 | delete(r.clientProviders, clientID) |
| 148 | 153 | misc.LogCredentialSeparator() |
| 149 | 154 | return |
| @@ -152,7 +157,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ · 2 unmodified lines | ||
| 152 | 157 | now := time.Now() |
| 153 | 158 | |
| 154 | 159 | oldModels, hadExisting := r.clientModels[clientID] |
| 155 | − oldProvider, _ := r.clientProviders[clientID] | |
| 160 | + oldProvider := r.clientProviders[clientID] | |
| 156 | 161 | providerChanged := oldProvider != provider |
| 157 | 162 | if !hadExisting { |
| 158 | 163 | // Pure addition path. |
| @@ -161,6 +166,12 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ · 2 unmodified lines | ||
| 161 | 166 | r.addModelRegistration(modelID, provider, model, now) |
| 162 | 167 | } |
| 163 | 168 | r.clientModels[clientID] = append([]string(nil), rawModelIDs...) |
| 169 | + // Store client's own model infos | |
| 170 | + clientInfos := make(map[string]*ModelInfo, len(newModels)) | |
| 171 | + for id, m := range newModels { | |
| 172 | + clientInfos[id] = cloneModelInfo(m) | |
| 173 | + } | |
| 174 | + r.clientModelInfos[clientID] = clientInfos | |
| 164 | 175 | if provider != "" { |
| 165 | 176 | r.clientProviders[clientID] = provider |
| 166 | 177 | } else { |
| @@ -287,6 +298,12 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ · 120 unmodified lines | ||
| 287 | 298 | if len(rawModelIDs) > 0 { |
| 288 | 299 | r.clientModels[clientID] = append([]string(nil), rawModelIDs...) |
| 289 | 300 | } |
| 301 | + // Update client's own model infos | |
| 302 | + clientInfos := make(map[string]*ModelInfo, len(newModels)) | |
| 303 | + for id, m := range newModels { | |
| 304 | + clientInfos[id] = cloneModelInfo(m) | |
| 305 | + } | |
| 306 | + r.clientModelInfos[clientID] = clientInfos | |
| 290 | 307 | if provider != "" { |
| 291 | 308 | r.clientProviders[clientID] = provider |
| 292 | 309 | } else { |
| @@ -436,6 +453,7 @@ func (r *ModelRegistry) unregisterClientInternal(clientID string) { · 143 unmodified lines | ||
| 436 | 453 | } |
| 437 | 454 | |
| 438 | 455 | delete(r.clientModels, clientID) |
| 456 | + delete(r.clientModelInfos, clientID) | |
| 439 | 457 | if hasProvider { |
| 440 | 458 | delete(r.clientProviders, clientID) |
| 441 | 459 | } |
| @@ -887,6 +905,9 @@ func (r *ModelRegistry) GetModelsForClient(clientID string) []*ModelInfo { · 445 unmodified lines | ||
| 887 | 905 | return nil |
| 888 | 906 | } |
| 889 | 907 | |
| 908 | + // Try to use client-specific model infos first | |
| 909 | + clientInfos := r.clientModelInfos[clientID] | |
| 910 | + | |
| 890 | 911 | seen := make(map[string]struct{}) |
| 891 | 912 | result := make([]*ModelInfo, 0, len(modelIDs)) |
| 892 | 913 | for _, modelID := range modelIDs { |
| @@ -894,6 +915,15 @@ func (r *ModelRegistry) GetModelsForClient(clientID string) []*ModelInfo { · 1 unmodified line | ||
| 894 | 915 | continue |
| 895 | 916 | } |
| 896 | 917 | seen[modelID] = struct{}{} |
| 918 | + | |
| 919 | + // Prefer client's own model info to preserve original type/owned_by | |
| 920 | + if clientInfos != nil { | |
| 921 | + if info, ok := clientInfos[modelID]; ok && info != nil { | |
| 922 | + result = append(result, info) | |
| 923 | + continue | |
| 924 | + } | |
| 925 | + } | |
| 926 | + // Fallback to global registry (for backwards compatibility) | |
| 897 | 927 | if reg, ok := r.models[modelID]; ok && reg.Info != nil { |
| 898 | 928 | result = append(result, reg.Info) |
| 899 | 929 | } |