homelab-frontend/internal/proxy/llm_metrics_test.go

package proxy

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"
	"time"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)

// exportMetrics renders the Prometheus text exposition for m.
func exportMetrics(m *observability.Metrics) string {
	return observability.NewPrometheusExporter(m).Export()
}

// requireSeries asserts the exporter emitted a sample line with the exact value.
func requireSeries(t *testing.T, out, want string) {
	t.Helper()
	for _, line := range strings.Split(out, "\n") {
		if strings.TrimSpace(line) == want {
			return
		}
	}
	t.Errorf("missing metrics series %q\n--- got ---\n%s", want, out)
}

// TestResponseWriterWrapperRecordsSSEMetrics guards the regression where the
// LLM latency wrapper existed but was never constructed, leaving
// llm_ttft_seconds / llm_itl_seconds / llm_tokens_total as HELP-only stubs.
func TestResponseWriterWrapperRecordsSSEMetrics(t *testing.T) {
	m := observability.NewMetrics()

	rec := httptest.NewRecorder()
	rec.Header().Set("Content-Type", "text/event-stream")

	w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))

	if _, err := w.Write([]byte("data: {\"token\":\"a\"}\n\n")); err != nil {
		t.Fatalf("write: %v", err)
	}
	time.Sleep(5 * time.Millisecond)
	if _, err := w.Write([]byte("data: {\"token\":\"b\"}\n\n")); err != nil {
		t.Fatalf("write: %v", err)
	}
	// The terminator is not a generated token.
	if _, err := w.Write([]byte("data: [DONE]\n\n")); err != nil {
		t.Fatalf("write: %v", err)
	}

	out := exportMetrics(m)
	requireSeries(t, out, `llm_ttft_seconds_count{model="reasoning"} 1`)
	requireSeries(t, out, `llm_itl_seconds_count{model="reasoning"} 1`)
	requireSeries(t, out, `llm_tokens_total{model="reasoning"} 2`)
}

// TestResponseWriterWrapperIgnoresNonStreaming verifies a single-shot JSON
// response does not produce a bogus TTFT sample equal to full request latency.
func TestResponseWriterWrapperIgnoresNonStreaming(t *testing.T) {
	m := observability.NewMetrics()

	rec := httptest.NewRecorder()
	rec.Header().Set("Content-Type", "application/json")

	w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))
	if _, err := w.Write([]byte(`{"choices":[]}`)); err != nil {
		t.Fatalf("write: %v", err)
	}

	out := exportMetrics(m)
	if strings.Contains(out, `llm_ttft_seconds_count{model="reasoning"}`) {
		t.Errorf("non-streaming response recorded TTFT\n--- got ---\n%s", out)
	}
	if strings.Contains(out, `llm_tokens_total{model="reasoning"}`) {
		t.Errorf("non-streaming response recorded tokens\n--- got ---\n%s", out)
	}
}

// TestResponseWriterWrapperPreservesStreaming verifies the wrapper stays
// transparent: bytes reach the client and Flush still works, so wrapping the
// writer cannot silently break SSE (issue #33).
func TestResponseWriterWrapperPreservesStreaming(t *testing.T) {
	m := observability.NewMetrics()

	rec := httptest.NewRecorder()
	rec.Header().Set("Content-Type", "text/event-stream")

	w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now())
	payload := "data: {\"token\":\"a\"}\n\n"
	if _, err := w.Write([]byte(payload)); err != nil {
		t.Fatalf("write: %v", err)
	}
	w.Flush()

	if got := rec.Body.String(); got != payload {
		t.Errorf("body = %q, want %q", got, payload)
	}
	if !rec.Flushed {
		t.Error("Flush did not reach the underlying writer")
	}
}

// TestHandlerRecordsLLMMetricsForStreaming is the end-to-end guard: it drives a
// real streaming request through the Handler, so it fails if the wrapper is ever
// constructed nowhere (the original defect) rather than merely being broken.
func TestHandlerRecordsLLMMetricsForStreaming(t *testing.T) {
	// Unique model name keeps the observations isolated in the global registry.
	model := fmt.Sprintf("metrics-test-%d", time.Now().UnixNano())

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		rc := http.NewResponseController(w)
		for _, tok := range []string{"a", "b", "c"} {
			time.Sleep(2 * time.Millisecond) // guarantee distinct ITL samples
			fmt.Fprintf(w, "data: {\"token\":%q}\n\n", tok)
			if err := rc.Flush(); err != nil {
				return
			}
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
		_ = rc.Flush()
	}))
	defer upstream.Close()

	cfg := &config.Config{
		Models: map[string]*config.ModelUpstream{
			model: {
				Name:    model,
				Address: strings.TrimPrefix(upstream.URL, "http://"),
				Path:    "/v1/chat/completions",
			},
		},
		Routes: make(map[string]*config.Route),
	}

	handler := New(cfg)
	defer handler.Close()
	srv := httptest.NewServer(handler)
	defer srv.Close()

	body := fmt.Sprintf(`{"model":%q,"stream":true,"messages":[{"role":"user","content":"hi"}]}`, model)
	resp, err := http.Post(srv.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, want 200", resp.StatusCode)
	}
	if _, err := io.ReadAll(resp.Body); err != nil {
		t.Fatalf("read body: %v", err)
	}

	out := exportMetrics(observability.Global())
	requireSeries(t, out, fmt.Sprintf(`llm_ttft_seconds_count{model=%q} 1`, model))
	requireSeries(t, out, fmt.Sprintf(`llm_tokens_total{model=%q} 3`, model))

	itlCount := regexp.MustCompile(fmt.Sprintf(`llm_itl_seconds_count\{model=%q\} (\d+)`, model))
	match := itlCount.FindStringSubmatch(out)
	if match == nil {
		t.Errorf("no llm_itl_seconds samples recorded for %s\n--- got ---\n%s", model, out)
	} else if match[1] == "0" {
		t.Errorf("llm_itl_seconds_count = 0, want > 0")
	}
}