feat(ai-security-guard): structured x_higress deny response, error-path metrics, and AI logging (#3894)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: rinfx <yucheng.lxr@alibaba-inc.com>
2026-06-07 11:47:30 +08:00 · 2026-05-29 10:45:10 +08:00
parent 385f8d8b4e
commit c21a38e783
14 changed files with 2181 additions and 195 deletions
--- a/plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/text/openai.go
+++ b/plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/text/openai.go
@@ -2,7 +2,6 @@ package text

 import (
 	"encoding/json"
-	"fmt"
 	"net/http"
 	"strings"
 	"time"
@@ -67,6 +66,8 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 	hasMasked := false
 	maskedContent := []byte(content)
 	sessionID, _ := utils.GenerateHexID(20)
+	currentSubmissionIndex := 0
+	currentImageSubmissionIndex := 0
 	var singleCall func()
 	var singleCallForImage func()
 	// prevContentIndex tracks the start of the current chunk for masking replacement
@@ -74,6 +75,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 	callback := func(statusCode int, responseHeaders http.Header, responseBody []byte) {
 		log.Info(string(responseBody))
 		if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
+			cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
 			proxywasm.ResumeHttpRequest()
 			return
 		}
@@ -81,6 +83,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 		err := json.Unmarshal(responseBody, &response)
 		if err != nil {
 			log.Errorf("%+v", err)
+			cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
 			proxywasm.ResumeHttpRequest()
 			return
 		}
@@ -97,26 +100,17 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 					if replaceErr != nil {
 						log.Errorf("failed to replace request body content, falling back to block: %v", replaceErr)
 						// Fall back to block to prevent leaking sensitive data
-						denyMessage := cfg.DefaultDenyMessage
-						if config.DenyMessage != "" {
-							denyMessage = config.DenyMessage
-						}
-						marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
-						if config.ProtocolOriginal {
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
-						} else if gjson.GetBytes(body, "stream").Bool() {
-							randomID := utils.GenerateRandomChatID()
-							jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
-						} else {
-							randomID := utils.GenerateRandomChatID()
-							jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
+						if sendErr := cfg.SendFallbackDenyResponse(config, gjson.GetBytes(body, "stream").Bool()); sendErr != nil {
+							log.Errorf("failed to build deny response body: %v", sendErr)
+							cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
+							proxywasm.ResumeHttpRequest()
+							return
 						}
 						ctx.DontReadResponseBody()
 						config.IncrementCounter("ai_sec_request_deny", 1)
 						ctx.SetUserAttribute("safecheck_status", "reqeust deny")
-						ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+						cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
+						cfg.WriteGuardrailLog(ctx)
 						return
 					}
 					proxywasm.ReplaceHttpRequestBody(newBody)
@@ -125,10 +119,13 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 				} else {
 					ctx.SetUserAttribute("safecheck_status", "request pass")
 				}
-				ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+			}
+			cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultPass)
+			if contentIndex >= len(maskedContent) {
 				if len(images) > 0 && config.CheckRequestImage {
 					singleCallForImage()
 				} else {
+					cfg.WriteGuardrailLog(ctx)
 					proxywasm.ResumeHttpRequest()
 				}
 			} else {
@@ -140,6 +137,30 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 			if desensitization == "" {
 				proxywasm.LogInfof("safecheck_action_source=mask_fallback_to_block, reason=empty_desensitization")
 				log.Warnf("desensitization content is empty, falling back to block logic")
+				// Keep this fallback separate from RiskBlock: legacy reuses the
+				// original deny body in content, while structured emits an empty
+				// fallback guardrail object.
+				isStream := gjson.GetBytes(body, "stream").Bool()
+				var sendErr error
+				if !config.ProtocolOriginal && config.OpenAIDenyResponseFormat != cfg.OpenAIDenyResponseFormatStructured {
+					sendErr = cfg.SendDenyResponse(config, response, consumer, isStream)
+				} else {
+					sendErr = cfg.SendFallbackDenyResponse(config, isStream)
+				}
+				if sendErr != nil {
+					log.Errorf("failed to build deny response body: %v", sendErr)
+					cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
+					proxywasm.ResumeHttpRequest()
+					return
+				}
+				ctx.DontReadResponseBody()
+				config.IncrementCounter("ai_sec_request_deny", 1)
+				endTime := time.Now().UnixMilli()
+				ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
+				ctx.SetUserAttribute("safecheck_status", "reqeust deny")
+				cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
+				cfg.WriteGuardrailLog(ctx)
+				return
 			} else {
 				// Replace only the current chunk portion in maskedContent
 				chunkStart := prevContentIndex
@@ -156,28 +177,19 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 					if replaceErr != nil {
 						log.Errorf("failed to replace request body content, falling back to block: %v", replaceErr)
 						// Fall back to block to prevent leaking sensitive data
-						denyMessage := cfg.DefaultDenyMessage
-						if config.DenyMessage != "" {
-							denyMessage = config.DenyMessage
-						}
-						marshalledDenyMessage := wrapper.MarshalStr(denyMessage)
-						if config.ProtocolOriginal {
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(marshalledDenyMessage), -1)
-						} else if gjson.GetBytes(body, "stream").Bool() {
-							randomID := utils.GenerateRandomChatID()
-							jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
-						} else {
-							randomID := utils.GenerateRandomChatID()
-							jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
-							proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
+						if sendErr := cfg.SendFallbackDenyResponse(config, gjson.GetBytes(body, "stream").Bool()); sendErr != nil {
+							log.Errorf("failed to build deny response body: %v", sendErr)
+							cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
+							proxywasm.ResumeHttpRequest()
+							return
 						}
 						ctx.DontReadResponseBody()
 						config.IncrementCounter("ai_sec_request_deny", 1)
 						endTime := time.Now().UnixMilli()
 						ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
 						ctx.SetUserAttribute("safecheck_status", "reqeust deny")
-						ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+						cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
+						cfg.WriteGuardrailLog(ctx)
 						return
 					}
 					proxywasm.ReplaceHttpRequestBody(newBody)
@@ -185,52 +197,41 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 					endTime := time.Now().UnixMilli()
 					ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
 					ctx.SetUserAttribute("safecheck_status", "request mask")
-					ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+					cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultMask)
 					if len(images) > 0 && config.CheckRequestImage {
 						singleCallForImage()
 					} else {
+						cfg.WriteGuardrailLog(ctx)
 						proxywasm.ResumeHttpRequest()
 					}
 				} else {
+					cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultMask)
 					singleCall()
 				}
 				return
 			}
-			// Fall through to block logic when desensitization is empty
-			fallthrough
 		case cfg.RiskBlock:
-			denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
-			if err != nil {
+			if err := cfg.SendDenyResponse(config, response, consumer, gjson.GetBytes(body, "stream").Bool()); err != nil {
 				log.Errorf("failed to build deny response body: %v", err)
+				cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
 				proxywasm.ResumeHttpRequest()
 				return
 			}
-			if config.ProtocolOriginal {
-				proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
-			} else if gjson.GetBytes(body, "stream").Bool() {
-				randomID := utils.GenerateRandomChatID()
-				marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
-				jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
-				proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
-			} else {
-				randomID := utils.GenerateRandomChatID()
-				marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
-				jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
-				proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
-			}
 			ctx.DontReadResponseBody()
 			config.IncrementCounter("ai_sec_request_deny", 1)
 			endTime := time.Now().UnixMilli()
 			ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
 			ctx.SetUserAttribute("safecheck_status", "reqeust deny")
-			if response.Data.Advice != nil {
+			if len(response.Data.Result) > 0 {
 				ctx.SetUserAttribute("safecheck_riskLabel", response.Data.Result[0].Label)
 				ctx.SetUserAttribute("safecheck_riskWords", response.Data.Result[0].RiskWords)
 			}
-			ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+			cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
+			cfg.WriteGuardrailLog(ctx)
 		}
 	}
 	singleCall = func() {
+		currentSubmissionIndex = cfg.BeginGuardrailSubmissionEvent(ctx, cfg.GuardrailPhaseRequest, cfg.GuardrailModalityText)
 		prevContentIndex = contentIndex
 		var nextContentIndex int
 		if contentIndex+cfg.LengthLimit >= len(maskedContent) {
@@ -245,6 +246,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 		err := config.Client.Post(path, headers, body, callback, config.Timeout)
 		if err != nil {
 			log.Errorf("failed call the safe check service: %v", err)
+			cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, nil, startTime)
 			proxywasm.ResumeHttpRequest()
 		}
 	}
@@ -253,6 +255,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 		imageIndex += 1
 		log.Info(string(responseBody))
 		if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
+			cfg.MarkGuardrailRequestError(ctx, currentImageSubmissionIndex, responseBody, startTime)
 			if imageIndex < len(images) {
 				singleCallForImage()
 			} else {
@@ -264,6 +267,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 		err := json.Unmarshal(responseBody, &response)
 		if err != nil {
 			log.Errorf("%+v", err)
+			cfg.MarkGuardrailRequestError(ctx, currentImageSubmissionIndex, responseBody, startTime)
 			if imageIndex < len(images) {
 				singleCallForImage()
 			} else {
@@ -276,7 +280,10 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 			if imageIndex >= len(images) {
 				ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
 				ctx.SetUserAttribute("safecheck_status", "request pass")
-				ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+			}
+			cfg.CompleteGuardrailSubmissionEvent(ctx, currentImageSubmissionIndex, responseBody, cfg.GuardrailResultPass)
+			if imageIndex >= len(images) {
+				cfg.WriteGuardrailLog(ctx)
 				proxywasm.ResumeHttpRequest()
 			} else {
 				singleCallForImage()
@@ -284,36 +291,25 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 			return
 		}

-		denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
-		if err != nil {
+		if err := cfg.SendDenyResponse(config, response, consumer, gjson.GetBytes(body, "stream").Bool()); err != nil {
 			log.Errorf("failed to build deny response body: %v", err)
+			cfg.MarkGuardrailRequestError(ctx, currentImageSubmissionIndex, responseBody, startTime)
 			proxywasm.ResumeHttpRequest()
 			return
 		}
-		if config.ProtocolOriginal {
-			proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, denyBody, -1)
-		} else if gjson.GetBytes(body, "stream").Bool() {
-			randomID := utils.GenerateRandomChatID()
-			marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
-			jsonData := []byte(fmt.Sprintf(cfg.OpenAIStreamResponseFormat, randomID, marshalledDenyMessage, randomID))
-			proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "text/event-stream;charset=UTF-8"}}, jsonData, -1)
-		} else {
-			randomID := utils.GenerateRandomChatID()
-			marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
-			jsonData := []byte(fmt.Sprintf(cfg.OpenAIResponseFormat, randomID, marshalledDenyMessage))
-			proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, jsonData, -1)
-		}
 		ctx.DontReadResponseBody()
 		config.IncrementCounter("ai_sec_request_deny", 1)
 		ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
 		ctx.SetUserAttribute("safecheck_status", "reqeust deny")
-		if response.Data.Advice != nil {
+		if len(response.Data.Result) > 0 {
 			ctx.SetUserAttribute("safecheck_riskLabel", response.Data.Result[0].Label)
 			ctx.SetUserAttribute("safecheck_riskWords", response.Data.Result[0].RiskWords)
 		}
-		ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
+		cfg.CompleteGuardrailSubmissionEvent(ctx, currentImageSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
+		cfg.WriteGuardrailLog(ctx)
 	}
 	singleCallForImage = func() {
+		currentImageSubmissionIndex = cfg.BeginGuardrailSubmissionEvent(ctx, cfg.GuardrailPhaseRequest, cfg.GuardrailModalityImage)
 		img := images[imageIndex]
 		imgUrl := ""
 		imgBase64 := ""
@@ -326,6 +322,7 @@ func HandleTextGenerationRequestBody(ctx wrapper.HttpContext, config cfg.AISecur
 		err := config.Client.Post(path, headers, body, callbackForImage, config.Timeout)
 		if err != nil {
 			log.Errorf("failed call the safe check service: %v", err)
+			cfg.MarkGuardrailRequestError(ctx, currentImageSubmissionIndex, nil, startTime)
 			proxywasm.ResumeHttpRequest()
 		}
 	}