From 49d1167c8e92d0cc93e912ac8368e6bee3b6b474 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Thu, 3 Sep 2026 12:49:33 +0200 Subject: [PATCH 1/2] feat(ai-guard): apply sensitive data redaction from evaluate responses --- .../custom/aiguard/MessageWriter.java | 30 +- .../aiguard/MessageWriterTest.groovy | 34 + dd-java-agent/agent-aiguard/build.gradle | 2 + .../com/datadog/aiguard/AIGuardInternal.java | 62 +- .../com/datadog/aiguard/MessageRedactor.java | 356 ++++++++++ .../aiguard/AIGuardInternalTests.groovy | 153 +++-- .../aiguard/AIGuardInternalRedactionTest.java | 300 ++++++++ .../datadog/aiguard/MessageRedactorTest.java | 643 ++++++++++++++++++ .../datadog/trace/api/aiguard/AIGuard.java | 89 ++- .../trace/api/aiguard/noop/NoOpEvaluator.java | 9 +- .../trace/api/config/AIGuardConfig.java | 2 + .../trace/api/aiguard/AIGuardTest.java | 90 +++ .../main/java/datadog/trace/api/Config.java | 16 + .../api/telemetry/WafMetricCollector.java | 97 ++- .../telemetry/WafMetricCollectorTest.groovy | 45 +- .../trace/api/ConfigAIGuardRedactionTest.java | 56 ++ metadata/supported-configurations.json | 8 + 17 files changed, 1868 insertions(+), 124 deletions(-) create mode 100644 dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java create mode 100644 dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java create mode 100644 dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java create mode 100644 internal-api/src/test/java/datadog/trace/api/ConfigAIGuardRedactionTest.java diff --git a/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java b/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java index e8876b86df3..55669be9afa 100644 --- a/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java +++ b/communication/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java @@ -13,12 +13,14 @@ public class MessageWriter implements ValueWriter { public void write( final AIGuard.Message value, final Writable writable, final EncodingCache encodingCache) { final int[] size = {0}; - final boolean hasRole = isNotBlank(value.getRole(), size); - final boolean hasToolCallId = isNotBlank(value.getToolCallId(), size); - final boolean hasToolCalls = isNotEmpty(value.getToolCalls(), size); + final boolean hasRole = present(Strings.isNotBlank(value.getRole()), size); + final boolean hasToolCallId = present(Strings.isNotBlank(value.getToolCallId()), size); + final boolean hasToolCalls = present(isNotEmpty(value.getToolCalls()), size); - final boolean hasContentParts = isNotEmpty(value.getContentParts(), size); - final boolean hasContentString = !hasContentParts && isNotBlank(value.getContent(), size); + final boolean hasContentParts = present(isNotEmpty(value.getContentParts()), size); + // An empty content string is still written: "" is what the redaction remove strategy leaves + // behind, and dropping it would be indistinguishable from a message that never had content. + final boolean hasContentString = present(!hasContentParts && value.getContent() != null, size); writable.startMap(size[0]); writeString(hasRole, "role", value.getRole(), writable, encodingCache); @@ -83,19 +85,15 @@ private static void writeToolCallArray( } } - private static boolean isNotBlank(final String value, final int[] nonBlankCount) { - final boolean hasText = Strings.isNotBlank(value); - if (hasText) { - nonBlankCount[0]++; + /** Counts a field towards the map size when it is present, and reports whether it is. */ + private static boolean present(final boolean present, final int[] fieldCount) { + if (present) { + fieldCount[0]++; } - return hasText; + return present; } - private static boolean isNotEmpty(final List value, final int[] nonEmptyCount) { - final boolean nonEmpty = value != null && !value.isEmpty(); - if (nonEmpty) { - nonEmptyCount[0]++; - } - return nonEmpty; + private static boolean isNotEmpty(final List value) { + return value != null && !value.isEmpty(); } } diff --git a/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy b/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy index 45191a2fc8f..fb55fc0e0a8 100644 --- a/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy +++ b/communication/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy @@ -224,6 +224,40 @@ class MessageWriterTest extends DDSpecification { } } + void 'test write message with empty content'() { + given: + // The redaction "remove" strategy replaces content with "", which must stay visible as an + // empty string rather than being dropped like a message that never carried content. + final message = AIGuard.Message.message('user', '') + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringValueMap(unpacker.unpackValue()) + assert value.size() == 2 + assert value.role == 'user' + assert value.content == '' + } + } + + void 'test write message without content'() { + given: + final message = AIGuard.Message.message('user', (String) null) + + when: + writer.writeObject(message, encodingCache) + + then: + try (final unpacker = MessagePack.newDefaultUnpacker(buffer.slice())) { + final value = asStringValueMap(unpacker.unpackValue()) + assert value.size() == 1 + assert value.role == 'user' + assert !value.containsKey('content') + } + } + void 'test backward compatibility with string content'() { given: final message = AIGuard.Message.message('user', 'Plain text message') diff --git a/dd-java-agent/agent-aiguard/build.gradle b/dd-java-agent/agent-aiguard/build.gradle index ad2a97bbf93..7047ed3d34c 100644 --- a/dd-java-agent/agent-aiguard/build.gradle +++ b/dd-java-agent/agent-aiguard/build.gradle @@ -22,6 +22,8 @@ dependencies { implementation project(':communication') testImplementation project(':utils:test-utils') + testImplementation libs.bundles.junit5 + testImplementation libs.bundles.mockito testImplementation('org.skyscreamer:jsonassert:1.5.3') testImplementation('com.fasterxml.jackson.core:jackson-databind:2.20.0') } diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index d921be448d8..cefa5feb4ee 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -28,6 +28,7 @@ import datadog.trace.api.aiguard.noop.NoOpEvaluator; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.telemetry.WafMetricCollector; +import datadog.trace.api.telemetry.WafMetricCollector.AIGuardRedaction; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData; @@ -51,6 +52,8 @@ import okhttp3.Response; import okhttp3.ResponseBody; import okio.BufferedSink; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Concrete implementation of the SDK used to interact with the AIGuard REST API. @@ -60,6 +63,8 @@ */ public class AIGuardInternal implements Evaluator { + private static final Logger log = LoggerFactory.getLogger(AIGuardInternal.class); + public static class BadConfigurationException extends RuntimeException { public BadConfigurationException(final String message) { super(message); @@ -72,6 +77,9 @@ public BadConfigurationException(final String message) { static final String ACTION_TAG = "ai_guard.action"; static final String REASON_TAG = "ai_guard.reason"; static final String BLOCKED_TAG = "ai_guard.blocked"; + static final String REDACTED_TAG = "ai_guard.redacted"; + + static final String RESPONSE_REDACTION_REPLACEMENTS = "redaction_replacements"; static final String META_STRUCT_TAG = "ai_guard"; static final String META_STRUCT_MESSAGES = "messages"; @@ -128,6 +136,7 @@ static void uninstall() { private final OkHttpClient client; private final Map meta; private final Map headers; + private final MessageRedactor redactor; AIGuardInternal(final HttpUrl url, final Map headers, final OkHttpClient client) { this.url = url; @@ -136,13 +145,17 @@ static void uninstall() { this.moshi = new Moshi.Builder().add(new AIGuardFactory()).build(); final Config config = Config.get(); this.meta = mapOf("service", config.getServiceName(), "env", config.getEnv()); + this.redactor = + config.isAiGuardRedactionEnabled() + ? new MessageRedactor.DefaultRedactor() + : new MessageRedactor.NoOp(); } /** * Creates a deep copy of the messages before storing them in the metastruct to avoid concurrent * modifications prior to trace serialization. */ - private static List messagesForMetaStruct(List messages) { + private static List messagesForMetaStruct(final List messages) { final Config config = Config.get(); final int size = Math.min(messages.size(), config.getAiGuardMaxMessagesLength()); if (size < messages.size()) { @@ -191,6 +204,32 @@ private static List messagesForMetaStruct(List messages) { return result; } + /** + * Applies the redaction requested by the AI Guard service and reports the outcome on the span. + * + *

This runs before the blocking decision on purpose: a blocked evaluation still reports its + * conversation through the meta struct, and that report must be redacted too. The {@link + * AIGuardAbortError} raised on that path deliberately carries no messages. + * + * @return the telemetry state, {@link AIGuardRedaction#DISABLED} when the kill switch is off, in + * which case no {@code ai_guard.redacted} tag is attached either + */ + private AIGuardRedaction reportRedaction( + final AgentSpan span, final MessageRedactor.Result redaction) { + if (!redactor.enabled()) { + // No tag at all, so an absent tag ("redaction is off") stays distinguishable from a false + // one ("redaction is on and nothing was redacted"). + return AIGuardRedaction.DISABLED; + } + span.setTag(REDACTED_TAG, redaction.redacted()); + if (redaction.skipped > 0) { + log.debug( + "AI Guard skipped {} redaction replacement(s) that could not be applied", + redaction.skipped); + } + return redaction.redacted() ? AIGuardRedaction.APPLIED : AIGuardRedaction.NOT_APPLIED; + } + private static boolean isToolCall(final Message message) { return message.getToolCalls() != null || message.getToolCallId() != null; } @@ -283,6 +322,7 @@ public Evaluation evaluate(final List messages, final Options options) // sure client IP tags were populated. copyAnomalyDetectionTags(span, localRootSpan); } + List finalMessages = messages; try (final ContextScope scope = tracer.activateSpan(span)) { final Message last = messages.get(messages.size() - 1); if (isToolCall(last)) { @@ -295,7 +335,6 @@ public Evaluation evaluate(final List messages, final Options options) span.setTag(TARGET_TAG, "prompt"); } final Map metaStruct = new HashMap<>(2); - metaStruct.put(META_STRUCT_MESSAGES, messagesForMetaStruct(messages)); span.setMetaStruct(META_STRUCT_TAG, metaStruct); final Request.Builder request = new Request.Builder() @@ -329,14 +368,29 @@ public Evaluation evaluate(final List messages, final Options options) if (sdsFindings != null && !sdsFindings.isEmpty()) { metaStruct.put(META_STRUCT_SDS, sdsFindings); } + final Object rawReplacements = result.get(RESPONSE_REDACTION_REPLACEMENTS); + final MessageRedactor.Result redaction = + redactor.redact( + messages, rawReplacements instanceof List ? (List) rawReplacements : null); + final AIGuardRedaction redactionState = reportRedaction(span, redaction); + finalMessages = redaction.messages; final boolean shouldBlock = isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW; - WafMetricCollector.get().aiGuardRequest(action, shouldBlock); + WafMetricCollector.get().aiGuardRequest(action, shouldBlock, redactionState); if (shouldBlock) { span.setTag(BLOCKED_TAG, true); throw new AIGuardAbortError(action, reason, tags, tagProbs, sdsFindings); } - return new Evaluation(action, reason, tags, tagProbs, sdsFindings); + return new Evaluation( + action, + reason, + tags, + tagProbs, + sdsFindings, + redaction.messages, + redaction.replacements); + } finally { + metaStruct.put(META_STRUCT_MESSAGES, messagesForMetaStruct(finalMessages)); } } catch (AIGuardAbortError e) { span.addThrowable(e); diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java new file mode 100644 index 00000000000..da7dea9d96a --- /dev/null +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java @@ -0,0 +1,356 @@ +package com.datadog.aiguard; + +import datadog.trace.api.aiguard.AIGuard.ContentPart; +import datadog.trace.api.aiguard.AIGuard.Message; +import datadog.trace.api.aiguard.AIGuard.ToolCall; +import datadog.trace.api.aiguard.AIGuard.ToolCall.Function; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import javax.annotation.Nullable; + +public interface MessageRedactor { + + /** Outcome of a redaction pass. */ + final class Result { + + /** The redacted messages, or the very same list that was passed in when nothing was applied. */ + final List messages; + + /** + * The {@code {path, replacement}} entries that were actually overwritten, in the order the + * service returned them. Entries skipped fail-safe never appear here, so this list describes + * exactly the transformation {@link #messages} underwent, and it is what the SDK hands back + * through {@code Evaluation.getRedactionReplacements()}. + */ + final List> replacements; + + /** Number of paths successfully overwritten. */ + final int applied; + + /** Number of entries skipped fail-safe (unresolvable, non-string, missing or conflicting). */ + final int skipped; + + private Result( + final List messages, + final List> replacements, + final int skipped) { + this.messages = messages; + this.replacements = replacements; + this.applied = replacements.size(); + this.skipped = skipped; + } + + /** Whether at least one replacement was applied. */ + boolean redacted() { + return applied > 0; + } + + private static Result nothingApplied(final List messages, final int skipped) { + return new Result(messages, Collections.emptyList(), skipped); + } + } + + Result redact(final List messages, @Nullable final List replacements); + + /** + * Whether redaction is active. When {@code false} the {@code ai_guard.redacted} tag is not + * reported at all, so an absent tag means "redaction is off", which stays distinguishable from a + * {@code false} one meaning "redaction is on and nothing was redacted". + */ + boolean enabled(); + + class NoOp implements MessageRedactor { + + @Override + public Result redact(final List messages, final @Nullable List replacements) { + return Result.nothingApplied(messages, 0); + } + + @Override + public boolean enabled() { + return false; + } + } + + /** + * Applies the {@code redaction_replacements} returned by the AI Guard service to a message list. + * + *

The service returns the fully redacted string for each affected path, so this class + * never slices, concatenates, or reasons about offsets and string encodings: it resolves a path + * to a single string and overwrites it verbatim. Placeholder selection and redaction strategy are + * resolved server side and already baked into each replacement. + * + *

Two properties matter to callers: + * + *

    + *
  • Copy on write. The caller's list and messages are never mutated. Only + * the messages that are actually redacted are rebuilt; the rest are shared with the input + * list. When nothing is applied, {@link Result#messages} is the very same reference that + * was passed in, so {@code result.messages != messages} tells the caller whether anything + * changed. + *
  • Never throws. A malformed response must not break the caller's control + * flow, so unresolvable paths, non-string targets, and missing or conflicting replacements + * are skipped and counted in {@link Result#skipped}. + *
+ */ + class DefaultRedactor implements MessageRedactor { + + @Override + public boolean enabled() { + return true; + } + + /** + * Matches a single path segment, e.g. {@code messages[1]} or {@code function}. Kept verbatim + * from the cross-tracer specification, so every tracer tokenizes paths identically. + */ + private static final Pattern SEGMENT = + Pattern.compile("\\A([A-Za-z0-9_]+)(?:\\[([0-9]+)\\])?\\z"); + + /** No supported target is deeper than {@code messages[i].tool_calls[k].function.arguments}. */ + private static final int MAX_SEGMENTS = 4; + + private static final int NO_INDEX = -1; + private static final int INVALID_INDEX = -2; + + /** Longest index we bother parsing; anything longer cannot address a real list. */ + private static final int MAX_INDEX_DIGITS = 9; + + private static Map entry(final String path, final String replacement) { + final Map entry = new LinkedHashMap<>(4); + entry.put("path", path); + entry.put("replacement", replacement); + return Collections.unmodifiableMap(entry); + } + + /** + * Overwrites every path in {@code replacements} with its replacement string. + * + * @param messages the evaluated messages, never mutated + * @param replacements the raw {@code redaction_replacements} array from the response + * @return the redaction outcome, holding {@code messages} itself when nothing was applied + */ + @Override + public Result redact(final List messages, @Nullable final List replacements) { + if (messages == null + || messages.isEmpty() + || replacements == null + || replacements.isEmpty()) { + return Result.nothingApplied(messages, 0); + } + + int skipped = 0; + + // Collect one authoritative replacement per path, dropping the ones the backend contradicts + // itself on. Insertion ordered so that a malformed response yields reproducible counters. + final Map byPath = new LinkedHashMap<>(); + Set conflicting = null; + for (final Object entry : replacements) { + if (!(entry instanceof Map)) { + skipped++; + continue; + } + final Map map = (Map) entry; + final Object path = map.get("path"); + final Object replacement = map.get("replacement"); + // An empty replacement is legitimate: "" is a supported placeholder meaning "remove". + if (!(path instanceof String) + || ((String) path).isEmpty() + || !(replacement instanceof String)) { + skipped++; + continue; + } + final String previous = byPath.put((String) path, (String) replacement); + if (previous != null && !previous.equals(replacement)) { + // Conflicting replacements for one path: skip it rather than guess which one wins. + if (conflicting == null) { + conflicting = new HashSet<>(2); + } + conflicting.add((String) path); + } + } + if (conflicting != null) { + for (final String path : conflicting) { + byPath.remove(path); + skipped++; + } + } + if (byPath.isEmpty()) { + return Result.nothingApplied(messages, skipped); + } + + final String[] names = new String[MAX_SEGMENTS]; + final int[] indices = new int[MAX_SEGMENTS]; + List working = null; + final List> applied = new ArrayList<>(byPath.size()); + + for (final Map.Entry entry : byPath.entrySet()) { + final int count = parseSegments(entry.getKey(), names, indices); + // Paths are rooted at the evaluated array, so they always start with `messages[i]`. + if (count < 2 + || indices[0] == NO_INDEX + || !"messages".equals(names[0]) + || indices[0] >= messages.size()) { + skipped++; + continue; + } + final int index = indices[0]; + final Message current = working == null ? messages.get(index) : working.get(index); + final Message updated = apply(current, names, indices, count, entry.getValue()); + if (updated == null) { + skipped++; + continue; + } + if (working == null) { + working = new ArrayList<>(messages); + } + working.set(index, updated); + applied.add(entry(entry.getKey(), entry.getValue())); + } + + if (working == null) { + return Result.nothingApplied(messages, skipped); + } + return new Result(working, Collections.unmodifiableList(applied), skipped); + } + + /** + * Rebuilds {@code message} with {@code replacement} written at the target the path resolves to. + * + * @return the rebuilt message, or {@code null} when the path does not resolve to a writable + * string, in which case the caller skips it fail-safe + */ + @Nullable + private static Message apply( + final Message message, + final String[] names, + final int[] indices, + final int count, + final String replacement) { + + if ("content".equals(names[1])) { + if (indices[1] == NO_INDEX) { + // messages[i].content + if (count != 2 || message.getContentParts() != null || message.getContent() == null) { + // A message holding content parts resolves to a list, not a string: skip. + return null; + } + return withContent(message, replacement); + } + // messages[i].content[j].text + if (count != 3 || !"text".equals(names[2]) || indices[2] != NO_INDEX) { + return null; + } + final List parts = message.getContentParts(); + if (parts == null || indices[1] >= parts.size()) { + return null; + } + final ContentPart part = parts.get(indices[1]); + // Only text parts are redactable; image locators are out of scope. + if (part.getType() != ContentPart.Type.TEXT) { + return null; + } + final List updated = new ArrayList<>(parts); + updated.set(indices[1], ContentPart.text(replacement)); + return withContentParts(message, updated); + } + + if ("tool_calls".equals(names[1])) { + // messages[i].tool_calls[k].function.arguments + if (indices[1] == NO_INDEX + || count != 4 + || !"function".equals(names[2]) + || indices[2] != NO_INDEX + || !"arguments".equals(names[3]) + || indices[3] != NO_INDEX) { + return null; + } + final List toolCalls = message.getToolCalls(); + if (toolCalls == null || indices[1] >= toolCalls.size()) { + return null; + } + final ToolCall toolCall = toolCalls.get(indices[1]); + final Function function = toolCall.getFunction(); + if (function == null || function.getArguments() == null) { + return null; + } + final List updated = new ArrayList<>(toolCalls); + updated.set( + indices[1], + new ToolCall(toolCall.getId(), new Function(function.getName(), replacement))); + return withToolCalls(message, updated); + } + + return null; + } + + private static Message withContent(final Message message, final String content) { + return new Message( + message.getRole(), content, message.getToolCalls(), message.getToolCallId()); + } + + private static Message withContentParts( + final Message message, final List contentParts) { + return new Message( + message.getRole(), contentParts, message.getToolCalls(), message.getToolCallId()); + } + + private static Message withToolCalls(final Message message, final List toolCalls) { + final List contentParts = message.getContentParts(); + // A message carries either content parts or a content string, never both; preserve whichever. + return contentParts != null + ? new Message(message.getRole(), contentParts, toolCalls, message.getToolCallId()) + : new Message( + message.getRole(), message.getContent(), toolCalls, message.getToolCallId()); + } + + /** + * Splits a path on {@code .} and matches every segment against {@link #SEGMENT}, filling {@code + * names} and {@code indices}. + * + * @return the number of segments, or {@code -1} when any segment is malformed or the path is + * deeper than any supported target + */ + private static int parseSegments(final String path, final String[] names, final int[] indices) { + int count = 0; + int start = 0; + while (true) { + if (count == MAX_SEGMENTS) { + return -1; + } + final int dot = path.indexOf('.', start); + final String segment = dot < 0 ? path.substring(start) : path.substring(start, dot); + final Matcher matcher = SEGMENT.matcher(segment); + if (!matcher.matches()) { + return -1; + } + final int index = parseIndex(matcher.group(2)); + if (index == INVALID_INDEX) { + return -1; + } + names[count] = matcher.group(1); + indices[count] = index; + count++; + if (dot < 0) { + return count; + } + start = dot + 1; + } + } + + private static int parseIndex(@Nullable final String index) { + if (index == null) { + return NO_INDEX; + } + // The regex guarantees digits only, so the sole remaining hazard is overflow. + return index.length() > MAX_INDEX_DIGITS ? INVALID_INDEX : Integer.parseInt(index); + } + } +} diff --git a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy index 114e83da3f1..22e1069d39b 100644 --- a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy +++ b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy @@ -225,7 +225,8 @@ class AIGuardInternalTests extends DDSpecification { eval.tagProbabilities == suite.tagProbabilities eval.sds == [] } - assertTelemetry('ai_guard.requests', "action:$suite.action", "block:$throwAbortError", 'error:false') + // no redaction_replacements in these responses, and redaction is on by default + assertTelemetry('requests', "action:$suite.action", "block:$throwAbortError", 'error:false', 'redacted:false') where: suite << TestSuite.build() @@ -390,7 +391,7 @@ class AIGuardInternalTests extends DDSpecification { final exception = thrown(AIGuard.AIGuardClientError) exception.errors == errors 1 * span.addThrowable(_ as AIGuard.AIGuardClientError) - assertTelemetry('ai_guard.requests', 'error:true') + assertTelemetry('requests', 'error:true') } void 'test evaluate with invalid JSON'() { @@ -403,7 +404,7 @@ class AIGuardInternalTests extends DDSpecification { then: thrown(AIGuard.AIGuardClientError) 1 * span.addThrowable(_ as AIGuard.AIGuardClientError) - assertTelemetry('ai_guard.requests', 'error:true') + assertTelemetry('requests', 'error:true') } void 'test evaluate with missing action'() { @@ -416,7 +417,7 @@ class AIGuardInternalTests extends DDSpecification { then: thrown(AIGuard.AIGuardClientError) 1 * span.addThrowable(_ as AIGuard.AIGuardClientError) - assertTelemetry('ai_guard.requests', 'error:true') + assertTelemetry('requests', 'error:true') } void 'test evaluate with non JSON response'() { @@ -429,7 +430,7 @@ class AIGuardInternalTests extends DDSpecification { then: thrown(AIGuard.AIGuardClientError) 1 * span.addThrowable(_ as AIGuard.AIGuardClientError) - assertTelemetry('ai_guard.requests', 'error:true') + assertTelemetry('requests', 'error:true') } void 'test evaluate with empty response'() { @@ -442,11 +443,12 @@ class AIGuardInternalTests extends DDSpecification { then: thrown(AIGuard.AIGuardClientError) 1 * span.addThrowable(_ as AIGuard.AIGuardClientError) - assertTelemetry('ai_guard.requests', 'error:true') + assertTelemetry('requests', 'error:true') } void 'test message length truncation'() { given: + Map receivedMeta = null final maxMessages = Config.get().getAiGuardMaxMessagesLength() final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) final messages = (0..maxMessages) @@ -458,15 +460,18 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final received = (List) it[1].messages - assert received.size() == maxMessages - assert received.size() < messages.size() + receivedMeta = it[1] as Map + return span } - assertTelemetry('ai_guard.truncated', 'type:messages') + final received = (List) receivedMeta.messages + assert received.size() == maxMessages + assert received.size() < messages.size() + assertTelemetry('truncated', 'type:messages') } void 'test message content truncation'() { given: + Map receivedMeta = null final maxContent = Config.get().getAiGuardMaxContentSize() final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine']]]) final message = AIGuard.Message.message("user", (0..maxContent).collect { 'A' }.join()) @@ -476,13 +481,14 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final received = (List) it[1].messages - received.last().with { - assert it.content.length() == maxContent - assert it.content.length() < message.content.length() - } + receivedMeta = it[1] as Map + return span } - assertTelemetry('ai_guard.truncated', 'type:content') + final received = (List) receivedMeta.messages + final truncated = received.last() + assert truncated.content.length() == maxContent + assert truncated.content.length() < message.content.length() + assertTelemetry('truncated', 'type:content') } void 'test no messages'() { @@ -651,7 +657,7 @@ class AIGuardInternalTests extends DDSpecification { drain() } final filtered = metrics.findAll { - it.namespace == 'appsec' + it.namespace == 'ai_guard' && it.metricName == metric && it.tags == tags.toList() } @@ -710,6 +716,7 @@ class AIGuardInternalTests extends DDSpecification { void 'test JSON serialization with text content parts'() { given: + Map receivedMeta = null final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final messages = [AIGuard.Message.message('user', [AIGuard.ContentPart.text('Hello world')])] @@ -718,18 +725,19 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages.size() == 1 - assert receivedMessages[0].contentParts.size() == 1 - assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT - assert receivedMessages[0].contentParts[0].text == 'Hello world' + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 1 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[0].text == 'Hello world' } void 'test JSON serialization with image_url content parts'() { given: + Map receivedMeta = null final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final messages = [ AIGuard.Message.message('user', [AIGuard.ContentPart.imageUrl('https://example.com/image.jpg')]) @@ -740,18 +748,19 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages.size() == 1 - assert receivedMessages[0].contentParts.size() == 1 - assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.IMAGE_URL - assert receivedMessages[0].contentParts[0].imageUrl.url == 'https://example.com/image.jpg' + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 1 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[0].imageUrl.url == 'https://example.com/image.jpg' } void 'test JSON serialization with mixed content parts'() { given: + Map receivedMeta = null final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final messages = [ AIGuard.Message.message('user', [ @@ -766,22 +775,23 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages.size() == 1 - assert receivedMessages[0].contentParts.size() == 3 - assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT - assert receivedMessages[0].contentParts[0].text == 'Describe this image:' - assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL - assert receivedMessages[0].contentParts[1].imageUrl.url == 'https://example.com/image.jpg' - assert receivedMessages[0].contentParts[2].type == AIGuard.ContentPart.Type.TEXT - assert receivedMessages[0].contentParts[2].text == 'What do you see?' + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].contentParts.size() == 3 + assert receivedMessages[0].contentParts[0].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[0].text == 'Describe this image:' + assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[1].imageUrl.url == 'https://example.com/image.jpg' + assert receivedMessages[0].contentParts[2].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[2].text == 'What do you see?' } void 'test content parts order is preserved'() { given: + Map receivedMeta = null final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final parts = (0..4).collect { it % 2 == 0 ? AIGuard.ContentPart.text("Text $it") : AIGuard.ContentPart.imageUrl("https://example.com/image${it}.jpg") @@ -793,24 +803,25 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages[0].contentParts.size() == 5 - (0..4).each { i -> - if (i % 2 == 0) { - assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.TEXT - assert receivedMessages[0].contentParts[i].text == "Text $i" - } else { - assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.IMAGE_URL - assert receivedMessages[0].contentParts[i].imageUrl.url == "https://example.com/image${i}.jpg" - } - } + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages[0].contentParts.size() == 5 + (0..4).each { i -> + if (i % 2 == 0) { + assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.TEXT + assert receivedMessages[0].contentParts[i].text == "Text $i" + } else { + assert receivedMessages[0].contentParts[i].type == AIGuard.ContentPart.Type.IMAGE_URL + assert receivedMessages[0].contentParts[i].imageUrl.url == "https://example.com/image${i}.jpg" + } + } } void 'test content part text truncation'() { given: + Map receivedMeta = null final maxContent = Config.get().getAiGuardMaxContentSize() final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final longText = (0..maxContent).collect { 'A' }.join() @@ -823,19 +834,20 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages[0].contentParts.size() == 2 - assert receivedMessages[0].contentParts[0].text.length() == maxContent - assert receivedMessages[0].contentParts[0].text.length() < longText.length() - assert receivedMessages[0].contentParts[1].text == 'Short text' + receivedMeta = it[1] as Map return span } - assertTelemetry('ai_guard.truncated', 'type:content') + final receivedMessages = receivedMeta.messages as List + assert receivedMessages[0].contentParts.size() == 2 + assert receivedMessages[0].contentParts[0].text.length() == maxContent + assert receivedMessages[0].contentParts[0].text.length() < longText.length() + assert receivedMessages[0].contentParts[1].text == 'Short text' + assertTelemetry('truncated', 'type:content') } void 'test content part image_url not truncated even with long data URI'() { given: + Map receivedMeta = null final maxContent = Config.get().getAiGuardMaxContentSize() final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) // Create a very long data URI (longer than max content size) @@ -852,15 +864,15 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages[0].contentParts.size() == 2 - assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL - // Image URL should NOT be truncated - assert receivedMessages[0].contentParts[1].imageUrl.url == longDataUri - assert receivedMessages[0].contentParts[1].imageUrl.url.length() > maxContent + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages[0].contentParts.size() == 2 + assert receivedMessages[0].contentParts[1].type == AIGuard.ContentPart.Type.IMAGE_URL + // Image URL should NOT be truncated + assert receivedMessages[0].contentParts[1].imageUrl.url == longDataUri + assert receivedMessages[0].contentParts[1].imageUrl.url.length() > maxContent } void 'test adapter serializes content parts'() { @@ -887,6 +899,7 @@ class AIGuardInternalTests extends DDSpecification { void 'test backward compatibility with string content'() { given: + Map receivedMeta = null final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Good']]]) final messages = [AIGuard.Message.message('user', 'Hello world')] @@ -895,13 +908,13 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { - final meta = it[1] as Map - final receivedMessages = meta.messages as List - assert receivedMessages.size() == 1 - assert receivedMessages[0].content == 'Hello world' - assert receivedMessages[0].contentParts == null + receivedMeta = it[1] as Map return span } + final receivedMessages = receivedMeta.messages as List + assert receivedMessages.size() == 1 + assert receivedMessages[0].content == 'Hello world' + assert receivedMessages[0].contentParts == null } private static class TestSuite { diff --git a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java new file mode 100644 index 00000000000..8314da498f5 --- /dev/null +++ b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java @@ -0,0 +1,300 @@ +package com.datadog.aiguard; + +import static com.datadog.aiguard.AIGuardInternal.META_STRUCT_MESSAGES; +import static com.datadog.aiguard.AIGuardInternal.META_STRUCT_TAG; +import static com.datadog.aiguard.AIGuardInternal.REDACTED_TAG; +import static java.util.Arrays.asList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError; +import datadog.trace.api.aiguard.AIGuard.Evaluation; +import datadog.trace.api.aiguard.AIGuard.Message; +import datadog.trace.api.aiguard.AIGuard.Options; +import datadog.trace.api.telemetry.MetricCollector; +import datadog.trace.api.telemetry.WafMetricCollector; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.test.junit.utils.config.WithConfigExtension; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import okhttp3.Call; +import okhttp3.HttpUrl; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +/** Covers how {@link AIGuardInternal} wires the backend redaction decision into an evaluation. */ +@ExtendWith(WithConfigExtension.class) +class AIGuardInternalRedactionTest { + + private static final HttpUrl URL = + HttpUrl.parse("https://app.datadoghq.com/api/v2/ai-guard/evaluate"); + + private static final String SENSITIVE = "My SSN is 123-45-6789"; + private static final String REDACTED = "My SSN is "; + + private static final List MESSAGES = + asList( + Message.message("system", "You are a helpful assistant."), + Message.message("user", SENSITIVE)); + + private AgentSpan span; + private AgentTracer.TracerAPI originalTracer; + private Map metaStruct; + + @BeforeEach + void setUp() { + originalTracer = AgentTracer.get(); + + span = mock(AgentSpan.class); + lenient().when(span.getLocalRootSpan()).thenReturn(mock(AgentSpan.class)); + doAnswer( + invocation -> { + metaStruct = invocation.getArgument(1); + return span; + }) + .when(span) + .setMetaStruct(eq(META_STRUCT_TAG), any()); + + final AgentTracer.SpanBuilder builder = mock(AgentTracer.SpanBuilder.class, RETURNS_DEEP_STUBS); + lenient().when(builder.start()).thenReturn(span); + final AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class); + lenient().when(tracer.buildSpan(anyString(), anyString())).thenReturn(builder); + lenient().when(tracer.activateSpan(any())).thenReturn(mock(AgentScope.class)); + AgentTracer.forceRegister(tracer); + } + + @AfterEach + void tearDown() { + AgentTracer.forceRegister(originalTracer); + AIGuardInternal.uninstall(); + } + + @Test + void redactsMessagesAndRewritesMetaStruct() { + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + assertEquals(REDACTED, evaluation.getMessages().get(1).getContent()); + // the untouched message is carried over as is + assertEquals("You are a helpful assistant.", evaluation.getMessages().get(0).getContent()); + assertEquals(REDACTED, metaStructMessages().get(1).getContent()); + verify(span).setTag(REDACTED_TAG, true); + // the caller's list is left alone + assertEquals(SENSITIVE, MESSAGES.get(1).getContent()); + } + + @Test + void reportsNotRedactedWhenResponseCarriesNoReplacements() { + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, null)); + + assertSame(MESSAGES, evaluation.getMessages()); + assertEquals(SENSITIVE, metaStructMessages().get(1).getContent()); + verify(span).setTag(REDACTED_TAG, false); + } + + @Test + void reportsNotRedactedWhenReplacementsAreEmpty() { + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, "[]")); + + assertSame(MESSAGES, evaluation.getMessages()); + verify(span).setTag(REDACTED_TAG, false); + } + + @Test + void reportsNotRedactedWhenEveryReplacementIsSkipped() { + final String unresolvable = "[{\"path\":\"messages[9].content\",\"replacement\":\"nope\"}]"; + + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, unresolvable)); + + assertSame(MESSAGES, evaluation.getMessages()); + assertEquals(SENSITIVE, metaStructMessages().get(1).getContent()); + verify(span).setTag(REDACTED_TAG, false); + } + + @Test + void killSwitchLeavesMessagesUntouchedAndEmitsNoTag() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_REDACTION_ENABLED", "false", false); + + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + assertSame(MESSAGES, evaluation.getMessages()); + assertEquals(SENSITIVE, metaStructMessages().get(1).getContent()); + // an absent tag means "redaction is off", which is distinct from a false one + verify(span, never()).setTag(eq(REDACTED_TAG), anyBoolean()); + } + + @Test + void redactsMetaStructOnBlockedEvaluation() { + final AIGuardAbortError error = + assertThrows( + AIGuardAbortError.class, + () -> evaluate(responseWith("DENY", true, replacements(REDACTED)))); + + // reporting is still redacted on the blocked path + assertEquals(REDACTED, metaStructMessages().get(1).getContent()); + verify(span).setTag(REDACTED_TAG, true); + + // the abort error must not carry the conversation, redacted or otherwise + assertFalse(errorText(error).contains(SENSITIVE)); + assertFalse(errorText(error).contains(REDACTED)); + } + + /** + * The meta struct messages are built exactly once per evaluation, in the {@code finally} that + * closes the request, so one evaluation reports one truncation no matter how much redaction it + * applied. + */ + @Test + void reportsContentTruncationOnlyOnceWhenRedacting() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_MAX_CONTENT_SIZE", "8", false); + drainTelemetry(); + + evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + assertEquals(1, truncationMetrics()); + } + + @Test + void reportsRedactedTrueInTelemetryWhenSomethingWasRedacted() { + drainTelemetry(); + + evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + assertTrue(requestTags().contains("redacted:true")); + } + + @Test + void reportsRedactedFalseInTelemetryWhenNothingWasRedacted() { + drainTelemetry(); + + evaluate(responseWith("ALLOW", false, null)); + + assertTrue(requestTags().contains("redacted:false")); + } + + @Test + void reportsNoRedactedTagInTelemetryWhenKillSwitchIsOff() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_REDACTION_ENABLED", "false", false); + drainTelemetry(); + + evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + final List tags = requestTags(); + // absence of the tag is what tells a consumer that redaction is off + assertFalse(tags.contains("redacted:true")); + assertFalse(tags.contains("redacted:false")); + assertTrue(tags.contains("error:false")); + } + + private static List requestTags() { + WafMetricCollector.get().prepareMetrics(); + for (final MetricCollector.Metric metric : WafMetricCollector.get().drain()) { + if ("requests".equals(metric.metricName)) { + return metric.tags; + } + } + throw new AssertionError("no ai_guard requests metric was reported"); + } + + private static void drainTelemetry() { + WafMetricCollector.get().prepareMetrics(); + WafMetricCollector.get().drain(); + } + + private static long truncationMetrics() { + WafMetricCollector.get().prepareMetrics(); + long count = 0; + for (final MetricCollector.Metric metric : WafMetricCollector.get().drain()) { + if ("truncated".equals(metric.metricName)) { + count += metric.value.longValue(); + } + } + return count; + } + + @Test + void doesNotFailWhenReplacementsAreMalformed() { + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, "\"not-an-array\"")); + + assertSame(MESSAGES, evaluation.getMessages()); + verify(span).setTag(REDACTED_TAG, false); + } + + private static String errorText(final AIGuardAbortError error) { + return String.valueOf(error.getMessage()) + error + String.valueOf(error.getReason()); + } + + @SuppressWarnings("unchecked") + private List metaStructMessages() { + assertNotNull(metaStruct, "meta struct was never attached to the span"); + return (List) metaStruct.get(META_STRUCT_MESSAGES); + } + + private static String replacements(final String replacement) { + return "[{\"path\":\"messages[1].content\",\"replacement\":\"" + replacement + "\"}]"; + } + + private static String responseWith( + final String action, final boolean blocking, final String redactionReplacements) { + final StringBuilder attributes = new StringBuilder(); + attributes + .append("{\"action\":\"") + .append(action) + .append("\",\"reason\":\"a reason\",\"is_blocking_enabled\":") + .append(blocking) + .append(",\"tags\":[]"); + if (redactionReplacements != null) { + attributes.append(",\"redaction_replacements\":").append(redactionReplacements); + } + attributes.append('}'); + return "{\"data\":{\"id\":\"1\",\"type\":\"evaluations\",\"attributes\":" + attributes + "}}"; + } + + private Evaluation evaluate(final String responseBody) { + final OkHttpClient client = mock(OkHttpClient.class); + final Call call = mock(Call.class); + try { + lenient().when(call.execute()).thenReturn(response(responseBody)); + } catch (final IOException e) { + throw new AssertionError(e); + } + lenient().when(client.newCall(any(Request.class))).thenReturn(call); + return new AIGuardInternal(URL, Collections.emptyMap(), client) + .evaluate(MESSAGES, new Options().block(true)); + } + + private static Response response(final String json) { + return new Response.Builder() + .protocol(Protocol.HTTP_1_1) + .message("ok") + .request(new Request.Builder().url(URL).build()) + .code(200) + .body(ResponseBody.create(MediaType.parse("application/json"), json)) + .build(); + } +} diff --git a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java new file mode 100644 index 00000000000..e81a19d4214 --- /dev/null +++ b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java @@ -0,0 +1,643 @@ +package com.datadog.aiguard; + +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.api.aiguard.AIGuard.ContentPart; +import datadog.trace.api.aiguard.AIGuard.Message; +import datadog.trace.api.aiguard.AIGuard.ToolCall; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +class MessageRedactorTest { + + private static final MessageRedactor REDACTOR = new MessageRedactor.DefaultRedactor(); + + private static Map replacement(final Object path, final Object replacement) { + final Map entry = new HashMap<>(2); + entry.put("path", path); + entry.put("replacement", replacement); + return entry; + } + + private static Map applied(final String path, final String replacement) { + final Map entry = new LinkedHashMap<>(2); + entry.put("path", path); + entry.put("replacement", replacement); + return entry; + } + + private static List messages(final Message... messages) { + return new ArrayList<>(asList(messages)); + } + + @Nested + class ContentString { + + @Test + void redactsPlainContent() { + final List messages = + messages( + Message.message("system", "You are a helpful assistant."), + Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[1].content", "My SSN is "))); + + assertTrue(result.redacted()); + assertEquals(1, result.applied); + assertEquals(0, result.skipped); + assertEquals("My SSN is ", result.messages.get(1).getContent()); + // the untouched message is shared, not rebuilt + assertSame(messages.get(0), result.messages.get(0)); + } + + @Test + void neverMutatesTheCallersMessages() { + final Message original = Message.message("user", "My SSN is 123-45-6789"); + final List messages = messages(original); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content", "My SSN is "))); + + assertNotSame(messages, result.messages); + assertEquals("My SSN is 123-45-6789", original.getContent()); + assertEquals("My SSN is 123-45-6789", messages.get(0).getContent()); + } + + @Test + void preservesRoleAndToolCallId() { + final List messages = messages(Message.tool("call_1", "Account 000123456789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content", "Account "))); + + final Message redacted = result.messages.get(0); + assertEquals("tool", redacted.getRole()); + assertEquals("call_1", redacted.getToolCallId()); + assertEquals("Account ", redacted.getContent()); + } + + @Test + void appliesEmptyReplacementBecauseItMeansRemoval() { + final List messages = messages(Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact(messages, singletonList(replacement("messages[0].content", ""))); + + assertTrue(result.redacted()); + assertEquals(0, result.skipped); + assertEquals("", result.messages.get(0).getContent()); + } + + @Test + void copiesReplacementVerbatimIncludingAstralPlaneCharacters() { + final String replacement = "๐ŸŽญ ๐Ÿ‘จโ€๐Ÿ‘ฉโ€๐Ÿ‘งโ€๐Ÿ‘ฆ cafรฉ"; + final List messages = + messages(Message.message("user", "๐ŸŽญ 123-45-6789 ๐Ÿ‘จโ€๐Ÿ‘ฉโ€๐Ÿ‘งโ€๐Ÿ‘ฆ cafรฉ")); + + final MessageRedactor.Result result = + REDACTOR.redact(messages, singletonList(replacement("messages[0].content", replacement))); + + assertEquals(replacement, result.messages.get(0).getContent()); + } + } + + @Nested + class ContentParts { + + @Test + void redactsTextContentPart() { + final List messages = + messages( + Message.message( + "user", + asList( + ContentPart.text("here is my card 4111111111111111"), + ContentPart.imageUrl("https://example.com/image.jpg")))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement("messages[0].content[0].text", "here is my card "))); + + assertTrue(result.redacted()); + final List parts = result.messages.get(0).getContentParts(); + assertEquals("here is my card ", parts.get(0).getText()); + // the image part is untouched + assertEquals("https://example.com/image.jpg", parts.get(1).getImageUrl().getUrl()); + } + + @Test + void skipsImageUrlPart() { + final List messages = + messages( + Message.message( + "user", singletonList(ContentPart.imageUrl("https://example.com/image.jpg")))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content[0].text", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void skipsContentPartPathWhenMessageHoldsAPlainString() { + // messages[0].content[0].text has nothing to resolve against + final List messages = messages(Message.message("user", "plain content")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content[0].text", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void skipsContentPathWhenMessageHoldsContentParts() { + // messages[0].content resolves to a list, not a string + final List messages = + messages(Message.message("user", singletonList(ContentPart.text("secret")))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + } + + @Nested + class ToolCallArguments { + + @Test + void redactsToolCallArguments() { + final List messages = + messages( + Message.assistant( + ToolCall.toolCall("call_1", "send_email", "{\"ssn\":\"123-45-6789\"}"))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement( + "messages[0].tool_calls[0].function.arguments", "{\"ssn\":\"\"}"))); + + assertTrue(result.redacted()); + final ToolCall toolCall = result.messages.get(0).getToolCalls().get(0); + assertEquals("call_1", toolCall.getId()); + assertEquals("send_email", toolCall.getFunction().getName()); + assertEquals("{\"ssn\":\"\"}", toolCall.getFunction().getArguments()); + } + + @Test + void skipsToolCallPathWhenMessageHasNoToolCalls() { + final List messages = messages(Message.message("user", "no tools here")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement("messages[0].tool_calls[0].function.arguments", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void skipsToolCallWhoseFunctionCarriesNoArguments() { + final List messages = + messages( + Message.assistant(new ToolCall("call_1", new ToolCall.Function("send_email", null)))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement("messages[0].tool_calls[0].function.arguments", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void redactsOnlyTheTargetedToolCall() { + final List messages = + messages( + Message.assistant( + ToolCall.toolCall("call_1", "a", "{\"x\":1}"), + ToolCall.toolCall("call_2", "b", "{\"ssn\":\"123-45-6789\"}"))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement( + "messages[0].tool_calls[1].function.arguments", "{\"ssn\":\"\"}"))); + + final List toolCalls = result.messages.get(0).getToolCalls(); + assertEquals("{\"x\":1}", toolCalls.get(0).getFunction().getArguments()); + assertEquals("{\"ssn\":\"\"}", toolCalls.get(1).getFunction().getArguments()); + } + + @Test + void preservesContentPartsWhenRedactingToolCallOnTheSameMessage() { + final Message message = + new Message( + "assistant", + singletonList(ContentPart.text("calling a tool")), + singletonList(ToolCall.toolCall("call_1", "send", "{\"ssn\":\"123-45-6789\"}")), + null); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages(message), + singletonList( + replacement( + "messages[0].tool_calls[0].function.arguments", "{\"ssn\":\"\"}"))); + + final Message redacted = result.messages.get(0); + assertNull(redacted.getContent()); + assertEquals("calling a tool", redacted.getContentParts().get(0).getText()); + assertEquals( + "{\"ssn\":\"\"}", redacted.getToolCalls().get(0).getFunction().getArguments()); + } + } + + @Nested + class MultipleReplacements { + + @Test + void appliesReplacementsAcrossSeveralMessages() { + final List messages = + messages( + Message.message("system", "contact ops@acme.io"), + Message.message("user", "My SSN is 123-45-6789"), + Message.message("assistant", "nothing sensitive here")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "contact "), + replacement("messages[1].content", "My SSN is "))); + + assertEquals(2, result.applied); + assertEquals("contact ", result.messages.get(0).getContent()); + assertEquals("My SSN is ", result.messages.get(1).getContent()); + assertSame(messages.get(2), result.messages.get(2)); + } + + @Test + void composesTwoReplacementsTargetingTheSameMessage() { + final Message message = + new Message( + "assistant", + asList( + ContentPart.text("card 4111111111111111"), ContentPart.text("ssn 123-45-6789")), + singletonList(ToolCall.toolCall("call_1", "send", "{\"email\":\"a@b.io\"}")), + null); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages(message), + asList( + replacement("messages[0].content[0].text", "card "), + replacement("messages[0].content[1].text", "ssn "), + replacement( + "messages[0].tool_calls[0].function.arguments", + "{\"email\":\"\"}"))); + + assertEquals(3, result.applied); + assertEquals(0, result.skipped); + final Message redacted = result.messages.get(0); + assertEquals("card ", redacted.getContentParts().get(0).getText()); + assertEquals("ssn ", redacted.getContentParts().get(1).getText()); + assertEquals( + "{\"email\":\"\"}", + redacted.getToolCalls().get(0).getFunction().getArguments()); + } + + @Test + void appliesSurvivorsWhenOneEntryIsUnresolvable() { + final List messages = messages(Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "My SSN is "), + replacement("messages[9].content", "never applied"))); + + assertEquals(1, result.applied); + assertEquals(1, result.skipped); + assertEquals("My SSN is ", result.messages.get(0).getContent()); + } + } + + @Nested + class AppliedReplacements { + + @Test + void reportsOnlyTheEntriesThatWereApplied() { + final List messages = messages(Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "My SSN is "), + replacement("messages[9].content", "never applied"))); + + assertEquals( + singletonList(applied("messages[0].content", "My SSN is ")), + result.replacements); + } + + @Test + void keepsAnEmptyReplacementWhichMeansRemove() { + final List messages = messages(Message.message("user", "SSN 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact(messages, singletonList(replacement("messages[0].content", ""))); + + assertEquals(singletonList(applied("messages[0].content", "")), result.replacements); + } + + @Test + void reportsNothingWhenNothingWasApplied() { + final List messages = messages(Message.message("user", "hello")); + + assertTrue(REDACTOR.redact(messages, null).replacements.isEmpty()); + assertTrue(new MessageRedactor.NoOp().redact(messages, null).replacements.isEmpty()); + } + } + + @Nested + class MalformedResponses { + + @Test + void returnsSameListWhenReplacementsAreNull() { + final List messages = messages(Message.message("user", "hello")); + + final MessageRedactor.Result result = REDACTOR.redact(messages, null); + + assertSame(messages, result.messages); + assertFalse(result.redacted()); + assertEquals(0, result.skipped); + } + + @Test + void returnsSameListWhenReplacementsAreEmpty() { + final List messages = messages(Message.message("user", "hello")); + + final MessageRedactor.Result result = REDACTOR.redact(messages, Collections.emptyList()); + + assertSame(messages, result.messages); + assertFalse(result.redacted()); + } + + @Test + void returnsSameListWhenThereAreNoMessages() { + final MessageRedactor.Result nullMessages = + REDACTOR.redact(null, singletonList(replacement("messages[0].content", ""))); + assertNull(nullMessages.messages); + assertFalse(nullMessages.redacted()); + + final List empty = Collections.emptyList(); + final MessageRedactor.Result noMessages = + REDACTOR.redact(empty, singletonList(replacement("messages[0].content", ""))); + assertSame(empty, noMessages.messages); + assertFalse(noMessages.redacted()); + } + + @Test + void skipsEveryPathTheServiceContradictsItselfOn() { + final List messages = + messages( + Message.message("system", "contact ops@acme.io"), + Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "first"), + replacement("messages[0].content", "second"), + replacement("messages[1].content", "third"), + replacement("messages[1].content", "fourth"))); + + assertFalse(result.redacted()); + assertEquals(2, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void skipsConflictingReplacementsForTheSamePath() { + final List messages = messages(Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "first"), + replacement("messages[0].content", "second"))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + assertEquals("My SSN is 123-45-6789", messages.get(0).getContent()); + } + + @Test + void acceptsDuplicateReplacementsCarryingTheSameValue() { + final List messages = messages(Message.message("user", "My SSN is 123-45-6789")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement("messages[0].content", "My SSN is "), + replacement("messages[0].content", "My SSN is "))); + + assertEquals(1, result.applied); + assertEquals(0, result.skipped); + assertEquals("My SSN is ", result.messages.get(0).getContent()); + } + + @Test + void skipsEntriesThatAreNotObjects() { + final List messages = messages(Message.message("user", "hello")); + + final MessageRedactor.Result result = REDACTOR.redact(messages, asList("not an object", 42)); + + assertFalse(result.redacted()); + assertEquals(2, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void skipsEntriesWithMissingOrNonStringFields() { + final List messages = messages(Message.message("user", "hello")); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + asList( + replacement(null, ""), + replacement("messages[0].content", null), + replacement("", ""), + replacement(42, ""), + replacement("messages[0].content", 42), + new HashMap())); + + assertFalse(result.redacted()); + assertEquals(6, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void neverThrowsOnAToolCallWithoutAFunction() { + final List messages = + messages( + new Message( + "assistant", (String) null, singletonList(new ToolCall("call_1", null)), null)); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, + singletonList( + replacement("messages[0].tool_calls[0].function.arguments", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + } + + @Test + void skipsContentPathOnAMessageWithoutContent() { + final List messages = + messages(Message.assistant(ToolCall.toolCall("call_1", "send", "{}"))); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[0].content", ""))); + + assertFalse(result.redacted()); + assertEquals(1, result.skipped); + } + } + + @Nested + class PathGrammar { + + @ParameterizedTest + @ValueSource( + strings = { + "", // empty path + "content", // not rooted at messages + "messages", // resolves to the list itself + "messages.content", // missing root index + "messages[0]", // resolves to a message, not a string + "messages[0].foo", // unknown field + "messages[0].content.text", // missing part index + "messages[0].content[0]", // resolves to a content part + "messages[0].content[0].foo", // unknown content part field + "messages[0].content[0].text.extra", // too deep + "messages[0].content[1].image_url.url", // image locator, out of scope + "messages[0].tool_calls.function.arguments", // missing tool call index + "messages[0].tool_calls[0].function", // resolves to a function + "messages[0].tool_calls[0].function.name", // not a redactable target + "messages[0].tool_calls[0].arguments", // missing function hop + "messages[-1].content", // negative index rejected by the segment regex + "messages[0].content ", // trailing space, whole segment must match + " messages[0].content", // leading space + "messages[0]..content", // empty segment + "messages[0].content[", // unbalanced bracket + "messages[0].con-tent", // hyphen is not a word character + "messages[99999999999].content", // index overflows + "messages[0].tool_calls[0].function.arguments.extra", // deeper than any target + "conversation[0].content", // rooted at something other than messages + "messages[0].content[0].text[0]", // an index on the text field itself + "messages[0].content[9].text", // content part index out of range + "messages[0].tool_calls[0].foo.arguments", // the hop is not `function` + "messages[0].tool_calls[0].function[0].arguments", // an index on `function` + "messages[0].tool_calls[0].function.arguments[0]", // an index on `arguments` + "messages[0].tool_calls[9].function.arguments" // tool call index out of range + }) + void skipsUnsupportedPaths(final String path) { + final List messages = + messages( + new Message( + "assistant", + asList( + ContentPart.text("text"), ContentPart.imageUrl("https://example.com/i.jpg")), + singletonList(ToolCall.toolCall("call_1", "send", "{}")), + null)); + + final MessageRedactor.Result result = + REDACTOR.redact(messages, singletonList(replacement(path, ""))); + + assertFalse(result.redacted(), "expected path to be skipped: " + path); + assertEquals(1, result.skipped); + assertSame(messages, result.messages); + } + + @Test + void acceptsMultiDigitIndexes() { + final List messages = new ArrayList<>(); + for (int i = 0; i < 12; i++) { + messages.add(Message.message("user", "message " + i)); + } + + final MessageRedactor.Result result = + REDACTOR.redact( + messages, singletonList(replacement("messages[11].content", ""))); + + assertTrue(result.redacted()); + assertEquals("", result.messages.get(11).getContent()); + } + + @Test + void toleratesUnknownKeysOnAReplacementEntry() { + final Map entry = new LinkedHashMap<>(); + entry.put("path", "messages[0].content"); + entry.put("replacement", ""); + entry.put("category", "ssn"); + + final MessageRedactor.Result result = + REDACTOR.redact( + messages(Message.message("user", "My SSN is 123-45-6789")), singletonList(entry)); + + assertTrue(result.redacted()); + assertEquals("", result.messages.get(0).getContent()); + } + } +} diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java index f8111224ed7..f18cc0bf4b7 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java @@ -49,7 +49,9 @@ public static Evaluation evaluate(final List messages) { * * @param messages the collection of messages to evaluate (prompts, responses, tool calls, etc.) * @param options configuration options for the evaluation process - * @return an {@link Evaluation} containing the security decision and reasoning + * @return an {@link Evaluation} containing the security decision, the reasoning, and the + * evaluated messages (redacted when redaction was applied, see {@link + * Evaluation#getMessages()}) * @throws AIGuardAbortError if the evaluation action is not ALLOW (DENY or ABORT) and blocking is * enabled * @throws AIGuardClientError if there are client-side errors communicating with the AIGuard REST @@ -169,27 +171,81 @@ public static class Evaluation { final List tags; final Map tagProbs; final List sds; + final List messages; + final List> redactionReplacements; /** - * Creates a new evaluation result. + * Creates a new evaluation result carrying no messages. * * @param action the recommended action for the evaluated content * @param reason human-readable explanation for the decision * @param tags list of tags associated with the evaluation (e.g. indirect-prompt-injection) * @param tagProbs map of tags associated to their probability * @param sds list of Sensitive Data Scanner findings + * @deprecated use {@link #Evaluation(Action, String, List, Map, List, List)} instead, so the + * evaluated messages are available to the caller. */ + @Deprecated public Evaluation( final Action action, final String reason, final List tags, final Map tagProbs, final List sds) { + this(action, reason, tags, tagProbs, sds, Collections.emptyList()); + } + + /** + * Creates a new evaluation result carrying no redaction replacements. + * + * @param action the recommended action for the evaluated content + * @param reason human-readable explanation for the decision + * @param tags list of tags associated with the evaluation (e.g. indirect-prompt-injection) + * @param tagProbs map of tags associated to their probability + * @param sds list of Sensitive Data Scanner findings + * @param messages the evaluated messages, redacted when redaction was applied + * @deprecated use {@link #Evaluation(Action, String, List, Map, List, List, List)} instead, so + * the applied redaction replacements are available to the caller. + */ + @Deprecated + public Evaluation( + final Action action, + final String reason, + final List tags, + final Map tagProbs, + final List sds, + final List messages) { + this(action, reason, tags, tagProbs, sds, messages, Collections.emptyList()); + } + + /** + * Creates a new evaluation result. + * + * @param action the recommended action for the evaluated content + * @param reason human-readable explanation for the decision + * @param tags list of tags associated with the evaluation (e.g. indirect-prompt-injection) + * @param tagProbs map of tags associated to their probability + * @param sds list of Sensitive Data Scanner findings + * @param messages the evaluated messages, redacted when redaction was applied + * @param redactionReplacements the redactions that produced {@code messages}, one {@code {path, + * replacement}} entry per rewritten path + */ + public Evaluation( + final Action action, + final String reason, + final List tags, + final Map tagProbs, + final List sds, + final List messages, + final List> redactionReplacements) { this.action = action; this.reason = reason; this.tags = tags; this.tagProbs = tagProbs; this.sds = sds != null ? sds : Collections.emptyList(); + this.messages = messages != null ? messages : Collections.emptyList(); + this.redactionReplacements = + redactionReplacements != null ? redactionReplacements : Collections.emptyList(); } /** @@ -236,6 +292,35 @@ public Map getTagProbabilities() { public List getSds() { return sds; } + + /** + * Returns the evaluated messages, with sensitive data redacted whenever the AIGuard service + * requested redaction and redaction is enabled locally. + * + *

When nothing was redacted, this is the very same list that was passed to {@link + * AIGuard#evaluate(List, Options)}. The caller's list is never mutated. + * + * @return the evaluated messages, redacted when redaction was applied + */ + public List getMessages() { + return messages; + } + + /** + * Returns the redactions that were applied to produce {@link #getMessages()}. + * + *

Each entry is a {@code {path, replacement}} pair addressing one rewritten string in the + * evaluated conversation, e.g. {@code messages[1].content} or {@code + * messages[2].tool_calls[0].function.arguments}. Only the replacements that were actually + * applied are reported: entries the AI Guard service returned but that could not be resolved + * are skipped fail-safe and never surface here, and the list is empty when redaction is + * disabled locally. + * + * @return the applied redaction replacements, empty when nothing was redacted + */ + public List> getRedactionReplacements() { + return redactionReplacements; + } } /** diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java index 4389eba0c41..95dce0803c7 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java @@ -14,6 +14,13 @@ public final class NoOpEvaluator implements Evaluator { @Override public Evaluation evaluate(final List messages, final Options options) { - return new Evaluation(ALLOW, "AI Guard is not enabled", emptyList(), emptyMap(), emptyList()); + return new Evaluation( + ALLOW, + "AI Guard is not enabled", + emptyList(), + emptyMap(), + emptyList(), + messages, + emptyList()); } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/AIGuardConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/AIGuardConfig.java index 2d685c4a098..57d5714a5f3 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/AIGuardConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/AIGuardConfig.java @@ -7,8 +7,10 @@ public final class AIGuardConfig { public static final String AI_GUARD_TIMEOUT = "ai_guard.timeout"; public static final String AI_GUARD_MAX_CONTENT_SIZE = "ai_guard.max-content-size"; public static final String AI_GUARD_MAX_MESSAGES_LENGTH = "ai_guard.max-messages-length"; + public static final String AI_GUARD_REDACTION_ENABLED = "ai_guard.redaction.enabled"; public static final boolean DEFAULT_AI_GUARD_ENABLED = false; + public static final boolean DEFAULT_AI_GUARD_REDACTION_ENABLED = true; public static final int DEFAULT_AI_GUARD_TIMEOUT = 10_000; public static final int DEFAULT_AI_GUARD_MAX_CONTENT_SIZE = 512 * 1024; public static final int DEFAULT_AI_GUARD_MAX_MESSAGES_LENGTH = 16; diff --git a/dd-trace-api/src/test/java/datadog/trace/api/aiguard/AIGuardTest.java b/dd-trace-api/src/test/java/datadog/trace/api/aiguard/AIGuardTest.java index 0d44d81a975..6676c698a4e 100644 --- a/dd-trace-api/src/test/java/datadog/trace/api/aiguard/AIGuardTest.java +++ b/dd-trace-api/src/test/java/datadog/trace/api/aiguard/AIGuardTest.java @@ -4,11 +4,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; import org.junit.jupiter.api.Test; class AIGuardTest { @@ -76,6 +78,94 @@ void testNoopImplementation() { assertEquals(ALLOW, evaluation.getAction()); assertEquals("AI Guard is not enabled", evaluation.getReason()); + // nothing is redacted, so the very same list is handed back + assertSame(messages, evaluation.getMessages()); + } + + @Test + void testNoopImplementationWithoutMessages() { + // Evaluation normalises a null message list, so the no-op evaluator need not guard for it. + AIGuard.Evaluation evaluation = AIGuard.evaluate(null); + + assertEquals(ALLOW, evaluation.getAction()); + assertNotNull(evaluation.getMessages()); + assertTrue(evaluation.getMessages().isEmpty()); + } + + @Test + @SuppressWarnings("deprecation") + void testEvaluationCarriesMessages() { + List messages = + Collections.singletonList(AIGuard.Message.message("user", "My SSN is ")); + + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, "No rule match.", Collections.emptyList(), null, null, messages); + + assertSame(messages, evaluation.getMessages()); + } + + @Test + @SuppressWarnings("deprecation") + void testEvaluationWithoutMessagesReturnsEmptyList() { + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, "No rule match.", Collections.emptyList(), null, null); + + assertNotNull(evaluation.getMessages()); + assertTrue(evaluation.getMessages().isEmpty()); + } + + @Test + @SuppressWarnings("deprecation") + void testEvaluationWithNullMessagesReturnsEmptyList() { + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, "No rule match.", Collections.emptyList(), null, null, null); + + assertNotNull(evaluation.getMessages()); + assertTrue(evaluation.getMessages().isEmpty()); + } + + @Test + void testEvaluationCarriesRedactionReplacements() { + List messages = + Collections.singletonList(AIGuard.Message.message("user", "My SSN is ")); + List> replacements = + Collections.singletonList(Collections.singletonMap("path", "messages[0].content")); + + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, + "No rule match.", + Collections.emptyList(), + null, + null, + messages, + replacements); + + assertSame(replacements, evaluation.getRedactionReplacements()); + } + + @Test + @SuppressWarnings("deprecation") + void testEvaluationWithoutReplacementsReturnsEmptyList() { + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, "No rule match.", Collections.emptyList(), null, null, null); + + assertNotNull(evaluation.getRedactionReplacements()); + assertTrue(evaluation.getRedactionReplacements().isEmpty()); + } + + @Test + void testEvaluationWithNullReplacementsReturnsEmptyList() { + AIGuard.Evaluation evaluation = + new AIGuard.Evaluation( + ALLOW, "No rule match.", Collections.emptyList(), null, null, null, null); + + assertNotNull(evaluation.getRedactionReplacements()); + assertTrue(evaluation.getRedactionReplacements().isEmpty()); } @Test diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 18bd4ae31fd..a3388360e9c 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -213,10 +213,12 @@ import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_ENDPOINT; import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_MAX_CONTENT_SIZE; import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_MAX_MESSAGES_LENGTH; +import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_REDACTION_ENABLED; import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_TIMEOUT; import static datadog.trace.api.config.AIGuardConfig.DEFAULT_AI_GUARD_ENABLED; import static datadog.trace.api.config.AIGuardConfig.DEFAULT_AI_GUARD_MAX_CONTENT_SIZE; import static datadog.trace.api.config.AIGuardConfig.DEFAULT_AI_GUARD_MAX_MESSAGES_LENGTH; +import static datadog.trace.api.config.AIGuardConfig.DEFAULT_AI_GUARD_REDACTION_ENABLED; import static datadog.trace.api.config.AIGuardConfig.DEFAULT_AI_GUARD_TIMEOUT; import static datadog.trace.api.config.AppSecConfig.API_SECURITY_DOWNSTREAM_BODY_ANALYSIS_SAMPLE_RATE; import static datadog.trace.api.config.AppSecConfig.API_SECURITY_DOWNSTREAM_REQUEST_ANALYSIS_SAMPLE_RATE; @@ -1463,6 +1465,7 @@ public static String getHostName() { private final int aiGuardTimeout; private final int aiGuardMaxMessagesLength; private final int aiGuardMaxContentSize; + private final boolean aiGuardRedactionEnabled; static { // Bind telemetry collector to config module before initializing ConfigProvider @@ -3440,6 +3443,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment()) this.aiGuardMaxMessagesLength = configProvider.getInteger( AI_GUARD_MAX_MESSAGES_LENGTH, DEFAULT_AI_GUARD_MAX_MESSAGES_LENGTH); + this.aiGuardRedactionEnabled = + configProvider.getBoolean(AI_GUARD_REDACTION_ENABLED, DEFAULT_AI_GUARD_REDACTION_ENABLED); log.debug("New instance: {}", this); } @@ -6227,6 +6232,15 @@ public int getAiGuardTimeout() { return aiGuardTimeout; } + /** + * Global kill-switch for AI Guard sensitive data redaction. When {@code false}, the tracer never + * applies the redaction requested by the AI Guard service, even when the evaluation response asks + * for it. + */ + public boolean isAiGuardRedactionEnabled() { + return aiGuardRedactionEnabled; + } + private Set getSettingsSetFromEnvironment( String name, Function mapper, boolean splitOnWS) { final String value = configProvider.getString(name, ""); @@ -6991,6 +7005,8 @@ public String toString() { + aiGuardEnabled + ", aiGuardEndpoint=" + aiGuardEndpoint + + ", aiGuardRedactionEnabled=" + + aiGuardRedactionEnabled + ", logsOtelExporter=" + logsOtelExporter + ", logsOtelInterval=" diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java index 5d510f23a1a..9d1480047db 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/WafMetricCollector.java @@ -34,6 +34,16 @@ private WafMetricCollector() { private static final String NAMESPACE = "appsec"; + /** + * AI Guard metrics live in their own telemetry namespace. The namespace and the metric name are + * reported as separate fields and joined downstream, so {@code ai_guard} + {@code requests} is + * what surfaces the {@code ai_guard.requests} metric the AI Guard RFC specifies. + */ + private static final String AI_GUARD_NAMESPACE = "ai_guard"; + + /** Hoisted because {@link Enum#values()} clones its backing array on every call. */ + private static final AIGuardRedaction[] REDACTION_VALUES = AIGuardRedaction.values(); + private static final BlockingQueue rawMetricsQueue = new ArrayBlockingQueue<>(RAW_QUEUE_SIZE); @@ -64,7 +74,10 @@ private WafMetricCollector() { private static final AtomicInteger wafConfigErrorCounter = new AtomicInteger(); private static final AtomicInteger contextClosedRaceCounter = new AtomicInteger(); private static final AtomicLongArray aiGuardRequests = - new AtomicLongArray(AIGuard.Action.values().length * 2); // 3 actions * block + new AtomicLongArray( + AIGuard.Action.values().length + * 2 + * REDACTION_VALUES.length); // actions * block * redaction state private static final AtomicInteger aiGuardErrors = new AtomicInteger(); private static final AtomicLongArray aiGuardTruncated = new AtomicLongArray(AIGuardTruncationType.values().length); @@ -246,8 +259,14 @@ public void appSecSdkEvent(final LoginEvent event, final LoginVersion version) { appSecSdkEventQueue.incrementAndGet(index); } - public void aiGuardRequest(final AIGuard.Action action, final boolean block) { - aiGuardRequests.incrementAndGet(action.ordinal() * 2 + (block ? 1 : 0)); + public void aiGuardRequest( + final AIGuard.Action action, final boolean block, final AIGuardRedaction redaction) { + aiGuardRequests.incrementAndGet(aiGuardRequestIndex(action, block, redaction)); + } + + private static int aiGuardRequestIndex( + final AIGuard.Action action, final boolean block, final AIGuardRedaction redaction) { + return (action.ordinal() * 2 + (block ? 1 : 0)) * REDACTION_VALUES.length + redaction.ordinal(); } public void aiGuardError() { @@ -521,17 +540,18 @@ public void prepareMetrics() { } // AI Guard successful requests + aiGuardSuccesses: for (final AIGuard.Action action : AIGuard.Action.values()) { - final long blocked = aiGuardRequests.getAndSet(action.ordinal() * 2 + 1, 0); - if (blocked > 0) { - if (!rawMetricsQueue.offer(AIGuardRequests.success(blocked, action, true))) { - break; - } - } - final long nonBlocked = aiGuardRequests.getAndSet(action.ordinal() * 2, 0); - if (nonBlocked > 0) { - if (!rawMetricsQueue.offer(AIGuardRequests.success(nonBlocked, action, false))) { - break; + for (int blockFlag = 1; blockFlag >= 0; blockFlag--) { + final boolean block = blockFlag == 1; + for (final AIGuardRedaction redaction : REDACTION_VALUES) { + final long count = + aiGuardRequests.getAndSet(aiGuardRequestIndex(action, block, redaction), 0); + if (count > 0) { + if (!rawMetricsQueue.offer(AIGuardRequests.success(count, action, block, redaction))) { + break aiGuardSuccesses; + } + } } } } @@ -589,7 +609,11 @@ public void prepareMetrics() { public abstract static class WafMetric extends MetricCollector.Metric { public WafMetric(String metricName, long counter, String... tags) { - super(NAMESPACE, true, metricName, "count", counter, tags); + this(NAMESPACE, metricName, counter, tags); + } + + protected WafMetric(String namespace, String metricName, long counter, String... tags) { + super(namespace, true, metricName, "count", counter, tags); } } @@ -823,14 +847,33 @@ public WafInputTruncated(final long counter, final int bitfield) { } } - public static class AIGuardRequests extends WafMetric { + /** Base class for the metrics reported under the {@code ai_guard} namespace. */ + public abstract static class AIGuardMetric extends WafMetric { + protected AIGuardMetric(final String metricName, final long counter, final String... tags) { + super(AI_GUARD_NAMESPACE, metricName, counter, tags); + } + } + + public static class AIGuardRequests extends AIGuardMetric { private AIGuardRequests(final long count, final String... tags) { - super("ai_guard.requests", count, tags); + super("requests", count, tags); } public static AIGuardRequests success( - final long count, final AIGuard.Action action, final boolean block) { - return new AIGuardRequests(count, "action:" + action, "block:" + block, "error:false"); + final long count, + final AIGuard.Action action, + final boolean block, + final AIGuardRedaction redaction) { + if (redaction == AIGuardRedaction.DISABLED) { + // No redacted tag at all, so its absence stays distinguishable from a false value. + return new AIGuardRequests(count, "action:" + action, "block:" + block, "error:false"); + } + return new AIGuardRequests( + count, + "action:" + action, + "block:" + block, + "error:false", + "redacted:" + (redaction == AIGuardRedaction.APPLIED)); } public static AIGuardRequests error(final long count) { @@ -838,9 +881,9 @@ public static AIGuardRequests error(final long count) { } } - public static class AIGuardTruncated extends WafMetric { + public static class AIGuardTruncated extends AIGuardMetric { public AIGuardTruncated(final long count, final AIGuardTruncationType type) { - super("ai_guard.truncated", count, "type:" + type.tagValue); + super("truncated", count, "type:" + type.tagValue); } } @@ -862,6 +905,20 @@ public ApiSecurityRequestNoSchema(final long counter, final String framework) { } } + /** + * Whether an evaluation redacted anything, as reported by the {@code redacted} tag on {@code + * ai_guard.requests}. {@link #DISABLED} reports no tag at all, so an absent tag means "redaction + * is off" and stays distinguishable from {@code redacted:false}. + */ + public enum AIGuardRedaction { + /** Redaction is disabled locally, so nothing was even attempted. */ + DISABLED, + /** Redaction is enabled and at least one replacement was applied. */ + APPLIED, + /** Redaction is enabled but nothing was redacted. */ + NOT_APPLIED + } + public enum AIGuardTruncationType { MESSAGES("messages"), CONTENT("content"); diff --git a/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy index d35a4d6816d..e28ee7bdbff 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/api/telemetry/WafMetricCollectorTest.groovy @@ -1,6 +1,9 @@ package datadog.trace.api.telemetry import static datadog.trace.api.aiguard.AIGuard.Action.ABORT +import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardRedaction.APPLIED +import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardRedaction.DISABLED +import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardRedaction.NOT_APPLIED import static datadog.trace.api.aiguard.AIGuard.Action.ALLOW import static datadog.trace.api.aiguard.AIGuard.Action.DENY import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardTruncationType.CONTENT @@ -533,19 +536,19 @@ class WafMetricCollectorTest extends DDSpecification { final collector = WafMetricCollector.get() when: - collector.aiGuardRequest(action, block) + collector.aiGuardRequest(action, block, NOT_APPLIED) then: collector.prepareMetrics() final metrics = collector.drain() - final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.requests' } + final configErrorMetrics = metrics.findAll { it.metricName == 'requests' } final metric = configErrorMetrics[0] metric.type == 'count' - metric.metricName == 'ai_guard.requests' - metric.namespace == 'appsec' + metric.metricName == 'requests' + metric.namespace == 'ai_guard' metric.value == 1 - metric.tags.toSet() == ["action:${action.name()}", "block:${block}", 'error:false'].toSet() + metric.tags.toSet() == ["action:${action.name()}", "block:${block}", 'error:false', 'redacted:false'].toSet() where: action | block @@ -557,6 +560,26 @@ class WafMetricCollectorTest extends DDSpecification { ABORT | false } + void 'test ai guard redaction telemetry tag'() { + given: + final collector = WafMetricCollector.get() + + when: + collector.aiGuardRequest(ALLOW, false, redaction) + + then: + collector.prepareMetrics() + final metric = collector.drain().find { it.metricName == 'requests' } + metric.tags.toSet() == expectedTags.toSet() + + where: + redaction | expectedTags + APPLIED | ['action:ALLOW', 'block:false', 'error:false', 'redacted:true'] + NOT_APPLIED | ['action:ALLOW', 'block:false', 'error:false', 'redacted:false'] + // the kill switch reports no redacted tag at all + DISABLED | ['action:ALLOW', 'block:false', 'error:false'] + } + void 'test ai guard error'() { given: final collector = WafMetricCollector.get() @@ -567,12 +590,12 @@ class WafMetricCollectorTest extends DDSpecification { then: collector.prepareMetrics() final metrics = collector.drain() - final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.requests' } + final configErrorMetrics = metrics.findAll { it.metricName == 'requests' } final metric = configErrorMetrics[0] metric.type == 'count' - metric.metricName == 'ai_guard.requests' - metric.namespace == 'appsec' + metric.metricName == 'requests' + metric.namespace == 'ai_guard' metric.value == 1 metric.tags.toSet() == ['error:true'].toSet() } @@ -587,12 +610,12 @@ class WafMetricCollectorTest extends DDSpecification { then: collector.prepareMetrics() final metrics = collector.drain() - final configErrorMetrics = metrics.findAll { it.metricName == 'ai_guard.truncated' } + final configErrorMetrics = metrics.findAll { it.metricName == 'truncated' } final metric = configErrorMetrics[0] metric.type == 'count' - metric.metricName == 'ai_guard.truncated' - metric.namespace == 'appsec' + metric.metricName == 'truncated' + metric.namespace == 'ai_guard' metric.value == 1 metric.tags.toSet() == ["type:${type.tagValue}"].toSet() diff --git a/internal-api/src/test/java/datadog/trace/api/ConfigAIGuardRedactionTest.java b/internal-api/src/test/java/datadog/trace/api/ConfigAIGuardRedactionTest.java new file mode 100644 index 00000000000..338b2e30145 --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/ConfigAIGuardRedactionTest.java @@ -0,0 +1,56 @@ +package datadog.trace.api; + +import static datadog.trace.api.config.AIGuardConfig.AI_GUARD_REDACTION_ENABLED; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.test.junit.utils.config.WithConfigExtension; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +@ExtendWith(WithConfigExtension.class) +class ConfigAIGuardRedactionTest { + + @Test + void redactionIsEnabledByDefault() { + assertTrue(Config.get().isAiGuardRedactionEnabled()); + } + + @Test + void killSwitchDisablesRedactionViaSystemProperty() { + WithConfigExtension.injectSysConfig(AI_GUARD_REDACTION_ENABLED, "false"); + + assertFalse(Config.get().isAiGuardRedactionEnabled()); + } + + @Test + void killSwitchDisablesRedactionViaEnvironmentVariable() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_REDACTION_ENABLED", "false", false); + + assertFalse(Config.get().isAiGuardRedactionEnabled()); + } + + @Test + void killSwitchAcceptsNumericAndMixedCaseValues() { + WithConfigExtension.injectSysConfig(AI_GUARD_REDACTION_ENABLED, "0"); + assertFalse(Config.get().isAiGuardRedactionEnabled()); + + WithConfigExtension.injectSysConfig(AI_GUARD_REDACTION_ENABLED, "FALSE"); + assertFalse(Config.get().isAiGuardRedactionEnabled()); + + WithConfigExtension.injectSysConfig(AI_GUARD_REDACTION_ENABLED, "True"); + assertTrue(Config.get().isAiGuardRedactionEnabled()); + } + + /** + * Invalid boolean values resolve to {@code false} rather than to the configured default: see the + * backward-compatibility branch in {@code ConfigProvider#get}. A typo therefore turns redaction + * off, even though it defaults to on. + */ + @Test + void unparseableValueDisablesRedaction() { + WithConfigExtension.injectSysConfig(AI_GUARD_REDACTION_ENABLED, "not-a-boolean"); + + assertFalse(Config.get().isAiGuardRedactionEnabled()); + } +} diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 6c3fe354b68..e237cee5dac 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -97,6 +97,14 @@ "aliases": [] } ], + "DD_AI_GUARD_REDACTION_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "true", + "aliases": [] + } + ], "DD_AI_GUARD_TIMEOUT": [ { "version": "A", From d8d7fb7985cb0dc0df61ebcbde103f91dff9f4e7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Thu, 3 Sep 2026 14:45:57 +0200 Subject: [PATCH 2/2] Apply PR feedback --- .../com/datadog/aiguard/AIGuardInternal.java | 34 ++++++++----- .../com/datadog/aiguard/MessageRedactor.java | 32 +++---------- .../aiguard/AIGuardInternalRedactionTest.java | 48 +++++++++++++++++++ .../datadog/aiguard/MessageRedactorTest.java | 19 +++----- .../datadog/trace/api/aiguard/AIGuard.java | 29 ++++++----- 5 files changed, 97 insertions(+), 65 deletions(-) diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index cefa5feb4ee..b0e898e8eab 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -211,8 +211,12 @@ private static List messagesForMetaStruct(final List messages) * conversation through the meta struct, and that report must be redacted too. The {@link * AIGuardAbortError} raised on that path deliberately carries no messages. * + *

The {@code ai_guard.redacted} tag is set to {@code false} before the request is issued and + * only raised here, so an evaluation that fails before this point still reports that nothing was + * redacted rather than looking like the kill switch is off. + * * @return the telemetry state, {@link AIGuardRedaction#DISABLED} when the kill switch is off, in - * which case no {@code ai_guard.redacted} tag is attached either + * which case no {@code ai_guard.redacted} tag is attached at all */ private AIGuardRedaction reportRedaction( final AgentSpan span, final MessageRedactor.Result redaction) { @@ -221,13 +225,17 @@ private AIGuardRedaction reportRedaction( // one ("redaction is on and nothing was redacted"). return AIGuardRedaction.DISABLED; } - span.setTag(REDACTED_TAG, redaction.redacted()); if (redaction.skipped > 0) { log.debug( "AI Guard skipped {} redaction replacement(s) that could not be applied", redaction.skipped); } - return redaction.redacted() ? AIGuardRedaction.APPLIED : AIGuardRedaction.NOT_APPLIED; + if (!redaction.redacted()) { + // The tag was already set to false before the request; nothing to correct. + return AIGuardRedaction.NOT_APPLIED; + } + span.setTag(REDACTED_TAG, true); + return AIGuardRedaction.APPLIED; } private static boolean isToolCall(final Message message) { @@ -334,6 +342,11 @@ public Evaluation evaluate(final List messages, final Options options) } else { span.setTag(TARGET_TAG, "prompt"); } + if (redactor.enabled()) { + // Reported before the request goes out so an evaluation that fails, and therefore redacts + // nothing, still says so. An absent tag stays reserved for the kill switch being off. + span.setTag(REDACTED_TAG, false); + } final Map metaStruct = new HashMap<>(2); span.setMetaStruct(META_STRUCT_TAG, metaStruct); final Request.Builder request = @@ -369,9 +382,10 @@ public Evaluation evaluate(final List messages, final Options options) metaStruct.put(META_STRUCT_SDS, sdsFindings); } final Object rawReplacements = result.get(RESPONSE_REDACTION_REPLACEMENTS); - final MessageRedactor.Result redaction = - redactor.redact( - messages, rawReplacements instanceof List ? (List) rawReplacements : null); + // Reported back to the caller verbatim, including entries redaction could not apply. + final List redactionReplacements = + rawReplacements instanceof List ? (List) rawReplacements : null; + final MessageRedactor.Result redaction = redactor.redact(messages, redactionReplacements); final AIGuardRedaction redactionState = reportRedaction(span, redaction); finalMessages = redaction.messages; final boolean shouldBlock = @@ -382,13 +396,7 @@ public Evaluation evaluate(final List messages, final Options options) throw new AIGuardAbortError(action, reason, tags, tagProbs, sdsFindings); } return new Evaluation( - action, - reason, - tags, - tagProbs, - sdsFindings, - redaction.messages, - redaction.replacements); + action, reason, tags, tagProbs, sdsFindings, redaction.messages, redactionReplacements); } finally { metaStruct.put(META_STRUCT_MESSAGES, messagesForMetaStruct(finalMessages)); } diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java index da7dea9d96a..933294e36ba 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/MessageRedactor.java @@ -5,7 +5,6 @@ import datadog.trace.api.aiguard.AIGuard.ToolCall; import datadog.trace.api.aiguard.AIGuard.ToolCall.Function; import java.util.ArrayList; -import java.util.Collections; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -23,27 +22,15 @@ final class Result { /** The redacted messages, or the very same list that was passed in when nothing was applied. */ final List messages; - /** - * The {@code {path, replacement}} entries that were actually overwritten, in the order the - * service returned them. Entries skipped fail-safe never appear here, so this list describes - * exactly the transformation {@link #messages} underwent, and it is what the SDK hands back - * through {@code Evaluation.getRedactionReplacements()}. - */ - final List> replacements; - /** Number of paths successfully overwritten. */ final int applied; /** Number of entries skipped fail-safe (unresolvable, non-string, missing or conflicting). */ final int skipped; - private Result( - final List messages, - final List> replacements, - final int skipped) { + private Result(final List messages, final int applied, final int skipped) { this.messages = messages; - this.replacements = replacements; - this.applied = replacements.size(); + this.applied = applied; this.skipped = skipped; } @@ -53,7 +40,7 @@ boolean redacted() { } private static Result nothingApplied(final List messages, final int skipped) { - return new Result(messages, Collections.emptyList(), skipped); + return new Result(messages, 0, skipped); } } @@ -123,13 +110,6 @@ public boolean enabled() { /** Longest index we bother parsing; anything longer cannot address a real list. */ private static final int MAX_INDEX_DIGITS = 9; - private static Map entry(final String path, final String replacement) { - final Map entry = new LinkedHashMap<>(4); - entry.put("path", path); - entry.put("replacement", replacement); - return Collections.unmodifiableMap(entry); - } - /** * Overwrites every path in {@code replacements} with its replacement string. * @@ -189,7 +169,7 @@ public Result redact(final List messages, @Nullable final List repla final String[] names = new String[MAX_SEGMENTS]; final int[] indices = new int[MAX_SEGMENTS]; List working = null; - final List> applied = new ArrayList<>(byPath.size()); + int applied = 0; for (final Map.Entry entry : byPath.entrySet()) { final int count = parseSegments(entry.getKey(), names, indices); @@ -212,13 +192,13 @@ public Result redact(final List messages, @Nullable final List repla working = new ArrayList<>(messages); } working.set(index, updated); - applied.add(entry(entry.getKey(), entry.getValue())); + applied++; } if (working == null) { return Result.nothingApplied(messages, skipped); } - return new Result(working, Collections.unmodifiableList(applied), skipped); + return new Result(working, applied, skipped); } /** diff --git a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java index 8314da498f5..f5a75ac6f47 100644 --- a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java +++ b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/AIGuardInternalRedactionTest.java @@ -22,6 +22,7 @@ import static org.mockito.Mockito.verify; import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError; +import datadog.trace.api.aiguard.AIGuard.AIGuardClientError; import datadog.trace.api.aiguard.AIGuard.Evaluation; import datadog.trace.api.aiguard.AIGuard.Message; import datadog.trace.api.aiguard.AIGuard.Options; @@ -237,6 +238,53 @@ private static long truncationMetrics() { return count; } + @Test + void reportsTheServiceReplacementsVerbatimIncludingUnappliedEntries() { + final String replacements = + "[{\"path\":\"messages[1].content\",\"replacement\":\"" + + REDACTED + + "\"},{\"path\":\"messages[9].content\",\"replacement\":\"never applied\"}]"; + + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, replacements)); + + assertEquals(REDACTED, evaluation.getMessages().get(1).getContent()); + // messages[9] cannot resolve and is skipped fail-safe, but the service asked for it, so it is + // still reported back to the caller + assertEquals(2, evaluation.getRedactionReplacements().size()); + } + + @Test + void reportsTheServiceReplacementsEvenWhenTheKillSwitchIsOff() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_REDACTION_ENABLED", "false", false); + + final Evaluation evaluation = evaluate(responseWith("ALLOW", false, replacements(REDACTED))); + + // nothing was applied, but the service's request is still visible to the caller + assertSame(MESSAGES, evaluation.getMessages()); + assertEquals(1, evaluation.getRedactionReplacements().size()); + } + + @Test + void reportsNotRedactedWhenTheEvaluationFails() { + // no action field, so the response is rejected before any redaction work happens + assertThrows( + AIGuardClientError.class, () -> evaluate("{\"data\":{\"attributes\":{\"tags\":[]}}}")); + + // a failed evaluation redacted nothing, and must not look like the kill switch is off + verify(span).setTag(REDACTED_TAG, false); + verify(span, never()).setTag(REDACTED_TAG, true); + } + + @Test + void emitsNoTagWhenTheEvaluationFailsAndTheKillSwitchIsOff() { + WithConfigExtension.injectEnvConfig("DD_AI_GUARD_REDACTION_ENABLED", "false", false); + + assertThrows( + AIGuardClientError.class, () -> evaluate("{\"data\":{\"attributes\":{\"tags\":[]}}}")); + + verify(span, never()).setTag(eq(REDACTED_TAG), anyBoolean()); + } + @Test void doesNotFailWhenReplacementsAreMalformed() { final Evaluation evaluation = evaluate(responseWith("ALLOW", false, "\"not-an-array\"")); diff --git a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java index e81a19d4214..1734cfc4cf6 100644 --- a/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java +++ b/dd-java-agent/agent-aiguard/src/test/java/com/datadog/aiguard/MessageRedactorTest.java @@ -34,13 +34,6 @@ private static Map replacement(final Object path, final Object r return entry; } - private static Map applied(final String path, final String replacement) { - final Map entry = new LinkedHashMap<>(2); - entry.put("path", path); - entry.put("replacement", replacement); - return entry; - } - private static List messages(final Message... messages) { return new ArrayList<>(asList(messages)); } @@ -377,9 +370,8 @@ void reportsOnlyTheEntriesThatWereApplied() { replacement("messages[0].content", "My SSN is "), replacement("messages[9].content", "never applied"))); - assertEquals( - singletonList(applied("messages[0].content", "My SSN is ")), - result.replacements); + assertEquals(1, result.applied); + assertEquals(1, result.skipped); } @Test @@ -389,15 +381,16 @@ void keepsAnEmptyReplacementWhichMeansRemove() { final MessageRedactor.Result result = REDACTOR.redact(messages, singletonList(replacement("messages[0].content", ""))); - assertEquals(singletonList(applied("messages[0].content", "")), result.replacements); + assertEquals(1, result.applied); + assertEquals("", result.messages.get(0).getContent()); } @Test void reportsNothingWhenNothingWasApplied() { final List messages = messages(Message.message("user", "hello")); - assertTrue(REDACTOR.redact(messages, null).replacements.isEmpty()); - assertTrue(new MessageRedactor.NoOp().redact(messages, null).replacements.isEmpty()); + assertEquals(0, REDACTOR.redact(messages, null).applied); + assertEquals(0, new MessageRedactor.NoOp().redact(messages, null).applied); } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java index f18cc0bf4b7..0054655b6f7 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java @@ -172,7 +172,7 @@ public static class Evaluation { final Map tagProbs; final List sds; final List messages; - final List> redactionReplacements; + final List redactionReplacements; /** * Creates a new evaluation result carrying no messages. @@ -227,8 +227,8 @@ public Evaluation( * @param tagProbs map of tags associated to their probability * @param sds list of Sensitive Data Scanner findings * @param messages the evaluated messages, redacted when redaction was applied - * @param redactionReplacements the redactions that produced {@code messages}, one {@code {path, - * replacement}} entry per rewritten path + * @param redactionReplacements the {@code redaction_replacements} array as returned by the + * AIGuard service */ public Evaluation( final Action action, @@ -237,7 +237,7 @@ public Evaluation( final Map tagProbs, final List sds, final List messages, - final List> redactionReplacements) { + final List redactionReplacements) { this.action = action; this.reason = reason; this.tags = tags; @@ -307,18 +307,21 @@ public List getMessages() { } /** - * Returns the redactions that were applied to produce {@link #getMessages()}. + * Returns the redaction replacements the AIGuard service asked for, exactly as it returned + * them. * - *

Each entry is a {@code {path, replacement}} pair addressing one rewritten string in the - * evaluated conversation, e.g. {@code messages[1].content} or {@code - * messages[2].tool_calls[0].function.arguments}. Only the replacements that were actually - * applied are reported: entries the AI Guard service returned but that could not be resolved - * are skipped fail-safe and never surface here, and the list is empty when redaction is - * disabled locally. + *

Each entry is a {@code {path, replacement}} pair addressing one string in the evaluated + * conversation, e.g. {@code messages[1].content} or {@code + * messages[2].tool_calls[0].function.arguments}. * - * @return the applied redaction replacements, empty when nothing was redacted + *

This is detection metadata, not a record of what the tracer did: entries the tracer could + * not resolve are skipped fail-safe but still reported here, and the whole array is reported + * even when redaction is disabled locally, in which case none of it was applied. Use {@link + * #getMessages()} for the conversation as it actually stands. + * + * @return the service's redaction replacements, empty when it asked for none */ - public List> getRedactionReplacements() { + public List getRedactionReplacements() { return redactionReplacements; } }