llmgw/internal/router/router_test.go

package router

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// catModel builds a minimal catalog entry for router tests, which care only
// about names, aliases and fallbacks.
func catModel(name string, aliases []string, fallbacks ...string) catalog.Model {
	return catalog.Model{
		Name: name, Source: "/models/" + name, Aliases: aliases, Fallbacks: fallbacks,
		Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}},
	}
}

// fakeEngine serves the minimum an endpoint needs: the metric contract plus a
// scripted chat response.
type fakeEngine struct {
	srv   *httptest.Server
	calls atomic.Int64
	// status, when non-zero, is returned instead of a completion.
	status atomic.Int64
	// last is the body of the last completion asked for.
	last atomic.Value
	// traceparent is the one the last completion came with.
	traceparent atomic.Value
}

func newFakeEngine(t *testing.T, model string) *fakeEngine {
	t.Helper()
	f := &fakeEngine{}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n")
		fmt.Fprintf(w, "# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n")
		fmt.Fprintf(w, "# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0.1\n")
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		f.calls.Add(1)
		f.traceparent.Store(r.Header.Get("traceparent"))
		if b, err := io.ReadAll(r.Body); err == nil {
			f.last.Store(b)
		}
		if code := int(f.status.Load()); code != 0 {
			http.Error(w, "unavailable", code)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model)
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

func newTestRouter(t *testing.T, cfg config.Server) *Router {
	t.Helper()
	rt, err := New(cfg, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	rt.Start(t.Context())
	t.Cleanup(rt.Stop)
	return rt
}

func post(t *testing.T, h http.Handler, body any, hdr map[string]string) *httptest.ResponseRecorder {
	t.Helper()
	b, err := json.Marshal(body)
	if err != nil {
		t.Fatal(err)
	}
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	for k, v := range hdr {
		req.Header.Set(k, v)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func baseConfig(eps ...config.Endpoint) config.Server {
	c := config.Default()
	c.Router.ScrapeInterval = 50 * time.Millisecond
	c.Router.StaleAfter = 5 * time.Second
	c.Endpoints = eps
	return c
}

func TestRoutesToHealthyEndpoint(t *testing.T) {
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{
		ID: "a", URL: e.srv.URL, Models: []string{"m"},
	}))

	rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if e.calls.Load() != 1 {
		t.Errorf("engine received %d calls", e.calls.Load())
	}
}

func TestRetriesOnAnotherEndpoint(t *testing.T) {
	bad := newFakeEngine(t, "m")
	bad.status.Store(http.StatusServiceUnavailable)
	good := newFakeEngine(t, "m")

	rt := newTestRouter(t, baseConfig(
		config.Endpoint{ID: "bad", URL: bad.srv.URL, Models: []string{"m"}},
		config.Endpoint{ID: "good", URL: good.srv.URL, Models: []string{"m"}},
	))

	// With two endpoints and two retries, the bad one must not sink the request.
	rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if good.calls.Load() == 0 {
		t.Error("the healthy endpoint was never tried")
	}
}

func TestNoRetryAfterBytesReachClient(t *testing.T) {
	// An endpoint that starts streaming then dies must not be retried: splicing
	// two responses together is worse than a short stream (plan §7.7).
	var served atomic.Int64
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		io.WriteString(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n"+
			"# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n"+
			"# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0\n")
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
		served.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		f := w.(http.Flusher)
		io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		f.Flush()
		// Then hang up mid-stream.
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()

	rt := newTestRouter(t, baseConfig(
		config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}},
		config.Endpoint{ID: "b", URL: srv.URL, Models: []string{"m"}},
	))

	rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}, "stream": true}, nil)
	if !strings.Contains(rec.Body.String(), `"content":"a"`) {
		t.Fatalf("partial stream not delivered: %q", rec.Body.String())
	}
	if n := served.Load(); n != 1 {
		t.Errorf("request was sent %d times; it must not be retried once streaming began", n)
	}
}

func TestFallbackModel(t *testing.T) {
	primary := newFakeEngine(t, "primary")
	primary.status.Store(http.StatusServiceUnavailable)
	backup := newFakeEngine(t, "backup")

	cfg := baseConfig(
		config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}},
		config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
	)
	cfg.Models = []catalog.Model{catModel("primary", nil, "backup")}
	rt := newTestRouter(t, cfg)

	rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if backup.calls.Load() == 0 {
		t.Fatal("the fallback model was never tried")
	}
	// The client asked for one model and got another; they must be told.
	if got := rec.Header().Get("x-llmgw-served-model"); got != "backup" {
		t.Errorf("x-llmgw-served-model=%q, expected \"backup\"", got)
	}
}

func TestAliasResolution(t *testing.T) {
	e := newFakeEngine(t, "llama-3.3-70b")
	cfg := baseConfig(config.Endpoint{
		ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"},
	})
	cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})}
	rt := newTestRouter(t, cfg)

	rec := post(t, rt.Handler(), map[string]any{"model": "llama", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("alias was not resolved: %d %s", rec.Code, rec.Body.String())
	}
}

func TestUnknownModelIs503(t *testing.T) {
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))

	rec := post(t, rt.Handler(), map[string]any{"model": "nope", "messages": []any{}}, nil)
	if rec.Code != http.StatusServiceUnavailable {
		t.Errorf("status %d, expected 503", rec.Code)
	}
	var body map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
		t.Fatalf("error body is not JSON: %q", rec.Body.String())
	}
	if _, ok := body["error"]; !ok {
		t.Errorf("error body is not OpenAI-shaped: %v", body)
	}
}

func TestMissingModelIs400(t *testing.T) {
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))

	rec := post(t, rt.Handler(), map[string]any{"messages": []any{}}, nil)
	if rec.Code != http.StatusBadRequest {
		t.Errorf("status %d, expected 400", rec.Code)
	}
}

func TestModelsEndpointListsAliases(t *testing.T) {
	e := newFakeEngine(t, "llama-3.3-70b")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}})
	cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})}
	rt := newTestRouter(t, cfg)

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)

	var out struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
		t.Fatal(err)
	}
	ids := map[string]bool{}
	for _, d := range out.Data {
		ids[d.ID] = true
	}
	if !ids["llama-3.3-70b"] || !ids["llama"] {
		t.Errorf("both the model and its alias should be listed: %v", ids)
	}
}

func TestModelsEndpointMarksSimulatedModels(t *testing.T) {
	// A client picking a model from this list has nothing else to go on. Both
	// entries are OpenAI-shaped and only owned_by distinguishes the one that
	// returns placeholder text, so it has to be right for aliases too.
	e := newFakeEngine(t, "real")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"real", "fake"}})
	sim := catModel("fake", []string{"fake-alias"})
	sim.Profiles = []catalog.Profile{{Name: "p", Engine: "sim", GPUClass: "h100-80g", TP: 1}}
	cfg.Models = []catalog.Model{catModel("real", []string{"real-alias"}), sim}
	rt := newTestRouter(t, cfg)

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)

	var out struct {
		Data []struct {
			ID      string `json:"id"`
			OwnedBy string `json:"owned_by"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
		t.Fatal(err)
	}
	owner := map[string]string{}
	for _, d := range out.Data {
		owner[d.ID] = d.OwnedBy
	}
	for _, name := range []string{"fake", "fake-alias"} {
		if owner[name] != "llmgw-sim" {
			t.Errorf("%s: owned_by = %q, want llmgw-sim", name, owner[name])
		}
	}
	for _, name := range []string{"real", "real-alias"} {
		if owner[name] != "llmgw" {
			t.Errorf("%s: owned_by = %q, want llmgw", name, owner[name])
		}
	}
}

func TestInjectIncludeUsagePreservesFields(t *testing.T) {
	// The one place the router rewrites a request body; nothing may be lost.
	in := []byte(`{"model":"m","stream":true,"temperature":0.7,"tools":[{"x":1}]}`)
	out, ok := injectIncludeUsage(in)
	if !ok {
		t.Fatal("injection declined a well-formed body")
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatal(err)
	}
	for _, k := range []string{"model", "stream", "temperature", "tools"} {
		if _, present := m[k]; !present {
			t.Errorf("field %q was dropped", k)
		}
	}
	opts, _ := m["stream_options"].(map[string]any)
	if opts["include_usage"] != true {
		t.Errorf("include_usage not set: %v", m["stream_options"])
	}
}

func TestInjectIncludeUsageKeepsExistingOptions(t *testing.T) {
	in := []byte(`{"model":"m","stream":true,"stream_options":{"something":"else"}}`)
	out, ok := injectIncludeUsage(in)
	if !ok {
		t.Fatal("injection declined")
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatal(err)
	}
	opts, ok := m["stream_options"].(map[string]any)
	if !ok {
		t.Fatalf("stream_options is not an object: %v", m["stream_options"])
	}
	if opts["something"] != "else" {
		t.Error("an existing stream_options field was dropped")
	}
	if opts["include_usage"] != true {
		t.Error("include_usage not set")
	}
}

func TestPromptPrefixIsBounded(t *testing.T) {
	long := strings.Repeat("x", prefixBytes*3)
	b := inferenceBody{Messages: []struct {
		Role    string          `json:"role"`
		Content json.RawMessage `json:"content"`
	}{{Role: "user", Content: json.RawMessage(`"` + long + `"`)}}}

	if got := len(promptPrefix(&b)); got > prefixBytes {
		t.Errorf("prefix is %d bytes, over the %d cap", got, prefixBytes)
	}
}

func TestSessionHeaderFallsBackToUser(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/", nil)
	if got := sessionID(req, "x-session-id", "user-42"); got != "user-42" {
		t.Errorf("expected the user field to be used, got %q", got)
	}
	req.Header.Set("x-session-id", "sess-1")
	if got := sessionID(req, "x-session-id", "user-42"); got != "sess-1" {
		t.Errorf("the header should win, got %q", got)
	}
}

func TestUnknownModelErrorNamesWhatIsAvailable(t *testing.T) {
	// Reported by a user who copied the curl from the account page and sent
	// its placeholder unchanged. "no endpoint serves \"MODEL\"" is true and
	// useless: it reads the same for a typo, a placeholder and a real outage.
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))

	rec := post(t, rt.Handler(), map[string]any{"model": "MODEL", "messages": []any{}}, nil)
	if rec.Code != http.StatusServiceUnavailable {
		t.Fatalf("status %d, expected 503", rec.Code)
	}
	var body struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
		t.Fatalf("error body is not JSON: %q", rec.Body.String())
	}
	if !strings.Contains(body.Error.Message, `"MODEL"`) {
		t.Errorf("the message does not repeat what was asked for: %q", body.Error.Message)
	}
	if !strings.Contains(body.Error.Message, "m") ||
		!strings.Contains(body.Error.Message, "available") {
		t.Errorf("the message does not say what is available: %q", body.Error.Message)
	}
}

func TestTheErrorDoesNotBecomeAPage(t *testing.T) {
	// A large catalog must not turn one error into a wall of names.
	var many []string
	for i := range 40 {
		many = append(many, fmt.Sprintf("model-%02d", i))
	}
	msg := noEndpointFor("nope", many)
	if !strings.Contains(msg, "and 30 more") {
		t.Errorf("the list was not capped: %q", msg)
	}
	if strings.Contains(msg, "model-11") {
		t.Errorf("more than the cap was listed: %q", msg)
	}
}

func TestTheErrorIsHonestWhenNothingIsServed(t *testing.T) {
	// Listing nothing as "available models are " would read as a truncated
	// message rather than an empty cluster.
	msg := noEndpointFor("m", nil)
	if !strings.Contains(msg, "no model is currently being served") {
		t.Errorf("an empty fleet is not explained: %q", msg)
	}
}

func TestAnOverLongRequestIsRefusedBeforeItReachesAnEngine(t *testing.T) {
	// The whole point of checking here: an over-long prompt used to be
	// admitted, scheduled, given a replica and a slot, and refused by the
	// engine — paying the cost of the request to learn something knowable from
	// the request itself.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	m := catModel("m", nil)
	m.ContextLength = 2048
	cfg.Models = []catalog.Model{m}
	rt := newTestRouter(t, cfg)

	before := e.calls.Load()
	rec := post(t, rt.Handler(), map[string]any{
		"model":    "m",
		"messages": []any{map[string]any{"role": "user", "content": strings.Repeat("x", 40_000)}},
	}, nil)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status %d, expected 400; body %s", rec.Code, rec.Body.String())
	}
	if got := e.calls.Load(); got != before {
		t.Errorf("the engine was called %d times for a request that could not fit", got-before)
	}
	if !strings.Contains(rec.Body.String(), "2048") {
		t.Errorf("the error does not name the limit: %s", rec.Body.String())
	}
}

func TestModelsEndpointPublishesContextAndCapabilities(t *testing.T) {
	// A client sizing its prompts had nothing to read, and a client deciding
	// whether to send tools had nothing either.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	m := catModel("m", nil)
	m.ContextLength, m.MaxOutput = 32768, 4096
	m.Capabilities = catalog.Capabilities{Tools: true}
	cfg.Models = []catalog.Model{m}
	rt := newTestRouter(t, cfg)

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)

	var out struct {
		Data []struct {
			ID           string          `json:"id"`
			Context      int             `json:"context_length"`
			MaxOutput    int             `json:"max_output_tokens"`
			Capabilities map[string]bool `json:"capabilities"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
		t.Fatal(err)
	}
	var found bool
	for _, d := range out.Data {
		if d.ID != "m" {
			continue
		}
		found = true
		if d.Context != 32768 || d.MaxOutput != 4096 {
			t.Errorf("context %d, max output %d", d.Context, d.MaxOutput)
		}
		if !d.Capabilities["tools"] || d.Capabilities["vision"] {
			t.Errorf("capabilities wrong: %v", d.Capabilities)
		}
	}
	if !found {
		t.Fatal("the model was not listed")
	}
}

func TestAModelWithNoContextPublishesNoContextField(t *testing.T) {
	// Zero would read as "no context", which is worse than saying nothing.
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))

	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if strings.Contains(rec.Body.String(), "context_length") {
		t.Errorf("an unknown context was published as a number: %s", rec.Body.String())
	}
}

type fakeActivator struct {
	onDemand map[string]bool
	asked    []string
}

func (f *fakeActivator) Want(model string) bool {
	f.asked = append(f.asked, model)
	return f.onDemand[model]
}

func TestALoadingModelIsToldToRetryRatherThanRefused(t *testing.T) {
	// "no endpoint serves this" would be a lie while a model is starting, and
	// the difference matters: a client told to retry succeeds, one told the
	// model does not exist gives up.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	// A very short wait, because the point here is what happens when it runs
	// out rather than how long it is.
	cfg.Router.ActivationWait = 50 * time.Millisecond
	rt := newTestRouter(t, cfg)
	rt.SetActivator(&fakeActivator{onDemand: map[string]bool{"sleepy": true}})

	rec := post(t, rt.Handler(), map[string]any{"model": "sleepy", "messages": []any{}}, nil)
	if rec.Code != http.StatusServiceUnavailable {
		t.Fatalf("status %d, expected 503", rec.Code)
	}
	if got := rec.Header().Get("Retry-After"); got == "" {
		t.Error("no Retry-After; a well-behaved client backs off as it would from a permanent failure")
	}
	if !strings.Contains(rec.Body.String(), "loading") {
		t.Errorf("the client is not told what is happening: %s", rec.Body.String())
	}
}

func TestAModelThatIsSimplyUnknownStillSaysSo(t *testing.T) {
	// The other half. On-demand must not turn every typo into "loading",
	// because then nothing is ever wrong and the useful error disappears.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	rt := newTestRouter(t, cfg)
	rt.SetActivator(&fakeActivator{onDemand: map[string]bool{}})

	rec := post(t, rt.Handler(), map[string]any{"model": "typo", "messages": []any{}}, nil)
	if rec.Code != http.StatusServiceUnavailable {
		t.Fatalf("status %d", rec.Code)
	}
	if rec.Header().Get("Retry-After") != "" {
		t.Error("an unknown model was presented as retryable")
	}
	if !strings.Contains(rec.Body.String(), "available models") {
		t.Errorf("the useful error was lost: %s", rec.Body.String())
	}
}

func TestAServedModelNeverWaitsForActivation(t *testing.T) {
	// A model that is up must not pay a millisecond for the feature.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	cfg.Router.ActivationWait = 10 * time.Second
	rt := newTestRouter(t, cfg)
	act := &fakeActivator{onDemand: map[string]bool{"m": true}}
	rt.SetActivator(act)

	start := time.Now()
	rec := post(t, rt.Handler(), map[string]any{"model": "m", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if elapsed := time.Since(start); elapsed > time.Second {
		t.Errorf("a served request took %s; activation is on the hot path", elapsed)
	}
	if len(act.asked) != 0 {
		t.Errorf("the activator was consulted for a model that is already serving: %v", act.asked)
	}
}

func TestWithoutAnActivatorNothingChanges(t *testing.T) {
	// Every deployment before this existed, and a router running without a
	// control plane.
	e := newFakeEngine(t, "m")
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}}))

	rec := post(t, rt.Handler(), map[string]any{"model": "nope", "messages": []any{}}, nil)
	if rec.Code != http.StatusServiceUnavailable {
		t.Errorf("status %d", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "available models") {
		t.Errorf("body = %s", rec.Body.String())
	}
}

func TestTheServedModelHeaderArrivesBeforeTheBody(t *testing.T) {
	// It used to be set in finish, which runs after the proxy has written the
	// response — so it announced a fallback to nobody. Nothing failed, because
	// the only thing asserting it was a comment.
	// Nothing serves the primary at all, which is the only condition that
	// justifies moving to another model — an upstream returning 500 is a
	// failure the fallback would share, and dispatch rightly declines to try.
	live := newFakeEngine(t, "backup")
	cfg := baseConfig(config.Endpoint{ID: "b", URL: live.srv.URL, Models: []string{"backup"}})
	m := catModel("primary", nil, "backup")
	cfg.Models = []catalog.Model{m, catModel("backup", nil)}
	rt := newTestRouter(t, cfg)

	rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if got := rec.Header().Get("x-llmgw-served-model"); got != "backup" {
		t.Errorf("served-model header = %q, want backup", got)
	}
}

type unservableActivator struct {
	fakeActivator
	why map[string]string
}

func (u *unservableActivator) Unservable(model string) (string, bool) {
	w, ok := u.why[model]
	return w, ok
}

func TestAModelThatCannotComeUpIsSaidSoAtOnce(t *testing.T) {
	// A model that fits on no node, or failed every time it started, was
	// answered "loading; retry shortly" — and a client that retried did so
	// forever.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	cfg.Router.ActivationWait = 10 * time.Second
	rt := newTestRouter(t, cfg)
	rt.SetActivator(&unservableActivator{
		fakeActivator: fakeActivator{onDemand: map[string]bool{"phi-4": true}},
		why:           map[string]string{"phi-4": "phi-4 cannot be placed on any node: needs 15 GB"},
	})

	start := time.Now()
	rec := post(t, rt.Handler(), map[string]any{"model": "phi-4", "messages": []any{}}, nil)
	if took := time.Since(start); took > 2*time.Second {
		t.Errorf("the request waited %s for a model that will not load", took)
	}
	if rec.Code != http.StatusServiceUnavailable || rec.Header().Get("Retry-After") != "" {
		t.Errorf("status %d, Retry-After %q", rec.Code, rec.Header().Get("Retry-After"))
	}
	if !strings.Contains(rec.Body.String(), "needs 15 GB") {
		t.Errorf("the reason was not given: %s", rec.Body.String())
	}
}

func TestAMadeUpModelNameDoesNotBecomeAMetricLabel(t *testing.T) {
	// Counted as sent, every request for a nonsense model minted a new series.
	e := newFakeEngine(t, "m")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"m"}})
	cfg.Models = []catalog.Model{{Name: "m", Aliases: []string{"mm"}}}
	rt := newTestRouter(t, cfg)

	if got := rt.metricModel("mm"); got != "m" {
		t.Errorf("an alias is counted as %q, not the model it names", got)
	}
	for i := 0; i < 20; i++ {
		post(t, rt.Handler(), map[string]any{"model": fmt.Sprintf("junk-%d", i), "messages": []any{}}, nil)
	}
	// Twenty names, one series: all of them "unknown".
	reg := prometheus.NewRegistry()
	reg.MustRegister(rt.m.requests)
	mfs, err := reg.Gather()
	if err != nil {
		t.Fatal(err)
	}
	series := 0
	for _, mf := range mfs {
		series += len(mf.GetMetric())
	}
	if series != 1 {
		t.Errorf("%d series after twenty made-up names", series)
	}
}

func TestAFallbackTheKeyMayNotUseIsNotServed(t *testing.T) {
	primary := newFakeEngine(t, "primary")
	primary.status.Store(http.StatusServiceUnavailable)
	backup := newFakeEngine(t, "backup")
	cfg := baseConfig(
		config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}},
		config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}},
	)
	cfg.Models = []catalog.Model{catModel("primary", nil, "backup"), catModel("backup", nil)}
	rt := newTestRouter(t, cfg)
	asKey := func(allowed ...string) http.Handler {
		p := auth.Principal{KeyID: "k1", Limits: store.Effective{AllowedModels: allowed}}
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			rt.Handler().ServeHTTP(w, r.WithContext(auth.WithPrincipal(r.Context(), p)))
		})
	}

	rec := post(t, asKey("primary"), map[string]any{"model": "primary", "messages": []any{}}, nil)
	if rec.Code == http.StatusOK || backup.calls.Load() != 0 {
		t.Fatalf("a key allowed only primary was served by backup: %d", rec.Code)
	}
	rec = post(t, asKey("primary", "backup"), map[string]any{"model": "primary", "messages": []any{}}, nil)
	if rec.Code != http.StatusOK || rec.Header().Get("x-llmgw-served-model") != "backup" {
		t.Fatalf("a key allowed both was not given the fallback: %d", rec.Code)
	}
}

func TestModelsEndpointListsOnlyWhatTheKeyMayUse(t *testing.T) {
	e := newFakeEngine(t, "real")
	cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"real", "other"}})
	cfg.Models = []catalog.Model{catModel("real", []string{"real-alias"}), catModel("other", nil)}
	rt := newTestRouter(t, cfg)

	p := auth.Principal{Limits: store.Effective{AllowedModels: []string{"real"}}}
	req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
	req = req.WithContext(auth.WithPrincipal(req.Context(), p))
	rec := httptest.NewRecorder()
	rt.handleModels(rec, req)

	var out struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil {
		t.Fatal(err)
	}
	var ids []string
	for _, d := range out.Data {
		ids = append(ids, d.ID)
	}
	if !slices.Equal(ids, []string{"real", "real-alias"}) {
		t.Errorf("listed %v; the key may use real and its alias only", ids)
	}
}