165 lines
6.9 KiB
Go
165 lines
6.9 KiB
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)
|
||
}
|
||
// v2.3: llama-server's own shape for this error, so a client handles crossbar's refusal the
|
||
// way it handles the server's (Boxmaker keys on error.type; the error JSON must come first).
|
||
if !strings.HasPrefix(body, `{"error":`) {
|
||
t.Errorf("body must start with the error object: %s", body)
|
||
}
|
||
var e struct {
|
||
Error struct {
|
||
Code int `json:"code"`
|
||
Type string `json:"type"`
|
||
Message string `json:"message"`
|
||
NPromptTokens float64 `json:"n_prompt_tokens"`
|
||
NCtx float64 `json:"n_ctx"`
|
||
} `json:"error"`
|
||
}
|
||
if err := json.Unmarshal([]byte(body), &e); err != nil || e.Error.Code != 400 || e.Error.Type != "exceed_context_size_error" || e.Error.Message != "prompt too large" {
|
||
t.Fatalf("body = %s, want {\"error\":{\"code\":400,\"type\":\"exceed_context_size_error\",\"message\":\"prompt too large\",…}}", body)
|
||
}
|
||
if est := e.Error.NPromptTokens; est < 8000 || est > 13000 {
|
||
t.Errorf("n_prompt_tokens = %v, want roughly 10000 tokens", est)
|
||
}
|
||
if max := e.Error.NCtx; max != 4096 {
|
||
t.Errorf("n_ctx = %v, want the largest per-slot context among the route's hosts (4096)", max)
|
||
}
|
||
if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "application/json") {
|
||
t.Errorf("Content-Type = %q, want application/json", ct)
|
||
}
|
||
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))
|
||
}
|
||
}
|