Support DD_APM_TRACING_ENABLED=false for AI Guard traces

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<Message> 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.

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')
     }

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);
       }

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<String, String> ptags = span.spanContext().getPropagationTags().createTagMap();
+    assertEquals("-13", ptags.get("_dd.p.dm"), "_dd.p.dm must be -13 (AI Guard decision maker)");
+  }
+}

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";