fix(ai-statistics): skip empty streaming model values (#3851)

Signed-off-by: jianwei.wjw <jianwei.wjw@alibaba-inc.com>
2026-06-05 10:37:28 +08:00 · 2026-06-02 12:10:46 +08:00
parent 52c99eb27d
commit a45b1dde6c
4 changed files with 180 additions and 2 deletions
--- a/plugins/wasm-go/extensions/ai-statistics/main_test.go
+++ b/plugins/wasm-go/extensions/ai-statistics/main_test.go
@@ -21,6 +21,7 @@ import (

 	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
 	"github.com/higress-group/wasm-go/pkg/test"
+	"github.com/higress-group/wasm-go/pkg/wrapper"
 	"github.com/stretchr/testify/require"
 )

@@ -110,6 +111,33 @@ var streamingBodyConfig = func() json.RawMessage {
 	return data
 }()

+var streamingModelExtractionConfig = func() json.RawMessage {
+	data, _ := json.Marshal(map[string]interface{}{
+		"attributes": []map[string]interface{}{
+			{
+				"key":                   "first_model",
+				"value_source":          "response_streaming_body",
+				"value":                 "model",
+				"rule":                  "first",
+				"apply_to_log":          true,
+				"apply_to_span":         false,
+				"as_separate_log_field": false,
+			},
+			{
+				"key":                   "replace_model",
+				"value_source":          "response_streaming_body",
+				"value":                 "model",
+				"rule":                  "replace",
+				"apply_to_log":          true,
+				"apply_to_span":         false,
+				"as_separate_log_field": false,
+			},
+		},
+		"disable_openai_usage": true,
+	})
+	return data
+}()
+
 // 测试配置：请求体属性配置
 var requestBodyConfig = func() json.RawMessage {
 	data, _ := json.Marshal(map[string]interface{}{
@@ -467,6 +495,17 @@ func TestOnHttpResponseHeaders(t *testing.T) {
 	})
 }

+func getAILogAttributes(t *testing.T, host test.TestHost) map[string]interface{} {
+	t.Helper()
+
+	raw, err := host.GetProperty([]string{wrapper.AILogKey})
+	require.NoError(t, err)
+
+	var attrs map[string]interface{}
+	require.NoError(t, json.Unmarshal([]byte(wrapper.UnmarshalStr(`"`+string(raw)+`"`)), &attrs))
+	return attrs
+}
+
 func TestOnHttpStreamingBody(t *testing.T) {
 	test.RunTest(t, func(t *testing.T) {
 		// 测试流式响应体处理
@@ -511,6 +550,68 @@ func TestOnHttpStreamingBody(t *testing.T) {
 			host.CompleteHttp()
 		})

+		t.Run("streaming first and replace skip empty model chunks", func(t *testing.T) {
+			host, status := test.NewTestHost(streamingModelExtractionConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			action := host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+			})
+			require.Equal(t, types.ActionContinue, action)
+
+			action = host.CallOnHttpResponseHeaders([][2]string{
+				{":status", "200"},
+				{"content-type", "text/event-stream"},
+			})
+			require.Equal(t, types.ActionContinue, action)
+
+			action = host.CallOnHttpStreamingResponseBody([]byte("data: {\"model\":\"\"}\n\n"), false)
+			require.Equal(t, types.ActionContinue, action)
+			action = host.CallOnHttpStreamingResponseBody([]byte("data: {\"model\":null}\n\n"), false)
+			require.Equal(t, types.ActionContinue, action)
+			action = host.CallOnHttpStreamingResponseBody([]byte("data: {\"model\":\"gpt-4o\"}\n\n"), true)
+			require.Equal(t, types.ActionContinue, action)
+
+			attrs := getAILogAttributes(t, host)
+			require.Equal(t, "gpt-4o", attrs["first_model"])
+			require.Equal(t, "gpt-4o", attrs["replace_model"])
+
+			host.CompleteHttp()
+		})
+
+		t.Run("streaming first and replace return nil when model path is missing", func(t *testing.T) {
+			host, status := test.NewTestHost(streamingModelExtractionConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			action := host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+			})
+			require.Equal(t, types.ActionContinue, action)
+
+			action = host.CallOnHttpResponseHeaders([][2]string{
+				{":status", "200"},
+				{"content-type", "text/event-stream"},
+			})
+			require.Equal(t, types.ActionContinue, action)
+
+			action = host.CallOnHttpStreamingResponseBody([]byte("data: {\"choices\":[]}\n\n"), false)
+			require.Equal(t, types.ActionContinue, action)
+			action = host.CallOnHttpStreamingResponseBody([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n"), true)
+			require.Equal(t, types.ActionContinue, action)
+
+			attrs := getAILogAttributes(t, host)
+			require.Nil(t, attrs["first_model"])
+			require.Nil(t, attrs["replace_model"])
+
+			host.CompleteHttp()
+		})
+
 		// 测试不包含 token 统计的流式响应体处理
 		t.Run("streaming body without token usage", func(t *testing.T) {
 			host, status := test.NewTestHost(streamingBodyConfig)
@@ -1452,6 +1553,58 @@ func TestSessionIdExtraction(t *testing.T) {
 	})
 }

+// TestExtractStreamingBodyByJsonPath 单独测试流式响应 body 的 JSONPath 提取规则
+func TestExtractStreamingBodyByJsonPath(t *testing.T) {
+	t.Run("first skips empty string chunk", func(t *testing.T) {
+		// Azure/OpenAI 兼容流可能先返回带空 model 的过滤结果 chunk，后续 chunk 才有真实模型名。
+		chunks := []byte(`data: {"choices":[],"created":0,"id":"","model":"","object":""}
+
+data: {"choices":[{"delta":{"content":""}}],"created":1777444731,"id":"chatcmpl-1","model":"gpt-5.4-2026-03-05","object":"chat.completion.chunk"}`)
+
+		value := extractStreamingBodyByJsonPath(chunks, "model", RuleFirst)
+
+		require.Equal(t, "gpt-5.4-2026-03-05", value)
+	})
+
+	t.Run("replace skips trailing empty string chunk", func(t *testing.T) {
+		chunks := []byte(`data: {"model":"gpt-4o"}
+
+data: {"model":""}`)
+
+		value := extractStreamingBodyByJsonPath(chunks, "model", RuleReplace)
+
+		require.Equal(t, "gpt-4o", value)
+	})
+
+	t.Run("first returns nil when path is missing in all chunks", func(t *testing.T) {
+		chunks := []byte(`data: {"choices":[]}
+
+data: {"choices":[{"delta":{"content":"hello"}}]}`)
+
+		value := extractStreamingBodyByJsonPath(chunks, "model", RuleFirst)
+
+		require.Nil(t, value)
+	})
+
+	t.Run("first skips explicit null chunk", func(t *testing.T) {
+		chunks := []byte(`data: {"model":null}
+
+data: {"model":"gpt-4o"}`)
+
+		value := extractStreamingBodyByJsonPath(chunks, "model", RuleFirst)
+
+		require.Equal(t, "gpt-4o", value)
+	})
+
+	t.Run("zero and false remain valid values", func(t *testing.T) {
+		numberValue := extractStreamingBodyByJsonPath([]byte(`data: {"usage":{"total_tokens":0}}`), "usage.total_tokens", RuleFirst)
+		boolValue := extractStreamingBodyByJsonPath([]byte(`data: {"filtered":false}`), "filtered", RuleFirst)
+
+		require.Equal(t, float64(0), numberValue)
+		require.Equal(t, false, boolValue)
+	})
+}
+
 // TestExtractStreamingToolCalls 单独测试 extractStreamingToolCalls 函数
 func TestExtractStreamingToolCalls(t *testing.T) {
 	t.Run("single tool call assembly", func(t *testing.T) {