From 42d77deb22eb59d60eeaf638f669e892db3ee55e Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Wed, 8 Jul 2026 14:27:04 +0200 Subject: [PATCH 1/2] Support DD_APM_TRACING_ENABLED=false for AI Guard traces Co-Authored-By: Claude Sonnet 4.6 --- .../com/datadog/aiguard/AIGuardInternal.java | 4 +- .../datadog/trace/core/TraceCollector.java | 7 +- .../common/sampling/AIGuardSamplingTest.java | 113 ++++++++++++++++++ .../bootstrap/instrumentation/api/Tags.java | 3 + 4 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java 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 11ede90b51c..912eec8b4cc 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 @@ -71,7 +71,7 @@ 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 EVENT_TAG = "ai_guard.event"; + static final String META_STRUCT_TAG = "ai_guard"; static final String META_STRUCT_MESSAGES = "messages"; static final String META_STRUCT_CATEGORIES = "attack_categories"; @@ -275,7 +275,7 @@ public Evaluation evaluate(final List messages, final Options options) final AgentSpan localRootSpan = span.getLocalRootSpan(); if (localRootSpan != null) { localRootSpan.setTag(Tags.AI_GUARD_KEEP, true); - localRootSpan.setTag(EVENT_TAG, true); + localRootSpan.setTag(Tags.AI_GUARD_EVENT, true); applyClientIpTags(localRootSpan); // copyAnomalyDetectionTags MUST run after applyClientIpTags, to make // sure client IP tags were populated. diff --git a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java index 2d97e895fb0..40652664319 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java @@ -6,6 +6,7 @@ import datadog.trace.api.sampling.PrioritySampling; import datadog.trace.api.time.TimeSource; import datadog.trace.bootstrap.instrumentation.api.AgentTraceCollector; +import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.common.sampling.PrioritySampler; import java.util.concurrent.atomic.AtomicLongFieldUpdater; import javax.annotation.Nonnull; @@ -65,11 +66,13 @@ public void setSamplingPriorityIfNecessary() { DDSpan rootSpan = getRootSpan(); if (traceConfig.sampler instanceof PrioritySampler && rootSpan != null) { // Ignore the force-keep priority in the absence of propagated _dd.p.ts span tag marked for - // ASM. + // ASM. AI Guard traces (ai_guard.event=true) preserve the USER_KEEP priority set by + // ai_guard.keep, bypassing the sampler override. if ((!Config.get().isApmTracingEnabled() && !ProductTraceSource.isProductMarked( rootSpan.spanContext().getPropagationTags().getTraceSource(), - ProductTraceSource.ASM)) + ProductTraceSource.ASM) + && !Boolean.TRUE.equals(rootSpan.getTag(Tags.AI_GUARD_EVENT))) || rootSpan.spanContext().getSamplingPriority() == PrioritySampling.UNSET) { ((PrioritySampler) traceConfig.sampler).setSamplingPriority(rootSpan); } diff --git a/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java b/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java new file mode 100644 index 00000000000..fd72a968ee3 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/common/sampling/AIGuardSamplingTest.java @@ -0,0 +1,113 @@ +package datadog.trace.common.sampling; + +import static datadog.trace.api.config.AppSecConfig.APPSEC_ENABLED; +import static datadog.trace.api.config.GeneralConfig.APM_TRACING_ENABLED; +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP; +import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.CoreTracer; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.DDSpan; +import datadog.trace.junit.utils.config.WithConfig; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Verifies that AI Guard traces are kept regardless of APM/ASM configuration. */ +public class AIGuardSamplingTest extends DDCoreJavaSpecification { + + private ListWriter writer; + private CoreTracer tracer; + + @BeforeEach + void setUp() { + writer = new ListWriter(); + tracer = tracerBuilder().writer(writer).build(); + } + + @Test + void aiGuardTraceIsKeptWhenApmEnabled() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + void aiGuardTraceIsKeptWhenApmAndAsmDisabled() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void aiGuardTraceIsKeptInAsmStandaloneMode() throws Exception { + DDSpan span = (DDSpan) tracer.buildSpan("datadog", "operation").start(); + span.setTag(Tags.AI_GUARD_KEEP, true); + span.setTag(Tags.AI_GUARD_EVENT, true); + span.finish(); + writer.waitForTraces(1); + assertEquals(USER_KEEP, (int) span.getSamplingPriority()); + assertDecisionMakerIsAiGuard(span); + } + + /** ai_guard.event alone (without AI_GUARD_KEEP) must not bypass the sampler. */ + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void aiGuardEventTagAloneDoesNotBypassSampler() throws Exception { + // consume the first allowed slot + DDSpan first = (DDSpan) tracer.buildSpan("datadog", "op").start(); + first.finish(); + writer.waitForTraces(1); + + // second trace sets ai_guard.event but NOT AI_GUARD_KEEP — should still be dropped + DDSpan second = (DDSpan) tracer.buildSpan("datadog", "op").start(); + second.setTag(Tags.AI_GUARD_EVENT, true); + second.finish(); + writer.waitForTraces(2); + + assertEquals(SAMPLER_DROP, (int) second.getSamplingPriority()); + } + + /** + * AI Guard traces must be kept even when AsmStandaloneSampler has exhausted its rate-limit slot. + */ + @Test + @WithConfig(key = APM_TRACING_ENABLED, value = "false") + @WithConfig(key = APPSEC_ENABLED, value = "true") + void aiGuardTraceIsKeptAfterRateLimitSlotIsExhausted() throws Exception { + // consume the first allowed slot + DDSpan first = (DDSpan) tracer.buildSpan("datadog", "op").start(); + first.finish(); + writer.waitForTraces(1); + + // AI Guard trace must still be force-kept even though the rate-limit slot is gone + DDSpan aiGuard = (DDSpan) tracer.buildSpan("datadog", "op").start(); + aiGuard.setTag(Tags.AI_GUARD_KEEP, true); + aiGuard.setTag(Tags.AI_GUARD_EVENT, true); + aiGuard.finish(); + writer.waitForTraces(2); + + assertEquals(USER_KEEP, (int) aiGuard.getSamplingPriority()); + assertDecisionMakerIsAiGuard(aiGuard); + } + + private static void assertDecisionMakerIsAiGuard(DDSpan span) { + Map ptags = span.spanContext().getPropagationTags().createTagMap(); + assertEquals("-13", ptags.get("_dd.p.dm"), "_dd.p.dm must be -13 (AI Guard decision maker)"); + } +} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java index 14496e8b243..782e63fa1ec 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java @@ -177,6 +177,9 @@ public class Tags { /** AI Guard force tracer to keep the trace */ public static final String AI_GUARD_KEEP = "ai_guard.keep"; + /** Marks a span as originating from an AI Guard evaluation */ + public static final String AI_GUARD_EVENT = "ai_guard.event"; + public static final String PROPAGATED_TRACE_SOURCE = "_dd.p.ts"; public static final String PROPAGATED_DEBUG = "_dd.p.debug"; From 1e9ccb8e0b91d2b32790e6404d88c3a1d730835c Mon Sep 17 00:00:00 2001 From: Clara Poncet Date: Wed, 8 Jul 2026 14:39:52 +0200 Subject: [PATCH 2/2] Fix remaining EVENT_TAG reference in Groovy test --- .../test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e355b2eec3c..0cbd610e043 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 @@ -192,7 +192,7 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target) 1 * localRootSpan.setTag(Tags.AI_GUARD_KEEP, true) - 1 * localRootSpan.setTag(AIGuardInternal.EVENT_TAG, true) + 1 * localRootSpan.setTag(Tags.AI_GUARD_EVENT, true) if (suite.target == 'tool') { 1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc') }