diff --git a/.needle-predispatch-sha b/.needle-predispatch-sha index 8920927..570cbff 100644 --- a/.needle-predispatch-sha +++ b/.needle-predispatch-sha @@ -1 +1 @@ -15ce7dc54ccf06d09bdba2e964f47aaaf7964184 +3c2482734b9dd8a1bae7cc90c6aaf3f1eff2d091 diff --git a/proxy/helpers_test.go b/proxy/helpers_test.go index d6601fe..2595e72 100644 --- a/proxy/helpers_test.go +++ b/proxy/helpers_test.go @@ -702,6 +702,10 @@ func TestHelperFunctions(t *testing.T) { }) } +// ============================================================================ +// Mock Upstream Tests +// ============================================================================ + // TestMockUpstreamBasic tests basic mock upstream functionality func TestMockUpstreamBasic(t *testing.T) { t.Run("success scenario", func(t *testing.T) { diff --git a/proxy/response_validation_test.go b/proxy/response_validation_test.go new file mode 100644 index 0000000..4f12353 --- /dev/null +++ b/proxy/response_validation_test.go @@ -0,0 +1,663 @@ +package main + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus/testutil" +) + +// TestResponseValidation_EmptyJSONNonStreaming tests that empty JSON body on 2xx responses +// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway. +func TestResponseValidation_EmptyJSONNonStreaming(t *testing.T) { + testCases := []struct { + name string + maxRetries int + wantCalls int32 // Expected upstream calls: 1 initial + maxRetries + }{ + { + name: "MAX_RETRIES=1", + maxRetries: 1, + wantCalls: 2, // 1 initial + 1 retry + }, + { + name: "MAX_RETRIES=2", + maxRetries: 2, + wantCalls: 3, // 1 initial + 2 retries + }, + { + name: "MAX_RETRIES=5", + maxRetries: 5, + wantCalls: 6, // 1 initial + 5 retries + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with empty body every time + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.WriteHeader(http.StatusOK) + // Empty body - write nothing + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + start := time.Now() + handler.ServeHTTP(w, req) + duration := time.Since(start) + + // Verify 502 Bad Gateway is returned after exhausting retries + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + // Verify exact upstream call count: 1 initial + maxRetries retries + actualCalls := upstreamCallCount.Load() + if actualCalls != tc.wantCalls { + t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d", + tc.wantCalls, tc.maxRetries, actualCalls) + } + + // Verify exponential backoff took reasonable time + // With maxRetries=N, backoff durations are: 1s, 2s, 4s, ... + // For maxRetries=1: ~1s delay + // For maxRetries=2: ~1s+2s=3s delay + // For maxRetries=5: ~1s+2s+4s+8s+16s=31s delay + // We allow 20% tolerance for test execution variance + minExpectedDuration := time.Duration(1< 0 && duration < minExpectedDuration { + t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration) + } + + t.Logf("PASS: Empty JSON body triggered %d retries, %d upstream calls in %v, returned 502", + tc.maxRetries, actualCalls, duration) + }) + } +} + +// TestResponseValidation_InvalidJSONNonStreaming tests that invalid JSON body on 2xx responses +// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway. +func TestResponseValidation_InvalidJSONNonStreaming(t *testing.T) { + testCases := []struct { + name string + maxRetries int + wantCalls int32 + body string + }{ + { + name: "MAX_RETRIES=1 with malformed JSON", + maxRetries: 1, + wantCalls: 2, + body: "{invalid json", + }, + { + name: "MAX_RETRIES=2 with truncated JSON", + maxRetries: 2, + wantCalls: 3, + body: `{"id": "msg_123", "content":`, + }, + { + name: "MAX_RETRIES=3 with garbage data", + maxRetries: 3, + wantCalls: 4, + body: "not json at all!!!", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with invalid JSON every time + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.WriteHeader(http.StatusOK) + w.Write([]byte(tc.body)) + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + start := time.Now() + handler.ServeHTTP(w, req) + duration := time.Since(start) + + // Verify 502 Bad Gateway is returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + // Verify exact upstream call count + actualCalls := upstreamCallCount.Load() + if actualCalls != tc.wantCalls { + t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d", + tc.wantCalls, tc.maxRetries, actualCalls) + } + + t.Logf("PASS: Invalid JSON body triggered %d retries, %d upstream calls in %v, returned 502", + tc.maxRetries, actualCalls, duration) + }) + } +} + +// TestResponseValidation_EmptyStreaming tests that streaming responses opening with zero bytes +// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway. +func TestResponseValidation_EmptyStreaming(t *testing.T) { + testCases := []struct { + name string + maxRetries int + wantCalls int32 + }{ + { + name: "MAX_RETRIES=1", + maxRetries: 1, + wantCalls: 2, + }, + { + name: "MAX_RETRIES=2", + maxRetries: 2, + wantCalls: 3, + }, + { + name: "MAX_RETRIES=4", + maxRetries: 4, + wantCalls: 5, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with empty streaming response + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Empty streaming - write no data + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + start := time.Now() + handler.ServeHTTP(w, req) + duration := time.Since(start) + + // Verify 502 Bad Gateway is returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + // Verify exact upstream call count + actualCalls := upstreamCallCount.Load() + if actualCalls != tc.wantCalls { + t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d", + tc.wantCalls, tc.maxRetries, actualCalls) + } + + // Verify exponential backoff took reasonable time + minExpectedDuration := time.Duration(1< 0 && duration < minExpectedDuration { + t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration) + } + + t.Logf("PASS: Empty streaming triggered %d retries, %d upstream calls in %v, returned 502", + tc.maxRetries, actualCalls, duration) + }) + } +} + +// TestResponseValidation_MetricsEmptyJSON tests that zai_proxy_upstream_errors_total +// metric increments with correct error_type="truncated_response" for empty JSON retries. +func TestResponseValidation_MetricsEmptyJSON(t *testing.T) { + // Reset metrics before test + retryAttempts.Reset() + upstreamErrors.Reset() + + maxRetries := 2 + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with empty body + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.WriteHeader(http.StatusOK) + // Empty body + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify zai_proxy_upstream_errors_total metric increments + // Each failed attempt (except the final one that returns 502) increments the metric + // With maxRetries=2, we have: initial attempt fails (error count 1), retry 1 fails (error count 2), retry 2 fails (error count 3, then returns 502) + // So we expect 3 error increments + expectedErrorCount := float64(maxRetries + 1) + errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test") + errorCount := testutil.ToFloat64(errorMetric) + if errorCount != expectedErrorCount { + t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f", + expectedErrorCount, errorCount) + } + + // Verify zai_proxy_retry_attempts_total metric increments + // Each retry increments the retry attempts metric + // With maxRetries=2, we expect 2 retry attempts (not 3 - the final attempt doesn't retry, it returns 502) + expectedRetryCount := float64(maxRetries) + retryMetric := retryAttempts.WithLabelValues("truncated_response", "test") + retryCount := testutil.ToFloat64(retryMetric) + if retryCount != expectedRetryCount { + t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f", + expectedRetryCount, retryCount) + } + + // Verify upstream call count + expectedCalls := int32(maxRetries + 1) + actualCalls := upstreamCallCount.Load() + if actualCalls != expectedCalls { + t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls) + } + + // Verify 502 returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d", + errorCount, retryCount, actualCalls, w.Code) +} + +// TestResponseValidation_MetricsInvalidJSON tests that zai_proxy_upstream_errors_total +// metric increments with correct error_type="truncated_response" for invalid JSON retries. +func TestResponseValidation_MetricsInvalidJSON(t *testing.T) { + // Reset metrics before test + retryAttempts.Reset() + upstreamErrors.Reset() + + maxRetries := 3 + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with invalid JSON + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.WriteHeader(http.StatusOK) + w.Write([]byte("{this is not valid json")) + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify upstream errors metric + expectedErrorCount := float64(maxRetries + 1) + errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test") + errorCount := testutil.ToFloat64(errorMetric) + if errorCount != expectedErrorCount { + t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f", + expectedErrorCount, errorCount) + } + + // Verify retry attempts metric + expectedRetryCount := float64(maxRetries) + retryMetric := retryAttempts.WithLabelValues("truncated_response", "test") + retryCount := testutil.ToFloat64(retryMetric) + if retryCount != expectedRetryCount { + t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f", + expectedRetryCount, retryCount) + } + + // Verify upstream call count + expectedCalls := int32(maxRetries + 1) + actualCalls := upstreamCallCount.Load() + if actualCalls != expectedCalls { + t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls) + } + + // Verify 502 returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d", + errorCount, retryCount, actualCalls, w.Code) +} + +// TestResponseValidation_MetricsEmptyStreaming tests that zai_proxy_upstream_errors_total +// metric increments with correct error_type="empty_streaming" for empty streaming retries. +func TestResponseValidation_MetricsEmptyStreaming(t *testing.T) { + // Reset metrics before test + retryAttempts.Reset() + upstreamErrors.Reset() + + maxRetries := 2 + var upstreamCallCount atomic.Int32 + + // Create upstream that returns 200 with empty streaming response + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Empty streaming response + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify upstream errors metric with error_type="empty_streaming" + expectedErrorCount := float64(maxRetries + 1) + errorMetric := upstreamErrors.WithLabelValues("empty_streaming", "test") + errorCount := testutil.ToFloat64(errorMetric) + if errorCount != expectedErrorCount { + t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"empty_streaming\"}} = %f, got %f", + expectedErrorCount, errorCount) + } + + // Verify retry attempts metric with reason="empty_streaming" + expectedRetryCount := float64(maxRetries) + retryMetric := retryAttempts.WithLabelValues("empty_streaming", "test") + retryCount := testutil.ToFloat64(retryMetric) + if retryCount != expectedRetryCount { + t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"empty_streaming\"}} = %f, got %f", + expectedRetryCount, retryCount) + } + + // Verify upstream call count + expectedCalls := int32(maxRetries + 1) + actualCalls := upstreamCallCount.Load() + if actualCalls != expectedCalls { + t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls) + } + + // Verify 502 returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + + t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d", + errorCount, retryCount, actualCalls, w.Code) +} + +// TestResponseValidation_ValidResponsePassesThrough tests that valid responses +// (both non-streaming and streaming) pass through without triggering retries. +func TestResponseValidation_ValidResponsePassesThrough(t *testing.T) { + t.Run("valid non-streaming response", func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "msg_test123", + "type": "message", + "role": "assistant", + "content": []map[string]string{{"type": "text", "text": "Hello!"}}, + "usage": map[string]int{"input_tokens": 10, "output_tokens": 5}, + }) + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, 3) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify success response returned + if w.Code != http.StatusOK { + t.Errorf("Expected status 200 OK, got %d", w.Code) + } + + // Verify response body is valid JSON + body := w.Body.String() + if !json.Valid([]byte(body)) { + t.Error("Response body is not valid JSON") + } + + // Verify exactly 1 upstream call (no retries) + if upstreamCallCount.Load() != 1 { + t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load()) + } + + t.Logf("PASS: Valid non-streaming response passed through without retry") + }) + + t.Run("valid streaming response", func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Send some SSE data + w.Write([]byte("data: {\"type\":\"message_start\"}\n\n")) + w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n")) + w.Write([]byte("data: [DONE]\n\n")) + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, 3) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify success response returned + if w.Code != http.StatusOK { + t.Errorf("Expected status 200 OK, got %d", w.Code) + } + + // Verify exactly 1 upstream call (no retries) + if upstreamCallCount.Load() != 1 { + t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load()) + } + + t.Logf("PASS: Valid streaming response passed through without retry") + }) +} + +// TestResponseValidation_VariousInvalidBodies tests various forms of invalid response bodies +// to ensure they all trigger retries correctly. +func TestResponseValidation_VariousInvalidBodies(t *testing.T) { + testCases := []struct { + name string + body string + valid bool + }{ + { + name: "empty body", + body: "", + valid: false, + }, + { + name: "single opening brace", + body: "{", + valid: false, + }, + { + name: "unclosed object", + body: `{"key": "value"`, + valid: false, + }, + { + name: "unclosed array", + body: `["item1", "item2"`, + valid: false, + }, + { + name: "random text", + body: "this is not json at all", + valid: false, + }, + { + name: "partial JSON with trailing comma", + body: `{"id": "msg123", "content": "text",}`, + valid: false, + }, + { + name: "valid JSON object", + body: `{"id": "msg123", "content": "hello"}`, + valid: true, + }, + { + name: "valid JSON array", + body: `[{"id": "msg123"}, {"id": "msg456"}]`, + valid: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var upstreamCallCount atomic.Int32 + maxRetries := 1 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.WriteHeader(http.StatusOK) + if len(tc.body) > 0 { + w.Write([]byte(tc.body)) + } + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + if tc.valid { + // Valid JSON should pass through + if w.Code != http.StatusOK { + t.Errorf("Expected status 200 OK for valid JSON, got %d", w.Code) + } + if upstreamCallCount.Load() != 1 { + t.Errorf("Expected 1 upstream call for valid JSON (no retry), got %d", upstreamCallCount.Load()) + } + } else { + // Invalid JSON should trigger retry and return 502 + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway for invalid JSON, got %d", w.Code) + } + expectedCalls := int32(maxRetries + 1) + if upstreamCallCount.Load() != expectedCalls { + t.Errorf("Expected %d upstream calls for invalid JSON (1 initial + %d retries), got %d", + expectedCalls, maxRetries, upstreamCallCount.Load()) + } + } + + t.Logf("PASS: %s - status=%d, calls=%d", tc.name, w.Code, upstreamCallCount.Load()) + }) + } +} + +// TestResponseValidation_StreamingZeroBytesPeek tests that the streaming zero-byte detection +// works by peeking at the first chunk of the response body. +func TestResponseValidation_StreamingZeroBytesPeek(t *testing.T) { + t.Run("zero bytes on first peek", func(t *testing.T) { + var upstreamCallCount atomic.Int32 + maxRetries := 1 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Close immediately without writing anything + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, maxRetries) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify retry was triggered and 502 returned + if w.Code != http.StatusBadGateway { + t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code) + } + expectedCalls := int32(maxRetries + 1) + if upstreamCallCount.Load() != expectedCalls { + t.Errorf("Expected %d upstream calls, got %d", expectedCalls, upstreamCallCount.Load()) + } + + t.Logf("PASS: Zero bytes on first peek triggered retry, %d calls, status 502", upstreamCallCount.Load()) + }) + + t.Run("non-zero bytes on first peek (no retry)", func(t *testing.T) { + var upstreamCallCount atomic.Int32 + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCallCount.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // Write at least one byte immediately + w.Write([]byte("d")) + })) + defer upstream.Close() + + handler := createTestProxyHandler(t, upstream.URL, 3) + + req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody())) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + handler.ServeHTTP(w, req) + + // Verify no retry - response passes through + if w.Code != http.StatusOK { + t.Errorf("Expected status 200 OK, got %d", w.Code) + } + if upstreamCallCount.Load() != 1 { + t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load()) + } + + t.Logf("PASS: Non-zero bytes on first peek, no retry, status 200") + }) +}