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/lease/lease.go

@@ -194,6 +194,29 @@ func (t *Table) Acquire(k Key, candidates []string, now time.Time) (host string,
 	return host, false, nil
 }
 
+// Move relocates k to host, dropping any existing lease for k (on any host). It re-leases k onto
+// host, deletes the old row, and records a ctx event carrying the old and new hosts. It returns an
+// error only from the persister, rolling back the in-memory lease on save failure.
+func (t *Table) Move(k Key, host string, now time.Time) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	var from string
+	if l, exists := t.leases[k]; exists {
+		from = l.Host
+		delete(t.leases, k)
+		_ = t.p.DeleteLease(k.Route, k.FP, k.Model)
+	}
+	l := &Lease{k, host, store.Active, now, now}
+	t.leases[k] = l
+	if err := t.save(l); err != nil {
+		delete(t.leases, k)
+		return fmt.Errorf("lease: %w", err)
+	}
+	t.event(now, k, store.ReasonCtx, from, host)
+	return nil
+}
+
 // Candidates records hosts as seen for route (idempotent), so Pin can accept a host the route
 // is configured for before any request has used it. cmd/crossbar calls it for every route at
 // start; the admin handler calls it before Pin.

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)
+}

internal/proxy/ctxguard_test.go

@@ -0,0 +1,148 @@
+package proxy_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"git.wntrmute.dev/kyle/crossbar/internal/proxy"
+)
+
+// ctxUpstream is a fake router that reports a context size in /props and echoes completions.
+func ctxUpstream(t *testing.T, name string, nCtx, slots int) *upstream {
+	u := &upstream{name: name}
+	mux := http.NewServeMux()
+	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"status":"ok"}`) })
+	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"data":[{"id":"shared"}]}`) })
+	mux.HandleFunc("/props", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, `{"default_generation_settings":{"n_ctx":%d},"total_slots":%d}`, nCtx, slots)
+	})
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		u.hits.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
+	})
+	u.srv = httptest.NewServer(mux)
+	t.Cleanup(u.srv.Close)
+	return u
+}
+
+const ctxHosts = `
+listen = "127.0.0.1:1"
+queue_max = 2
+[hosts.small]
+base_url = %q
+weight = 10.0
+models = { "shared" = { parallel = 2 } }
+[hosts.big]
+base_url = %q
+weight = 1.0
+models = { "shared" = { parallel = 1 } }
+[routes.r]
+hosts = ["small", "big"]
+default_model = "shared"
+`
+
+// bodyOfTokens builds a chat body whose byte size implies roughly n tokens under the guard's
+// estimate (bytes/4 × 1.2): n tokens ≈ 3.33 n bytes ≈ 2n/3 five-byte words.
+func bodyOfTokens(n int) string {
+	text := strings.Repeat("word ", n*2/3)
+	return fmt.Sprintf(`{"model":"shared","stream":false,"messages":[{"role":"user","content":"%s"}]}`, text)
+}
+
+func TestOversizedPromptMovesToAHostWhereItFits(t *testing.T) {
+	small := ctxUpstream(t, "small", 8192, 2) // 4096 per slot
+	big := ctxUpstream(t, "big", 131072, 1)   // 131072 per slot
+	r := newRig(t, ctxHosts, small, big)
+	// A small prompt starts on `small` (weight 10).
+	resp := r.post("/r/v1/chat/completions", bodyOfTokens(100))
+	drain(resp)
+	if resp.Header.Get(proxy.HostHeader) != "small" {
+		t.Fatalf("small prompt went to %q, want small", resp.Header.Get(proxy.HostHeader))
+	}
+	// A new conversation with ~10k tokens does not fit small's 4096-token slot: it must be
+	// placed on big, with the reason visible in a header.
+	resp = r.post("/r/v1/chat/completions", bodyOfTokens(10000))
+	drain(resp)
+	if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "big" {
+		t.Fatalf("oversized prompt: %d from %q, want 200 from big", resp.StatusCode, resp.Header.Get(proxy.HostHeader))
+	}
+	if got := resp.Header.Get(proxy.CtxHeader); !strings.HasPrefix(got, "moved") {
+		t.Errorf("%s = %q, want moved:… ", proxy.CtxHeader, got)
+	}
+}
+
+func TestOversizedPromptWithNoFitIs400(t *testing.T) {
+	small := ctxUpstream(t, "small", 8192, 2)
+	tiny := ctxUpstream(t, "big", 4096, 2) // also too small
+	r := newRig(t, ctxHosts, small, tiny)
+	resp := r.post("/r/v1/chat/completions", bodyOfTokens(10000))
+	body := drain(resp)
+	if resp.StatusCode != http.StatusBadRequest {
+		t.Fatalf("status %d body %s, want 400", resp.StatusCode, body)
+	}
+	var e map[string]any
+	if err := json.Unmarshal([]byte(body), &e); err != nil || e["error"] != "prompt too large" {
+		t.Fatalf("body = %s, want error 'prompt too large'", body)
+	}
+	if est, _ := e["estimate"].(float64); est < 8000 || est > 13000 {
+		t.Errorf("estimate = %v, want roughly 10000 tokens", e["estimate"])
+	}
+	if max, _ := e["max"].(float64); max != 4096 {
+		t.Errorf("max = %v, want the largest per-slot context among the route's hosts (4096)", e["max"])
+	}
+	if small.hits.Load()+tiny.hits.Load() != 0 {
+		t.Errorf("a refused prompt must not reach any upstream")
+	}
+}
+
+func TestUnknownContextNeverBlocks(t *testing.T) {
+	// /props missing on both hosts: NCtx 0 means "unknown", and the guard must stay out of the way.
+	alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+	r := newRig(t, twoHosts, alpha, beta)
+	resp := r.post("/r/v1/chat/completions", bodyOfTokens(50000))
+	drain(resp)
+	if resp.StatusCode != 200 || resp.Header.Get(proxy.CtxHeader) != "" {
+		t.Errorf("unknown context: %d %q, want 200 and no ctx header", resp.StatusCode, resp.Header.Get(proxy.CtxHeader))
+	}
+}
+
+// grow appends later turns to a conversation body without touching its system prompt or first
+// user message, so the fingerprint — and therefore the lease — stays the same.
+func grow(body string, words int) string {
+	turn := `,{"role":"assistant","content":"ok"},{"role":"user","content":"` + strings.Repeat("x ", words) + `"}`
+	return strings.Replace(body, `]}`, turn+`]}`, 1)
+}
+
+func TestStickyLeaseSurvivesGrowthUntilItDoesNotFit(t *testing.T) {
+	small := ctxUpstream(t, "small", 8192, 2)
+	big := ctxUpstream(t, "big", 131072, 1)
+	r := newRig(t, ctxHosts, small, big)
+	body := bodyOfTokens(100)
+	resp := r.post("/r/v1/chat/completions", body)
+	drain(resp)
+	if resp.Header.Get(proxy.HostHeader) != "small" {
+		t.Fatal("setup: first turn must be on small")
+	}
+	// Same conversation, a later turn well under 4096 tokens: stays.
+	resp = r.post("/r/v1/chat/completions", grow(body, 500))
+	drain(resp)
+	if resp.Header.Get(proxy.HostHeader) != "small" || resp.Header.Get(proxy.LeaseHeader) != "reused" {
+		t.Errorf("turn 2: %q %q, want small reused", resp.Header.Get(proxy.HostHeader), resp.Header.Get(proxy.LeaseHeader))
+	}
+	// A turn that outgrows the slot moves the lease — once — and the move is visible in the header.
+	huge := grow(body, 30000)
+	resp = r.post("/r/v1/chat/completions", huge)
+	drain(resp)
+	if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "big" || !strings.HasPrefix(resp.Header.Get(proxy.CtxHeader), "moved") {
+		t.Fatalf("outgrown turn: %d %q ctx=%q, want 200 from big with a moved header", resp.StatusCode, resp.Header.Get(proxy.HostHeader), resp.Header.Get(proxy.CtxHeader))
+	}
+	resp = r.post("/r/v1/chat/completions", huge)
+	drain(resp)
+	if resp.Header.Get(proxy.HostHeader) != "big" || resp.Header.Get(proxy.LeaseHeader) != "reused" {
+		t.Errorf("after the move the lease is on big: %q %q", resp.Header.Get(proxy.HostHeader), resp.Header.Get(proxy.LeaseHeader))
+	}
+}

internal/proxy/forward.go

@@ -14,8 +14,10 @@ import (
 )
 
 // forward builds the reverse proxy for one host, tees the response, records the accounting row, and
-// logs. leaseState is "new" or "reused"; waited is the time spent in the queue.
-func (p *Handler) forward(w http.ResponseWriter, r *http.Request, route, host, leaseState, rest, fp, model string, started time.Time, waited time.Duration) {
+// logs. leaseState is "new" or "reused"; waited is the time spent in the queue. ctxEst is the
+// prompt size the context guard estimated (0 when the guard did not run); ctxHeader is the
+// "moved:…<host>" header to set when the guard relocated the conversation.
+func (p *Handler) forward(w http.ResponseWriter, r *http.Request, route, host, leaseState, rest, fp, model string, started time.Time, waited time.Duration, ctxEst int, ctxHeader string) {
 	hostCfg, ok := p.cfg.Hosts[host]
 	if !ok {
 		p.writeError(w, http.StatusBadGateway, "upstream failed")
@@ -28,7 +30,7 @@ func (p *Handler) forward(w http.ResponseWriter, r *http.Request, route, host, l
 	}
 
 	rev := &forwardState{started: started}
-	rp := newReverseProxy(p.health, host, leaseState, target, rest, rev)
+	rp := newReverseProxy(p.health, host, leaseState, ctxHeader, target, rest, rev)
 	rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
 
 	// ServeHTTP unwinds with http.ErrAbortHandler when a client leaves mid-stream; recover so the
@@ -71,6 +73,7 @@ func (p *Handler) forward(w http.ResponseWriter, r *http.Request, route, host, l
 		"lease", leaseState,
 		"queued_ms", waited.Milliseconds(),
 		"fp", fp8,
+		"ctx_est", ctxEst,
 		"ms", total.Milliseconds(),
 	)
 }
@@ -163,7 +166,7 @@ func (p *Handler) writeRecord(req store.Request) {
 // original query string. It flushes after every write so long server-sent-event streams are not
 // buffered, tees the response for usage/timings, and marks the host down on any transport error
 // other than a client disconnect.
-func newReverseProxy(h Health, host, leaseState string, target *url.URL, rest string, rev *forwardState) *httputil.ReverseProxy {
+func newReverseProxy(h Health, host, leaseState, ctxHeader string, target *url.URL, rest string, rev *forwardState) *httputil.ReverseProxy {
 	return &httputil.ReverseProxy{
 		Rewrite: func(pr *httputil.ProxyRequest) {
 			pr.SetURL(target)
@@ -176,6 +179,9 @@ func newReverseProxy(h Health, host, leaseState string, target *url.URL, rest st
 		ModifyResponse: func(resp *http.Response) error {
 			resp.Header.Set(HostHeader, host)
 			resp.Header.Set(LeaseHeader, leaseState)
+			if ctxHeader != "" {
+				resp.Header.Set(CtxHeader, ctxHeader)
+			}
 			rev.ttfb = time.Now()
 			rev.streamed = strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
 			t := newTee(resp.Body, rev.streamed)

internal/proxy/proxy.go

@@ -28,6 +28,7 @@ const (
 	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.
@@ -213,7 +214,7 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			p.writeError(w, http.StatusServiceUnavailable, "no healthy host")
 			return
 		}
-		p.forward(w, r, route, name, "", rest, fp, model, started, 0)
+		p.forward(w, r, route, name, "", rest, fp, model, started, 0, 0, "")
 		return
 	}
 
@@ -263,5 +264,12 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	defer release()
 
-	p.forward(w, r, route, host, leaseState(reused), rest, fp, model, started, waited)
+	// 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)
 }

internal/store/store.go

@@ -30,6 +30,7 @@ const (
 	ReasonPin       = "pin"
 	ReasonRelease   = "release"
 	ReasonDrain     = "drain"
+	ReasonCtx       = "ctx"
 )
 
 // By selects the grouping column of a Usage query.