localmodelroutex/backend/internal/catalog/cache.go

package catalog

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/google/uuid"

	"git.govekar.net/LovroG05/localmodelroutex/internal/store"
)

// CatalogStore is the read side the cache depends on.
type CatalogStore interface {
	ListCatalog(ctx context.Context) ([]store.CatalogRow, error)
	GetSetting(ctx context.Context, key string, v any) (bool, error)
}

// Cache is the in-memory model library, rebuilt from the store on each sync.
type Cache struct {
	store CatalogStore

	mu       sync.RWMutex
	byAddr   map[string]*Entry
	byID     map[uuid.UUID]*Entry
	entries  []*Entry
	strategy Strategy
	// standby holds enabled, unblocked models the provider does not serve
	// right now. They are not servable, but a request for one can still be
	// routed to its backup.
	standby map[string]*Entry

	rr sync.Map // map[string]*atomic.Uint64, round-robin counters per query
}

// NewCache returns an empty cache; call Reload to populate it.
func NewCache(s CatalogStore) *Cache {
	return &Cache{store: s, byAddr: map[string]*Entry{}, byID: map[uuid.UUID]*Entry{}, standby: map[string]*Entry{}, strategy: StrategyRoundRobin}
}

// maxFallbackHops bounds a backup chain. A chain may legitimately loop (two
// models that back each other up), so the walk also tracks visited ids.
const maxFallbackHops = 8

// Reload rebuilds the cache (servable entries only) and routing strategy from
// the store.
func (c *Cache) Reload(ctx context.Context) error {
	rows, err := c.store.ListCatalog(ctx)
	if err != nil {
		return err
	}

	byAddr := make(map[string]*Entry, len(rows))
	byID := make(map[uuid.UUID]*Entry, len(rows))
	standby := map[string]*Entry{}
	entries := make([]*Entry, 0, len(rows))
	for _, r := range rows {
		if !r.Enabled || r.Blocked {
			continue
		}
		var overrides map[string]any
		if len(r.ParamOverrides) > 0 {
			_ = json.Unmarshal(r.ParamOverrides, &overrides)
		}
		e := &Entry{
			ID:                 r.ID,
			CreatedAt:          r.CreatedAt,
			ProviderID:         r.ProviderID,
			ProviderName:       r.ProviderName,
			ProviderType:       r.ProviderType,
			Slug:               r.Slug,
			Address:            r.ProviderName + "/" + r.Slug,
			DisplayName:        r.DisplayName,
			Capabilities:       r.Capabilities,
			ContextLength:      r.ContextLength,
			ParamOverrides:     overrides,
			InputPricePerMTok:  r.InputPricePerMTok,
			OutputPricePerMTok: r.OutputPricePerMTok,
			CachedPricePerMTok: r.CachedPricePerMTok,
			RPMLimit:           r.RPMLimit,
			TPMLimit:           r.TPMLimit,
			FallbackID:         r.FallbackModelID,
		}
		if !r.Available {
			standby[e.Address] = e
			continue
		}
		byAddr[e.Address] = e
		byID[e.ID] = e
		entries = append(entries, e)
	}

	strategy := StrategyRoundRobin
	var s string
	if ok, _ := c.store.GetSetting(ctx, "routing.strategy", &s); ok && s != "" {
		strategy = Strategy(s)
	}

	c.mu.Lock()
	c.byAddr = byAddr
	c.byID = byID
	c.standby = standby
	c.entries = entries
	c.strategy = strategy
	c.mu.Unlock()
	return nil
}

// Resolve maps a model address to the servable entry a request would reach
// first: the model itself, or its first live backup when the model is down.
func (c *Cache) Resolve(address string) (*Entry, error) {
	r, err := c.Route(address)
	if err != nil {
		return nil, err
	}
	return r.Targets[0], nil
}

// Route maps a model address to the entries that can serve it, in order.
// "any/<slug>" selects among matching providers using the configured strategy;
// the selected entry's backups follow it. A model that is enabled but not
// available routes straight to its backups.
func (c *Cache) Route(address string) (*Route, error) {
	provider, slug, ok := SplitAddress(address)
	if !ok {
		return nil, fmt.Errorf("invalid model address %q (want <provider>/<slug>)", address)
	}

	c.mu.RLock()
	defer c.mu.RUnlock()

	if strings.EqualFold(provider, AnyProvider) {
		cands := c.candidatesLocked(slug)
		if len(cands) == 0 {
			return nil, fmt.Errorf("no provider serves model %q", slug)
		}
		e := c.selectLocked(slug, cands)
		return &Route{Address: address, Primary: e.Address, Slug: e.Slug, Targets: c.chainLocked(e)}, nil
	}

	if e, ok := c.byAddr[address]; ok {
		return &Route{Address: address, Primary: address, Slug: e.Slug, Targets: c.chainLocked(e)}, nil
	}
	if e, ok := c.standby[address]; ok {
		if targets := c.chainLocked(e); len(targets) > 0 {
			return &Route{Address: address, Primary: address, Slug: e.Slug, Targets: targets}, nil
		}
	}
	return nil, fmt.Errorf("model %q not found", address)
}

// chainLocked returns e, when it is servable, followed by its live backups.
// Backups that are not servable are skipped and the walk continues through
// them, so a down backup does not hide the ones behind it.
func (c *Cache) chainLocked(e *Entry) []*Entry {
	var out []*Entry
	seen := map[uuid.UUID]bool{}
	for hops := 0; e != nil && hops <= maxFallbackHops && !seen[e.ID]; hops++ {
		seen[e.ID] = true
		if c.byID[e.ID] == e {
			out = append(out, e)
		}
		e = c.lookupLocked(e.FallbackID)
	}
	return out
}

// lookupLocked finds an enabled, unblocked entry by id, servable or standby.
func (c *Cache) lookupLocked(id *uuid.UUID) *Entry {
	if id == nil {
		return nil
	}
	if e, ok := c.byID[*id]; ok {
		return e
	}
	for _, e := range c.standby {
		if e.ID == *id {
			return e
		}
	}
	return nil
}

// candidatesLocked returns servable entries whose slug equals slug or ends with
// "/slug" (so "any/gpt-4o" also matches OpenRouter's "openai/gpt-4o").
func (c *Cache) candidatesLocked(slug string) []*Entry {
	suffix := "/" + slug
	var out []*Entry
	for _, e := range c.entries {
		if e.Slug == slug || strings.HasSuffix(e.Slug, suffix) {
			out = append(out, e)
		}
	}
	return out
}

func (c *Cache) selectLocked(key string, cands []*Entry) *Entry {
	if len(cands) == 1 {
		return cands[0]
	}
	switch c.strategy {
	case StrategyLowestCost:
		best := cands[0]
		for _, e := range cands[1:] {
			if cost(e) < cost(best) {
				best = e
			}
		}
		return best
	default: // round_robin; priority has no meaning for `any`, so it falls back here
		ctr, _ := c.rr.LoadOrStore(key, new(atomic.Uint64))
		n := ctr.(*atomic.Uint64).Add(1) - 1
		return cands[n%uint64(len(cands))]
	}
}

func cost(e *Entry) float64 { return e.InputPricePerMTok + e.OutputPricePerMTok }

// List returns a snapshot of the entries clients can call: the servable ones,
// plus unavailable models that have a live backup.
func (c *Cache) List() []*Entry {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]*Entry, len(c.entries), len(c.entries)+len(c.standby))
	copy(out, c.entries)
	for _, e := range c.standby {
		if len(c.chainLocked(e)) > 0 {
			out = append(out, e)
		}
	}
	return out
}

// Strategy returns the currently configured selection strategy.
func (c *Cache) Strategy() Strategy {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.strategy
}

// SplitAddress splits "<provider>/<slug>" on the first slash. The slug may
// itself contain slashes (e.g. OpenRouter's "openai/gpt-4o").
func SplitAddress(address string) (provider, slug string, ok bool) {
	i := strings.Index(address, "/")
	if i <= 0 || i == len(address)-1 {
		return "", "", false
	}
	return address[:i], address[i+1:], true
}