llmgw/internal/sim/server_test.go

package sim

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"
)

func newTestServer(t *testing.T, cfg Config) *httptest.Server {
	t.Helper()
	eng := New(cfg)
	srv := httptest.NewServer(NewServer(eng, nil))
	t.Cleanup(func() {
		srv.Close()
		eng.Close()
	})
	return srv
}

func postJSON(t *testing.T, url string, body any) *http.Response {
	t.Helper()
	b, err := json.Marshal(body)
	if err != nil {
		t.Fatal(err)
	}
	resp, err := http.Post(url, "application/json", bytes.NewReader(b)) //nolint:noctx // test
	if err != nil {
		t.Fatal(err)
	}
	return resp
}

func TestChatCompletionsBuffered(t *testing.T) {
	srv := newTestServer(t, testConfig())

	resp := postJSON(t, srv.URL+"/v1/chat/completions", map[string]any{
		"model":      "sim-model",
		"messages":   []map[string]string{{"role": "user", "content": "merhaba"}},
		"max_tokens": 7,
	})
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status %d", resp.StatusCode)
	}

	var out struct {
		Object  string `json:"object"`
		Choices []struct {
			Message struct {
				Role    string `json:"role"`
				Content string `json:"content"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage usage `json:"usage"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}

	if out.Object != "chat.completion" {
		t.Errorf("object=%q", out.Object)
	}
	if len(out.Choices) != 1 {
		t.Fatalf("choices=%d", len(out.Choices))
	}
	if n := len(strings.Fields(out.Choices[0].Message.Content)); n != 7 {
		t.Errorf("expected 7 tokens, got %d", n)
	}
	if out.Usage.CompletionTokens != 7 {
		t.Errorf("usage.completion_tokens=%d", out.Usage.CompletionTokens)
	}
	if out.Usage.TotalTokens != out.Usage.PromptTokens+out.Usage.CompletionTokens {
		t.Error("usage totals are inconsistent")
	}
}

// sseEvent is one decoded event from an SSE stream.
type sseEvent struct {
	raw  string
	data map[string]any
}

func readSSE(t *testing.T, body *bufio.Reader) []sseEvent {
	t.Helper()
	var out []sseEvent
	for {
		line, err := body.ReadString('\n')
		if err != nil {
			return out
		}
		line = strings.TrimRight(line, "\r\n")
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		payload := strings.TrimPrefix(line, "data: ")
		ev := sseEvent{raw: payload}
		if payload != "[DONE]" {
			if err := json.Unmarshal([]byte(payload), &ev.data); err != nil {
				t.Fatalf("could not decode SSE payload: %v (%q)", err, payload)
			}
		}
		out = append(out, ev)
	}
}

func TestChatCompletionsStreaming(t *testing.T) {
	srv := newTestServer(t, testConfig())

	resp := postJSON(t, srv.URL+"/v1/chat/completions", map[string]any{
		"model":          "sim-model",
		"messages":       []map[string]string{{"role": "user", "content": "merhaba"}},
		"max_tokens":     5,
		"stream":         true,
		"stream_options": map[string]any{"include_usage": true},
	})
	defer resp.Body.Close()

	if ct := resp.Header.Get("Content-Type"); ct != "text/event-stream" {
		t.Errorf("Content-Type=%q", ct)
	}

	events := readSSE(t, bufio.NewReader(resp.Body))
	if len(events) == 0 {
		t.Fatal("no SSE events received")
	}
	if last := events[len(events)-1]; last.raw != "[DONE]" {
		t.Errorf("stream did not end with [DONE]: %q", last.raw)
	}

	var contentTokens, usageChunks, finishChunks int
	for _, ev := range events {
		if ev.raw == "[DONE]" {
			continue
		}
		if _, ok := ev.data["usage"]; ok {
			usageChunks++
			continue
		}
		choices, _ := ev.data["choices"].([]any)
		if len(choices) == 0 {
			continue
		}
		c, _ := choices[0].(map[string]any)
		if delta, ok := c["delta"].(map[string]any); ok {
			if _, has := delta["content"]; has {
				contentTokens++
			}
		}
		if fr, ok := c["finish_reason"].(string); ok && fr == "length" {
			finishChunks++
		}
	}

	if contentTokens != 5 {
		t.Errorf("expected 5 content chunks, got %d", contentTokens)
	}
	if finishChunks != 1 {
		t.Errorf("expected 1 finish_reason chunk, got %d", finishChunks)
	}
	if usageChunks != 1 {
		t.Errorf("include_usage was requested but got %d usage chunks", usageChunks)
	}
}

func TestStreamingOmitsUsageUnlessRequested(t *testing.T) {
	srv := newTestServer(t, testConfig())

	resp := postJSON(t, srv.URL+"/v1/chat/completions", map[string]any{
		"messages":   []map[string]string{{"role": "user", "content": "x"}},
		"max_tokens": 2,
		"stream":     true,
	})
	defer resp.Body.Close()

	for _, ev := range readSSE(t, bufio.NewReader(resp.Body)) {
		if ev.raw == "[DONE]" {
			continue
		}
		if _, ok := ev.data["usage"]; ok {
			t.Error("usage chunk sent although it was not requested")
		}
	}
}

func TestClientDisconnectReachesEngine(t *testing.T) {
	// The basis of phase 1's exit criterion: when a client disconnects the GPU
	// must stop working on that request.
	cfg := testConfig()
	cfg.DecodeTokensPerSec = 20
	eng := New(cfg)
	defer eng.Close()
	srv := httptest.NewServer(NewServer(eng, nil))
	defer srv.Close()

	ctx, cancel := context.WithCancel(context.Background())
	body, _ := json.Marshal(map[string]any{
		"messages":   []map[string]string{{"role": "user", "content": "merhaba"}},
		"max_tokens": 5000,
		"stream":     true,
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		srv.URL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}

	// Wait for the first token, then drop the client.
	waitFor(t, 2*time.Second, func() bool { return eng.Stats().Running == 1 })
	cancel()
	_ = resp.Body.Close()

	waitFor(t, 3*time.Second, func() bool {
		s := eng.Stats()
		return s.Running == 0 && s.KVCacheUsage == 0
	})
}

func TestCompletionsEndpoint(t *testing.T) {
	srv := newTestServer(t, testConfig())

	resp := postJSON(t, srv.URL+"/v1/completions", map[string]any{
		"model": "sim-model", "prompt": "merhaba", "max_tokens": 3,
	})
	defer resp.Body.Close()

	var out struct {
		Object  string `json:"object"`
		Choices []struct {
			Text string `json:"text"`
		} `json:"choices"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	if out.Object != "text_completion" {
		t.Errorf("object=%q", out.Object)
	}
	if n := len(strings.Fields(out.Choices[0].Text)); n != 3 {
		t.Errorf("expected 3 tokens, got %d", n)
	}
}

func TestModelsAndHealth(t *testing.T) {
	cfg := testConfig()
	cfg.Model = "llama-3.3-70b"
	srv := newTestServer(t, cfg)

	resp, err := http.Get(srv.URL + "/health") //nolint:noctx // test
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Errorf("/health status %d", resp.StatusCode)
	}

	resp, err = http.Get(srv.URL + "/v1/models") //nolint:noctx // test
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	var out struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	if len(out.Data) != 1 || out.Data[0].ID != "llama-3.3-70b" {
		t.Errorf("/v1/models did not list the expected model: %+v", out.Data)
	}
}

// TestMetricsParseWithPrometheus verifies the engine contract from plan §7.6.
// The router reads these metrics with a real Prometheus parser; if the format
// is wrong, routing goes blind.
func TestMetricsParseWithPrometheus(t *testing.T) {
	srv := newTestServer(t, testConfig())

	// Generate some traffic so the counters are non-zero.
	resp := postJSON(t, srv.URL+"/v1/chat/completions", map[string]any{
		"messages":   []map[string]string{{"role": "user", "content": "merhaba dunya"}},
		"max_tokens": 4,
	})
	resp.Body.Close()

	mresp, err := http.Get(srv.URL + "/metrics") //nolint:noctx // test
	if err != nil {
		t.Fatal(err)
	}
	defer mresp.Body.Close()

	// TextParser's zero value is invalid; the scheme must be given explicitly.
	// The legacy scheme is stricter, so this also checks our metric names obey
	// classic Prometheus rules (the ':' in vLLM names is legal there too).
	parser := expfmt.NewTextParser(model.LegacyValidation)
	families, err := parser.TextToMetricFamilies(mresp.Body)
	if err != nil {
		t.Fatalf("metrics failed to parse: %v", err)
	}

	// The three metrics plan §7.6 makes mandatory.
	for _, name := range []string{
		"vllm:num_requests_running",
		"vllm:num_requests_waiting",
		"vllm:kv_cache_usage_perc",
	} {
		if _, ok := families[name]; !ok {
			t.Errorf("mandatory metric missing: %s", name)
		}
	}
	// Optional metrics and histograms.
	for _, name := range []string{
		"vllm:cache_config_info",
		"vllm:prefix_cache_queries_total",
		"vllm:prefix_cache_hits_total",
		"vllm:time_to_first_token_seconds",
	} {
		if _, ok := families[name]; !ok {
			t.Errorf("metric missing: %s", name)
		}
	}

	if f, ok := families["vllm:generation_tokens_total"]; ok {
		if got := f.Metric[0].GetCounter().GetValue(); got != 4 {
			t.Errorf("generation_tokens_total=%v, expected 4", got)
		}
	} else {
		t.Error("vllm:generation_tokens_total is missing")
	}
}

// TestUsageReportsCachedTokens verifies the prefix cache's effect reaches the
// client. The router will not consume this, but it is the only end-to-end
// signal that the cache actually works.
func TestUsageReportsCachedTokens(t *testing.T) {
	srv := newTestServer(t, testConfig())

	body := map[string]any{
		"messages":   []map[string]string{{"role": "user", "content": strings.Repeat("shared prefix. ", 100)}},
		"max_tokens": 3,
	}

	// The first request populates the cache.
	resp := postJSON(t, srv.URL+"/v1/chat/completions", body)
	resp.Body.Close()

	// The second one must reuse that prefix.
	resp = postJSON(t, srv.URL+"/v1/chat/completions", body)
	defer resp.Body.Close()

	var out struct {
		Usage usage `json:"usage"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	if out.Usage.PromptTokensDetails.CachedTokens == 0 {
		t.Error("cached_tokens is 0 on the second request; prefix cache is not reported")
	}
	if out.Usage.PromptTokensDetails.CachedTokens > out.Usage.PromptTokens {
		t.Errorf("cached_tokens (%d) exceeds prompt_tokens (%d)",
			out.Usage.PromptTokensDetails.CachedTokens, out.Usage.PromptTokens)
	}
}