steve/llama-swap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

proxy,ui-svelte: add request/response capturing (#508)

Browse Source 
Add saving request and response headers and bodies that go through
llama-swap in memory.

- captureBuffer added to configuration. Captures are enabled by default.
- 5MB of memory is allocated for req/response captures in a ring buffer.
Setting captureBuffer to 0 will disable captures.
- UI elements to view captured data added to Activity page. Includes
some
QOL features like json formatting and recombining SSE chat streams
- capture saving is done at the byte level and has minimal impact on
llama-swap performance

Fixes #464 
Ref #503
This commit is contained in:
Benson Wong
2026-02-07 15:40:01 -08:00
committed by GitHub
parent 7eef5defb8
commit b5fde8eb6d
14 changed files with 971 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/metrics_monitor_test.go
+254 -29
View File
@@ -18,7 +18,7 @@ import (
func TestMetricsMonitor_AddMetrics(t *testing.T) {
t.Run("adds metrics and assigns ID", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
metric := TokenMetrics{
Model: "test-model",
@@ -37,7 +37,7 @@ func TestMetricsMonitor_AddMetrics(t *testing.T) {
})
t.Run("increments ID for each metric", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
for i := 0; i < 5; i++ {
mm.addMetrics(TokenMetrics{Model: "model"})
@@ -51,7 +51,7 @@ func TestMetricsMonitor_AddMetrics(t *testing.T) {
})
t.Run("respects max metrics limit", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 3)
mm := newMetricsMonitor(testLogger, 3, 0)
// Add 5 metrics
for i := 0; i < 5; i++ {
@@ -71,7 +71,7 @@ func TestMetricsMonitor_AddMetrics(t *testing.T) {
})
t.Run("emits TokenMetricsEvent", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
receivedEvent := make(chan TokenMetricsEvent, 1)
cancel := event.On(func(e TokenMetricsEvent) {
@@ -101,14 +101,14 @@ func TestMetricsMonitor_AddMetrics(t *testing.T) {
func TestMetricsMonitor_GetMetrics(t *testing.T) {
t.Run("returns empty slice when no metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
metrics := mm.getMetrics()
assert.NotNil(t, metrics)
assert.Equal(t, 0, len(metrics))
})
t.Run("returns copy of metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
mm.addMetrics(TokenMetrics{Model: "model1"})
mm.addMetrics(TokenMetrics{Model: "model2"})
@@ -128,7 +128,7 @@ func TestMetricsMonitor_GetMetrics(t *testing.T) {
func TestMetricsMonitor_GetMetricsJSON(t *testing.T) {
t.Run("returns valid JSON for empty metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
jsonData, err := mm.getMetricsJSON()
assert.NoError(t, err)
assert.NotNil(t, jsonData)
@@ -140,7 +140,7 @@ func TestMetricsMonitor_GetMetricsJSON(t *testing.T) {
})
t.Run("returns valid JSON with metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
mm.addMetrics(TokenMetrics{
Model: "model1",
InputTokens: 100,
@@ -168,7 +168,7 @@ func TestMetricsMonitor_GetMetricsJSON(t *testing.T) {
func TestMetricsMonitor_WrapHandler(t *testing.T) {
t.Run("successful non-streaming request with usage data", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{
"usage": {
@@ -199,7 +199,7 @@ func TestMetricsMonitor_WrapHandler(t *testing.T) {
})
t.Run("successful request with timings data", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{
"timings": {
@@ -239,7 +239,7 @@ func TestMetricsMonitor_WrapHandler(t *testing.T) {
})
t.Run("streaming request with SSE format", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
// Note: SSE format requires proper line breaks - each data line followed by blank line
responseBody := `data: {"choices":[{"text":"Hello"}]}
@@ -275,7 +275,7 @@ data: [DONE]
})
t.Run("non-OK status code does not record metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
w.WriteHeader(http.StatusBadRequest)
@@ -295,7 +295,7 @@ data: [DONE]
})
t.Run("empty response body records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
w.WriteHeader(http.StatusOK)
@@ -317,7 +317,7 @@ data: [DONE]
})
t.Run("invalid JSON records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
w.Header().Set("Content-Type", "application/json")
@@ -341,7 +341,7 @@ data: [DONE]
})
t.Run("next handler error is propagated", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
expectedErr := assert.AnError
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
@@ -360,7 +360,7 @@ data: [DONE]
})
t.Run("response without usage or timings records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{"result": "ok"}`
@@ -437,7 +437,7 @@ func TestMetricsMonitor_ResponseBodyCopier(t *testing.T) {
func TestMetricsMonitor_Concurrent(t *testing.T) {
t.Run("concurrent addMetrics is safe", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 1000)
mm := newMetricsMonitor(testLogger, 1000, 0)
var wg sync.WaitGroup
numGoroutines := 10
@@ -464,7 +464,7 @@ func TestMetricsMonitor_Concurrent(t *testing.T) {
})
t.Run("concurrent reads and writes are safe", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 100)
mm := newMetricsMonitor(testLogger, 100, 0)
done := make(chan bool)
@@ -502,7 +502,7 @@ func TestMetricsMonitor_Concurrent(t *testing.T) {
func TestMetricsMonitor_ParseMetrics(t *testing.T) {
t.Run("prefers timings over usage data", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
// Timings should take precedence over usage
responseBody := `{
@@ -542,7 +542,7 @@ func TestMetricsMonitor_ParseMetrics(t *testing.T) {
})
t.Run("handles missing cache_n in timings", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{
"timings": {
@@ -577,7 +577,7 @@ func TestMetricsMonitor_ParseMetrics(t *testing.T) {
func TestMetricsMonitor_StreamingResponse(t *testing.T) {
t.Run("finds metrics in last valid SSE data", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
// Metrics should be found in the last data line before [DONE]
responseBody := `data: {"choices":[{"text":"First"}]}
@@ -611,7 +611,7 @@ data: [DONE]
})
t.Run("handles streaming with no valid JSON records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `data: not json
@@ -641,7 +641,7 @@ data: [DONE]
})
t.Run("handles empty streaming response records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := ``
@@ -669,7 +669,7 @@ data: [DONE]
// Benchmark tests
func BenchmarkMetricsMonitor_AddMetrics(b *testing.B) {
mm := newMetricsMonitor(testLogger, 1000)
mm := newMetricsMonitor(testLogger, 1000, 0)
metric := TokenMetrics{
Model: "test-model",
@@ -690,7 +690,7 @@ func BenchmarkMetricsMonitor_AddMetrics(b *testing.B) {
func BenchmarkMetricsMonitor_AddMetrics_SmallBuffer(b *testing.B) {
// Test performance with a smaller buffer where wrapping occurs more frequently
mm := newMetricsMonitor(testLogger, 100)
mm := newMetricsMonitor(testLogger, 100, 0)
metric := TokenMetrics{
Model: "test-model",
@@ -711,7 +711,7 @@ func BenchmarkMetricsMonitor_AddMetrics_SmallBuffer(b *testing.B) {
func TestMetricsMonitor_WrapHandler_Compression(t *testing.T) {
t.Run("gzip encoded response", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{"usage": {"prompt_tokens": 100, "completion_tokens": 50}}`
@@ -745,7 +745,7 @@ func TestMetricsMonitor_WrapHandler_Compression(t *testing.T) {
})
t.Run("deflate encoded response", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{"usage": {"prompt_tokens": 200, "completion_tokens": 75}}`
@@ -779,7 +779,7 @@ func TestMetricsMonitor_WrapHandler_Compression(t *testing.T) {
})
t.Run("invalid gzip data records minimal metrics", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
// Invalid compressed data
invalidData := []byte("this is not gzip data")
@@ -807,7 +807,7 @@ func TestMetricsMonitor_WrapHandler_Compression(t *testing.T) {
})
t.Run("unknown encoding treated as uncompressed", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10)
mm := newMetricsMonitor(testLogger, 10, 0)
responseBody := `{"usage": {"prompt_tokens": 300, "completion_tokens": 100}}`
@@ -832,3 +832,228 @@ func TestMetricsMonitor_WrapHandler_Compression(t *testing.T) {
assert.Equal(t, 100, metrics[0].OutputTokens)
})
}
func TestReqRespCapture_Size(t *testing.T) {
t.Run("calculates size correctly", func(t *testing.T) {
capture := ReqRespCapture{
ID: 1,
ReqPath: "/v1/chat/completions", // 20 bytes
ReqHeaders: map[string]string{
"Content-Type": "application/json", // 12 + 16 = 28
},
ReqBody: []byte("request body"), // 12 bytes
RespHeaders: map[string]string{
"X-Test": "value", // 6 + 5 = 11
},
RespBody: []byte("response body"), // 13 bytes
}
// Expected: 20 + 12 + 13 + 28 + 11 = 84
assert.Equal(t, 84, capture.Size())
})
t.Run("handles empty capture", func(t *testing.T) {
capture := ReqRespCapture{}
assert.Equal(t, 0, capture.Size())
})
}
func TestMetricsMonitor_AddCapture(t *testing.T) {
t.Run("does nothing when captures disabled", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 0)
capture := ReqRespCapture{
ID: 0,
ReqBody: []byte("test"),
}
mm.addCapture(capture)
// Should not store capture
assert.Nil(t, mm.getCaptureByID(0))
})
t.Run("adds capture when enabled", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
capture := ReqRespCapture{
ID: 0,
ReqBody: []byte("test request"),
RespBody: []byte("test response"),
}
mm.addCapture(capture)
retrieved := mm.getCaptureByID(0)
assert.NotNil(t, retrieved)
assert.Equal(t, 0, retrieved.ID)
assert.Equal(t, []byte("test request"), retrieved.ReqBody)
assert.Equal(t, []byte("test response"), retrieved.RespBody)
})
t.Run("evicts oldest when exceeding max size", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
mm.maxCaptureSize = 100 // Set small limit for test
// Add captures that will exceed the limit
capture1 := ReqRespCapture{ID: 0, ReqBody: make([]byte, 40)}
capture2 := ReqRespCapture{ID: 1, ReqBody: make([]byte, 40)}
capture3 := ReqRespCapture{ID: 2, ReqBody: make([]byte, 40)}
mm.addCapture(capture1)
mm.addCapture(capture2)
// Adding capture3 should evict capture1
mm.addCapture(capture3)
assert.Nil(t, mm.getCaptureByID(0), "capture 0 should be evicted")
assert.NotNil(t, mm.getCaptureByID(1), "capture 1 should exist")
assert.NotNil(t, mm.getCaptureByID(2), "capture 2 should exist")
})
t.Run("skips capture larger than max size", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
mm.maxCaptureSize = 100
// Add a capture larger than max
largeCapture := ReqRespCapture{ID: 0, ReqBody: make([]byte, 200)}
mm.addCapture(largeCapture)
assert.Nil(t, mm.getCaptureByID(0), "oversized capture should not be stored")
})
}
func TestMetricsMonitor_GetCaptureByID(t *testing.T) {
t.Run("returns nil for non-existent ID", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
assert.Nil(t, mm.getCaptureByID(999))
})
t.Run("returns capture by ID", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
capture := ReqRespCapture{
ID: 42,
ReqBody: []byte("test"),
}
mm.addCapture(capture)
retrieved := mm.getCaptureByID(42)
assert.NotNil(t, retrieved)
assert.Equal(t, 42, retrieved.ID)
})
}
func TestRedactHeaders(t *testing.T) {
t.Run("redacts sensitive headers", func(t *testing.T) {
headers := map[string]string{
"Authorization": "Bearer secret-token",
"Proxy-Authorization": "Basic creds",
"Cookie": "session=abc123",
"Set-Cookie": "session=xyz789",
"X-Api-Key": "sk-12345",
"Content-Type": "application/json",
"X-Custom": "safe-value",
}
redactHeaders(headers)
assert.Equal(t, "[REDACTED]", headers["Authorization"])
assert.Equal(t, "[REDACTED]", headers["Proxy-Authorization"])
assert.Equal(t, "[REDACTED]", headers["Cookie"])
assert.Equal(t, "[REDACTED]", headers["Set-Cookie"])
assert.Equal(t, "[REDACTED]", headers["X-Api-Key"])
assert.Equal(t, "application/json", headers["Content-Type"])
assert.Equal(t, "safe-value", headers["X-Custom"])
})
t.Run("handles mixed case header names", func(t *testing.T) {
headers := map[string]string{
"authorization": "Bearer token",
"COOKIE": "session=abc",
"x-api-key": "key123",
}
redactHeaders(headers)
assert.Equal(t, "[REDACTED]", headers["authorization"])
assert.Equal(t, "[REDACTED]", headers["COOKIE"])
assert.Equal(t, "[REDACTED]", headers["x-api-key"])
})
t.Run("handles empty headers", func(t *testing.T) {
headers := map[string]string{}
redactHeaders(headers)
assert.Empty(t, headers)
})
}
func TestMetricsMonitor_WrapHandler_Capture(t *testing.T) {
t.Run("captures request and response when enabled", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 5)
requestBody := `{"model": "test", "prompt": "hello"}`
responseBody := `{"usage": {"prompt_tokens": 100, "completion_tokens": 50}}`
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
w.Header().Set("Content-Type", "application/json")
w.Header().Set("X-Custom", "header-value")
w.WriteHeader(http.StatusOK)
w.Write([]byte(responseBody))
return nil
}
req := httptest.NewRequest("POST", "/test", bytes.NewBufferString(requestBody))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer secret")
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
err := mm.wrapHandler("test-model", ginCtx.Writer, req, nextHandler)
assert.NoError(t, err)
// Check metric was recorded
metrics := mm.getMetrics()
assert.Equal(t, 1, len(metrics))
metricID := metrics[0].ID
// Check capture was stored with same ID
capture := mm.getCaptureByID(metricID)
assert.NotNil(t, capture)
assert.Equal(t, metricID, capture.ID)
assert.Equal(t, []byte(requestBody), capture.ReqBody)
assert.Equal(t, []byte(responseBody), capture.RespBody)
assert.Equal(t, "/test", capture.ReqPath)
assert.Equal(t, "application/json", capture.ReqHeaders["Content-Type"])
assert.Equal(t, "[REDACTED]", capture.ReqHeaders["Authorization"])
assert.Equal(t, "application/json", capture.RespHeaders["Content-Type"])
assert.Equal(t, "header-value", capture.RespHeaders["X-Custom"])
})
t.Run("does not capture when disabled", func(t *testing.T) {
mm := newMetricsMonitor(testLogger, 10, 0)
requestBody := `{"model": "test"}`
responseBody := `{"usage": {"prompt_tokens": 100, "completion_tokens": 50}}`
nextHandler := func(modelID string, w http.ResponseWriter, r *http.Request) error {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
w.Write([]byte(responseBody))
return nil
}
req := httptest.NewRequest("POST", "/test", bytes.NewBufferString(requestBody))
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
err := mm.wrapHandler("test-model", ginCtx.Writer, req, nextHandler)
assert.NoError(t, err)
// Metrics should still be recorded
metrics := mm.getMetrics()
assert.Equal(t, 1, len(metrics))
// But no capture
capture := mm.getCaptureByID(metrics[0].ID)
assert.Nil(t, capture)
})
}