higress/plugins/wasm-go/extensions/ai-security-guard/config/ai_log.go

package config

import (
	"strings"
	"time"

	"github.com/higress-group/wasm-go/pkg/wrapper"
	"github.com/tidwall/gjson"
)

const (
	SafecheckRequestsKey   = "safecheck_requests"
	SafecheckRequestIDsKey = "safecheck_request_ids"
	SafecheckRequestIDKey  = "safecheck_request_id"

	GuardrailPhaseRequest  = "request"
	GuardrailPhaseResponse = "response"

	GuardrailModalityText  = "text"
	GuardrailModalityImage = "image"
	GuardrailModalityMCP   = "mcp"

	GuardrailResultPass  = "pass"
	GuardrailResultDeny  = "deny"
	GuardrailResultMask  = "mask"
	GuardrailResultError = "error"
)

type GuardrailSubmissionEvent struct {
	RequestID string `json:"requestId,omitempty"`
	Phase     string `json:"phase"`
	Modality  string `json:"modality"`
	Result    string `json:"result"`
}

// BeginGuardrailSubmissionEvent appends a placeholder event so append order matches
// the current serial submission order. The event is flushed only after completion.
func BeginGuardrailSubmissionEvent(ctx wrapper.HttpContext, phase, modality string) int {
	events := getGuardrailSubmissionEvents(ctx)
	events = append(events, GuardrailSubmissionEvent{
		Phase:    phase,
		Modality: modality,
	})
	setGuardrailSubmissionEvents(ctx, events)
	return len(events) - 1
}

func CompleteGuardrailSubmissionEvent(ctx wrapper.HttpContext, index int, responseBody []byte, result string) {
	CompleteGuardrailSubmissionEventWithRequestID(ctx, index, ExtractValidRequestID(responseBody), result)
}

func CompleteGuardrailSubmissionEventWithRequestID(ctx wrapper.HttpContext, index int, requestID, result string) {
	events := getGuardrailSubmissionEvents(ctx)
	if index < 0 || index >= len(events) {
		return
	}
	events[index].Result = result
	if requestID != "" {
		events[index].RequestID = requestID
	}
	setGuardrailSubmissionEvents(ctx, events)
}

// WriteGuardrailLog writes current guardrail-related user attributes to the AI log.
// Call after submission events are updated; Complete* does not flush the log.
func WriteGuardrailLog(ctx wrapper.HttpContext) {
	ctx.WriteUserAttributeToLogWithKey(wrapper.AILogKey)
}

// MarkGuardrailRequestError finalizes a request-phase safecheck submission that failed
// upstream (HTTP, unmarshal, downstream build, or dispatch error). It overwrites the
// legacy safecheck_status with "request error" so consumers that only watch the single
// status field do not see a stale "request pass" left over from a prior chunk, and
// records safecheck_request_rt so latency metrics cover failures too.
func MarkGuardrailRequestError(ctx wrapper.HttpContext, index int, responseBody []byte, startTime int64) {
	ctx.SetUserAttribute("safecheck_request_rt", time.Now().UnixMilli()-startTime)
	ctx.SetUserAttribute("safecheck_status", "request error")
	CompleteGuardrailSubmissionEvent(ctx, index, responseBody, GuardrailResultError)
	WriteGuardrailLog(ctx)
}

// MarkGuardrailResponseError is the response-phase counterpart of MarkGuardrailRequestError.
func MarkGuardrailResponseError(ctx wrapper.HttpContext, index int, responseBody []byte, startTime int64) {
	ctx.SetUserAttribute("safecheck_response_rt", time.Now().UnixMilli()-startTime)
	ctx.SetUserAttribute("safecheck_status", "response error")
	CompleteGuardrailSubmissionEvent(ctx, index, responseBody, GuardrailResultError)
	WriteGuardrailLog(ctx)
}

func ExtractValidRequestID(responseBody []byte) string {
	if len(responseBody) == 0 {
		return ""
	}
	requestID := gjson.GetBytes(responseBody, "RequestId")
	if !requestID.Exists() || requestID.Type != gjson.String {
		return ""
	}
	trimmed := strings.TrimSpace(requestID.String())
	if trimmed == "" {
		return ""
	}
	return trimmed
}

func getGuardrailSubmissionEvents(ctx wrapper.HttpContext) []GuardrailSubmissionEvent {
	events, ok := ctx.GetUserAttribute(SafecheckRequestsKey).([]GuardrailSubmissionEvent)
	if !ok || events == nil {
		return []GuardrailSubmissionEvent{}
	}
	return events
}

func setGuardrailSubmissionEvents(ctx wrapper.HttpContext, events []GuardrailSubmissionEvent) {
	ctx.SetUserAttribute(SafecheckRequestsKey, events)

	requestIDs := make([]string, 0, len(events))
	for _, event := range events {
		if event.RequestID != "" {
			requestIDs = append(requestIDs, event.RequestID)
		}
	}
	ctx.SetUserAttribute(SafecheckRequestIDsKey, requestIDs)
	if len(requestIDs) > 0 {
		ctx.SetUserAttribute(SafecheckRequestIDKey, requestIDs[len(requestIDs)-1])
	}
}