ModelRouter/internal/provider/registry.go

// Package provider also provides the top-level registry that owns all sources,
// routes model requests (explicit or AUTO), and supports hot reload.
package provider

import (
	"context"
	"sort"
	"strings"
	"sync"
	"time"
)

// Registry holds all configured providers and routes model requests.
type Registry struct {
	mu        sync.RWMutex
	providers []*Provider
	byModel   map[string][]*Provider // modelID -> providers (one per source)
	defaultM  string                 // default model id ("" means AUTO)
}

func NewRegistry(providers []*Provider, defaultModel string) *Registry {
	r := &Registry{byModel: map[string][]*Provider{}, defaultM: defaultModel}
	r.set(providers)
	return r
}

// Replace atomically swaps the provider set (hot reload).
func (r *Registry) Replace(providers []*Provider) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.set(providers)
}

func (r *Registry) set(providers []*Provider) {
	r.providers = providers
	m := map[string][]*Provider{}
	for _, p := range providers {
		for _, model := range p.Models() {
			key := strings.ToLower(model)
			m[key] = append(m[key], p)
		}
	}
	r.byModel = m
}

func (r *Registry) Providers() []*Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]*Provider, len(r.providers))
	copy(out, r.providers)
	return out
}

func (r *Registry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.providers)
}

// ModelList returns all exposed model ids (chat + image).
func (r *Registry) ModelList() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	seen := map[string]bool{}
	for _, p := range r.providers {
		for _, m := range p.Models() {
			if !seen[m] {
				seen[m] = true
			}
		}
	}
	out := make([]string, 0, len(seen))
	for m := range seen {
		out = append(out, m)
	}
	sort.Strings(out)
	return out
}

// Resolve returns the provider (or providers) serving a requested model,
// owning no AUTO scheduling logic anymore: AUTO chat scheduling is driven by
// the scheduler chain built from the runtime rules (see core/SaveAutoRules
// and scheduler.Chain).
//
// model "" or "AUTO" -> every provider in configured order. Used only by the
// image path (which then filters to image-capable sources) and tool-call
// anchoring; chat AUTO requests go through the chain instead.
// Otherwise the owning provider; "source-model"/"source:model"/"source/model"
// pinning resolves first; an unknown model resolves to nil (gateway answers
// 404) instead of silently falling back to the AUTO chain.
func (r *Registry) Resolve(model string) []*Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()

	model = strings.TrimSpace(model)
	if model == "" || strings.EqualFold(model, "AUTO") {
		out := make([]*Provider, len(r.providers))
		copy(out, r.providers)
		return out
	}

	// "source-model" / "source:model" / "source/model" pinning — disambiguates
	// duplicate model ids across sources.
	if p := r.ResolvePinned(model); p != nil {
		return []*Provider{p}
	}
	// explicit model — may be served by multiple sources
	if ps, ok := r.byModel[strings.ToLower(model)]; ok {
		out := make([]*Provider, len(ps))
		copy(out, ps)
		return out
	}
	return nil
}

// EffectiveModel strips a "source-model" / "source:model" / "source/model"
// pinning prefix and returns the bare model id when that source serves it;
// otherwise it returns the input unchanged.
func (r *Registry) EffectiveModel(model string) string {
	sep := strings.IndexAny(model, "-:/")
	if sep < 1 || sep == len(model)-1 {
		return model
	}
	src, m := model[:sep], model[sep+1:]
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, p := range r.providers {
		if strings.EqualFold(p.Name(), src) && p.ModelByID(m) != nil {
			return m
		}
	}
	return model
}

// ResolvePinned resolves "source-model" / "source:model" / "source/model" to
// the exact source, or nil if the source does not serve that model.
func (r *Registry) ResolvePinned(model string) *Provider {
	sep := strings.IndexAny(model, "-:/")
	if sep < 1 || sep == len(model)-1 {
		return nil
	}
	src, m := model[:sep], model[sep+1:]
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, p := range r.providers {
		if !strings.EqualFold(p.Name(), src) {
			continue
		}
		if p.ModelByID(m) != nil {
			return p
		}
	}
	return nil
}

// ProviderForModel returns the provider owning the model id (nil if unknown).
// When the same model exists on multiple sources, returns the first (source
// config order). Callers that need a specific source use ProviderForSlot.
func (r *Registry) ProviderForModel(model string) *Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	ps, ok := r.byModel[strings.ToLower(strings.TrimSpace(model))]
	if !ok || len(ps) == 0 {
		return nil
	}
	return ps[0]
}

// ProviderForSlot returns the provider for a (model, source) slot. When
// source is empty it behaves like ProviderForModel (first owner of the model id);
// when source is set it returns only that exact source (nil if the source
// does not serve the model).
func (r *Registry) ProviderForSlot(model, source string) *Provider {
	model = strings.ToLower(strings.TrimSpace(model))
	source = strings.TrimSpace(source)
	r.mu.RLock()
	defer r.mu.RUnlock()
	if source == "" {
		ps, ok := r.byModel[model]
		if !ok || len(ps) == 0 {
			return nil
		}
		return ps[0]
	}
	for _, p := range r.providers {
		if !strings.EqualFold(p.Name(), source) {
			continue
		}
		for _, m := range p.cfg.Models {
			if strings.EqualFold(m.ID, model) {
				return p
			}
		}
		return nil
	}
	return nil
}

// ModelKindInfo carries the per-model fields the chat test UI needs to tell a
// chat model from an image one.
type ModelKindInfo struct {
	ID   string `json:"id"`
	Kind string `json:"kind"` // "chat" (default) | "image"
}

// ModelStatus is a web-UI friendly snapshot per source.
type SourceStatus struct {
	Name          string   `json:"name"`
	Adapter       string   `json:"adapter"`
	BaseURL       string   `json:"base_url"`
	Models        []string `json:"models"`
	// ModelKinds parallels Models: Models[i] is ModelKinds[i].ID.
	//
	// It exists because a flat []string cannot answer "is gpt-image-2
	// reachable from /api/chat?" — it cannot, and offering it there produced
	// a guaranteed failure with no explanation. Models stays unchanged for
	// callers that only need names; this is purely additive and omitted when
	// empty, so an older UI keeps working.
	ModelKinds    []ModelKindInfo `json:"model_kinds,omitempty"`
	Available     bool     `json:"available"`
	Healthy       bool     `json:"healthy"`
	MaxConcurrent int      `json:"max_concurrent"`
	LiveAvailable bool     `json:"live_available"`
	LastError     string   `json:"last_error,omitempty"`
	LastChecked   int64    `json:"last_checked,omitempty"`
	FailCount     int      `json:"fail_count,omitempty"`
	BackoffUntil  int64    `json:"backoff_until,omitempty"`
	Permanent     bool     `json:"permanent,omitempty"`
	// RecentOK / RecentErr count real gateway requests served by this source
	// within the last 300s (drives the status column so a source that is
	// actually serving traffic can never be shown as down just because a
	// probe was rate-limited by the upstream).
	RecentOK  int64 `json:"recent_ok,omitempty"`
	RecentErr int64 `json:"recent_err,omitempty"`
	// AvgFirstByteMs is the mean time-to-first-byte (ms) for recently
	// successful requests served by this source (0 when unmeasured).
	AvgFirstByteMs int64 `json:"avg_first_byte_ms,omitempty"`
	// AvgTokPerS is the aggregate completion throughput (tokens/s) for
	// recently successful requests served by this source (0 when no data).
	AvgTokPerS int64 `json:"avg_tok_per_s,omitempty"`
}

// ProbeAll runs a live reachability check for every provider (in parallel).
func (r *Registry) ProbeAll(ctx context.Context) {
	r.mu.RLock()
	providers := append([]*Provider(nil), r.providers...)
	r.mu.RUnlock()
	var wg sync.WaitGroup
	for _, p := range providers {
		wg.Add(1)
		go func(p *Provider) {
			defer wg.Done()
			probeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
			defer cancel()
			p.Probe(probeCtx)
		}(p)
	}
	wg.Wait()
}

func (r *Registry) Status() []SourceStatus {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]SourceStatus, 0, len(r.providers))
	for _, p := range r.providers {
		live, lastErr, lastAt := p.LastProbe()
		fails, until, perm := p.HealthInfo()
		backoffUntil := int64(0)
		if !until.IsZero() {
			backoffUntil = until.Unix()
		}
		s := SourceStatus{
			Name:          p.Name(),
			Adapter:       p.Adapter(),
			BaseURL:       p.Config().BaseURL,
			Models:        p.Models(),
			ModelKinds:    p.ModelsWithKind(),
			Available:     p.Available(),
			Healthy:       p.Available(),
			MaxConcurrent: p.MaxConcurrent(),
			LiveAvailable: live,
			LastError:     lastErr,
			LastChecked:   lastAt,
			FailCount:     fails,
			BackoffUntil:  backoffUntil,
			Permanent:     perm,
		}
		out = append(out, s)
	}
	return out
}