crossbar/internal/proxy/ctxguard.go

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 4a: no host fits. Before refusing, ask a waker to rouse a candidate
	// whose context may grow when it comes up; a woken host takes the lease.
	if p.waker != nil {
		for _, name := range hosts {
			s, ok := p.health.Get(name)
			if !ok {
				continue
			}
			// A healthy host does not need waking; only a down host might grow
			// a larger context when it comes up.
			if s.Healthy {
				continue
			}
			// A host with a known per-slot context smaller than the estimate
			// cannot serve it no matter how it wakes.
			if psc := s.PerSlotCtx(); psc != 0 && psc < estimate {
				continue
			}
			if p.cfg.Hosts[name].Wake == nil || !p.waker.Wake(r.Context(), name) {
				continue
			}
			if err := p.leases.Move(lease.Key{Route: route, FP: fp, Model: model}, name, time.Now()); err != nil {
				p.writeError(w, http.StatusBadGateway, "upstream failed")
				return host, "", estimate, true
			}
			return name, movedHeader(host, name), 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)
}