v2.3 plan: clients that manage their own slots (control calls, route affinity, queue = false, route listeners, llama-server ctx error); given tests

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-09-25 17:36:59 -07:00
parent 4c6415819e
commit 3518e84dd7
16 changed files with 1655 additions and 0 deletions
@@ -0,0 +1,176 @@
+package proxy_test
+
+// v2.3 task 01: control-plane requests. A client that manages its own slots (Boxmaker) polls
+// /slots, reads /props, tokenizes and steers a running completion through
+// /v1/chat/completions/control. Those calls follow the route's lease like any other request but
+// must never wait for, or take, a slot: /control is sent while the client's own stream holds one.
+
+import (
+	"context"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+)
+
+// controlClient gives every control call a short deadline: a call that queues behind a full host
+// is the bug, and it must fail the test rather than hang it.
+var controlClient = &http.Client{Timeout: 2 * time.Second}
+
+func (r *rig) do(method, path, body string) *http.Response {
+	r.t.Helper()
+	var rd *strings.Reader
+	if body != "" {
+		rd = strings.NewReader(body)
+	}
+	var req *http.Request
+	var err error
+	if rd != nil {
+		req, err = http.NewRequest(method, r.front.URL+path, rd)
+		req.Header.Set("Content-Type", "application/json")
+	} else {
+		req, err = http.NewRequest(method, r.front.URL+path, nil)
+	}
+	if err != nil {
+		r.t.Fatal(err)
+	}
+	resp, err := controlClient.Do(req)
+	if err != nil {
+		r.t.Fatalf("%s %s: %v", method, path, err)
+	}
+	return resp
+}
+
+func (r *rig) rows(route string) int64 {
+	r.t.Helper()
+	counts, err := r.store.StatusCounts(time.Time{})
+	if err != nil {
+		r.t.Fatal(err)
+	}
+	var n int64
+	for _, c := range counts {
+		if c.Route == route {
+			n += c.Count
+		}
+	}
+	return n
+}
+
+// A GET names its model in the query string: /slots?model=alpha-only must reach the host that
+// has alpha-only loaded, not whichever host the route's default model would pick.
+func TestGetModelComesFromTheQuery(t *testing.T) {
+	alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+	r := newRig(t, twoHosts, alpha, beta)
+
+	resp := r.do(http.MethodGet, "/r/slots?model=alpha-only", "")
+	drain(resp)
+	if resp.StatusCode != 200 || resp.Header.Get("X-Crossbar-Host") != "alpha" {
+		t.Fatalf("GET /r/slots?model=alpha-only: %d on %q, want 200 on alpha", resp.StatusCode, resp.Header.Get("X-Crossbar-Host"))
+	}
+	if got := alpha.lastReq(); got.method != "GET" || got.path != "/slots?model=alpha-only" {
+		t.Errorf("alpha saw %s %s, want GET /slots?model=alpha-only", got.method, got.path)
+	}
+	// The same for beta-only, so a lucky default cannot pass the test.
+	resp = r.do(http.MethodGet, "/r/slots?model=beta-only", "")
+	drain(resp)
+	if resp.Header.Get("X-Crossbar-Host") != "beta" {
+		t.Errorf("GET /r/slots?model=beta-only went to %q, want beta", resp.Header.Get("X-Crossbar-Host"))
+	}
+}
+
+// /slots and /tokenize are proxied; the per-slot actions under /slots/ (save, restore, erase)
+// are not.
+func TestControlPathsAllowed(t *testing.T) {
+	alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+	r := newRig(t, twoHosts, alpha, beta)
+	for _, tc := range []struct {
+		method, path, body string
+		want               int
+	}{
+		{http.MethodGet, "/r/slots", "", 200},
+		{http.MethodGet, "/r/slots?model=shared", "", 200},
+		{http.MethodPost, "/r/tokenize", `{"model":"shared","content":"hello"}`, 200},
+		{http.MethodPost, "/r/v1/chat/completions/control", `{"id":"chatcmpl-1","action":"reasoning_end","model":"shared"}`, 200},
+		{http.MethodGet, "/r/slots/0", "", 404},
+		{http.MethodPost, "/r/slots/0?action=erase", "", 404},
+		{http.MethodPost, "/r/slots/0?action=save", `{"filename":"x"}`, 404},
+	} {
+		resp := r.do(tc.method, tc.path, tc.body)
+		body := drain(resp)
+		if resp.StatusCode != tc.want {
+			t.Errorf("%s %s: %d %s, want %d", tc.method, tc.path, resp.StatusCode, body, tc.want)
+		}
+	}
+}
+
+// With every slot on both hosts taken and the queue full, control-plane calls still go straight
+// through: no 503, no wait, no slot taken, no accounting row.
+func TestControlRequestsNeverTakeASlot(t *testing.T) {
+	alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+	r := newRig(t, twoHosts, alpha, beta)
+
+	// Take every "shared" slot (parallel 2 on each host) and the one queue place per host.
+	var releases []func()
+	for _, host := range []string{"alpha", "beta"} {
+		for i := 0; i < 2; i++ {
+			rel, _, err := r.lim.Acquire(context.Background(), host, "shared")
+			if err != nil {
+				t.Fatal(err)
+			}
+			releases = append(releases, rel)
+		}
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		go func() { _, _, _ = r.lim.Acquire(ctx, host, "shared") }()
+		waitUntil(t, func() bool { return r.lim.Queued(host, "shared") == 1 })
+	}
+	defer func() {
+		for _, rel := range releases {
+			rel()
+		}
+	}()
+
+	for _, tc := range []struct{ method, path, body string }{
+		{http.MethodGet, "/r/slots?model=shared", ""},
+		{http.MethodGet, "/r/props?model=shared", ""},
+		{http.MethodHead, "/r/props?model=shared", ""},
+		{http.MethodGet, "/r/v1/models", ""},
+		{http.MethodPost, "/r/tokenize", `{"model":"shared","content":"hello"}`},
+		{http.MethodPost, "/r/v1/chat/completions/control", `{"id":"chatcmpl-1","action":"reasoning_end","model":"shared"}`},
+	} {
+		resp := r.do(tc.method, tc.path, tc.body)
+		body := drain(resp)
+		if resp.StatusCode != 200 {
+			t.Errorf("%s %s with the host full: %d %s, want 200", tc.method, tc.path, resp.StatusCode, body)
+		}
+	}
+	for _, host := range []string{"alpha", "beta"} {
+		if n := r.lim.InFlight(host, "shared"); n != 2 {
+			t.Errorf("%s in flight = %d after control calls, want 2 (control takes no slot)", host, n)
+		}
+	}
+	time.Sleep(100 * time.Millisecond) // a row is written after the answer; give a stray one time to land
+	if n := r.rows("r"); n != 0 {
+		t.Errorf("control calls wrote %d accounting rows, want 0", n)
+	}
+
+	// A chat completion on the same full route still queues or is refused as before: the bypass
+	// is for control calls only.
+	resp := r.do(http.MethodPost, "/r/v1/chat/completions", conversation(1, 1))
+	drain(resp)
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Errorf("chat on a full route: %d, want 503 (queue full)", resp.StatusCode)
+	}
+}
+
+// A chat completion is not a control call just because its path starts the same way.
+func TestChatIsNotControl(t *testing.T) {
+	alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
+	r := newRig(t, twoHosts, alpha, beta)
+	resp := r.do(http.MethodPost, "/r/v1/chat/completions", conversation(1, 1))
+	drain(resp)
+	if resp.StatusCode != 200 {
+		t.Fatalf("chat: %d", resp.StatusCode)
+	}
+	waitUntil(t, func() bool { return r.rows("r") == 1 }) // the row lands just after the answer
+}