llmgw/internal/router/router_test.go

package router

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

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

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

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, _ *http.Request) {
		f.calls.Add(1)
		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 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)
	}
}