diff --git a/internal/proxy/llm_metrics.go b/internal/proxy/llm_metrics.go index bddc7c7..29c1fb1 100644 --- a/internal/proxy/llm_metrics.go +++ b/internal/proxy/llm_metrics.go @@ -3,7 +3,6 @@ package proxy import ( "bufio" "fmt" - "io" "net" "net/http" "strings" @@ -12,87 +11,37 @@ 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 - metrics *observability.Metrics - model string - startTime time.Time - firstByteTime time.Time - lastWriteTime time.Time - ttftRecorded bool + writer http.ResponseWriter + statusCode int + metrics *observability.Metrics + model string + startTime time.Time + 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, - model: model, - metrics: metrics, - startTime: startTime, - statusCode: 200, + writer: w, + model: model, + metrics: metrics, + startTime: startTime, + statusCode: 200, } } @@ -101,29 +50,48 @@ func (w *ResponseWriterWrapper) Header() http.Header { return w.writer.Header() } +// 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 +} + +// 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) { - // Record TTFT on first write - if !w.ttftRecorded && len(b) > 0 { + if len(b) > 0 && w.isStreaming() { now := time.Now() - ttft := now.Sub(w.startTime).Milliseconds() - w.metrics.RecordTTFT(w.model, ttft) - w.ttftRecorded = true - w.firstByteTime = now - w.lastWriteTime = now - } - - // Record ITL for subsequent writes (for streaming) - if w.ttftRecorded && len(b) > 0 { - 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 { - w.metrics.RecordITL(w.model, itl) - } + 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 +} diff --git a/internal/proxy/llm_metrics_test.go b/internal/proxy/llm_metrics_test.go new file mode 100644 index 0000000..69d40d7 --- /dev/null +++ b/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") + } +} diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index c3d9972..a2de895 100644 --- a/internal/proxy/proxy.go +++ b/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) }