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