llmgw/internal/router/federation.go

package router

import (
	"context"
	"net/http"
	"slices"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// Federation, as the router sees it: requests sent to another llmgw cluster
// say which clusters they have passed through, and a request that comes back
// to one of them is refused rather than sent round again.

const (
	// hopHeader lists the cluster ids a request has passed through, oldest
	// first: "3f9a…, 81c0…". Each gateway adds its own when it sends the
	// request to a peer.
	hopHeader = "x-llmgw-hop"
	// clusterHeader says this gateway's cluster id, on /v1/models, which is
	// where a peer learns it.
	clusterHeader = "x-llmgw-cluster"
	// maxHops bounds a chain of peers whatever their ids say: a backstop
	// for two clusters given the same id by a restore of one's store into
	// the other, where the loop check cannot tell them apart.
	maxHops = 4
)

// federationID is this cluster as its peers know it: its id, in the hop
// header, and its name, in the tags a peer may be sent.
type federationID struct {
	id, name string
}

// SetFederationIdentity says which cluster this gateway is. Without it,
// requests are never refused as loops and peers are sent no hop of this one.
func (r *Router) SetFederationIdentity(id, name string) {
	r.federation.Store(&federationID{id: id, name: name})
}

func (r *Router) clusterID() federationID {
	if f := r.federation.Load(); f != nil {
		return *f
	}
	return federationID{}
}

// hops reads the clusters a request has passed through.
func hops(h http.Header) []string {
	var out []string
	for _, v := range h.Values(hopHeader) {
		for _, p := range strings.Split(v, ",") {
			if p = strings.TrimSpace(p); p != "" {
				out = append(out, p)
			}
		}
	}
	return out
}

// refuseLoops answers 508 to a request that has already passed through this
// cluster, or through more clusters than any sensible chain of peers has.
//
// Before anything else, authentication included: such a request was sent by
// a peer that had no capacity, and answering it at once is what lets that
// peer move on to its next fallback instead of waiting out a queue here.
func (r *Router) refuseLoops(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.Header.Get(hopHeader) == "" || !strings.HasPrefix(req.URL.Path, "/v1/") {
			next.ServeHTTP(w, req)
			return
		}
		seen := hops(req.Header)
		self := r.clusterID().id
		switch {
		case self != "" && slices.Contains(seen, self):
			writeErrorCode(w, http.StatusLoopDetected, "invalid_request_error", "federation_loop",
				"this request has already passed through this cluster; it is not sent round again")
		case len(seen) >= maxHops:
			writeErrorCode(w, http.StatusLoopDetected, "invalid_request_error", "federation_loop",
				"this request has passed through too many federated clusters")
		default:
			next.ServeHTTP(w, req)
			return
		}
		r.m.observeRequest("", "", http.StatusLoopDetected, 0, nil, 0)
	})
}

// withoutVisited drops the endpoints of peers a request has already passed
// through: sending it there would only be refused as a loop, after a round
// trip to the other site.
func withoutVisited(eps []endpoint.Snapshot, h http.Header) []endpoint.Snapshot {
	seen := hops(h)
	if len(seen) == 0 {
		return eps
	}
	return slices.DeleteFunc(slices.Clone(eps), func(e endpoint.Snapshot) bool {
		return e.Peer != "" && e.PeerCluster != "" && slices.Contains(seen, e.PeerCluster)
	})
}

// peerHeaders are what a request sent to a peer carries beyond what any
// upstream is sent: the hop header with this cluster added, and, when the
// peer is set to receive them, tags naming this cluster and the project.
//
// The peer sees one key, this cluster's, so without the tags every request
// from here is one tenant to it. They are off by default because they tell
// another site who in this one is asking.
func (r *Router) peerHeaders(ctx context.Context, ep endpoint.Snapshot, h http.Header) http.Header {
	out := http.Header{}
	self := r.clusterID()
	chain := hops(h)
	if self.id != "" {
		chain = append(chain, self.id)
	}
	if len(chain) > 0 {
		out.Set(hopHeader, strings.Join(chain, ", "))
	}
	if ep.PeerTags {
		tags := []string{"federated_from=" + tagValue(self.name)}
		if p, ok := auth.FromContext(ctx); ok {
			for _, kv := range [][2]string{{"org", p.OrgID}, {"team", p.TeamID}, {"project", p.ProjectID}} {
				if kv[1] != "" {
					tags = append(tags, kv[0]+"="+tagValue(kv[1]))
				}
			}
		}
		out.Set(tagsHeader, strings.Join(tags, ","))
	}
	return out
}

// tagValue keeps a value inside one tag: x-llmgw-tags is comma-separated.
func tagValue(v string) string {
	v = strings.NewReplacer(",", "_", "=", "_").Replace(v)
	if len(v) > maxTagValue {
		v = v[:maxTagValue]
	}
	return v
}

// localEndpoints counts the endpoints that are not a federated peer's: what
// this cluster can serve with itself.
func (r *Router) localEndpoints() int {
	n := 0
	for _, s := range r.reg.Snapshots("") {
		if s.Peer == "" {
			n++
		}
	}
	return n
}