crossbar/docs/plans/v2/_files/internal/proxy/ctxguard_test.go

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

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 (same first user message), later turn well under 4096: stays.
	longer := strings.Replace(body, `"content":"`, `"content":"`+strings.Repeat("x ", 500), 1)
	resp = r.post("/r/v1/chat/completions", longer)
	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 recorded as an event.
	huge := strings.Replace(body, `"content":"`, `"content":"`+strings.Repeat("x ", 30000), 1)
	resp = r.post("/r/v1/chat/completions", huge)
	drain(resp)
	if resp.StatusCode != 200 || resp.Header.Get(proxy.HostHeader) != "big" {
		t.Fatalf("outgrown turn: %d %q, want 200 from big", resp.StatusCode, resp.Header.Get(proxy.HostHeader))
	}
	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))
	}
}