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"
	"log/slog"
	"net/http"
	"slices"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"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/proxy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
)

// 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
	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.
	models map[string]catalog.Model
}

// 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,
	}
	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)
	}

	models := map[string]catalog.Model{}
	for _, m := range cfg.Models {
		models[m.Name] = m
		for _, a := range m.Aliases {
			models[a] = m
		}
	}

	r := &Router{
		cfg:      cfg.Router,
		log:      log,
		reg:      reg,
		profiles: profiles,
		models:   models,
		m:        newMetrics(),
		px: proxy.New(proxy.Config{
			DialTimeout:           cfg.Router.Timeouts.Dial,
			ResponseHeaderTimeout: cfg.Router.Timeouts.ResponseHeader,
			FirstTokenTimeout:     cfg.Router.Timeouts.FirstToken,
			IdleTimeout:           cfg.Router.Timeouts.Idle,
			TotalTimeout:          cfg.Router.Timeouts.Total,
			MaxResponseBytes:      cfg.Router.MaxBodyBytes,
		}),
	}
	r.poller = scrape.NewPoller(reg, scrape.Config{
		Interval: cfg.Router.ScrapeInterval,
		Timeout:  cfg.Router.ScrapeTimeout,
	}, log)

	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
	}
	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,
		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
	}
	return admission.Saturation(eps, max(0, outstanding-dispatched), admission.Thresholds{
		TargetWaitingPerReplica: r.cfg.Scoring.TargetWaiting,
		KVCeiling:               r.cfg.Scoring.KVCeiling,
	}), 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, "/"),
			Models:   e.Models,
			Pool:     e.Pool,
			Weight:   e.Weight,
			Priority: e.Priority,
		})
	}
	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.adm != nil {
		r.adm.Start()
	}
}

// Stop halts scraping and admission.
func (r *Router) Stop() {
	r.poller.Stop()
	if r.adm != nil {
		r.adm.Close()
	}
}

// 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.
func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
	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.models[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
}

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)

	// 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 := append([]string{canonical}, entry.Fallbacks...)

	var last dispatchResult
	for i, model := range chain {
		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")

func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry catalog.Model) dispatchResult {
	candidates := r.reg.Snapshots(model)
	if len(candidates) == 0 {
		return dispatchResult{err: fmt.Errorf("%w: no endpoint serves %q", errNoCapacity, model)}
	}

	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,
		},
	}

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

	for range attempts {
		if len(candidates) == 0 {
			break
		}
		dec, err := profile.Schedule(sctx, candidates)
		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
		}

		out, perr := r.px.Do(ctx, w, proxy.Request{
			URL:                dec.Endpoint.URL + req.Path,
			Body:               req.UpstreamBody,
			Header:             req.Header,
			Stream:             req.Stream,
			SuppressUsageChunk: req.SuppressUsageChunk,
		})
		// A retryable failure is the endpoint's fault; a client disconnect is not.
		release(perr == nil || errors.Is(perr, proxy.ErrClientGone))

		res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr}
		r.m.observeScores(dec)

		if perr == nil || out.Wrote || !out.Retryable {
			return res
		}
		last = res
		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
}

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