crossbar/internal/proxy/control_test.go

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
}