Set sampling decision maker for AI Guard

Author: smola

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

@@ -12,7 +12,6 @@
 import com.squareup.moshi.Types;
 import datadog.communication.http.OkHttpUtils;
 import datadog.trace.api.Config;
-import datadog.trace.api.DDTags;
 import datadog.trace.api.aiguard.AIGuard;
 import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError;
 import datadog.trace.api.aiguard.AIGuard.AIGuardClientError;
@@ -29,6 +28,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
 import java.io.IOException;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Type;
@@ -221,7 +221,7 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
     final AgentSpan span = builder.start();
     final AgentSpan localRootSpan = span.getLocalRootSpan();
     if (localRootSpan != null) {
-      localRootSpan.setTag(DDTags.MANUAL_KEEP, true);
+      localRootSpan.setTag(Tags.AI_GUARD_KEEP, true);
     }
     try (final AgentScope scope = tracer.activateSpan(span)) {
       final Message last = messages.get(messages.size() - 1);

dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy

@@ -6,11 +6,11 @@ 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.DDTags
 import datadog.trace.api.aiguard.AIGuard
 import datadog.trace.api.telemetry.WafMetricCollector
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.test.util.DDSpecification
 import okhttp3.Call
 import okhttp3.HttpUrl
@@ -189,7 +189,7 @@ class AIGuardInternalTests extends DDSpecification {
 
     then:
     1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target)
-    1 * localRootSpan.setTag(DDTags.MANUAL_KEEP, true)
+    1 * localRootSpan.setTag(Tags.AI_GUARD_KEEP, true)
     if (suite.target == 'tool') {
       1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc')
     }

dd-trace-core/src/main/java/datadog/trace/core/taginterceptor/TagInterceptor.java

@@ -108,6 +108,7 @@ public boolean needsIntercept(String tag) {
       case DDTags.MANUAL_KEEP:
       case DDTags.MANUAL_DROP:
       case Tags.ASM_KEEP:
+      case Tags.AI_GUARD_KEEP:
       case Tags.SAMPLING_PRIORITY:
       case Tags.PROPAGATED_TRACE_SOURCE:
       case Tags.PROPAGATED_DEBUG:
@@ -155,6 +156,12 @@ public boolean interceptTag(DDSpanContext span, String tag, Object value) {
           return true;
         }
         return false;
+      case Tags.AI_GUARD_KEEP:
+        if (asBoolean(value)) {
+          span.forceKeep(SamplingMechanism.AI_GUARD);
+          return true;
+        }
+        return false;
       case Tags.SAMPLING_PRIORITY:
         return interceptSamplingPriority(span, value);
       case Tags.PROPAGATED_TRACE_SOURCE:

internal-api/src/main/java/datadog/trace/api/sampling/SamplingMechanism.java

@@ -37,6 +37,7 @@ public class SamplingMechanism {
 
   public static final byte REMOTE_USER_RULE = 11;
   public static final byte REMOTE_ADAPTIVE_RULE = 12;
+  public static final byte AI_GUARD = 13;
 
   /** Force override sampling decision from external source, like W3C traceparent. */
   public static final byte EXTERNAL_OVERRIDE = Byte.MIN_VALUE;

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

@@ -167,6 +167,9 @@ public class Tags {
   /** ASM force tracer to keep the trace */
   public static final String ASM_KEEP = "asm.keep";
 
+  /** AI Guard force tracer to keep the trace */
+  public static final String AI_GUARD_KEEP = "ai_guard.keep";
+
   public static final String PROPAGATED_TRACE_SOURCE = "_dd.p.ts";
   public static final String PROPAGATED_DEBUG = "_dd.p.debug";