fix(phase-2): add thread-safety to ModelRegistry and Router reload operations

This commit is contained in:
newkirk
2026-07-21 13:20:29 -04:00
parent 42b70621a1
commit 4027ed4351
3 changed files with 29 additions and 4 deletions
+11 -1
View File
@@ -1,6 +1,9 @@
package models package models
import "strings" import (
"strings"
"sync"
)
// CanonicalProviders lists the original model creators in priority order. // CanonicalProviders lists the original model creators in priority order.
// When a model name exists in multiple providers (e.g. deepseek-v4-pro in // When a model name exists in multiple providers (e.g. deepseek-v4-pro in
@@ -22,6 +25,7 @@ var CanonicalProviders = []string{
} }
type ModelRegistry struct { type ModelRegistry struct {
mu sync.RWMutex
Providers map[string]ProviderInfo `json:"-"` Providers map[string]ProviderInfo `json:"-"`
} }
@@ -173,6 +177,12 @@ func (r *ModelRegistry) findAllForwardFuzzy(modelID string) (*ModelMetadata, boo
// etc.) from overriding the original provider's authoritative pricing and // etc.) from overriding the original provider's authoritative pricing and
// limits. // limits.
func (r *ModelRegistry) FindModel(modelID string) *ModelMetadata { func (r *ModelRegistry) FindModel(modelID string) *ModelMetadata {
if r == nil {
return nil
}
r.mu.RLock()
defer r.mu.RUnlock()
// 1. Exact key match — canonical first, then all // 1. Exact key match — canonical first, then all
if m, ok := r.findInCanonical(modelID); ok { if m, ok := r.findInCanonical(modelID); ok {
return m return m
+14 -2
View File
@@ -5,6 +5,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"strings" "strings"
"sync"
"gophergate/internal/db" "gophergate/internal/db"
) )
@@ -32,6 +33,7 @@ type ClassifierFunc func(ctx context.Context, selectorModel, systemPrompt, userM
// Router resolves model groups to concrete models. // Router resolves model groups to concrete models.
type Router struct { type Router struct {
mu sync.RWMutex
groups map[string]db.ModelGroup groups map[string]db.ModelGroup
classify ClassifierFunc classify ClassifierFunc
} }
@@ -50,6 +52,8 @@ func New(groups []db.ModelGroup, classify ClassifierFunc) *Router {
// Groups returns all registered model group IDs. // Groups returns all registered model group IDs.
func (r *Router) Groups() []string { func (r *Router) Groups() []string {
r.mu.RLock()
defer r.mu.RUnlock()
ids := make([]string, 0, len(r.groups)) ids := make([]string, 0, len(r.groups))
for id := range r.groups { for id := range r.groups {
ids = append(ids, id) ids = append(ids, id)
@@ -59,13 +63,17 @@ func (r *Router) Groups() []string {
// IsGroup returns true if the model name is a group ID. // IsGroup returns true if the model name is a group ID.
func (r *Router) IsGroup(modelID string) bool { func (r *Router) IsGroup(modelID string) bool {
r.mu.RLock()
defer r.mu.RUnlock()
_, ok := r.groups[modelID] _, ok := r.groups[modelID]
return ok return ok
} }
// Route resolves a group to a concrete model. // Route resolves a group to a concrete model.
func (r *Router) Route(ctx context.Context, groupID string, routeCtx *RouteContext) (*Decision, error) { func (r *Router) Route(ctx context.Context, groupID string, routeCtx *RouteContext) (*Decision, error) {
r.mu.RLock()
group, ok := r.groups[groupID] group, ok := r.groups[groupID]
r.mu.RUnlock()
if !ok { if !ok {
return nil, fmt.Errorf("unknown model group: %s", groupID) return nil, fmt.Errorf("unknown model group: %s", groupID)
} }
@@ -133,8 +141,12 @@ func (r *Router) RouteToConcrete(ctx context.Context, modelID string, routeCtx *
// Reload replaces the group definitions without recreating the router. // Reload replaces the group definitions without recreating the router.
func (r *Router) Reload(groups []db.ModelGroup) { func (r *Router) Reload(groups []db.ModelGroup) {
r.groups = make(map[string]db.ModelGroup) newGroups := make(map[string]db.ModelGroup)
for _, g := range groups { for _, g := range groups {
r.groups[g.ID] = g newGroups[g.ID] = g
} }
r.mu.Lock()
r.groups = newGroups
r.mu.Unlock()
} }
+4 -1
View File
@@ -628,7 +628,10 @@ func (s *Server) handleChatCompletions(c *gin.Context) {
// Inject or cap max_tokens from model registry. // Inject or cap max_tokens from model registry.
s.registryMu.RLock() s.registryMu.RLock()
meta := s.registry.FindModel(modelID) var meta *models.ModelMetadata
if s.registry != nil {
meta = s.registry.FindModel(modelID)
}
s.registryMu.RUnlock() s.registryMu.RUnlock()
if meta != nil && meta.Limit != nil && meta.Limit.Output > 0 { if meta != nil && meta.Limit != nil && meta.Limit.Output > 0 {