internal/proxy/llm_metrics.go

@@ -3,7 +3,6 @@ package proxy
 import (
 	"bufio"
 	"fmt"
-	"io"
 	"net"
 	"net/http"
 	"strings"
@@ -12,68 +11,14 @@ import (
 	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
 )
 
-// LLMMetricsCapture wraps a response writer to capture TTFT and ITL metrics
-type LLMMetricsCapture struct {
-	writer              io.WriteCloser
-	model               string
-	metrics             *observability.Metrics
-	firstTokenTime      time.Time
-	lastTokenTime       time.Time
-	requestStartTime    time.Time
-	ttftRecorded        bool
-	tokenCount          int64
-	responseStartTime   time.Time
-}
-
-// NewLLMMetricsCapture creates a new metrics capture wrapper
-func NewLLMMetricsCapture(writer io.WriteCloser, model string, metrics *observability.Metrics, startTime time.Time) *LLMMetricsCapture {
-	return &LLMMetricsCapture{
-		writer:            writer,
-		model:             model,
-		metrics:           metrics,
-		requestStartTime:  startTime,
-		responseStartTime: time.Now(),
-	}
-}
-
-// Write intercepts writes to detect tokens and record metrics
-func (c *LLMMetricsCapture) Write(p []byte) (int, error) {
-	// Record first token time
-	if !c.ttftRecorded && len(p) > 0 {
-		now := time.Now()
-		ttft := now.Sub(c.requestStartTime).Milliseconds()
-		c.metrics.RecordTTFT(c.model, ttft)
-		c.ttftRecorded = true
-		c.firstTokenTime = now
-		c.lastTokenTime = now
-	}
-
-	// Count tokens in SSE stream (simple: count "data: " lines)
-	if c.ttftRecorded {
-		tokenCount := strings.Count(string(p), "data: ")
-		if tokenCount > 0 {
-			now := time.Now()
-			if !c.firstTokenTime.IsZero() && c.lastTokenTime != now {
-				itl := now.Sub(c.lastTokenTime).Milliseconds()
-				c.metrics.RecordITL(c.model, itl)
-			}
-			c.lastTokenTime = now
-			c.tokenCount += int64(tokenCount)
-		}
-	}
-
-	return c.writer.Write(p)
-}
-
-// Close records final metrics and closes writer
-func (c *LLMMetricsCapture) Close() error {
-	if c.tokenCount > 0 {
-		c.metrics.RecordTokenCount(c.model, c.tokenCount)
-	}
-	return c.writer.Close()
-}
-
-// ResponseWriterWrapper wraps http.ResponseWriter to capture metrics
+// ResponseWriterWrapper wraps an http.ResponseWriter to capture per-request LLM
+// latency metrics (TTFT, ITL, generated tokens) while a streaming response is
+// copied from the upstream by httputil.ReverseProxy.
+//
+// Only server-sent-event responses are measured. A non-streaming completion is
+// written in a single Write, so recording it would emit one "TTFT" sample equal
+// to the whole request latency and no ITL samples at all — numbers that look
+// plausible but mean nothing.
 type ResponseWriterWrapper struct {
 	writer        http.ResponseWriter
 	statusCode    int
@@ -83,9 +28,13 @@ type ResponseWriterWrapper struct {
 	firstByteTime time.Time
 	lastWriteTime time.Time
 	ttftRecorded  bool
+	streamChecked bool
+	streaming     bool
 }
 
-// NewResponseWriterWrapper creates a wrapper for response writer
+// NewResponseWriterWrapper creates a wrapper for response writer.
+// startTime must be the time the request entered the gateway, not the time this
+// wrapper was constructed, otherwise TTFT measures the wrong interval.
 func NewResponseWriterWrapper(w http.ResponseWriter, model string, metrics *observability.Metrics, startTime time.Time) *ResponseWriterWrapper {
 	return &ResponseWriterWrapper{
 		writer:     w,
@@ -101,29 +50,48 @@ func (w *ResponseWriterWrapper) Header() http.Header {
 	return w.writer.Header()
 }
 
-// Write implements http.ResponseWriter
-func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
-	// Record TTFT on first write
-	if !w.ttftRecorded && len(b) > 0 {
-		now := time.Now()
-		ttft := now.Sub(w.startTime).Milliseconds()
-		w.metrics.RecordTTFT(w.model, ttft)
-		w.ttftRecorded = true
-		w.firstByteTime = now
-		w.lastWriteTime = now
+// isStreaming reports whether this is an SSE response. Resolved once from the
+// Content-Type header, which ReverseProxy has already set before the first write.
+func (w *ResponseWriterWrapper) isStreaming() bool {
+	if !w.streamChecked {
+		w.streamChecked = true
+		w.streaming = strings.Contains(w.writer.Header().Get("Content-Type"), "text/event-stream")
+	}
+	return w.streaming
 }
 
-	// Record ITL for subsequent writes (for streaming)
-	if w.ttftRecorded && len(b) > 0 {
+// countSSETokens counts generated tokens in a chunk of an SSE stream. These
+// upstreams emit one `data:` event per token, so the event count is the token
+// count; the `[DONE]` terminator is not a token.
+func countSSETokens(b []byte) int64 {
+	var n int64
+	for _, line := range strings.Split(string(b), "\n") {
+		line = strings.TrimSpace(line)
+		if strings.HasPrefix(line, "data:") && !strings.HasPrefix(line, "data: [DONE]") {
+			n++
+		}
+	}
+	return n
+}
+
+// Write implements http.ResponseWriter
+func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
+	if len(b) > 0 && w.isStreaming() {
 		now := time.Now()
-		if !w.firstByteTime.IsZero() && w.lastWriteTime != now {
-			itl := now.Sub(w.lastWriteTime).Milliseconds()
-			// Only record if ITL > 0 (avoid recording same millisecond twice)
-			if itl > 0 {
+		if !w.ttftRecorded {
+			w.metrics.RecordTTFT(w.model, now.Sub(w.startTime).Milliseconds())
+			w.ttftRecorded = true
+			w.firstByteTime = now
+		} else if itl := now.Sub(w.lastWriteTime).Milliseconds(); itl > 0 {
+			// Only record positive ITL: a zero means two chunks landed in the
+			// same millisecond, which carries no timing information.
 			w.metrics.RecordITL(w.model, itl)
 		}
-		}
 		w.lastWriteTime = now
+
+		if n := countSSETokens(b); n > 0 {
+			w.metrics.RecordTokenCount(w.model, n)
+		}
 	}
 
 	return w.writer.Write(b)
@@ -150,4 +118,8 @@ func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return nil, nil, fmt.Errorf("response writer does not implement Hijacker")
 }
 
-
+// Unwrap exposes the underlying writer to http.ResponseController so optional
+// capabilities (Flush, Hijack, deadlines) keep working through the wrapper.
+func (w *ResponseWriterWrapper) Unwrap() http.ResponseWriter {
+	return w.writer
+}

internal/proxy/llm_metrics_test.go

@@ -0,0 +1,171 @@
+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")
+	}
+}

internal/proxy/proxy.go

@@ -314,6 +314,10 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d
 
 // ServeHTTP implements http.Handler.
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Captured before any work so TTFT measures from the moment the gateway
+	// accepted the request, not from when the upstream response is copied.
+	requestStart := time.Now()
+
 	if h.handleAuthEndpoint(w, r) {
 		return
 	}
@@ -358,7 +362,14 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// For streaming responses (SSE, chunked), disable buffering. Issue #33.
 	w.Header().Set("X-Accel-Buffering", "no")
 
-	proxy.ServeHTTP(w, r)
+	// Capture per-request LLM latency metrics (TTFT/ITL/tokens) while the
+	// response is copied from the upstream. Inert for non-SSE responses.
+	var respWriter http.ResponseWriter = w
+	if route.ModelName != "" {
+		respWriter = NewResponseWriterWrapper(w, route.ModelName, observability.Global(), requestStart)
+	}
+
+	proxy.ServeHTTP(respWriter, r)
 }