refactor: updates modelcatalog to memoise provider configs recomputing for allowed models / providers

framework/modelcatalog/datasheet/store.go

@@ -4,6 +4,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	bifrost "github.com/maximhq/bifrost/core"
@@ -67,6 +68,10 @@ type Store struct {
 	datasheetByProvider      map[schemas.ModelProvider][]string // rebuilt every reload
 	deprecatedByProvider     map[schemas.ModelProvider][]string // rebuilt every reload
 
+	// writeGen counts membership rebuilds; the composer's model→provider memo
+	// stamps it to detect staleness. Atomic so readers skip mu.
+	writeGen atomic.Uint64
+
 	// Overrides under their own mutex: writes here don't block pricing reads
 	// (the hot CalculateCost path takes mu.RLock and overridesMu.RLock
 	// independently and the orderings never invert).
@@ -500,4 +505,11 @@ func (s *Store) rebuildDatasheetViewUnsafe() {
 		slices.Sort(models)
 		s.deprecatedByProvider[provider] = models
 	}
+
+	s.writeGen.Add(1)
+}
+
+// WriteGen returns the monotonic count of membership rebuilds.
+func (s *Store) WriteGen() uint64 {
+	return s.writeGen.Load()
 }

framework/modelcatalog/keyconfig/store.go

@@ -18,6 +18,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	bifrost "github.com/maximhq/bifrost/core"
 	"github.com/maximhq/bifrost/core/schemas"
@@ -60,6 +61,10 @@ type Store struct {
 	mu      sync.RWMutex
 	entries map[schemas.ModelProvider]*providerState
 	logger  schemas.Logger
+
+	// writeGen counts entry mutations; the composer's model→provider memo
+	// stamps it to detect staleness. Atomic so readers skip mu.
+	writeGen atomic.Uint64
 }
 
 // New constructs an empty Store.
@@ -88,6 +93,12 @@ func (s *Store) Replace(snapshot map[schemas.ModelProvider][]schemas.Key) {
 		}
 	}
 	s.entries = next
+	s.writeGen.Add(1)
+}
+
+// WriteGen returns the monotonic count of entry mutations.
+func (s *Store) WriteGen() uint64 {
+	return s.writeGen.Load()
 }
 
 // SetProvider replaces the cached state for one provider. Call after a
@@ -98,16 +109,19 @@ func (s *Store) SetProvider(provider schemas.ModelProvider, keys []schemas.Key)
 	defer s.mu.Unlock()
 	if st == nil {
 		delete(s.entries, provider)
+		s.writeGen.Add(1)
 		return
 	}
 	s.entries[provider] = st
+	s.writeGen.Add(1)
 }
 
 // RemoveProvider drops all state for the provider. Call on provider delete.
 func (s *Store) RemoveProvider(provider schemas.ModelProvider) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	delete(s.entries, provider)
+	s.writeGen.Add(1)
 }
 
 // EntriesFor returns all per-key entries for the provider (including

framework/modelcatalog/live/store.go

@@ -18,6 +18,7 @@ package live
 import (
 	"slices"
 	"sync"
+	"sync/atomic"
 
 	bifrost "github.com/maximhq/bifrost/core"
 	"github.com/maximhq/bifrost/core/schemas"
@@ -50,6 +51,11 @@ type Store struct {
 	// look current again after the provider was re-added.
 	gen    map[schemas.ModelProvider]uint64
 	logger schemas.Logger
+
+	// writeGen counts entry mutations across all providers; the composer's
+	// model→provider memo stamps it to detect staleness. Distinct from gen,
+	// which is per-provider and guards in-flight fetch races.
+	writeGen atomic.Uint64
 }
 
 func New(logger schemas.Logger) *Store {
@@ -73,6 +79,12 @@ func (s *Store) Upsert(provider schemas.ModelProvider, keyID string, unfiltered
 	s.mu.Lock()
 	s.entries[k] = Entry{Models: cp}
 	s.mu.Unlock()
+	s.writeGen.Add(1)
+}
+
+// WriteGen returns the monotonic count of entry mutations across all providers.
+func (s *Store) WriteGen() uint64 {
+	return s.writeGen.Load()
 }
 
 // Generation returns the provider's current invalidation counter. Read it
@@ -102,6 +114,7 @@ func (s *Store) UpsertIfCurrent(provider schemas.ModelProvider, keyID string, un
 		return false
 	}
 	s.entries[k] = Entry{Models: cp}
+	s.writeGen.Add(1)
 	return true
 }
 
@@ -113,6 +126,7 @@ func (s *Store) UpsertIfCurrent(provider schemas.ModelProvider, keyID string, un
 // nothing cached yet, and that is the write this has to stop.
 func (s *Store) bumpLocked(provider schemas.ModelProvider) {
 	s.gen[provider]++
+	s.writeGen.Add(1)
 }
 
 // Invalidate drops both filtered and unfiltered entries for one key. Called

framework/modelcatalog/main.go

@@ -40,6 +40,12 @@ type ModelCatalog struct {
 	// loadCapabilities fills a capability cache miss.
 	loadCapabilities func(schemas.ModelProvider, string) (*schemas.ModelCapabilities, error)
 
+	// providerMemo caches GetProvidersForModel per model, stamped with
+	// catalogGeneration() at compute time; any store write invalidates every
+	// entry. Capped at providerMemoMaxEntries, flushed on overflow.
+	providerMemoMu sync.RWMutex
+	providerMemo   map[string]providerMemoEntry
+
 	// MCP library sync configuration (protected by syncMu)
 	mcpLibraryURL          string
 	mcpLibrarySyncInterval time.Duration
@@ -114,6 +120,7 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto
 		live:         live.New(logger),
 		keyconf:      keyconfig.New(logger),
 		capabilities: lrucache.New[*schemas.ModelCapabilities](capabilityCacheSize),
+		providerMemo: make(map[string]providerMemoEntry),
 		done:         make(chan struct{}),
 	}
 	mc.syncCtx, mc.syncCancel = context.WithCancel(ctx)

framework/modelcatalog/models.go

@@ -150,11 +150,63 @@ func (mc *ModelCatalog) GetDistinctBaseModelNames() []string {
 	return mc.datasheet.DistinctBaseModelNames()
 }
 
+// providerMemoEntry is a memoized GetProvidersForModel result and the
+// catalogGeneration it was computed under.
+type providerMemoEntry struct {
+	gen uint64
+	val []schemas.ModelProvider
+}
+
+// providerMemoMaxEntries bounds the memo; model strings are client-controlled
+// and would otherwise grow it without limit. On overflow the map is flushed
+// and hot entries repopulate on demand.
+const providerMemoMaxEntries = 4096
+
+// catalogGeneration sums the three stores' write counters. Every write bumps
+// exactly one counter by one, so the sum strictly increases and no two
+// catalog states share a value.
+func (mc *ModelCatalog) catalogGeneration() uint64 {
+	return mc.datasheet.WriteGen() + mc.keyconf.WriteGen() + mc.live.WriteGen()
+}
+
 // GetProvidersForModel returns every provider that can serve the model.
-// Composes across stores and applies the cross-provider special cases
-// (openrouter / vertex / groq-gpt / bedrock-claude) preserved verbatim from
-// the pre-refactor implementation.
+// Memoized per model (bounded, empty results uncached); any store write
+// invalidates (see catalogGeneration). Returns a fresh clone the caller
+// may mutate.
 func (mc *ModelCatalog) GetProvidersForModel(model string) []schemas.ModelProvider {
+	gen := mc.catalogGeneration()
+
+	mc.providerMemoMu.RLock()
+	entry, ok := mc.providerMemo[model]
+	mc.providerMemoMu.RUnlock()
+	if ok && entry.gen == gen {
+		return slices.Clone(entry.val)
+	}
+
+	val := mc.computeProvidersForModel(model)
+	// Unknown models resolve to nothing; skipping memoising
+	if len(val) == 0 {
+		return val
+	}
+
+	// Stamp with the generation read before compute: a write during compute
+	// leaves the entry stale and the next call recomputes.
+	mc.providerMemoMu.Lock()
+	if mc.providerMemo == nil {
+		mc.providerMemo = make(map[string]providerMemoEntry)
+	}
+	if _, exists := mc.providerMemo[model]; !exists && len(mc.providerMemo) >= providerMemoMaxEntries {
+		mc.providerMemo = make(map[string]providerMemoEntry)
+	}
+	mc.providerMemo[model] = providerMemoEntry{gen: gen, val: val}
+	mc.providerMemoMu.Unlock()
+
+	return slices.Clone(val)
+}
+
+// computeProvidersForModel composes the uncached answer across stores,
+// including the openrouter / vertex / groq-gpt / bedrock-claude special cases.
+func (mc *ModelCatalog) computeProvidersForModel(model string) []schemas.ModelProvider {
 	baseModel := mc.datasheet.BaseModelName(model)
 
 	providers := make([]schemas.ModelProvider, 0)
@@ -265,17 +317,25 @@ func (mc *ModelCatalog) IsModelAllowedForProvider(provider schemas.ModelProvider
 		return false
 	}
 
+	// Bare-name match needs no catalog access and covers most allowlists.
+	if slices.Contains(allowedModels, model) {
+		return true
+	}
+
+	// Only provider-prefixed entries ("openai/gpt-4o") need the provider
+	// catalog; build it once, and only when one exists.
+	if !slices.ContainsFunc(allowedModels, func(m string) bool { return strings.Contains(m, "/") }) {
+		return false
+	}
 	providerCatalogModels := mc.GetModelsForProvider(provider)
 	for _, allowedModel := range allowedModels {
-		if allowedModel == model {
-			return true
+		if !strings.Contains(allowedModel, "/") {
+			continue
 		}
-		if strings.Contains(allowedModel, "/") {
-			if slices.Contains(providerCatalogModels, allowedModel) {
-				_, modelPart := schemas.ParseModelString(allowedModel, "")
-				if modelPart == model {
-					return true
-				}
+		if slices.Contains(providerCatalogModels, allowedModel) {
+			_, modelPart := schemas.ParseModelString(allowedModel, "")
+			if modelPart == model {
+				return true
 			}
 		}
 	}