Skip to content

feat(ai-guard): apply sensitive data redaction from evaluate responses #12394

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
manuel-alvarez-alvarez wants to merge 2 commits into
base: master
Choose a base branch
from
Open

feat(ai-guard): apply sensitive data redaction from evaluate responses #12394

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
49d1167
feat(ai-guard): apply sensitive data redaction from evaluate responses
manuel-alvarez-alvarez Sep 3, 2026
d8d7fb7
Apply PR feedback
manuel-alvarez-alvarez Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 14 additions & 16 deletions ...ation/src/main/java/datadog/communication/serialization/custom/aiguard/MessageWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,14 @@ public class MessageWriter implements ValueWriter<AIGuard.Message> {
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);
Expand Down Expand Up @@ -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();
}
}
34 changes: 34 additions & 0 deletions ...tion/src/test/groovy/datadog/communication/serialization/aiguard/MessageWriterTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
2 changes: 2 additions & 0 deletions dd-java-agent/agent-aiguard/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
}
Expand Down
70 changes: 66 additions & 4 deletions dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.
Expand All @@ -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);
Expand All @@ -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";
Expand Down Expand Up @@ -128,6 +136,7 @@ static void uninstall() {
private final OkHttpClient client;
private final Map<String, String> meta;
private final Map<String, String> headers;
private final MessageRedactor redactor;

AIGuardInternal(final HttpUrl url, final Map<String, String> headers, final OkHttpClient client) {
this.url = url;
Expand All @@ -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<Message> messagesForMetaStruct(List<Message> messages) {
private static List<Message> messagesForMetaStruct(final List<Message> messages) {
final Config config = Config.get();
final int size = Math.min(messages.size(), config.getAiGuardMaxMessagesLength());
if (size < messages.size()) {
Expand Down Expand Up @@ -191,6 +204,40 @@ private static List<Message> messagesForMetaStruct(List<Message> messages) {
return result;
}

/**
* Applies the redaction requested by the AI Guard service and reports the outcome on the span.
*
* <p>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.
*
* <p>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 at all
*/
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;
}
if (redaction.skipped > 0) {

@IlyasShabi IlyasShabi Sep 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should also report redaction_error telemetry on failure

@manuel-alvarez-alvarez manuel-alvarez-alvarez Sep 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is it defined in the RFC? (I couldn't find it 😓)

@IlyasShabi IlyasShabi Sep 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was mentioned on the tracer recommended behavior

If two entries for the same path carry different replacement values (a backend bug), the tracer skips that path and records a telemetry error rather than guessing

But not listed on telemetry part.

@manuel-alvarez-alvarez manuel-alvarez-alvarez Sep 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we update the RFC or add and addendum?, just to set redaction_error as the tag for this particular case

log.debug(
"AI Guard skipped {} redaction replacement(s) that could not be applied",
redaction.skipped);
}
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) {
return message.getToolCalls() != null || message.getToolCallId() != null;
}
Expand Down Expand Up @@ -283,6 +330,7 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
// sure client IP tags were populated.
copyAnomalyDetectionTags(span, localRootSpan);
}
List<Message> finalMessages = messages;
try (final ContextScope scope = tracer.activateSpan(span)) {
final Message last = messages.get(messages.size() - 1);
if (isToolCall(last)) {
Expand All @@ -294,8 +342,12 @@ public Evaluation evaluate(final List<Message> 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<String, Object> metaStruct = new HashMap<>(2);
metaStruct.put(META_STRUCT_MESSAGES, messagesForMetaStruct(messages));
span.setMetaStruct(META_STRUCT_TAG, metaStruct);
final Request.Builder request =
new Request.Builder()
Expand Down Expand Up @@ -329,14 +381,24 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
if (sdsFindings != null && !sdsFindings.isEmpty()) {
metaStruct.put(META_STRUCT_SDS, sdsFindings);
}
final Object rawReplacements = result.get(RESPONSE_REDACTION_REPLACEMENTS);
// 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);
Comment thread
manuel-alvarez-alvarez marked this conversation as resolved.
Comment thread
manuel-alvarez-alvarez marked this conversation as resolved.
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, redactionReplacements);
} finally {
metaStruct.put(META_STRUCT_MESSAGES, messagesForMetaStruct(finalMessages));
}
} catch (AIGuardAbortError e) {
span.addThrowable(e);
Expand Down
Loading