Support DD_APM_TRACING_ENABLED=false for AI Guard traces

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -13,6 +13,7 @@
 import datadog.communication.http.OkHttpUtils;
 import datadog.context.ContextScope;
 import datadog.trace.api.Config;
+import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.aiguard.AIGuard;
 import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError;
 import datadog.trace.api.aiguard.AIGuard.AIGuardClientError;
@@ -71,7 +72,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 +276,8 @@ 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);
+      localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
       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

@@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.PropertyNamingStrategies
 import com.squareup.moshi.Moshi
 import datadog.common.version.VersionInfo
 import datadog.trace.api.Config
+import datadog.trace.api.ProductTraceSource
 import datadog.trace.api.aiguard.AIGuard
 import datadog.trace.api.gateway.RequestContext
 import datadog.trace.api.telemetry.WafMetricCollector
@@ -192,7 +193,8 @@ 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)
+    1 * localRootSpan.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD)
     if (suite.target == 'tool') {
       1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc')
     }

dd-trace-core/src/main/java/datadog/trace/core/TraceCollector.java

@@ -64,12 +64,12 @@ public void setSamplingPriorityIfNecessary() {
     // Locks inside DDSpanContext ensure the correct behavior in the race case
     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.
+      // Skip sampler override when _dd.p.ts is marked for ASM or AI Guard.
       if ((!Config.get().isApmTracingEnabled()
               && !ProductTraceSource.isProductMarked(
                   rootSpan.spanContext().getPropagationTags().getTraceSource(),
-                  ProductTraceSource.ASM))
+                  ProductTraceSource.ASM,
+                  ProductTraceSource.AI_GUARD))
           || 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,114 @@
+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.USER_KEEP;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import datadog.trace.api.ProductTraceSource;
+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.test.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.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
+    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.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
+    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.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
+    span.finish();
+    writer.waitForTraces(1);
+    assertEquals(USER_KEEP, (int) span.getSamplingPriority());
+    assertDecisionMakerIsAiGuard(span);
+  }
+
+  /** _dd.p.ts 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 propagatedTraceSourceAloneDoesNotBypassSampler() throws Exception {
+    DDSpan span = (DDSpan) tracer.buildSpan("datadog", "op").start();
+    span.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
+    span.finish();
+    writer.waitForTraces(1);
+
+    // _dd.p.ts alone must not force-keep the trace; only AI_GUARD_KEEP does that.
+    assertNotEquals(USER_KEEP, (int) span.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.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.AI_GUARD);
+    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)");
+    assertEquals("20", ptags.get("_dd.p.ts"), "_dd.p.ts must have AI Guard bit set (0x20)");
+  }
+}

dd-trace-core/src/test/java/datadog/trace/core/taginterceptor/TagInterceptorTest.java

@@ -389,6 +389,13 @@ static Stream<Arguments> setPrioritySamplingViaTagArguments() {
         arguments("asm.keep / 'true'", Tags.ASM_KEEP, "true", (int) PrioritySampling.USER_KEEP),
         arguments("asm.keep / 'false'", Tags.ASM_KEEP, "false", null),
         arguments("asm.keep / 'asdf'", Tags.ASM_KEEP, "asdf", null),
+        arguments(
+            "ai_guard.keep / true", Tags.AI_GUARD_KEEP, true, (int) PrioritySampling.USER_KEEP),
+        arguments("ai_guard.keep / false", Tags.AI_GUARD_KEEP, false, null),
+        arguments(
+            "ai_guard.keep / 'true'", Tags.AI_GUARD_KEEP, "true", (int) PrioritySampling.USER_KEEP),
+        arguments("ai_guard.keep / 'false'", Tags.AI_GUARD_KEEP, "false", null),
+        arguments("ai_guard.keep / 'asdf'", Tags.AI_GUARD_KEEP, "asdf", null),
         arguments(
             "sampling.priority / -1", Tags.SAMPLING_PRIORITY, -1, (int) PrioritySampling.USER_DROP),
         arguments(

internal-api/src/main/java/datadog/trace/api/ProductTraceSource.java

@@ -22,6 +22,7 @@ public class ProductTraceSource {
   public static final int DSM = 0x04;
   public static final int DJM = 0x08;
   public static final int DBM = 0x10;
+  public static final int AI_GUARD = 0x20;
 
   /** Updates the bitfield by setting the bit corresponding to a specific product. */
   public static int updateProduct(int bitfield, int product) {
@@ -30,7 +31,12 @@ public static int updateProduct(int bitfield, int product) {
 
   /** Checks if the bitfield is marked for a specific product. */
   public static boolean isProductMarked(final int bitfield, int product) {
-    return (bitfield & product) != 0; // Check if the bit is set
+    return (bitfield & product) != 0;
+  }
+
+  /** Checks if the bitfield is marked for either of the two given products. */
+  public static boolean isProductMarked(final int bitfield, int productA, int productB) {
+    return (bitfield & (productA | productB)) != 0;
   }
 
   /**

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java

@@ -183,6 +183,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";
 

internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy

@@ -28,6 +28,23 @@ class ProductTraceSourceTest extends DDSpecification {
     value         | product       | expected
     ProductTraceSource.ASM | ProductTraceSource.ASM | true
     ProductTraceSource.DSM | ProductTraceSource.ASM | false
+    ProductTraceSource.AI_GUARD | ProductTraceSource.AI_GUARD | true
+    ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
+  }
+
+  void 'test isProductMarked with two products'(){
+    when:
+    final result = ProductTraceSource.isProductMarked(value, productA, productB)
+
+    then:
+    result == expected
+
+    where:
+    value                   | productA               | productB                    | expected
+    ProductTraceSource.ASM      | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | true
+    ProductTraceSource.AI_GUARD | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | true
+    ProductTraceSource.DSM      | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
+    ProductTraceSource.UNSET    | ProductTraceSource.ASM | ProductTraceSource.AI_GUARD | false
   }
 
   void 'test getBitfieldHex'(){
@@ -41,6 +58,7 @@ class ProductTraceSourceTest extends DDSpecification {
     value         | expected
     ProductTraceSource.UNSET               | "00"
     ProductTraceSource.ASM | "02"
+    ProductTraceSource.AI_GUARD | "20"
   }
 
   void 'test parseBitfieldHex'(){
@@ -56,5 +74,6 @@ class ProductTraceSourceTest extends DDSpecification {
     null | ProductTraceSource.UNSET
     ""   | ProductTraceSource.UNSET
     "02" | ProductTraceSource.ASM
+    "20" | ProductTraceSource.AI_GUARD
   }
 }