llmgw/internal/router/router.go

// Package router wires the pieces into a working OpenAI-compatible gateway:
// resolve the model, pick an endpoint, proxy the stream, retry when it is safe.
//
// The sub-packages hold the reusable parts (scheduler, endpoint, scrape, proxy,
// breaker); this package holds the orchestration and the HTTP surface.
package router

import (
	"context"
	"errors"
	"fmt"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mail"
	"log/slog"
	"math"
	"net"
	"net/http"
	"net/netip"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"

	"gitea.mynodes.xyz/ilker/llmgw/internal/account"
	"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/web"
)

// Router serves the OpenAI API in front of a set of engine replicas.
type Router struct {
	cfg config.Router
	log *slog.Logger

	reg    *endpoint.Registry
	poller *scrape.Poller
	px     *proxy.Proxy
	// comfy asks ComfyUI replicas for pictures.
	comfy *comfy.Client
	m     *metrics
	// adm is nil when admission control is disabled.
	adm *admission.Controller

	profiles map[string]*scheduler.Profile
	// models maps every name a client may use — canonical names and aliases —
	// onto its catalog entry.
	//
	// Held atomically because the catalog is no longer fixed at startup: adding
	// a model used to mean restarting the control plane, and this is the
	// router's half of not having to. A request takes the map once and works
	// from it, so a swap mid-request is invisible rather than a race.
	models atomic.Pointer[map[string]catalog.Model]
	// adapters are the LoRA adapters by the name clients ask for them by.
	adapters atomic.Pointer[map[string]adapterRef]
	// canaryTargets maps each model that is another's canary to that one;
	// canaries counts how the two sides served.
	canaryTargets atomic.Pointer[map[string]string]
	canaries      canaryStats
	// slo keeps each model's recent token timings, for admission.
	slo sloTracker
	// batchWake has the batch worker look for work at once.
	batchWake chan struct{}
	// reqLogs is each project's keeping of requests, as last read.
	reqLogs requestLogs
	// answers are the answers projects reuse; nil reuses none.
	answers *answerCache

	// extraMetrics are collectors other components publish on /metrics.
	extraMetrics []prometheus.Collector

	// activator loads a model that is configured but not currently running.
	activator Activator

	// metricsHandler is built with the admin registry and reused by any
	// listener that should serve metrics without serving pprof.
	metricsHandler http.Handler

	// landing configures the public page, which the router serves but does not
	// otherwise care about.
	landing config.Landing

	// auth, store, acct, recorder and accounts are nil when authentication is
	// disabled.
	auth     *auth.Authenticator
	accounts *account.Service
	store    *storepkg.DB
	// usageRetention is how long request records are kept; zero is for ever.
	usageRetention time.Duration
	// isLeader says whether this process leads; nil is always.
	isLeader atomic.Pointer[func() bool]
	// federation is which cluster this is, to its peers; nil until said.
	federation atomic.Pointer[federationID]
	acct       *accounting.Accountant
	// guard judges requests and answers; nil judges nothing. guardLibrary
	// is the policies it was built from, for the pages that show them.
	guard        guard.Guard
	guardLibrary func() guard.Config
	// oidc is sign-in through an identity provider, and scim provisioning
	// by it.
	oidc config.OIDC
	// mail sends what the portal mails: a password reset link.
	mail mail.Config
	// stopped is the operators having stopped the service (llmgw stop).
	stopped atomic.Bool
	scim    config.SCIM

	// threads holds conversations and memory notes. It is store when auth is on,
	// its own file when only the threads block is configured, and nil when
	// nothing is kept. Separate from store because a non-nil store also switches
	// on the accounts, web pages and usage recording, none of which belong to a
	// gateway that only wants to remember things.
	threads *storepkg.DB
	// anonymous means threads has no project to scope to: everything shares the
	// one pool named by anonymousProject.
	anonymous bool
	// personalOff refuses personal use: the chat page then only offers the
	// projects a person is a member of.
	personalOff bool
	// personalDefaults is what personal use is held to when nobody wrote a
	// limit for the person, for the quota page to say so.
	personalDefaults *storepkg.Limits
	// llamaVersion is the llama.cpp build the nodes install, so the mirror page
	// can refuse weights that build cannot load before fetching them.
	llamaVersion string
	// ownsThreads is true when threads is not store and so must be closed here.
	ownsThreads bool
	recorder    *accounting.Recorder

	// prefix scores by prefix: the precise index for models with kv_events,
	// the approximate tree when router.prefix is on. Its tree is nil when it
	// is off.
	prefix *scheduler.PrefixScorer
	// lora knows which replicas hold which adapters that load on demand.
	lora *loraTracker
	// kv is the precise index (plan §7.5 tier 2), fed by the control plane
	// from the agents' KV cache events; kvClient makes its tokenize calls.
	kv       *kvindex.Index
	kvClient *http.Client

	// clusterView is nil until the control plane hands it over, and stays nil
	// in a router-only deployment.
	clusterView  web.Source
	clusterUsage web.UsageSource

	// mirror is nil unless a model mirror is configured, which leaves
	// /admin/mirror unmounted — the right page for a deployment with no mirror
	// is no page, not an empty one pointing at nothing.
	mirror web.MirrorAdmin
	// registry browses a model registry from the mirror page. Nil where the
	// deployment has no way out to the internet.
	registry web.Registry
	// catalogAdmin puts a mirrored model into service. Nil without a store,
	// where a model added at runtime would not survive a restart.
	catalogAdmin web.CatalogAdmin
	// history is the metrics store the dashboard charts read from. Nil where
	// none is configured, which leaves the charts off the page.
	history    web.History
	histWindow time.Duration

	// kbDefCache remembers each caller's default knowledge-base
	// collections for kbDefaultTTL.
	kbDefMu    sync.Mutex
	kbDefCache map[string]kbDefaultEntry

	// settingsCheck and settingsLocal mount the portal's settings page.
	settingsCheck func([]byte) error
	settingsLocal []string
}

// SetClusterView gives the router the data the /admin page renders.
//
// Set by the server command once the control plane exists, because the router
// is built first. Nil leaves /admin unmounted, which is what a router-only
// deployment should have: no cluster to show.
func (r *Router) SetClusterView(src web.Source, usage web.UsageSource) {
	r.clusterView, r.clusterUsage = src, usage
}

// SetMirror gives the router a model mirror to administer.
//
// Set by the server command when one is configured, for the same reason as the
// cluster view: the router is built before the things its admin pages show.
func (r *Router) SetMirror(m web.MirrorAdmin) { r.mirror = m }

// SetGuard has every request and answer judged by g.
func (r *Router) SetGuard(g guard.Guard) { r.guard = g }

// SetGuardLibrary is the policies the guard was built from, for the pages.
func (r *Router) SetGuardLibrary(c func() guard.Config) { r.guardLibrary = c }

// SetSettingsCheck enables the portal's settings page: check says whether a
// document would load as the cluster's settings beside this machine's own
// server.yaml, and local are the keys that file keeps.
func (r *Router) SetSettingsCheck(check func([]byte) error, local []string) {
	r.settingsCheck, r.settingsLocal = check, local
}

// SetStopped refuses, or again serves, every request for a model.
func (r *Router) SetStopped(stopped bool) { r.stopped.Store(stopped) }

// SetRegistry gives the mirror page somewhere to look models up.
func (r *Router) SetRegistry(reg web.Registry) { r.registry = reg }

// SetCatalogAdmin lets the mirror page put a model into service.
func (r *Router) SetCatalogAdmin(c web.CatalogAdmin) { r.catalogAdmin = c }

// SetHistory gives the dashboard somewhere to read the cluster's past.
func (r *Router) SetHistory(h web.History, window time.Duration) {
	r.history, r.histWindow = h, window
}

// Auth exposes the authenticator so ctl and the batch writer can reach it.
// Nil when authentication is off.
func (r *Router) Auth() *auth.Authenticator { return r.auth }

// Store exposes the key store, or nil when authentication is off.
func (r *Router) Store() *storepkg.DB { return r.store }

// New builds a router from configuration. Start must be called to begin scraping.
func New(cfg config.Server, log *slog.Logger) (*Router, error) {
	if log == nil {
		log = slog.Default()
	}

	reg := endpoint.NewRegistry(endpoint.Config{
		StaleAfter: cfg.Router.StaleAfter,
		Breaker: breaker.Config{
			FailureThreshold: cfg.Router.FailureThreshold,
			BaseCooldown:     cfg.Router.Cooldown,
		},
	})
	reg.Replace(toInfos(cfg.Endpoints))

	sc := scheduler.ProfileConfig{
		TargetWaiting:     cfg.Router.Scoring.TargetWaiting,
		TargetInflight:    cfg.Router.Scoring.TargetInflight,
		KVCeiling:         cfg.Router.Scoring.KVCeiling,
		SessionLoadFactor: cfg.Router.Scoring.SessionLoadFactor,
	}
	// The metrics come first because the prefix scorer reports into them, and
	// the decision metric is the whole difference between tuning this and
	// guessing at it.
	metrics := newMetrics()
	metrics.tenant.level = cfg.Router.TenantLabel
	metrics.tenant.setTagLabels(cfg.Router.TagLabels)
	// The scorer exists whether or not the tree does: which models use the
	// precise index is the catalog's to say, and the catalog changes while
	// the router runs. With neither in use it scores nothing and records
	// nothing.
	var tree *prefix.Tree
	if cfg.Router.Prefix.Enabled {
		tree = prefix.New(prefix.Config{
			MaxNodes: cfg.Router.Prefix.MaxNodes,
			MinMatch: cfg.Router.Prefix.MinMatch,
			TTL:      cfg.Router.Prefix.TTL,
		})
	}
	minBlocks := cfg.Router.KVIndex.MinBlocks
	if minBlocks <= 0 {
		minBlocks = 4
	}
	prefixScorer := scheduler.NewPrefixScorer(tree,
		prefix.Thresholds{
			MatchRatio: cfg.Router.Prefix.MatchRatio,
			BalanceAbs: cfg.Router.Prefix.BalanceAbs,
			BalanceRel: cfg.Router.Prefix.BalanceRel,
		},
		metrics.observePrefix).WithKV(minBlocks, metrics.observeKV)
	sc.Prefix = prefixScorer

	profiles := map[string]*scheduler.Profile{}
	for _, name := range []string{"chat", "batch", "long-context", "debug"} {
		profiles[name] = scheduler.Builtin(name, sc)
	}
	if _, ok := profiles[cfg.Router.DefaultProfile]; !ok {
		return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile)
	}

	// The portal's language for whoever neither chose one nor asked for one
	// it speaks. Checked here so a mistyped code stops the server starting
	// rather than quietly serving English.
	if err := web.SetDefaultLanguage(cfg.Landing.Language); err != nil {
		return nil, fmt.Errorf("landing.language: %w", err)
	}

	r := &Router{
		batchWake: make(chan struct{}, 1),
		cfg:       cfg.Router,
		landing:   cfg.Landing,
		log:       log,
		reg:       reg,
		profiles:  profiles,
		prefix:    prefixScorer,
		lora:      newLoraTracker(log, metrics.observeAdapter),
		kv:        kvindex.New(),
		kvClient:  newKVClient(),
		m:         metrics,
		comfy:     &comfy.Client{HTTP: &http.Client{Transport: &http.Transport{DialContext: (&net.Dialer{Timeout: cfg.Router.Timeouts.Dial}).DialContext}}},
		px: proxy.New(proxy.Config{
			DialTimeout:           cfg.Router.Timeouts.Dial,
			ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader,
			GenerationTimeout:     cfg.Router.Timeouts.Generation,
			FirstTokenTimeout:     cfg.Router.Timeouts.FirstToken,
			IdleTimeout:           cfg.Router.Timeouts.Idle,
			TotalTimeout:          cfg.Router.Timeouts.Total,
			MaxResponseBytes:      cfg.Router.MaxBodyBytes,
		}),
	}
	r.answers = newAnswerCache(cfg.Router.Cache.MaxBytes, cfg.Router.Cache.MaxEntryBytes,
		func(bytes int64, n int) {
			metrics.cacheBytes.Set(float64(bytes))
			metrics.cacheEntries.Set(float64(n))
		})
	r.poller = scrape.NewPoller(reg, scrape.Config{
		Interval: cfg.Router.ScrapeInterval,
		Timeout:  cfg.Router.ScrapeTimeout,
	}, log)

	if !cfg.Auth.Enabled && !loopbackOnly(cfg.Listen) {
		log.Warn("auth is off and the API listens beyond this machine: anyone who can reach "+
			cfg.Listen+" can use every model; set auth.enabled, or listen on 127.0.0.1",
			"listen", cfg.Listen)
	}
	if cfg.Auth.Enabled {
		st, err := storepkg.Open(context.Background(), cfg.Auth.Store)
		if err != nil {
			return nil, fmt.Errorf("auth store: %w", err)
		}
		r.store = st
		r.usageRetention = cfg.Auth.UsageRetention
		// Applied where no scope set one, so a project created without limits
		// is still metered. Zero leaves it unlimited, which was the old
		// behaviour and is now a choice rather than an omission.
		var defRPM, defTPM *int
		if cfg.Auth.DefaultRPM > 0 {
			defRPM = &cfg.Auth.DefaultRPM
		}
		if cfg.Auth.DefaultTPM > 0 {
			defTPM = &cfg.Auth.DefaultTPM
		}
		st.SetDefaultLimits(defRPM, defTPM)
		st.SetPersonalDefaults(personalLimits(cfg.Auth.Personal))
		r.personalOff = cfg.Auth.Personal.Disabled
		r.oidc = cfg.Auth.OIDC
		r.mail = cfg.Mail
		r.scim = cfg.Auth.SCIM
		r.personalDefaults = personalLimits(cfg.Auth.Personal)
		r.llamaVersion = cfg.WantedEngines()["llamacpp"]

		r.auth = auth.New(st, auth.Config{
			TTL:         cfg.Auth.CacheTTL,
			NegativeTTL: cfg.Auth.NegativeTTL,
		})

		prices := make(map[string]accounting.Price, len(cfg.Auth.Pricing))
		for name, p := range cfg.Auth.Pricing {
			prices[name] = accounting.Price{
				PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, PerImage: p.PerImage,
			}
		}
		r.recorder = accounting.NewRecorder(st, r.auth, accounting.RecorderConfig{
			FlushInterval: cfg.Auth.FlushInterval,
			BatchSize:     cfg.Auth.BatchSize,
			Capacity:      cfg.Auth.BufferSize,
		}, log)
		r.accounts = account.New(st, account.Config{
			SignupOpen: cfg.Auth.SignupOpen,
		}, log)

		r.acct = accounting.New(
			accounting.NewPricing(prices, accounting.Price{
				PromptPer1k:     cfg.Auth.DefaultPrice.PromptPer1k,
				CompletionPer1k: cfg.Auth.DefaultPrice.CompletionPer1k,
				PerImage:        cfg.Auth.DefaultPrice.PerImage,
			}),
			r.recorder, log)
		r.threads = r.store
	} else if cfg.Threads.Store != "" {
		st, err := storepkg.Open(context.Background(), cfg.Threads.Store)
		if err != nil {
			return nil, fmt.Errorf("threads store: %w", err)
		}
		if err := st.EnsureAnonymousProject(context.Background(), anonymousProject, cfg.Threads.RetainHours); err != nil {
			_ = st.Close()
			return nil, fmt.Errorf("threads store: %w", err)
		}
		r.threads, r.anonymous, r.ownsThreads = st, true, true
	}

	if cfg.Router.Admission.Enabled {
		adm, err := admission.New(admissionConfig(cfg.Router), r.saturation)
		if err != nil {
			return nil, fmt.Errorf("admission: %w", err)
		}
		r.adm = adm
	}

	// Seeded from configuration, not only pushed. A router running without a
	// control plane — a static endpoint list, which is a supported mode — has
	// nothing to push it a catalog, and would otherwise resolve no model at
	// all. A control plane pushes over the top of this at startup and on every
	// change after.
	if len(cfg.Models) > 0 {
		r.storeModels(cfg.Models)
	}
	return r, nil
}

func admissionConfig(rc config.Router) admission.Config {
	bands := make([]admission.Band, 0, len(rc.Admission.Bands))
	for _, b := range rc.Admission.Bands {
		bands = append(bands, admission.Band{
			Name:              b.Name,
			Priority:          b.Priority,
			SaturationCeiling: b.SaturationCeiling,
		})
	}
	return admission.Config{
		Bands:       bands,
		DefaultBand: rc.Admission.DefaultBand,
		Thresholds: admission.Thresholds{
			// Admission and the queue scorer must agree on what a healthy
			// buffer is, or the two layers fight each other.
			TargetWaitingPerReplica: rc.Scoring.TargetWaiting,
			KVCeiling:               rc.Scoring.KVCeiling,
		},
		QueueTTL:             rc.Admission.QueueTTL,
		QueueCapacityPerBand: rc.Admission.QueueCapacityPerBand,
		QueueCapacityTotal:   rc.Admission.QueueCapacityTotal,
		Interval:             rc.Admission.Interval,
	}
}

// saturation reports the pool pressure for a model, for the admission layer.
//
// outstanding is what admission has let through. It is compared against the
// registry's inflight counters rather than the scraped engine stats, because
// inflight updates the instant a request is dispatched while scraped stats lag
// by a whole interval. Comparing against scraped data would make every
// dispatched-but-unscraped request look like queue depth and throttle the
// steady state down to a few requests per scrape.
//
// What is left over — admitted but not yet dispatched — is the thundering herd:
// a burst that passed the check together before any of it became visible.
func (r *Router) saturation(model string, outstanding int) (float64, bool) {
	eps := r.reg.Snapshots(model)
	if len(eps) == 0 {
		return 0, false
	}
	dispatched := 0
	for _, e := range eps {
		dispatched += e.Inflight
	}
	sat := admission.Saturation(eps, max(0, outstanding-dispatched), admission.Thresholds{
		TargetWaitingPerReplica: r.cfg.Scoring.TargetWaiting,
		KVCeiling:               r.cfg.Scoring.KVCeiling,
	})
	// A model with an SLO is as saturated as it is past it: the queue and
	// the cache can look healthy while answers come too slowly, and that
	// is when the lower bands should wait.
	return max(sat, r.sloPressure(model)), true
}

func toInfos(eps []config.Endpoint) []endpoint.Info {
	out := make([]endpoint.Info, 0, len(eps))
	for _, e := range eps {
		out = append(out, endpoint.Info{
			ID:       e.ID,
			URL:      strings.TrimRight(e.URL, "/"),
			APIKey:   e.APIKey,
			Models:   e.Models,
			Engine:   e.Engine,
			Pool:     e.Pool,
			Weight:   e.Weight,
			Priority: e.Priority,
			Workflow: e.Workflow,
		})
	}
	return out
}

// Start begins scraping endpoint metrics and dispatching admissions.
// It returns immediately.
func (r *Router) Start(ctx context.Context) {
	r.poller.Sync(ctx)
	if r.lora != nil {
		go r.lora.run(ctx)
	}
	if r.adm != nil {
		r.adm.Start()
	}
	if r.recorder != nil {
		r.recorder.Start(ctx)
	}
	if r.acct != nil && r.store != nil {
		// Budgets are counted in memory; the store is what makes them survive a
		// restart. A failure here is not fatal — refusing to start because last
		// month's spend could not be read would be worse than starting with
		// counters that are low.
		if err := r.acct.RestoreBudgets(ctx, r.store, storepkg.PeriodMonth); err != nil {
			r.log.Warn("could not restore budgets; they start from zero this period", "err", err)
		}
	}
	if r.threads != nil {
		go r.pruneThreads(ctx)
	}
	if r.store != nil && r.usageRetention > 0 {
		go r.pruneUsage(ctx, r.usageRetention)
	}
	if r.store != nil {
		go r.runBatches(ctx)
		go r.pruneRequestBodies(ctx)
		go r.refreshTenants(ctx)
		go r.purgeSessions(ctx)
	}
}

// purgeSessions deletes sessions and key claims that can no longer be used,
// hourly and once at startup: refused already, but a table that only grows
// is its own problem, and a session row is a credential's hash kept for
// nothing.
func (r *Router) purgeSessions(ctx context.Context) {
	purge := func() {
		now := time.Now()
		if n, err := r.store.PurgeSessions(ctx, now); err != nil {
			r.log.Warn("could not remove expired sessions", "err", err)
		} else if n > 0 {
			r.log.Info("removed expired and idle sessions", "sessions", n)
		}
		if n, err := r.store.PurgeClaims(ctx, now); err != nil {
			r.log.Warn("could not remove used and expired key claims", "err", err)
		} else if n > 0 {
			r.log.Info("removed used and expired key claims", "claims", n)
		}
	}
	purge()
	t := time.NewTicker(time.Hour)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			purge()
		}
	}
}

// refreshTenants keeps the tenant names the request series are labelled with,
// and the budget gauges, a minute fresh.
func (r *Router) refreshTenants(ctx context.Context) {
	refresh := func() {
		if err := r.m.tenant.refresh(ctx, r.store, time.Now()); err != nil && ctx.Err() == nil {
			r.log.Warn("could not refresh tenant names and budget gauges", "err", err)
		}
	}
	refresh()
	t := time.NewTicker(time.Minute)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			refresh()
		}
	}
}

// pruneUsage deletes request records older than the retention, hourly and
// once at startup, as pruneThreads does. The five-minute sums stay, so no
// total or chart moves; what goes is the trail of each old request.
func (r *Router) pruneUsage(ctx context.Context, keep time.Duration) {
	prune := func() {
		n, err := r.store.PruneUsage(ctx, time.Now().Add(-keep))
		if err != nil {
			r.log.Warn("could not prune old request records", "err", err)
			return
		}
		if n > 0 {
			r.log.Info("removed request records past their retention", "rows", n, "retention", keep)
		}
	}
	prune()
	t := time.NewTicker(time.Hour)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			prune()
		}
	}
}

// pruneThreads enforces each project's retention window.
//
// Without something running this, "we keep conversations for 24 hours" is a
// comment rather than a promise — the rows would simply accumulate, and the
// number a project chose would describe nothing.
//
// Hourly, and once at startup. A gateway that has been down for a week should
// not wait another hour before honouring a window that expired during it, and
// an hour is fine granularity for a limit expressed in hours.
func (r *Router) pruneThreads(ctx context.Context) {
	prune := func() {
		n, err := r.threads.PruneThreads(ctx, time.Now())
		if err != nil {
			r.log.Warn("could not prune expired conversations", "err", err)
			return
		}
		if n > 0 {
			r.log.Info("removed conversations past their retention", "threads", n)
		}
		m, err := r.threads.PruneMemory(ctx, time.Now())
		if err != nil {
			r.log.Warn("could not prune expired memory notes", "err", err)
			return
		}
		if m > 0 {
			r.log.Info("removed memory notes past their retention", "notes", m)
		}
	}
	prune()

	t := time.NewTicker(time.Hour)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			prune()
		}
	}
}

// Stop halts scraping and admission.
func (r *Router) Stop() {
	r.poller.Stop()
	if r.adm != nil {
		r.adm.Close()
	}
	// Drain before closing the database, or the last batch has nowhere to go.
	if r.recorder != nil {
		r.recorder.Stop()
	}
	if r.store != nil {
		_ = r.store.Close()
	}
	if r.ownsThreads {
		_ = r.threads.Close()
	}
}

// SetLeader tells /leader how to know whether this process leads.
func (r *Router) SetLeader(f func() bool) { r.isLeader.Store(&f) }

// Registry exposes the endpoint registry, for tests and admin endpoints.
func (r *Router) Registry() *endpoint.Registry { return r.reg }

// SetEndpoints replaces the endpoint list and resynchronizes metric scraping.
//
// This is how the control plane feeds the router from phase 3 onwards, taking
// over from the static configuration list (plan §6.4). Endpoints that survive
// the update keep their live state — inflight counts, breaker state, drain rate
// — so a replica that was already serving does not start from cold just
// because another one moved.
// modelMap returns the name→model map in force.
//
// Never nil: a router built before any catalog arrived resolves nothing, which
// is the same answer it gave for an unknown name before.
func (r *Router) modelMap() map[string]catalog.Model {
	if m := r.models.Load(); m != nil {
		return *m
	}
	return nil
}

// SetCatalog replaces the models this router knows about.
//
// The catalog's counterpart to SetEndpoints, and pushed for the same reason:
// the control plane owns what should exist, the router needs to know, and
// neither has to be in the other's process for that to work.
func (r *Router) SetCatalog(c *catalog.Catalog) { r.storeModels(c.Models()) }

// storeModels makes models what names resolve to: each model by its name, its
// aliases and its adapters. One place for it, whether the models came from
// configuration at start or from the control plane since.
func (r *Router) storeModels(models []catalog.Model) {
	m := map[string]catalog.Model{}
	ads := map[string]adapterRef{}
	for _, entry := range models {
		m[entry.Name] = entry
		for _, a := range entry.Aliases {
			m[a] = entry
		}
		// An adapter resolves to the model it runs on, as an alias does, and
		// is applied to the request on the way there.
		for i, a := range entry.Adapters {
			m[a.Name] = entry
			ads[a.Name] = adapterRef{base: entry.Name, owner: a.Owner, index: i, scale: a.EffectiveScale(),
				limit: entry.MaxLoadedAdapters, path: engines.RuntimeLoRAPath(a.Name, a.Source)}
		}
	}
	r.models.Store(&m)
	r.adapters.Store(&ads)
	targets := map[string]string{}
	for _, entry := range models {
		if entry.Canary != nil {
			if t, ok := m[entry.Canary.Model]; ok {
				targets[t.Name] = entry.Name
			}
		}
	}
	r.canaryTargets.Store(&targets)
	if r.m != nil {
		r.m.token.setSLOs(models)
	}
}

// adapterRef is where an adapter runs and how it is applied.
type adapterRef struct {
	base string
	// owner is who sees it (catalog.Adapter.Owner); empty is everyone.
	owner string
	index int
	scale float64
	// limit is the model's max_loaded_adapters: above zero, the adapter
	// loads on demand, from path.
	limit int
	path  string
}

// onDemand reports whether the adapter is loaded when asked for rather than
// at start.
func (a adapterRef) onDemand() bool { return a.limit > 0 }

// adapter finds one by the name a client asked for.
func (r *Router) adapter(name string) (adapterRef, bool) {
	if p := r.adapters.Load(); p != nil {
		a, ok := (*p)[name]
		return a, ok
	}
	return adapterRef{}, false
}

// adapterHidden reports whether name is an adapter its owner keeps from this
// caller: one outside the part of the tree it belongs to, or one with no
// identity at all. Such a caller is told there is no such model.
func (r *Router) adapterHidden(p auth.Principal, authed bool, name string) bool {
	ad, ok := r.adapter(name)
	if !ok || ad.owner == "" {
		return false
	}
	if !authed {
		return true
	}
	scope, id, _ := strings.Cut(ad.owner, ":")
	switch scope {
	case "org":
		return p.OrgID != id
	case "team":
		return p.TeamID != id
	case "project":
		return p.ProjectID != id
	case "user":
		return p.UserID != id
	case "key":
		return p.KeyID != id
	}
	return true
}

// LastUsed reports when each endpoint was last sent a request, satisfying
// control.UsageSource so the control plane can tell a busy replica from one
// holding its cards out of habit.
func (r *Router) LastUsed() map[string]time.Time { return r.reg.LastUsed() }

// ModelLoad sums what each model's replicas report: how many there are, and
// requests running and waiting at their engines. Keyed by the catalog name,
// which the control plane publishes first among each endpoint's names.
func (r *Router) ModelLoad() map[string]reconciler.Load {
	out := map[string]reconciler.Load{}
	for _, s := range r.reg.Snapshots("") {
		if len(s.Models) == 0 {
			continue
		}
		l := out[s.Models[0]]
		l.Replicas++
		l.Running += s.Stats.Running
		l.Waiting += s.Stats.Waiting
		out[s.Models[0]] = l
	}
	return out
}

func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
	// A replica that has gone takes its cache with it. Leaving its entries
	// in the tree or the index would steer traffic towards something that is
	// not there, and since nothing reuses a replica id the entries would never
	// be corrected — only expire, minutes later, or never.
	before := map[string]struct{}{}
	for _, s := range r.reg.Snapshots("") {
		before[s.ID] = struct{}{}
	}
	for _, i := range infos {
		delete(before, i.ID)
	}
	for id := range before {
		r.prefix.Forget(id)
		r.kv.Forget(id)
	}

	if r.lora != nil {
		ids := make(map[string]struct{}, len(infos))
		for _, i := range infos {
			ids[i.ID] = struct{}{}
		}
		r.lora.keep(ids)
	}

	r.reg.Replace(infos)
	r.poller.Sync(ctx)
}

// --- model resolution ---

// resolve maps a client-supplied model name onto the catalog. Unknown names are
// passed through unchanged: an endpoint may serve a model the catalog does not
// list, and refusing it would be unhelpful.
func (r *Router) resolve(name string) (canonical string, m catalog.Model) {
	if e, ok := r.modelMap()[name]; ok {
		return e.Name, e
	}
	return name, catalog.Model{Name: name}
}

func (r *Router) profileFor(m catalog.Model, override string) *scheduler.Profile {
	for _, name := range []string{override, m.RoutingProfile, r.cfg.DefaultProfile} {
		if p, ok := r.profiles[name]; ok && name != "" {
			return p
		}
	}
	return r.profiles["chat"]
}

// --- dispatch ---

// attempt describes one upstream try, for logging and metrics.
type dispatchResult struct {
	endpointID string
	pool       string
	servedModel
	outcome proxy.Outcome
	err     error
	// guard is what the output filter decided, if there was one, and
	// judged whether there was.
	guard  guard.Verdict
	judged bool
	// canaryOf is the model whose canary this request was split by, canary
	// the candidate, and onCanary whether the candidate served it.
	canaryOf, canary string
	onCanary         bool
	// cached is an answer from the project's answer cache, which no model
	// was asked for.
	cached bool
}

type servedModel struct {
	model string
	// fallback is true when the served model is not the one that was requested.
	fallback bool
}

// dispatch resolves the model, then tries endpoints until one works, the client
// has been written to, or the attempt budget runs out.
//
// The retry rule (plan §7.7) is strict: retry only while nothing has reached the
// client. Once bytes are out, a retry would splice two different responses
// together, which is worse than the error the client would otherwise see.
func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parsedRequest) dispatchResult {
	canonical, entry := r.resolve(req.Model)

	// A canary takes its share before anything else, and falls back to the
	// stable side, not to a fallback, when it has nothing ready.
	// Not an adapter: the candidate's replicas do not run it under the
	// stable model's name.
	_, isAdapter := r.adapter(req.Model)
	if cn := entry.Canary; cn != nil && !isAdapter {
		onCanary := chooseVariant(req, cn)
		target, _ := r.resolve(cn.Model)
		if onCanary && len(r.reg.Snapshots(target)) == 0 {
			onCanary = false
		}
		if _, te := r.resolve(target); onCanary && (req.noStore || req.inCluster) && te.IsExternal() {
			onCanary = false
		}
		setVariantHeader(w.Header(), onCanary)
		if onCanary {
			// The engine is asked for the model by the name it serves.
			creq := *req
			creq.UpstreamBody = withModel(req.UpstreamBody, target)
			res := r.dispatchOne(ctx, w, &creq, target, entry)
			res.model, res.canaryOf, res.canary, res.onCanary = target, canonical, target, true
			if res.err == nil || res.outcome.Wrote || errors.Is(res.err, proxy.ErrClientGone) ||
				!errors.Is(res.err, errNoCapacity) {
				return res
			}
			setVariantHeader(w.Header(), false)
		}
		res := r.dispatchStable(ctx, w, req, canonical, entry)
		res.canaryOf, res.canary = canonical, target
		return res
	}
	return r.dispatchStable(ctx, w, req, canonical, entry)
}

// dispatchStable serves a model as asked, then its fallbacks.
func (r *Router) dispatchStable(ctx context.Context, w http.ResponseWriter, req *parsedRequest, canonical string, entry catalog.Model) dispatchResult {
	// The requested model first, then its fallbacks. A fallback is a different
	// model, so the client is told about it via x-llmgw-served-model.
	chain := []string{canonical}
	p, authed := auth.FromContext(ctx)
	for _, fb := range entry.Fallbacks {
		// A fallback is another model, and a key allowed only the one it
		// asked for is not served another because that one was full.
		if authed && !r.permits(p, fb) || r.isGuardModel(fb) {
			continue
		}
		// Nor sent outside the cluster under a policy that keeps nothing.
		if _, fe := r.resolve(fb); (req.noStore || req.inCluster) && fe.IsExternal() {
			continue
		}
		chain = append(chain, fb)
	}

	var last dispatchResult
	for i, model := range chain {
		if i > 0 {
			// Before dispatching, because dispatchOne writes the response and
			// a header set after that is a header nobody receives. It was set
			// in finish until now, which is after the proxy has already called
			// WriteHeader — so this has announced a fallback to nobody since it
			// was written.
			w.Header().Set("x-llmgw-served-model", model)
		}
		res := r.dispatchOne(ctx, w, req, model, entry)
		res.model = model
		res.fallback = i > 0

		if res.err == nil || res.outcome.Wrote || errors.Is(res.err, proxy.ErrClientGone) {
			return res
		}
		last = res
		// Only a total absence of capacity justifies moving to another model;
		// a bad request would fail identically there.
		if !errors.Is(res.err, errNoCapacity) {
			return res
		}
	}
	return last
}

var errNoCapacity = errors.New("no endpoint available")

// errLoading means the model is being started and the wait ran out.
//
// Distinct from having no capacity because the right client behaviour differs:
// this one succeeds on a retry, and the other does not.
var errLoading = errors.New("model loading")

func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry catalog.Model) dispatchResult {
	candidates := r.reg.Snapshots(model)
	if visited := withoutVisited(candidates, req.Header); len(visited) < len(candidates) {
		if len(visited) == 0 {
			return dispatchResult{err: fmt.Errorf("%w: %s is served by a federated cluster this request has already passed through",
				errNoCapacity, model)}
		}
		candidates = visited
	}
	if len(candidates) == 0 {
		// Nothing is serving it. That may be because nothing ever will, or
		// because it loads when asked and nobody has asked yet — and the second
		// is not an error, it is the feature working.
		if r.activate(ctx, model) {
			candidates = r.reg.Snapshots(model)
		}
	}
	if len(candidates) == 0 {
		if why, no := r.whyUnservable(model); no && r.activator != nil {
			// Not "loading; retry": it will not load, and a client told to
			// retry keeps retrying something that never arrives.
			return dispatchResult{err: fmt.Errorf("%w: %s", errNoCapacity, why)}
		}
		if r.activator != nil && r.activator.Want(model) {
			// Asked for, still loading. Saying "no endpoint serves this" would
			// be a lie, and a client told to retry will succeed where one told
			// the model does not exist gives up.
			return dispatchResult{err: fmt.Errorf("%w: %s", errLoading,
				"the model is loading; retry shortly")}
		}
		return dispatchResult{err: fmt.Errorf("%w: %s", errNoCapacity, noEndpointFor(model, slices.DeleteFunc(r.servedModels(), r.hiddenModel)))}
	}

	profile := r.profileFor(entry, req.ProfileOverride)
	sctx := &scheduler.Ctx{
		Now: time.Now(),
		Req: scheduler.Request{
			Model:        model,
			SessionID:    req.SessionID,
			PromptPrefix: req.PromptPrefix,
			PromptTokens: req.PromptTokens,
			MaxTokens:    req.MaxTokens,
		},
	}
	// An adapter asked for by name runs on its model's replicas — not on a
	// fallback to another model, which does not have it.
	ad, isAdapter := r.adapter(req.Model)
	isAdapter = isAdapter && ad.base == model
	onDemand := isAdapter && ad.onDemand() && r.lora != nil
	if onDemand {
		sctx.Req.Adapter = req.Model
		r.lora.syncUnknown(ctx, candidates, model)
	}
	// Only where a profile publishes events, and only for a profile with a
	// prefix scorer to read the answer: the batch profile does not care where
	// a prefix is cached, and should not pay a round trip to find out.
	if entry.KVIndex() && profile.Uses("prefix") {
		_, kspan := tracer().Start(ctx, "kv index")
		sctx.Req.KV = r.kvQuery(ctx, req, model, candidates)
		if sctx.Req.KV.Fallback != "" {
			kspan.SetAttributes(attribute.String("llmgw.kv_fallback", string(sctx.Req.KV.Fallback)))
		}
		kspan.End()
	}

	attempts := max(r.cfg.Retries+1, 1)
	var last dispatchResult

	for range attempts {
		if len(candidates) == 0 {
			break
		}
		if onDemand {
			sctx.Adapters = r.lora.view(candidates, req.Model, ad.limit)
		}
		_, sspan := tracer().Start(ctx, "schedule", trace.WithAttributes(
			attribute.String("llmgw.model", model), attribute.Int("llmgw.candidates", len(candidates))))
		dec, err := profile.Schedule(sctx, candidates)
		if err == nil {
			sspan.SetAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID))
		}
		endSpan(sspan, err)
		if err != nil {
			last = dispatchResult{err: fmt.Errorf("%w: %v", errNoCapacity, err)}
			break
		}

		release, granted := r.reg.Acquire(dec.Endpoint.ID)
		if !granted {
			// The breaker opened between scoring and dispatch; drop it and
			// score again rather than failing the request.
			candidates = without(candidates, dec.Endpoint.ID)
			continue
		}

		// The replica loads the adapter first if it does not hold it.
		// Counted against its inflight while it does, so the requests
		// arriving meanwhile see it busy.
		releaseAdapter := func() {}
		if onDemand {
			how, rel, lerr := r.lora.acquire(ctx, dec.Endpoint.Info, model, req.Model, ad.path, ad.limit)
			if lerr != nil {
				release(true)
				if ctx.Err() != nil {
					return dispatchResult{endpointID: dec.Endpoint.ID, err: proxy.ErrClientGone}
				}
				res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool,
					err: fmt.Errorf("adapter %s on %s: %w", req.Model, dec.Endpoint.ID, lerr)}
				if errors.Is(lerr, errAdapterLoad) {
					// Refused for what it is, it would be refused everywhere.
					return res
				}
				r.log.Warn("a replica could not load an adapter, trying another",
					"endpoint", dec.Endpoint.ID, "adapter", req.Model, "err", lerr)
				last = res
				candidates = without(candidates, dec.Endpoint.ID)
				continue
			}
			releaseAdapter = rel
			r.m.adapterAffinity.WithLabelValues(r.metricModel(model), how).Inc()
		}

		uctx, uspan := tracer().Start(ctx, "upstream", trace.WithSpanKind(trace.SpanKindClient),
			trace.WithAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID),
				attribute.String("url.full", dec.Endpoint.URL+req.Path)))
		body := req.UpstreamBody
		if dec.Endpoint.UpstreamModel != "" && req.ContentType == "" {
			body = withModel(body, dec.Endpoint.UpstreamModel)
		}
		// llama-server is asked for the model with the adapter applied to
		// this request; vLLM serves each adapter as a model of its own name,
		// which is what the client asked for already.
		if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" {
			body = withAdapter(body, model, ad.index, ad.scale)
		}
		// Passages share cached blocks only with requests that could have
		// been given the same ones: a prefix found warm would otherwise say
		// what another tenant asked about, by how fast its answer started.
		// vLLM keys its prefix cache by this salt; other engines are not
		// sent a field they may refuse.
		if req.kbSalt != "" && dec.Endpoint.Engine == "vllm" && req.ContentType == "" {
			body = withCacheSalt(body, req.kbSalt)
		}
		// What computes the answer, said before the proxy writes the headers.
		// A model can have replicas on cards and on CPUs at once, and the two
		// differ tenfold in speed: a reply that came slowly is otherwise
		// indistinguishable from a model that is slow everywhere. Set again on
		// a retry, which only happens before anything was written.
		if hw := hardwareOf(dec.Endpoint.Pool); hw != "" {
			w.Header().Set(hardwareHeader, hw)
		} else {
			w.Header().Del(hardwareHeader)
		}
		preq := proxy.Request{
			URL:                dec.Endpoint.URL + req.Path,
			APIKey:             dec.Endpoint.APIKey,
			Body:               body,
			Header:             withTraceparent(uctx, req.Header),
			Stream:             req.Stream,
			SuppressUsageChunk: req.SuppressUsageChunk,
			CaptureReply:       req.captureBytes(),
			ContentType:        req.ContentType,
		}
		if onDemand {
			// vLLM's answer for a model it does not have: the adapter was
			// unloaded behind this router's back, by another router or a
			// restart. Not relayed; loaded again and tried again.
			preq.RetryStatus = http.StatusNotFound
		}
		if dec.Endpoint.Peer != "" {
			preq.Extra = r.peerHeaders(ctx, dec.Endpoint, req.Header)
		}
		if req.Stream {
			gaps := r.m.token.gap.WithLabelValues(r.metricModel(model))
			preq.OnEvent = func(d time.Duration) { gaps.Observe(d.Seconds()) }
		}
		// A filter per attempt, for the model this attempt is served by.
		filter := r.outputFilter(uctx, req, model, &preq)
		var out proxy.Outcome
		var perr error
		if dec.Endpoint.Engine == endpoint.EngineComfyUI {
			out, perr = r.generateImages(uctx, w, req, dec.Endpoint.Info, body)
		} else {
			out, perr = r.px.Do(uctx, w, preq)
		}
		releaseAdapter()
		uspan.SetAttributes(attribute.Int("http.response.status_code", out.Status))
		endSpan(uspan, perr)
		adapterGone := onDemand && !out.Wrote && out.Status == http.StatusNotFound
		if adapterGone {
			r.lora.forget(dec.Endpoint.ID, req.Model)
		}
		// A retryable failure is the endpoint's fault; a client disconnect is
		// not, and nor is an adapter it was not known to have lost.
		release(perr == nil || errors.Is(perr, proxy.ErrClientGone) || adapterGone)

		res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr}
		if filter != nil {
			res.guard = filter.Verdict()
			res.judged = true
		}
		r.m.observeScores(dec)

		// Recorded after the decision, not during it: the scorer proposes and
		// the picker disposes, so recording the proposal would fill the tree
		// with replicas that were never sent anything — the one thing that
		// would make the estimate worse than no estimate at all.
		if r.prefix != nil {
			if !req.noStore {
				r.prefix.Observe(req.PromptPrefix, dec.Endpoint.ID)
			}
		}

		if perr == nil || out.Wrote || !out.Retryable {
			return res
		}
		last = res
		if adapterGone {
			// The replica is fine; the next try may load the adapter on it
			// again.
			r.log.Info("a replica no longer held an adapter; loading it again",
				"endpoint", dec.Endpoint.ID, "adapter", req.Model)
			continue
		}
		r.log.Warn("upstream attempt failed, trying another endpoint",
			"endpoint", dec.Endpoint.ID, "model", model, "err", perr)
		candidates = without(candidates, dec.Endpoint.ID)
	}

	if last.err == nil {
		last.err = fmt.Errorf("%w: every endpoint for %q failed", errNoCapacity, model)
	} else if !errors.Is(last.err, errNoCapacity) {
		last.err = fmt.Errorf("%w: %v", errNoCapacity, last.err)
	}
	return last
}

// hardwareHeader says what the replica that answered computes on: "cpu", or
// the GPU class, such as "6000-24g".
const hardwareHeader = "x-llmgw-hardware"

// hardwareOf reads it from a pool, which placement names fabric/class —
// "none/cpu", "none/6000-24g", "ib-a/h100-80g". Empty for an endpoint
// configured by hand, which has no pool and says nothing about its hardware.
func hardwareOf(pool string) string {
	i := strings.LastIndexByte(pool, '/')
	if i < 0 || i == len(pool)-1 {
		return ""
	}
	return pool[i+1:]
}

func without(eps []endpoint.Snapshot, id string) []endpoint.Snapshot {
	return slices.DeleteFunc(slices.Clone(eps), func(e endpoint.Snapshot) bool { return e.ID == id })
}

// personalLimits turns the personal block into limits, nil when it sets none.
func personalLimits(p config.Personal) *storepkg.Limits {
	var l storepkg.Limits
	if p.RPM > 0 {
		v := p.RPM
		l.RPM = &v
	}
	if p.TPM > 0 {
		v := p.TPM
		l.TPM = &v
	}
	if p.Budget > 0 {
		v := int64(math.Round(p.Budget * 1e6))
		l.Budget = &v
		l.BudgetPeriod = storepkg.PeriodMonth
		if p.Period == "day" {
			l.BudgetPeriod = storepkg.PeriodDay
		}
	}
	if l.RPM == nil && l.TPM == nil && l.Budget == nil {
		return nil
	}
	return &l
}

// loopbackOnly reports whether a listen address is reachable from this
// machine alone. An empty host (":8080") is every interface.
func loopbackOnly(listen string) bool {
	host, _, err := net.SplitHostPort(listen)
	if err != nil {
		return false
	}
	if host == "localhost" {
		return true
	}
	a, err := netip.ParseAddr(host)
	return err == nil && a.IsLoopback()
}