Proxy: move or refuse prompts that do not fit the leased host's context

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

internal/proxy/ctxguard.go

@@ -0,0 +1,159 @@
+package proxy
+
+import (
+	"encoding/json"
+	"net/http"
+	"time"
+
+	"git.wntrmute.dev/kyle/crossbar/internal/config"
+	"git.wntrmute.dev/kyle/crossbar/internal/lease"
+	"git.wntrmute.dev/kyle/crossbar/internal/store"
+)
+
+// movedHeader announces a context-driven move in the response header:
+// "moved: old><new". The client learns which host served it.
+func movedHeader(from, to string) string {
+	return "moved:" + from + "><" + to
+}
+
+// drainer is the drain flag the guard consults when rule 3 prefers a host not
+// being taken out of service. *health.Table (used in tests and production) does
+// not implement it, so the assertion is a no-op there; a host table with a
+// drain set would.
+type drainer interface {
+	Draining(string) bool
+}
+
+// ctxEstimate returns the prompt size the guard reasons about:
+// int(float64(len(body))/4*1.2), or 0 for a bodyless request (GET/HEAD). The
+// body was already read by peekModel and restored, so its length is known
+// without reading again.
+func ctxEstimate(r *http.Request) int {
+	if r.Method == http.MethodGet || r.Method == http.MethodHead {
+		return 0
+	}
+	if r.Body == nil || r.Body == http.NoBody {
+		return 0
+	}
+	return int(float64(r.ContentLength) / 4 * 1.2)
+}
+
+// guard runs the context-size guard's rules after a lease and a slot are held.
+// If the prompt does not fit the leased host's per-slot context, it moves the
+// conversation to a host where it fits (updating the lease) and returns that
+// host with a moved header, or answers 400 when no host fits. newHost is the
+// host to forward to (== host when nothing moved); done is true when the caller
+// must return without forwarding.
+func (p *Handler) guard(w http.ResponseWriter, r *http.Request, hosts []string, host, route, model, fp string, started time.Time) (string, string, int, bool) {
+	estimate := ctxEstimate(r)
+	if estimate == 0 {
+		return host, "", 0, false
+	}
+
+	// Rule 2: no move when the leased host has no context size or the estimate
+	// already fits its per-slot context.
+	psc := 0
+	if s, ok := p.health.Get(host); ok {
+		psc = s.PerSlotCtx()
+	}
+	if psc == 0 || estimate <= psc {
+		return host, "", 0, false
+	}
+
+	// Rule 3: move the conversation to a host where the prompt fits.
+	if newHost, ok := ctxFitHost(hosts, model, estimate, p.health, p.cfg); ok {
+		if err := p.leases.Move(lease.Key{Route: route, FP: fp, Model: model}, newHost, time.Now()); err != nil {
+			p.writeError(w, http.StatusBadGateway, "upstream failed")
+			return host, "", estimate, true
+		}
+		return newHost, movedHeader(host, newHost), estimate, false
+	}
+
+	// Rule 4: no host fits. Answer 400 with the estimate and the largest
+	// available per-slot context, and record the row.
+	p.refuseCtx(w, host, route, model, fp, started, estimate, largestSlotCtx(hosts, p.health))
+	return host, "", estimate, true
+}
+
+// ctxFitHost walks the route's ordered candidate hosts for rule 3: the first
+// healthy, non-draining host whose per-slot context fits the estimate,
+// preferring a host that has the model loaded, else one that can serve it. ok
+// is false when none fits.
+func ctxFitHost(hosts []string, model string, estimate int, h Health, cfg *config.Config) (string, bool) {
+	var drn drainer
+	if d, ok := h.(drainer); ok {
+		drn = d
+	}
+
+	// First pass: the model is resident and the per-slot context fits.
+	for _, name := range hosts {
+		s, ok := h.Get(name)
+		if !ok || !s.Healthy || isDraining(drn, name) {
+			continue
+		}
+		if s.PerSlotCtx() >= estimate && contains(s.Loaded, model) {
+			return name, true
+		}
+	}
+
+	// Second pass: the host is configured to serve the model and the per-slot
+	// context fits.
+	for _, name := range hosts {
+		s, ok := h.Get(name)
+		if !ok || !s.Healthy || isDraining(drn, name) {
+			continue
+		}
+		if s.PerSlotCtx() >= estimate && cfg.Serves(name, model) {
+			return name, true
+		}
+	}
+	return "", false
+}
+
+// largestSlotCtx is the largest PerSlotCtx across the route's healthy hosts, or
+// 0 when none is healthy or reports a context size.
+func largestSlotCtx(hosts []string, h Health) int {
+	best := 0
+	for _, name := range hosts {
+		s, ok := h.Get(name)
+		if !ok || !s.Healthy {
+			continue
+		}
+		if psc := s.PerSlotCtx(); psc > best {
+			best = psc
+		}
+	}
+	return best
+}
+
+// refuseCtx answers the 400 the guard's rule 4: the JSON body carries the
+// estimate and the largest available per-slot context, plus the error text. It
+// records the accounting row (status 400, Err "prompt too large") and never
+// marks the host down.
+func (p *Handler) refuseCtx(w http.ResponseWriter, host, route, model, fp string, started time.Time, estimate, maxSlot int) {
+	p.writeRecord(store.Request{
+		Route:   route,
+		FP:      fp,
+		Model:   model,
+		Host:    host,
+		Started: started,
+		TotalMs: time.Since(started).Milliseconds(),
+		Status:  http.StatusBadRequest,
+		Err:     "prompt too large",
+	})
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusBadRequest)
+	_ = json.NewEncoder(w).Encode(map[string]any{
+		"error":    "prompt too large",
+		"estimate": estimate,
+		"max":      maxSlot,
+	})
+}
+
+// isDraining reports whether a drain-capable host table marks name as draining.
+func isDraining(d drainer, name string) bool {
+	if d == nil {
+		return false
+	}
+	return d.Draining(name)
+}