feat(ai-statistics): add llm_failure_count metric for error response tracking (#3975)

Signed-off-by: ljb <1327941766@qq.com>
2026-07-24 05:10:34 +08:00 · 2026-07-17 19:16:16 +08:00
parent e01db53e9f
commit 2682c926fc
2 changed files with 312 additions and 4 deletions
--- 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
 	}
--- 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)
+		})
+	})
+}