diff --git a/plugins/wasm-go/extensions/ai-statistics/main.go b/plugins/wasm-go/extensions/ai-statistics/main.go index ce6a19aa0..07f5b1e0d 100644 --- a/plugins/wasm-go/extensions/ai-statistics/main.go +++ b/plugins/wasm-go/extensions/ai-statistics/main.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "regexp" + "strconv" "strings" "time" @@ -82,6 +83,7 @@ const ( LLMServiceDuration = "llm_service_duration" LLMDurationCount = "llm_duration_count" LLMStreamDurationCount = "llm_stream_duration_count" + LLMFailureCount = "llm_failure_count" ResponseType = "response_type" ChatID = "chat_id" ChatRound = "chat_round" @@ -820,6 +822,13 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat ctx.SetUserAttribute(ChatID, chatID.String()) } + // Track streaming errors across chunks — SSE failures often appear as + // data: {"error":{...}} before data: [DONE], so the last chunk alone is + // insufficient for error detection. + if !ctx.GetBoolContext("hasStreamError", false) && isErrorResponse(data) { + ctx.SetContext("hasStreamError", true) + } + // Get requestStartTime from http context requestStartTime, ok := ctx.GetContext(StatisticsRequestStartTime).(int64) if !ok { @@ -873,8 +882,13 @@ func onHttpStreamingBody(ctx wrapper.HttpContext, config AIStatisticsConfig, dat debugLogAiLog(ctx) _ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey) - // Write metrics - writeMetric(ctx, config) + // Write metrics — prefer the accumulated buffer for error detection + // so that errors split across multiple SSE chunks are not missed. + bodyForMetric := data + if config.shouldBufferStreamingBody && len(streamingBodyBuffer) > 0 { + bodyForMetric = streamingBodyBuffer + } + writeMetric(ctx, config, bodyForMetric) } return data } @@ -928,7 +942,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AIStatisticsConfig, body _ = ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey) // Write metrics - writeMetric(ctx, config) + writeMetric(ctx, config, body) return types.ActionContinue } @@ -1341,7 +1355,59 @@ func setSpanAttribute(key string, value interface{}) { } } -func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig) { +// isErrorResponse checks whether the LLM response indicates an error. +// Detects errors by: +// 1. Response body contains non-null "error" field at root level (OpenAI/Anthropic format). +// Handles both raw JSON and SSE "data: " prefixed chunks, including multi-event +// streaming buffers. +// 2. HTTP status code >= 400 as fallback when body is empty. +// +// Note: some providers (e.g. Anthropic streaming responses) emit {"error":""} +// even on success; an empty-string error is treated as not-an-error to avoid +// false positives. +func isErrorResponse(body []byte) bool { + if len(body) > 0 { + // SSE chunks are prefixed with "data: "; accumulated buffers contain + // multiple SSE events separated by \n\n. Split and check each event. + trimmed := bytes.TrimSpace(body) + if bytes.HasPrefix(trimmed, []byte("data: ")) { + for _, event := range bytes.Split(trimmed, []byte("\n\n")) { + jsonBody := bytes.TrimSpace(event) + if bytes.HasPrefix(jsonBody, []byte("data: ")) { + jsonBody = jsonBody[len("data: "):] + } + if hasErrorField(jsonBody) { + return true + } + } + return false + } + return hasErrorField(body) + } + // Fallback: check HTTP status code for errors with empty body (connection reset, timeout, etc.) + if len(body) == 0 { + if statusCode, err := proxywasm.GetHttpResponseHeader(":status"); err == nil { + if code, err := strconv.Atoi(statusCode); err == nil && code >= 400 { + return true + } + } + } + return false +} + +// hasErrorField checks whether a JSON body contains a non-null, non-empty-string "error" field. +func hasErrorField(jsonBody []byte) bool { + errorVal := gjson.GetBytes(jsonBody, "error") + if errorVal.Exists() && errorVal.Value() != nil { + if errorVal.Type == gjson.String && errorVal.String() == "" { + return false + } + return true + } + return false +} + +func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig, body []byte) { // Generate usage metrics var ok bool var route, cluster, model string @@ -1357,6 +1423,29 @@ func writeMetric(ctx wrapper.HttpContext, config AIStatisticsConfig) { return } + // Get model for metric label (may be empty for error responses) + modelStr := "-" + if m := ctx.GetUserAttribute(tokenusage.CtxKeyModel); m != nil { + if ms, ok := m.(string); ok { + modelStr = ms + } + } + // Fallback to request model for error responses where usage info is unavailable + if modelStr == "-" { + if rm, ok := ctx.GetContext(tokenusage.CtxKeyRequestModel).(string); ok && rm != "" { + modelStr = rm + } + } + + // Count failure before usage check, so error responses without usage info are still counted. + // For streaming, also check the hasStreamError flag set during onHttpStreamingBody. + // llm_failure_count is intentionally incremented regardless of disableOpenaiUsage, + // because error responses carry no usage info and operators still need the failure + // signal even when usage tracking is off. + if isErrorResponse(body) || ctx.GetBoolContext("hasStreamError", false) { + config.incrementCounter(generateMetricName(route, cluster, modelStr, consumer, LLMFailureCount), 1) + } + if config.disableOpenaiUsage { return } diff --git a/plugins/wasm-go/extensions/ai-statistics/main_test.go b/plugins/wasm-go/extensions/ai-statistics/main_test.go index 13dd363d2..e920dd77e 100644 --- a/plugins/wasm-go/extensions/ai-statistics/main_test.go +++ b/plugins/wasm-go/extensions/ai-statistics/main_test.go @@ -2167,3 +2167,222 @@ func TestConfigWithDefaultAttributes(t *testing.T) { }) }) } + +func TestIsErrorResponse(t *testing.T) { + // Test error body detection (OpenAI/Anthropic format) + t.Run("error object in body", func(t *testing.T) { + body := []byte(`{"error": {"type": "api_error", "message": "Unauthorized"}}`) + require.True(t, isErrorResponse(body)) + }) + + t.Run("error object with only code", func(t *testing.T) { + body := []byte(`{"error": {"code": "invalid_api_key"}}`) + require.True(t, isErrorResponse(body)) + }) + + t.Run("error string in body", func(t *testing.T) { + body := []byte(`{"error": "Something went wrong"}`) + require.True(t, isErrorResponse(body)) + }) + + t.Run("empty error string is not counted", func(t *testing.T) { + body := []byte(`{"error": ""}`) + require.False(t, isErrorResponse(body)) + }) + + t.Run("null error is not counted", func(t *testing.T) { + body := []byte(`{"error": null}`) + require.False(t, isErrorResponse(body)) + }) + + t.Run("no error field in body", func(t *testing.T) { + body := []byte(`{"choices": [{"message": {"content": "Hi"}}], "model": "gpt-4"}`) + require.False(t, isErrorResponse(body)) + }) + + t.Run("nested error field not at root", func(t *testing.T) { + body := []byte(`{"choices": [{"error": "x"}]}`) + require.False(t, isErrorResponse(body)) + }) + + // Empty body with HTTP status fallback is tested in TestIsErrorResponseWithHTTPStatus + // (requires wasm test environment for proxywasm.GetHttpResponseHeader) +} + +func TestIsErrorResponseWithHTTPStatus(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + t.Run("empty body with 401 status returns true", func(t *testing.T) { + host, status := test.NewTestHost(basicConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/api/chat"}, + {":method", "POST"}, + }) + host.CallOnHttpResponseHeaders([][2]string{ + {":status", "401"}, + {"content-type", "application/json"}, + }) + + require.True(t, isErrorResponse([]byte{})) + host.CompleteHttp() + }) + + t.Run("empty body with 200 status returns false", func(t *testing.T) { + host, status := test.NewTestHost(basicConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/api/chat"}, + {":method", "POST"}, + }) + host.CallOnHttpResponseHeaders([][2]string{ + {":status", "200"}, + {"content-type", "application/json"}, + }) + + require.False(t, isErrorResponse([]byte{})) + host.CompleteHttp() + }) + }) +} + +func TestFailureCountMetric(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + t.Run("error response body increments failure count only", func(t *testing.T) { + host, status := test.NewTestHost(basicConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.SetRouteName("api-v1") + host.SetClusterName("cluster-1") + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/api/chat"}, + {":method", "POST"}, + {"x-mse-consumer", "user1"}, + }) + + requestBody := []byte(`{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}`) + host.CallOnHttpRequestBody(requestBody) + + time.Sleep(10 * time.Millisecond) + + host.CallOnHttpResponseHeaders([][2]string{ + {":status", "401"}, + {"content-type", "application/json"}, + }) + errorBody := []byte(`{"error": {"type": "authentication_error", "message": "Invalid API key"}}`) + host.CallOnHttpResponseBody(errorBody) + + host.CompleteHttp() + + // Verify llm_failure_count is incremented + failureMetric := "route.api-v1.upstream.cluster-1.model.gpt-3.5-turbo.consumer.user1.metric.llm_failure_count" + failureValue, err := host.GetCounterMetric(failureMetric) + require.NoError(t, err) + require.Equal(t, uint64(1), failureValue) + + // Verify success metrics are NOT present (usage info unavailable in error response) + durationCountMetric := "route.api-v1.upstream.cluster-1.model.gpt-3.5-turbo.consumer.user1.metric.llm_duration_count" + _, err = host.GetCounterMetric(durationCountMetric) + require.Error(t, err, "llm_duration_count should not exist for error response without usage") + }) + + t.Run("success response does not increment failure count", func(t *testing.T) { + host, status := test.NewTestHost(basicConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.SetRouteName("api-v1") + host.SetClusterName("cluster-1") + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/api/chat"}, + {":method", "POST"}, + {"x-mse-consumer", "user1"}, + }) + + requestBody := []byte(`{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}`) + host.CallOnHttpRequestBody(requestBody) + + time.Sleep(10 * time.Millisecond) + + host.CallOnHttpResponseHeaders([][2]string{ + {":status", "200"}, + {"content-type", "application/json"}, + }) + responseBody := []byte(`{ + "choices": [{"message": {"content": "Hello!"}}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + "model": "gpt-4" + }`) + host.CallOnHttpResponseBody(responseBody) + + host.CompleteHttp() + + // Verify llm_failure_count is NOT present + failureMetric := "route.api-v1.upstream.cluster-1.model.gpt-4.consumer.user1.metric.llm_failure_count" + _, err := host.GetCounterMetric(failureMetric) + require.Error(t, err, "llm_failure_count should not exist for successful response") + + // Verify success metrics exist + durationCountMetric := "route.api-v1.upstream.cluster-1.model.gpt-4.consumer.user1.metric.llm_duration_count" + durationCountValue, err := host.GetCounterMetric(durationCountMetric) + require.NoError(t, err) + require.Equal(t, uint64(1), durationCountValue) + }) + }) +} + +func TestStreamingFailureCountMetric(t *testing.T) { + test.RunTest(t, func(t *testing.T) { + t.Run("sse error in middle chunk before done", func(t *testing.T) { + host, status := test.NewTestHost(streamingBodyConfig) + defer host.Reset() + require.Equal(t, types.OnPluginStartStatusOK, status) + + host.SetRouteName("api-v1") + host.SetClusterName("cluster-1") + + host.CallOnHttpRequestHeaders([][2]string{ + {":authority", "example.com"}, + {":path", "/v1/chat/completions"}, + {":method", "POST"}, + {"x-mse-consumer", "user1"}, + }) + + requestBody := []byte(`{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}`) + host.CallOnHttpRequestBody(requestBody) + + host.CallOnHttpResponseHeaders([][2]string{ + {":status", "200"}, + {"content-type", "text/event-stream"}, + }) + + // Simulate SSE stream: error appears in middle chunk, last chunk is [DONE] + chunk1 := []byte(`data: {"choices":[{"delta":{"content":"Hello"}}]}`) + host.CallOnHttpStreamingResponseBody(chunk1, false) + + chunk2 := []byte(`data: {"error":{"type":"api_error","message":"Something went wrong"}}`) + host.CallOnHttpStreamingResponseBody(chunk2, false) + + chunk3 := []byte(`data: [DONE]`) + host.CallOnHttpStreamingResponseBody(chunk3, true) + + host.CompleteHttp() + + // Verify llm_failure_count is incremented despite [DONE] being the last chunk + failureMetric := "route.api-v1.upstream.cluster-1.model.gpt-4.consumer.user1.metric.llm_failure_count" + failureValue, err := host.GetCounterMetric(failureMetric) + require.NoError(t, err) + require.Equal(t, uint64(1), failureValue) + }) + }) +}