llmgw/internal/control/federation/federation.go

// Package federation makes other llmgw clusters providers of this one.
//
// A peer is another llmgw gateway, reached with a key it issued to this
// cluster. Its models are offered here as external models named PEER/MODEL,
// kept in step with its own /v1/models, and sent requests only while its
// /ready answers. A local model names one as a fallback to send overflow to
// the other site.
//
// Always PEER/MODEL, never the peer's own name for a model, even when nothing
// here uses that name. A name that left the cluster or stayed depending on
// what the catalog happened to hold would move traffic between sites when
// somebody added an unrelated model; the prefix also says, on every usage row
// and every x-llmgw-served-model header, which site answered.
//
// Everything that keeps a request inside the cluster already applies: a peer
// model is an external model, so a project whose requests stay in the
// cluster, or a policy that keeps nothing, is never sent to one.
package federation

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"regexp"
	"slices"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// ClusterHeader is the response header in which a gateway says its own
// cluster id, on /v1/models. Federation reads a peer's from it, to know the
// peer in a request's hop header.
const ClusterHeader = "x-llmgw-cluster"

// Store is the persistence federation needs.
type Store interface {
	PutFederationPeer(ctx context.Context, p store.FederationPeer) error
	SetFederationPeerModels(ctx context.Context, name, cluster string, models json.RawMessage, at time.Time) error
	DeleteFederationPeer(ctx context.Context, name string) error
	FederationPeers(ctx context.Context) ([]store.FederationPeer, error)
}

// Catalog is what puts a changed peer list in force: the model service.
type Catalog interface {
	Refresh(ctx context.Context) error
}

// Config tunes federation.
type Config struct {
	// Self is this cluster's id. A peer that says the same is this cluster,
	// reached by another address, and is refused.
	Self string
	// Interval is how often each peer is checked; zero is 10s.
	Interval time.Duration
	// Timeout bounds one check of one peer; zero is 5s.
	Timeout time.Duration
	// Client makes the checks; nil is a plain one.
	Client *http.Client
}

// PeerModel is what a peer says about one model it offers.
type PeerModel struct {
	// ID is the peer's name for it, and what it is asked for.
	ID            string               `json:"id"`
	Description   string               `json:"description,omitempty"`
	ContextLength int                  `json:"context_length,omitempty"`
	MaxOutput     int                  `json:"max_output,omitempty"`
	Capabilities  catalog.Capabilities `json:"capabilities,omitempty"`
	// Withdrawn is a model the peer no longer offers, kept in the catalog
	// because a model here still falls back to it: removing it would make
	// that model invalid. Nothing is sent to it, and it goes when nothing
	// names it any more.
	Withdrawn bool `json:"withdrawn,omitempty"`
}

// Service keeps the peers, their health and their models.
type Service struct {
	cfg     Config
	store   Store
	catalog Catalog
	// current is the catalog in force, read to tell which of a peer's
	// models a local model names.
	current func() *catalog.Catalog
	log     *slog.Logger

	// op serialises what changes the peer set or a peer's models, so two
	// rebuilds cannot land in the wrong order.
	op sync.Mutex

	mu    sync.Mutex
	peers map[string]*peer
}

type peer struct {
	rec    store.FederationPeer
	models []PeerModel
	health Health
}

// Health is how a peer answered its last check.
type Health struct {
	Up bool `json:"up"`
	// Checked is when; zero before the first check.
	Checked time.Time `json:"checked,omitempty"`
	// Latency is how long the peer took to answer /ready: the network to
	// the other site and the peer's front, not a model's speed.
	Latency time.Duration `json:"latency_ns,omitempty"`
	// Error is why the peer is down, empty while it is up.
	Error string `json:"error,omitempty"`
}

// New builds the service. Load reads the peers kept in the store.
func New(cfg Config, st Store, cat Catalog, current func() *catalog.Catalog, log *slog.Logger) *Service {
	if cfg.Interval <= 0 {
		cfg.Interval = 10 * time.Second
	}
	if cfg.Timeout <= 0 {
		cfg.Timeout = 5 * time.Second
	}
	if cfg.Client == nil {
		cfg.Client = &http.Client{}
	}
	if log == nil {
		log = slog.Default()
	}
	return &Service{cfg: cfg, store: st, catalog: cat, current: current, log: log, peers: map[string]*peer{}}
}

// Self is this cluster's id.
func (s *Service) Self() string { return s.cfg.Self }

// Load reads the peers and the models they last offered from the store, so
// the catalog built at startup has them before any peer is asked again.
// Every peer reads as down until its first check.
func (s *Service) Load(ctx context.Context) error {
	recs, err := s.store.FederationPeers(ctx)
	if err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, r := range recs {
		var models []PeerModel
		if err := json.Unmarshal(r.Models, &models); err != nil {
			s.log.Warn("a federation peer's models could not be read; they come back at its next check",
				"peer", r.Name, "err", err)
		}
		s.peers[r.Name] = &peer{rec: r, models: models, health: Health{Error: "not checked yet"}}
	}
	return nil
}

// Models are the peers' models as catalog entries, withdrawn ones included:
// what the model service adds to the catalog.
func (s *Service) Models() []catalog.Model {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []catalog.Model
	for _, name := range s.namesLocked() {
		p := s.peers[name]
		for _, m := range p.models {
			out = append(out, catalogModel(p.rec, m))
		}
	}
	return out
}

func (s *Service) namesLocked() []string {
	names := make([]string, 0, len(s.peers))
	for n := range s.peers {
		names = append(names, n)
	}
	slices.Sort(names)
	return names
}

// catalogModel is a peer's model as an entry of this catalog.
func catalogModel(rec store.FederationPeer, m PeerModel) catalog.Model {
	desc := "on " + rec.Name
	if m.Description != "" {
		desc += ": " + m.Description
	}
	return catalog.Model{
		Name: rec.Name + "/" + m.ID, Source: "peer:" + rec.Name, Description: desc,
		ContextLength: m.ContextLength, MaxOutput: m.MaxOutput, Capabilities: m.Capabilities,
		External: &catalog.External{URL: rec.URL, Model: m.ID, APIKeyFile: rec.KeyFile, Peer: rec.Name},
	}
}

// PeerRoute says whether a peer's model may be sent requests now: the peer
// answered its last check and still offers it. The controller publishes the
// model's endpoint only while it does.
func (s *Service) PeerRoute(name, model string) (cluster string, sendTags, ok bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	p, found := s.peers[name]
	if !found || !p.health.Up {
		return "", false, false
	}
	for _, m := range p.models {
		if m.ID == model {
			return p.rec.Cluster, p.rec.SendTags, !m.Withdrawn
		}
	}
	return "", false, false
}

// peerName is what a peer may be called: the prefix of its models' names,
// and a tag value, so letters, digits and a few separators, and no slash.
var peerName = regexp.MustCompile(`^[a-z0-9][a-z0-9_.-]{0,31}$`)

// AddRequest adds a peer.
type AddRequest struct {
	// Name is what this cluster calls the peer, and the prefix of its
	// models: NAME/MODEL.
	Name string `json:"name"`
	// URL is the peer gateway's client address, without /v1.
	URL string `json:"url"`
	// KeyFile is a file on the control plane holding a key the peer issued
	// to this cluster (or env:NAME, vault:PATH#FIELD).
	KeyFile string `json:"key_file"`
	// SendTags makes requests sent to the peer carry x-llmgw-tags naming
	// this cluster and the sending project. Off unless asked: the peer
	// learns who in this cluster is asking.
	SendTags bool `json:"send_tags,omitempty"`
}

// Add checks a peer and keeps it, its models joining the catalog.
//
// Checked before anything is kept: a mistyped address or a key the peer
// refuses is said now, rather than found as a peer that is always down.
func (s *Service) Add(ctx context.Context, req AddRequest, by string) (Status, error) {
	req.Name = strings.TrimSpace(req.Name)
	req.URL = strings.TrimRight(strings.TrimSpace(req.URL), "/")
	req.URL = strings.TrimSuffix(req.URL, "/v1")
	if !peerName.MatchString(req.Name) {
		return Status{}, fmt.Errorf("peer name %q: 1 to 32 lower-case letters, digits, '_', '.' or '-', starting with a letter or digit", req.Name)
	}
	u, err := url.Parse(req.URL)
	if err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "" {
		return Status{}, fmt.Errorf("peer address %q is not an http(s) address", req.URL)
	}
	if req.KeyFile == "" {
		return Status{}, errors.New("give the file holding the key the peer issued to this cluster (--key-file)")
	}
	if cur := s.current(); cur != nil {
		for _, m := range cur.Models() {
			if !m.IsPeer() && strings.HasPrefix(m.Name, req.Name+"/") {
				return Status{}, fmt.Errorf("model %q of this cluster already starts with %s/; call the peer something else",
					m.Name, req.Name)
			}
		}
	}

	s.op.Lock()
	defer s.op.Unlock()

	rec := store.FederationPeer{Name: req.Name, URL: req.URL, KeyFile: req.KeyFile, SendTags: req.SendTags,
		CreatedAt: time.Now(), CreatedBy: by}
	s.mu.Lock()
	old, existed := s.peers[req.Name]
	if existed {
		rec.CreatedAt, rec.CreatedBy = old.rec.CreatedAt, old.rec.CreatedBy
	}
	s.mu.Unlock()

	res := s.check(ctx, rec)
	if res.err != nil {
		return Status{}, fmt.Errorf("peer %s: %w", req.Name, res.err)
	}
	if res.cluster != "" && res.cluster == s.cfg.Self {
		return Status{}, fmt.Errorf("%s is this cluster: its gateway answers with this cluster's own id", req.URL)
	}
	rec.Cluster, rec.SyncedAt = res.cluster, time.Now()
	models := res.models
	if existed {
		models = s.keepReferenced(old, res.models)
	}
	b, _ := json.Marshal(models)
	rec.Models = b
	if err := s.store.PutFederationPeer(ctx, rec); err != nil {
		return Status{}, err
	}
	s.mu.Lock()
	s.peers[req.Name] = &peer{rec: rec, models: models, health: res.health}
	s.mu.Unlock()
	if err := s.catalog.Refresh(ctx); err != nil {
		return Status{}, fmt.Errorf("the peer is kept, but the catalog could not be rebuilt: %w", err)
	}
	s.log.Info("federation peer added", "peer", req.Name, "url", req.URL, "models", len(models), "by", by)
	st, _ := s.Status(req.Name)
	return st, nil
}

// Remove takes a peer and its models away.
//
// Refused while a model here falls back to one of them, or has one as its
// canary, for the reason removing a model is: a fallback pointing at
// nothing is found during the outage it was meant to cover.
func (s *Service) Remove(ctx context.Context, name string) error {
	s.op.Lock()
	defer s.op.Unlock()

	s.mu.Lock()
	p, ok := s.peers[name]
	s.mu.Unlock()
	if !ok {
		return fmt.Errorf("no federation peer called %q: `llmgw ctl federation list` lists them", name)
	}
	var users []string
	if cur := s.current(); cur != nil {
		for _, m := range cur.Models() {
			if m.IsPeer() {
				continue
			}
			for _, ref := range references(m) {
				if strings.HasPrefix(ref, name+"/") {
					users = append(users, m.Name+" → "+ref)
				}
			}
		}
	}
	if len(users) > 0 {
		return fmt.Errorf("models here still use %s's: %s; change their fallbacks first (`llmgw ctl model fallbacks`)",
			name, strings.Join(users, ", "))
	}
	if err := s.store.DeleteFederationPeer(ctx, name); err != nil && !errors.Is(err, store.ErrNotFound) {
		return err
	}
	s.mu.Lock()
	delete(s.peers, name)
	s.mu.Unlock()
	if err := s.catalog.Refresh(ctx); err != nil {
		return fmt.Errorf("the peer is removed, but the catalog could not be rebuilt: %w", err)
	}
	s.log.Info("federation peer removed", "peer", name, "models", len(p.models))
	return nil
}

// references are the other models a model names: its fallbacks and its
// canary.
func references(m catalog.Model) []string {
	out := slices.Clone(m.Fallbacks)
	if m.Canary != nil {
		out = append(out, m.Canary.Model)
	}
	return out
}

// keepReferenced is what a peer offers now, plus those it offered before and
// no longer does that a model here still names, kept withdrawn.
func (s *Service) keepReferenced(old *peer, offered []PeerModel) []PeerModel {
	out := slices.Clone(offered)
	named := map[string]bool{}
	if cur := s.current(); cur != nil {
		for _, m := range cur.Models() {
			if m.IsPeer() {
				continue
			}
			for _, ref := range references(m) {
				named[ref] = true
			}
		}
	}
	for _, m := range old.models {
		if slices.ContainsFunc(offered, func(o PeerModel) bool { return o.ID == m.ID }) {
			continue
		}
		if named[old.rec.Name+"/"+m.ID] {
			m.Withdrawn = true
			out = append(out, m)
		}
	}
	return out
}

// Sync checks one peer now, rather than at the next interval.
func (s *Service) Sync(ctx context.Context, name string) (Status, error) {
	s.mu.Lock()
	p, ok := s.peers[name]
	var rec store.FederationPeer
	if ok {
		rec = p.rec
	}
	s.mu.Unlock()
	if !ok {
		return Status{}, fmt.Errorf("no federation peer called %q: `llmgw ctl federation list` lists them", name)
	}
	res := s.check(ctx, rec)
	s.apply(ctx, map[string]checkResult{name: res})
	st, _ := s.Status(name)
	return st, nil
}

// SyncAll checks every peer at once, and puts what changed in force.
func (s *Service) SyncAll(ctx context.Context) {
	s.mu.Lock()
	recs := make([]store.FederationPeer, 0, len(s.peers))
	for _, p := range s.peers {
		recs = append(recs, p.rec)
	}
	s.mu.Unlock()

	results := make(map[string]checkResult, len(recs))
	var mu sync.Mutex
	var wg sync.WaitGroup
	for _, rec := range recs {
		wg.Go(func() {
			res := s.check(ctx, rec)
			mu.Lock()
			results[rec.Name] = res
			mu.Unlock()
		})
	}
	wg.Wait()
	s.apply(ctx, results)
}

// apply records checks' results: each peer's health always, and its models
// when it said them, rebuilding the catalog once if any changed.
func (s *Service) apply(ctx context.Context, results map[string]checkResult) {
	s.op.Lock()
	defer s.op.Unlock()

	type change struct {
		name, cluster string
		models        []PeerModel
	}
	var changes []change
	s.mu.Lock()
	for name, res := range results {
		p, ok := s.peers[name]
		if !ok {
			continue // removed while it was being checked
		}
		if p.health.Up != res.health.Up {
			if res.health.Up {
				s.log.Info("federation peer is up", "peer", name, "latency", res.health.Latency)
			} else {
				s.log.Warn("federation peer is down; its models are not sent requests", "peer", name, "err", res.health.Error)
			}
		}
		p.health = res.health
		if res.err != nil {
			continue // its list is kept as it was: down is not empty
		}
		next := s.keepReferenced(p, res.models)
		if res.cluster != p.rec.Cluster || !slices.Equal(next, p.models) {
			p.models, p.rec.Cluster, p.rec.SyncedAt = next, res.cluster, time.Now()
			changes = append(changes, change{name, res.cluster, next})
		}
	}
	s.mu.Unlock()

	if len(changes) == 0 {
		return
	}
	for _, c := range changes {
		b, _ := json.Marshal(c.models)
		if err := s.store.SetFederationPeerModels(ctx, c.name, c.cluster, b, time.Now()); err != nil {
			s.log.Warn("could not keep a federation peer's models", "peer", c.name, "err", err)
		}
		s.log.Info("federation peer's models changed", "peer", c.name, "models", len(c.models))
	}
	if err := s.catalog.Refresh(ctx); err != nil {
		s.log.Warn("the catalog could not be rebuilt after a peer's models changed", "err", err)
	}
}

// checkResult is one check of one peer.
type checkResult struct {
	health  Health
	cluster string
	models  []PeerModel
	// err is set when the models could not be read; health says the peer
	// is down then too.
	err error
}

// check asks a peer whether it is ready, timed, and what it offers this
// cluster's key.
func (s *Service) check(ctx context.Context, rec store.FederationPeer) checkResult {
	ctx, cancel := context.WithTimeout(ctx, s.cfg.Timeout)
	defer cancel()
	res := checkResult{health: Health{Checked: time.Now()}}
	fail := func(err error) checkResult {
		res.health.Up, res.health.Error, res.err = false, err.Error(), err
		return res
	}

	// /ready first, and timed: it is open, cheap, and says whether the peer
	// has anything of its own to serve with.
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rec.URL+"/ready", nil)
	if err != nil {
		return fail(err)
	}
	start := time.Now()
	resp, err := s.cfg.Client.Do(req)
	if err != nil {
		return fail(fmt.Errorf("the peer did not answer: %w", err))
	}
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<12))
	resp.Body.Close()
	res.health.Latency = time.Since(start)
	if resp.StatusCode != http.StatusOK {
		return fail(fmt.Errorf("the peer is not ready (%s): it has nothing of its own to serve with", resp.Status))
	}

	key := ""
	if rec.KeyFile != "" {
		b, err := secrets.Read(ctx, rec.KeyFile)
		if err != nil {
			return fail(fmt.Errorf("the key for the peer could not be read: %w", err))
		}
		key = strings.TrimSpace(string(b))
	}
	req, err = http.NewRequestWithContext(ctx, http.MethodGet, rec.URL+"/v1/models", nil)
	if err != nil {
		return fail(err)
	}
	if key != "" {
		req.Header.Set("Authorization", "Bearer "+key)
	}
	resp, err = s.cfg.Client.Do(req)
	if err != nil {
		return fail(fmt.Errorf("the peer's models could not be read: %w", err))
	}
	defer resp.Body.Close()
	switch {
	case resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden:
		return fail(fmt.Errorf("the peer refused this cluster's key (%s): check the key file", resp.Status))
	case resp.StatusCode != http.StatusOK:
		return fail(fmt.Errorf("the peer's models could not be read: %s", resp.Status))
	}
	cluster := resp.Header.Get(ClusterHeader)
	if cluster == "" {
		return fail(errors.New("the peer does not say its cluster id: is it an llmgw gateway of this version or later?"))
	}
	var list struct {
		Data []struct {
			ID            string               `json:"id"`
			Description   string               `json:"description"`
			ContextLength int                  `json:"context_length"`
			MaxOutput     int                  `json:"max_output_tokens"`
			Capabilities  catalog.Capabilities `json:"capabilities"`
			AliasOf       string               `json:"alias_of"`
			Peer          string               `json:"peer"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(&list); err != nil {
		return fail(fmt.Errorf("the peer's model list is not what an llmgw gateway answers: %w", err))
	}
	for _, m := range list.Data {
		// Not an alias, which is a second name for a model already here,
		// nor a model the peer itself has from a peer of its own: federating
		// those would grow a name a hop, A/B/A/..., each time two clusters
		// read each other's lists.
		if m.AliasOf != "" || m.Peer != "" || !validModelID(m.ID) {
			continue
		}
		res.models = append(res.models, PeerModel{ID: m.ID, Description: m.Description,
			ContextLength: m.ContextLength, MaxOutput: m.MaxOutput, Capabilities: m.Capabilities})
	}
	slices.SortFunc(res.models, func(a, b PeerModel) int { return strings.Compare(a.ID, b.ID) })
	res.cluster = cluster
	res.health.Up = true
	return res
}

// validModelID is a model name this cluster can carry: printable, no spaces,
// and short enough to be a name.
func validModelID(id string) bool {
	if id == "" || len(id) > 200 {
		return false
	}
	return !strings.ContainsFunc(id, func(r rune) bool { return r <= ' ' || r == 0x7f })
}

// Run checks every peer at once each interval until ctx ends, the first
// time straight away.
func (s *Service) Run(ctx context.Context) {
	t := time.NewTicker(s.cfg.Interval)
	defer t.Stop()
	for {
		s.SyncAll(ctx)
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// Status is a peer as `llmgw ctl federation list` shows it.
type Status struct {
	Name     string `json:"name"`
	URL      string `json:"url"`
	KeyFile  string `json:"key_file,omitempty"`
	SendTags bool   `json:"send_tags"`
	// Cluster is the peer's own cluster id.
	Cluster string `json:"cluster,omitempty"`
	Health
	// LatencyMs is Latency in milliseconds, for reading.
	LatencyMs float64 `json:"latency_ms"`
	// SyncedAt is when its models last changed here.
	SyncedAt  time.Time     `json:"synced_at,omitempty"`
	Models    []ModelStatus `json:"models"`
	CreatedAt time.Time     `json:"created_at"`
	CreatedBy string        `json:"created_by,omitempty"`
}

// ModelStatus is one of a peer's models: its name here, the peer's own,
// and whether the peer has stopped offering it.
type ModelStatus struct {
	Name      string `json:"name"`
	Upstream  string `json:"upstream"`
	Withdrawn bool   `json:"withdrawn,omitempty"`
}

// Status is one peer; false when there is none of the name.
func (s *Service) Status(name string) (Status, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	p, ok := s.peers[name]
	if !ok {
		return Status{}, false
	}
	return statusOf(p), true
}

// List is every peer, by name.
func (s *Service) List() []Status {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Status, 0, len(s.peers))
	for _, n := range s.namesLocked() {
		out = append(out, statusOf(s.peers[n]))
	}
	return out
}

func statusOf(p *peer) Status {
	st := Status{Name: p.rec.Name, URL: p.rec.URL, KeyFile: p.rec.KeyFile, SendTags: p.rec.SendTags,
		Cluster: p.rec.Cluster, Health: p.health, LatencyMs: float64(p.health.Latency.Microseconds()) / 1000,
		SyncedAt: p.rec.SyncedAt, CreatedAt: p.rec.CreatedAt, CreatedBy: p.rec.CreatedBy,
		Models: make([]ModelStatus, 0, len(p.models))}
	for _, m := range p.models {
		st.Models = append(st.Models, ModelStatus{Name: p.rec.Name + "/" + m.ID, Upstream: m.ID, Withdrawn: m.Withdrawn})
	}
	return st
}