Merge pull request #536 from AoaoMH/feature/auth-model-check

Luis Pater committed 11 months agov6.6.12

Commit 5a75ef82 parents 71f788b 07279f81 file · +35 −5Browse files

Unified · Split

internal/registry/model_registry.go+35−5

 	models map[string]*ModelRegistration
 	// clientModels maps client ID to the models it provides
 	clientModels map[string][]string
+	// clientModelInfos maps client ID to a map of model ID -> ModelInfo
+	// This preserves the original model info provided by each client
+	clientModelInfos map[string]map[string]*ModelInfo
 	// clientProviders maps client ID to its provider identifier
 	clientProviders map[string]string
 	// mutex ensures thread-safe access to the registry
 func GetGlobalRegistry() *ModelRegistry {
 	registryOnce.Do(func() {
 		globalRegistry = &ModelRegistry{
-			models:          make(map[string]*ModelRegistration),
+			models:           make(map[string]*ModelRegistration),
-			clientModels:    make(map[string][]string),
+			clientModels:     make(map[string][]string),
-			clientProviders: make(map[string]string),
+			clientModelInfos: make(map[string]map[string]*ModelInfo),
-			mutex:           &sync.RWMutex{},
+			clientProviders:  make(map[string]string),
+			mutex:            &sync.RWMutex{},
 		}
 	})
 	return globalRegistry
 		// No models supplied; unregister existing client state if present.
 		r.unregisterClientInternal(clientID)
 		delete(r.clientModels, clientID)
+		delete(r.clientModelInfos, clientID)
 		delete(r.clientProviders, clientID)
 		misc.LogCredentialSeparator()
 		return
 	now := time.Now()
 	oldModels, hadExisting := r.clientModels[clientID]
-	oldProvider, _ := r.clientProviders[clientID]
+	oldProvider := r.clientProviders[clientID]
 	providerChanged := oldProvider != provider
 	if !hadExisting {
 		// Pure addition path.
 			r.addModelRegistration(modelID, provider, model, now)
 		}
 		r.clientModels[clientID] = append([]string(nil), rawModelIDs...)
+		// Store client's own model infos
+		clientInfos := make(map[string]*ModelInfo, len(newModels))
+		for id, m := range newModels {
+			clientInfos[id] = cloneModelInfo(m)
+		}
+		r.clientModelInfos[clientID] = clientInfos
 		if provider != "" {
 			r.clientProviders[clientID] = provider
 		} else {
 	if len(rawModelIDs) > 0 {
 		r.clientModels[clientID] = append([]string(nil), rawModelIDs...)
 	}
+	// Update client's own model infos
+	clientInfos := make(map[string]*ModelInfo, len(newModels))
+	for id, m := range newModels {
+		clientInfos[id] = cloneModelInfo(m)
+	}
+	r.clientModelInfos[clientID] = clientInfos
 	if provider != "" {
 		r.clientProviders[clientID] = provider
 	} else {
 	}
 	delete(r.clientModels, clientID)
+	delete(r.clientModelInfos, clientID)
 	if hasProvider {
 		delete(r.clientProviders, clientID)
 	}
 		return nil
 	}
+	// Try to use client-specific model infos first
+	clientInfos := r.clientModelInfos[clientID]
 	seen := make(map[string]struct{})
 	result := make([]*ModelInfo, 0, len(modelIDs))
 	for _, modelID := range modelIDs {
 			continue
 		}
 		seen[modelID] = struct{}{}
+		// Prefer client's own model info to preserve original type/owned_by
+		if clientInfos != nil {
+			if info, ok := clientInfos[modelID]; ok && info != nil {
+				result = append(result, info)
+				continue
+			}
+		}
+		// Fallback to global registry (for backwards compatibility)
 		if reg, ok := r.models[modelID]; ok && reg.Info != nil {
 			result = append(result, reg.Info)
 		}