jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-09 04:37:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
JianweiWang
2026-05-29 10:45:10 +08:00
committed by GitHub
parent 385f8d8b4e
commit c21a38e783
14 changed files with 2181 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
plugins/wasm-go/extensions/ai-security-guard/lvwang/multi_modal_guard/mcp/mcp.go
View File

@@ -43,10 +43,12 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
}
contentIndex := 0
sessionID, _ := utils.GenerateHexID(20)
currentSubmissionIndex := 0
var singleCall func()
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
}
@@ -54,6 +56,7 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
err := json.Unmarshal(responseBody, &response)
if err != nil {
log.Errorf("%+v", err)
cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
proxywasm.ResumeHttpRequest()
return
}
@@ -62,7 +65,10 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
endTime := time.Now().UnixMilli()
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
ctx.SetUserAttribute("safecheck_status", "request pass")
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
}
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultPass)
if contentIndex >= len(content) {
cfg.WriteGuardrailLog(ctx)
proxywasm.ResumeHttpRequest()
} else {
singleCall()
@@ -74,22 +80,25 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
endTime := time.Now().UnixMilli()
ctx.SetUserAttribute("safecheck_request_rt", endTime-startTime)
ctx.SetUserAttribute("safecheck_status", "request 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)
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
if err != nil {
log.Errorf("failed to build deny response body: %v", err)
cfg.MarkGuardrailRequestError(ctx, currentSubmissionIndex, responseBody, startTime)
proxywasm.ResumeHttpRequest()
return
}
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
cfg.WriteGuardrailLog(ctx)
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
denyResponse := fmt.Sprintf(DenyResponse, marshalledDenyMessage)
proxywasm.SendHttpResponse(uint32(config.DenyCode), [][2]string{{"content-type", "application/json"}}, []byte(denyResponse), -1)
}
singleCall = func() {
currentSubmissionIndex = cfg.BeginGuardrailSubmissionEvent(ctx, cfg.GuardrailPhaseRequest, cfg.GuardrailModalityMCP)
var nextContentIndex int
if contentIndex+cfg.LengthLimit >= len(content) {
nextContentIndex = len(content)
@@ -103,6 +112,7 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
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()
}
}
@@ -114,6 +124,7 @@ func HandleMcpRequestBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig, data []byte, endOfStream bool) []byte {
consumer, _ := ctx.GetContext("consumer").(string)
var frontBuffer []byte
currentSubmissionIndex := 0
var singleCall func()
callback := func(statusCode int, responseHeaders http.Header, responseBody []byte) {
defer func() {
@@ -122,6 +133,9 @@ func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecuri
}()
log.Info(string(responseBody))
if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
ctx.SetUserAttribute("safecheck_status", "response error")
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultError)
cfg.WriteGuardrailLog(ctx)
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
return
}
@@ -129,6 +143,9 @@ func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecuri
err := json.Unmarshal(responseBody, &response)
if err != nil {
log.Error("failed to unmarshal aliyun content security response at response phase")
ctx.SetUserAttribute("safecheck_status", "response error")
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultError)
cfg.WriteGuardrailLog(ctx)
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
return
}
@@ -136,13 +153,20 @@ func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecuri
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
if err != nil {
log.Errorf("failed to build deny response body: %v", err)
ctx.SetUserAttribute("safecheck_status", "response error")
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultError)
cfg.WriteGuardrailLog(ctx)
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
return
}
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
cfg.WriteGuardrailLog(ctx)
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
denySSEResponse := fmt.Sprintf(DenySSEResponse, marshalledDenyMessage)
proxywasm.InjectEncodedDataToFilterChain([]byte(denySSEResponse), true)
} else {
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultPass)
cfg.WriteGuardrailLog(ctx)
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
}
}
@@ -158,10 +182,14 @@ func HandleMcpStreamingResponseBody(ctx wrapper.HttpContext, config cfg.AISecuri
ctx.SetContext("during_call", true)
checkService := config.GetResponseCheckService(consumer)
sessionID, _ := utils.GenerateHexID(20)
currentSubmissionIndex = cfg.BeginGuardrailSubmissionEvent(ctx, cfg.GuardrailPhaseResponse, cfg.GuardrailModalityMCP)
path, headers, body := common.GenerateRequestForText(config, config.Action, checkService, msg, sessionID)
err := config.Client.Post(path, headers, body, callback, config.Timeout)
if err != nil {
log.Errorf("failed call the safe check service: %v", err)
ctx.SetUserAttribute("safecheck_status", "response error")
cfg.CompleteGuardrailSubmissionEventWithRequestID(ctx, currentSubmissionIndex, "", cfg.GuardrailResultError)
cfg.WriteGuardrailLog(ctx)
proxywasm.InjectEncodedDataToFilterChain(frontBuffer, false)
ctx.SetContext("during_call", false)
}
@@ -194,10 +222,12 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
}
contentIndex := 0
sessionID, _ := utils.GenerateHexID(20)
currentSubmissionIndex := 0
var singleCall func()
callback := func(statusCode int, responseHeaders http.Header, responseBody []byte) {
log.Info(string(responseBody))
if statusCode != 200 || gjson.GetBytes(responseBody, "Code").Int() != 200 {
cfg.MarkGuardrailResponseError(ctx, currentSubmissionIndex, responseBody, startTime)
proxywasm.ResumeHttpResponse()
return
}
@@ -205,6 +235,7 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
err := json.Unmarshal(responseBody, &response)
if err != nil {
log.Error("failed to unmarshal aliyun content security response at response phase")
cfg.MarkGuardrailResponseError(ctx, currentSubmissionIndex, responseBody, startTime)
proxywasm.ResumeHttpResponse()
return
}
@@ -213,7 +244,10 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
endTime := time.Now().UnixMilli()
ctx.SetUserAttribute("safecheck_response_rt", endTime-startTime)
ctx.SetUserAttribute("safecheck_status", "response pass")
ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
}
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultPass)
if contentIndex >= len(content) {
cfg.WriteGuardrailLog(ctx)
proxywasm.ResumeHttpResponse()
} else {
singleCall()
@@ -224,17 +258,19 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
endTime := time.Now().UnixMilli()
ctx.SetUserAttribute("safecheck_response_rt", endTime-startTime)
ctx.SetUserAttribute("safecheck_status", "response 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)
denyBody, err := cfg.BuildDenyResponseBody(response, config, consumer)
if err != nil {
log.Errorf("failed to build deny response body: %v", err)
cfg.MarkGuardrailResponseError(ctx, currentSubmissionIndex, responseBody, startTime)
proxywasm.ResumeHttpResponse()
return
}
cfg.CompleteGuardrailSubmissionEvent(ctx, currentSubmissionIndex, responseBody, cfg.GuardrailResultDeny)
cfg.WriteGuardrailLog(ctx)
marshalledDenyMessage := wrapper.MarshalStr(string(denyBody))
denyResponseBody := fmt.Sprintf(DenyResponse, marshalledDenyMessage)
proxywasm.RemoveHttpResponseHeader("content-length")
@@ -253,10 +289,12 @@ func HandleMcpResponseBody(ctx wrapper.HttpContext, config cfg.AISecurityConfig,
contentIndex = nextContentIndex
log.Debugf("current content piece: %s", contentPiece)
checkService := config.GetResponseCheckService(consumer)
currentSubmissionIndex = cfg.BeginGuardrailSubmissionEvent(ctx, cfg.GuardrailPhaseResponse, cfg.GuardrailModalityMCP)
path, headers, body := common.GenerateRequestForText(config, config.Action, checkService, contentPiece, sessionID)
err := config.Client.Post(path, headers, body, callback, config.Timeout)
if err != nil {
log.Errorf("failed call the safe check service: %v", err)
cfg.MarkGuardrailResponseError(ctx, currentSubmissionIndex, nil, startTime)
proxywasm.ResumeHttpResponse()
}
}