From 301c926d0816a3a1d3684001f8cb16afc7ccaaa0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Tue, 11 Nov 2025 15:12:33 +0100 Subject: [PATCH 1/4] chore: add the tags returned by the service to the ai_guard span --- .../com/datadog/aiguard/AIGuardInternal.java | 11 ++++++++++- .../aiguard/AIGuardInternalTests.groovy | 19 +++++++++++++++---- 2 files changed, 25 insertions(+), 5 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 fc2d2dafb64..eb2143b152f 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 @@ -224,8 +224,17 @@ public Evaluation evaluate(final List messages, final Options options) } final Action action = Action.valueOf(actionStr); final String reason = (String) result.get("reason"); + @SuppressWarnings("unchecked") + final List tags = (List) result.get("tags"); span.setTag(ACTION_TAG, action); - span.setTag(REASON_TAG, reason); + if (reason != null) { + span.setTag(REASON_TAG, reason); + } + if (tags != null) { + for (final String tag : tags) { + span.setTag("ai_guard.tag." + tag, true); + } + } final boolean shouldBlock = isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW; WafMetricCollector.get().aiGuardRequest(action, shouldBlock); 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 366b977a7a3..80426761c24 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 @@ -163,7 +163,7 @@ class AIGuardInternalTests extends DDSpecification { return mockResponse( request, 200, - [data: [attributes: [action: suite.action, reason: suite.reason, is_blocking_enabled: suite.blocking]]] + [data: [attributes: [action: suite.action, reason: suite.reason, tags: suite.tags ?: [], is_blocking_enabled: suite.blocking]]] ) } } @@ -193,6 +193,11 @@ class AIGuardInternalTests extends DDSpecification { if (throwAbortError) { 1 * span.addThrowable(_ as AIGuard.AIGuardAbortError) } + if (suite.tags) { + suite.tags.each { + 1 * span.setTag("ai_guard.tag.${it}", true) + } + } assertRequest(request, suite.messages) if (throwAbortError) { @@ -497,14 +502,16 @@ class AIGuardInternalTests extends DDSpecification { private static class TestSuite { private final AIGuard.Action action private final String reason + private final List tags private final boolean blocking private final String description private final String target private final List messages - TestSuite(AIGuard.Action action, String reason, boolean blocking, String description, String target, List messages) { + TestSuite(AIGuard.Action action, String reason, List tags, boolean blocking, String description, String target, List messages) { this.action = action this.reason = reason + this.tags = tags this.blocking = blocking this.description = description this.target = target @@ -512,7 +519,11 @@ class AIGuardInternalTests extends DDSpecification { } static List build() { - def actionValues = [[ALLOW, 'Go ahead'], [DENY, 'Nope'], [ABORT, 'Kill it with fire']] + def actionValues = [ + [ALLOW, 'Go ahead', []], + [DENY, 'Nope', ['deny_everything', 'test_deny']], + [ABORT, 'Kill it with fire', ['alarm_tag', 'abort_everything']] + ] def blockingValues = [true, false] def suiteValues = [ ['tool call', 'tool', TOOL_CALL], @@ -521,7 +532,7 @@ class AIGuardInternalTests extends DDSpecification { ] return combinations([actionValues, blockingValues, suiteValues] as Iterable) .collect { action, blocking, suite -> - new TestSuite(action[0], action[1], blocking, suite[0], suite[1], suite[2]) + new TestSuite(action[0], action[1], action[2], blocking, suite[0], suite[1], suite[2]) } } From e100df71c4e3f3a7dcf294ecc055034562b7dcc7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 12 Nov 2025 09:32:55 +0100 Subject: [PATCH 2/4] Rename tag name to tags --- .../src/main/java/com/datadog/aiguard/AIGuardInternal.java | 2 +- .../test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy | 2 +- 2 files changed, 2 insertions(+), 2 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 eb2143b152f..acd6d8b1ebb 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 @@ -232,7 +232,7 @@ public Evaluation evaluate(final List messages, final Options options) } if (tags != null) { for (final String tag : tags) { - span.setTag("ai_guard.tag." + tag, true); + span.setTag("ai_guard.tags." + tag, true); } } final boolean shouldBlock = 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 80426761c24..4a0f5f76b0c 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 @@ -195,7 +195,7 @@ class AIGuardInternalTests extends DDSpecification { } if (suite.tags) { suite.tags.each { - 1 * span.setTag("ai_guard.tag.${it}", true) + 1 * span.setTag("ai_guard.tags.${it}", true) } } From a6d1707fc5d168d3e88c612c386c8615ca295dac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 12 Nov 2025 14:44:51 +0100 Subject: [PATCH 3/4] Move matching tags to meta struct --- .../java/com/datadog/aiguard/AIGuardInternal.java | 14 ++++++-------- .../datadog/aiguard/AIGuardInternalTests.groovy | 13 ++++++++----- 2 files changed, 14 insertions(+), 13 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 acd6d8b1ebb..a20603f1a0e 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 @@ -4,7 +4,6 @@ import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardTruncationType.CONTENT; import static datadog.trace.api.telemetry.WafMetricCollector.AIGuardTruncationType.MESSAGES; import static datadog.trace.util.Strings.isBlank; -import static java.util.Collections.singletonMap; import com.squareup.moshi.JsonAdapter; import com.squareup.moshi.JsonReader; @@ -69,7 +68,8 @@ public BadConfigurationException(final String message) { static final String REASON_TAG = "ai_guard.reason"; static final String BLOCKED_TAG = "ai_guard.blocked"; static final String META_STRUCT_TAG = "ai_guard"; - static final String META_STRUCT_KEY = "messages"; + static final String META_STRUCT_MESSAGES = "messages"; + static final String META_STRUCT_MATCHING_RULES = "matching_rules"; public static void install() { final Config config = Config.get(); @@ -208,8 +208,8 @@ public Evaluation evaluate(final List messages, final Options options) } else { span.setTag(TARGET_TAG, "prompt"); } - final Map metaStruct = - singletonMap(META_STRUCT_KEY, messagesForMetaStruct(messages)); + 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() @@ -230,10 +230,8 @@ public Evaluation evaluate(final List messages, final Options options) if (reason != null) { span.setTag(REASON_TAG, reason); } - if (tags != null) { - for (final String tag : tags) { - span.setTag("ai_guard.tags." + tag, true); - } + if (tags != null && !tags.isEmpty()) { + metaStruct.put(META_STRUCT_MATCHING_RULES, tags); } final boolean shouldBlock = isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW; 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 4a0f5f76b0c..aadaf68d3d5 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 @@ -157,6 +157,7 @@ class AIGuardInternalTests extends DDSpecification { Request request = null Throwable error = null AIGuard.Evaluation eval = null + Map receivedMeta = null final throwAbortError = suite.blocking && suite.action != ALLOW final call = Mock(Call) { execute() >> { @@ -189,16 +190,18 @@ class AIGuardInternalTests extends DDSpecification { } 1 * span.setTag(AIGuardInternal.ACTION_TAG, suite.action) 1 * span.setTag(AIGuardInternal.REASON_TAG, suite.reason) - 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, [messages: suite.messages]) + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _ as Map) >> { + receivedMeta = it[1] as Map + return span + } if (throwAbortError) { 1 * span.addThrowable(_ as AIGuard.AIGuardAbortError) } + + receivedMeta.messages == suite.messages if (suite.tags) { - suite.tags.each { - 1 * span.setTag("ai_guard.tags.${it}", true) - } + receivedMeta.matching_rules == suite.tags } - assertRequest(request, suite.messages) if (throwAbortError) { error instanceof AIGuard.AIGuardAbortError From 61aa2a432f6241df7164fbd8de99113bfeb9f98a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 12 Nov 2025 14:47:17 +0100 Subject: [PATCH 4/4] Rename to attack_categories --- .../src/main/java/com/datadog/aiguard/AIGuardInternal.java | 4 ++-- .../groovy/com/datadog/aiguard/AIGuardInternalTests.groovy | 2 +- 2 files changed, 3 insertions(+), 3 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 a20603f1a0e..4c7604444b7 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 @@ -69,7 +69,7 @@ public BadConfigurationException(final String message) { static final String BLOCKED_TAG = "ai_guard.blocked"; static final String META_STRUCT_TAG = "ai_guard"; static final String META_STRUCT_MESSAGES = "messages"; - static final String META_STRUCT_MATCHING_RULES = "matching_rules"; + static final String META_STRUCT_CATEGORIES = "attack_categories"; public static void install() { final Config config = Config.get(); @@ -231,7 +231,7 @@ public Evaluation evaluate(final List messages, final Options options) span.setTag(REASON_TAG, reason); } if (tags != null && !tags.isEmpty()) { - metaStruct.put(META_STRUCT_MATCHING_RULES, tags); + metaStruct.put(META_STRUCT_CATEGORIES, tags); } final boolean shouldBlock = isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW; 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 aadaf68d3d5..77a27dbc61f 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 @@ -200,7 +200,7 @@ class AIGuardInternalTests extends DDSpecification { receivedMeta.messages == suite.messages if (suite.tags) { - receivedMeta.matching_rules == suite.tags + receivedMeta.attack_categories == suite.tags } assertRequest(request, suite.messages) if (throwAbortError) {