crossbar/internal/proxy/proxy.go

// Package proxy is the routing reverse proxy. It takes /{route}/v1/…, picks a host for the
// conversation from its ordered list using a lease table (or the health table alone), queues per
// (host, model), forwards the request streaming the answer back as it arrives, tees the response to
// read usage/timings, marks a host down when a forward fails, and records one accounting row per
// request.
package proxy

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"git.wntrmute.dev/kyle/crossbar/internal/config"
	"git.wntrmute.dev/kyle/crossbar/internal/fingerprint"
	"git.wntrmute.dev/kyle/crossbar/internal/health"
	"git.wntrmute.dev/kyle/crossbar/internal/lease"
	"git.wntrmute.dev/kyle/crossbar/internal/limiter"
	"git.wntrmute.dev/kyle/crossbar/internal/store"
)

const (
	MaxBody     = 16 << 20
	HostHeader  = "X-Crossbar-Host"
	LeaseHeader = "X-Crossbar-Lease" // "new" or "reused"
	RouteHeader = "X-Crossbar-Route" // client may name the route here instead of the path
	CtxHeader   = "X-Crossbar-Ctx"   // "moved:<old>new" when the context was relocated
)

// errBodyTooLarge is returned when a request body exceeds MaxBody during the model peek.
var errBodyTooLarge = errors.New("body too large")

// Health is what the proxy needs from the health table.
type Health interface {
	Get(name string) (health.Status, bool)
	MarkDown(name, reason string)
}

// Recorder is what the proxy needs to write an accounting row; *store.Store satisfies it.
type Recorder interface {
	RecordRequest(store.Request) error
}

// Waker rouses a sleeping host. *wake.Waker satisfies it.
type Waker interface {
	Wake(ctx context.Context, host string) bool
}

// Handler forwards requests for a route to one of the route's healthy hosts, choosing by lease when
// one is configured and by health alone otherwise.
type Handler struct {
	cfg    *config.Config
	health Health
	leases *lease.Table
	lim    *limiter.Limiter
	rec    Recorder
	log    *slog.Logger
	waker  Waker
}

// SetWaker installs the waker the consults when a route has no healthy host left. A nil waker
// (the default) leaves the ErrNoHost answer as it was in v0: a plain 503.
func (p *Handler) SetWaker(w Waker) {
	p.waker = w
}

// New builds a Handler. A nil logger becomes slog.Default(). With a nil lease table it behaves like
// the v0 proxy: first healthy host, no queueing, no recording; nil limiter and recorder are likewise
// no-ops.
func New(cfg *config.Config, h Health, leases *lease.Table, lim *limiter.Limiter, rec Recorder, log *slog.Logger) *Handler {
	if log == nil {
		log = slog.Default()
	}
	return &Handler{cfg: cfg, health: h, leases: leases, lim: lim, rec: rec, log: log}
}

// SplitRoute takes the first path segment as the route. "/a/v1/x" -> ("a", "/v1/x", true); "/a" and
// "/a/" -> ("a", "/", true); "/", "//x", "", "noslash/v1" -> ("", "", false). The query string, if
// present, is kept in rest.
func SplitRoute(path string) (route, rest string, ok bool) {
	if path == "" || path[0] != '/' {
		return "", "", false
	}
	after := path[1:]
	slash := strings.IndexByte(after, '/')
	if slash == -1 {
		if after == "" {
			return "", "", false
		}
		return after, "/", true
	}
	if after[:slash] == "" {
		return "", "", false
	}
	return after[:slash], after[slash:], true
}

// Choose returns the first host in order that is healthy and lists model in Loaded; failing that,
// the first healthy host. ok is false if none. model may be "".
func Choose(hosts []string, model string, h Health) (string, bool) {
	for _, name := range hosts {
		s, ok := h.Get(name)
		if !ok || !s.Healthy || !hasModel(s.Loaded, model) {
			continue
		}
		return name, true
	}
	for _, name := range hosts {
		s, ok := h.Get(name)
		if ok && s.Healthy {
			return name, true
		}
	}
	return "", false
}

func hasModel(loaded []string, model string) bool {
	if model == "" {
		return false
	}
	for _, m := range loaded {
		if m == model {
			return true
		}
	}
	return false
}

// allowedPath reports whether rest may be proxied: under /v1/, or the admin and control-plane paths.
func allowedPath(rest string) bool {
	return strings.HasPrefix(rest, "/v1/") || rest == "/health" || rest == "/props" || rest == "/slots" || rest == "/tokenize"
}

// route resolves the route name and the upstream path (rest) from the request, honouring the
// optional route header. code is non-zero when the request must be answered; msg is the JSON error
// text for that code.
func (p *Handler) route(r *http.Request) (route, rest string, code int, msg string) {
	hdr := r.Header.Get(RouteHeader)
	if hdr != "" {
		rest := r.URL.Path
		// A path that also carries a (different) route name is a client mistake: the header is the
		// operator's intent, but the path disagrees. Compare concrete names.
		if rname, _, ok := SplitRoute(rest); ok {
			if _, _, rok := p.cfg.Route(rname); rok && rname != hdr {
				return "", "", http.StatusBadRequest, "conflicting route"
			}
		}
		if _, _, ok := p.cfg.Route(hdr); !ok {
			return "", "", http.StatusNotFound, "unknown route"
		}
		if !allowedPath(rest) {
			return "", "", http.StatusNotFound, "not found"
		}
		return hdr, rest, 0, ""
	}
	route, rest, ok := SplitRoute(r.URL.Path)
	if !ok {
		return "", "", http.StatusBadRequest, "missing route"
	}
	if _, _, ok := p.cfg.Route(route); !ok {
		return "", "", http.StatusNotFound, "unknown route"
	}
	if !allowedPath(rest) {
		return "", "", http.StatusNotFound, "not found"
	}
	return route, rest, 0, ""
}

// peekModel reads a non-GET/HEAD body up to MaxBody+1 bytes, restores it on the request, and returns
// the top-level "model" and the body itself (for fingerprinting). A non-JSON body or one without a
// model gives "". A body larger than MaxBody returns errBodyTooLarge.
func peekModel(r *http.Request) (string, []byte, error) {
	if r.Method == http.MethodGet || r.Method == http.MethodHead {
		return "", nil, nil
	}
	if r.Body == nil || r.Body == http.NoBody {
		return "", nil, nil
	}
	body, err := io.ReadAll(io.LimitReader(r.Body, MaxBody+1))
	if err != nil {
		return "", nil, err
	}
	if len(body) > MaxBody {
		return "", nil, errBodyTooLarge
	}
	r.Body = io.NopCloser(bytes.NewReader(body))
	r.ContentLength = int64(len(body))

	var req struct {
		Model string `json:"model"`
	}
	_ = json.Unmarshal(body, &req)
	return req.Model, body, nil
}

// ServeHTTP routes, then serves the request over the shared flow below: fingerprint, lease a host,
// queue per (host, model), forward with streaming, tee the response for usage/timings, and record
// one accounting row. Every error answer is JSON {"error":"…"}.
func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	route, rest, code, msg := p.route(r)
	if code != 0 {
		p.writeError(w, code, msg)
		return
	}
	p.serve(w, r, route, rest)
}

// ForRoute serves route alone, for its dedicated listener: every request there is this route, the
// whole path passed upstream as it is (no route segment is taken from it). A request from an unknown
// route is 404; a different X-Crossbar-Route header is a 400 (an equal one is ignored); a prefixed
// path or a path outside the allowed set is 404. Then it runs the same serve flow as ServeHTTP.
func (p *Handler) ForRoute(route string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if _, _, ok := p.cfg.Route(route); !ok {
			p.writeError(w, http.StatusNotFound, "unknown route")
			return
		}
		if hdr := r.Header.Get(RouteHeader); hdr != "" && hdr != route {
			p.writeError(w, http.StatusBadRequest, "conflicting route")
			return
		}
		rest := r.URL.Path
		if !allowedPath(rest) {
			p.writeError(w, http.StatusNotFound, "not found")
			return
		}
		p.serve(w, r, route, rest)
	})
}

// serve fingerprints, leases a host, queues per (host, model), forwards with streaming, tees the
// response for usage/timings, and records one accounting row. Both ServeHTTP and ForRoute reach this
// after they have settled the route and the upstream path (rest).
func (p *Handler) serve(w http.ResponseWriter, r *http.Request, route, rest string) {
	routeCfg, _, _ := p.cfg.Route(route)
	isControl := isControlCall(r.Method, rest)

	model, body, err := peekModel(r)
	if err != nil {
		p.writeError(w, http.StatusRequestEntityTooLarge, "body too large")
		return
	}
	model = resolveModel(model, r, routeCfg)
	fp := fingerprint.Of(body)
	// A route-affinity route puts every request (chat or control) on one lease per model, so the
	// lease key's fingerprint is "" for all of them; the real fingerprint is kept for the row below.
	leaseFP := fp
	if routeCfg.PerRoute() {
		leaseFP = ""
	}
	started := time.Now()

	// v0 compatibility path: no lease table, no limiter, no recording.
	if p.leases == nil {
		name, ok := Choose(routeCfg.Hosts, model, p.health)
		if !ok {
			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
			return
		}
		p.forward(w, r, route, name, "", rest, fp, model, started, 0, 0, "", isControl)
		return
	}

	// Lease. The route's ordered host list is the candidate set.
	host, reused, err := p.leases.Acquire(lease.Key{Route: route, FP: leaseFP, Model: model}, routeCfg.Hosts, time.Now())
	if err != nil {
		switch {
		case errors.Is(err, lease.ErrNoHost):
			// No host healthy. Ask a waker to rouse a sleeping one; it answers
			// (served or 503) when it has had a turn, else falls through to the
			// plain 503.
			if p.waker != nil && p.wakeOnErrNoHost(w, r, route, routeCfg, rest, model, fp, started, lease.Key{Route: route, FP: leaseFP, Model: model}) {
				return
			}
			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
		case errors.Is(err, lease.ErrPinnedDown):
			p.writeError(w, http.StatusServiceUnavailable, "pinned host down")
		default:
			p.writeError(w, http.StatusBadGateway, "upstream failed")
		}
		return
	}

	// Slot, context guard and forward, holding the slot for the leased host.
	p.serveLeased(w, r, route, routeCfg, rest, model, fp, started, host, reused, isControl)
}

// serveLeased queues the request against the leased host's limiter, runs the
// context guard, and forwards. The slot is held for the originally leased host
// even if the guard relocates the lease: the guard already moved it.
func (p *Handler) serveLeased(w http.ResponseWriter, r *http.Request, route string, routeCfg config.Route, rest, model, fp string, started time.Time, host string, reused bool, control bool) {
	// A control-plane call follows the lease but takes no slot and runs no
	// context guard: it is sent beside its own stream, so it must never wait
	// for or hold a slot. forward writes no row and logs at Debug for it.
	if control {
		p.forward(w, r, route, host, leaseState(reused), rest, fp, model, started, 0, 0, "", true)
		return
	}
	// Slot or track. A queue = false route leaves queueing to the client's own
	// llama-server slot: crossbar only counts the request on the host, never
	// holding it or refusing it.
	var waited time.Duration
	var release func()
	if routeCfg.Queues() {
		var err error
		release, waited, err = p.lim.Acquire(r.Context(), host, model)
		if err != nil {
			if errors.Is(err, limiter.ErrQueueFull) {
				p.writeRecord(store.Request{
					Route:   route,
					FP:      fp,
					Model:   model,
					Host:    host,
					Started: started,
					TotalMs: time.Since(started).Milliseconds(),
					Status:  http.StatusServiceUnavailable,
					Err:     "queue full",
				})
				p.writeError(w, http.StatusServiceUnavailable, "queue full")
				return
			}
			p.log.Warn("request", "route", route, "host", host, "method", r.Method, "path", rest, "status", 499)
			p.writeRecord(store.Request{
				Route:   route,
				FP:      fp,
				Model:   model,
				Host:    host,
				Started: started,
				TotalMs: time.Since(started).Milliseconds(),
				Status:  499,
				Err:     "client cancelled while queued",
			})
			return
		}
	} else {
		release = p.lim.Track(host, model)
	}
	defer release()

	// Context guard: if the prompt does not fit the leased host's per-slot
	// context, move the conversation to a host where it fits, else answer 400.
	now := time.Now()
	host, header, _, done := p.guard(w, r, routeCfg.Hosts, host, route, model, fp, started)
	if done {
		return
	}
	p.forward(w, r, route, host, leaseState(reused), rest, fp, model, now, waited, 0, header, false)
}

// wakeOnErrNoHost answers the request when no host was healthy. It asks, in
// route order, each candidate with a wake target to rouse itself; a host that
// wakes is leased once more and then served. When none wakes, it answers 503
// with the hosts it tried. It returns true when the request has been answered.
func (p *Handler) wakeOnErrNoHost(w http.ResponseWriter, r *http.Request, route string, routeCfg config.Route, rest, model, fp string, started time.Time, key lease.Key) bool {
	var tried []string
	for _, name := range routeCfg.Hosts {
		if p.cfg.Hosts[name].Wake == nil {
			continue
		}
		tried = append(tried, name)
		if !p.waker.Wake(r.Context(), name) {
			continue
		}
		if newHost, _, err := p.leases.Acquire(key, routeCfg.Hosts, time.Now()); err == nil {
			p.serveLeased(w, r, route, routeCfg, rest, model, fp, started, newHost, true, isControlCall(r.Method, rest))
			return true
		}
	}
	p.writeNoHealthyHost(w, route, model, fp, started, tried)
	return true
}