Force manual keep of traces for AI Guard

smola · smola · commit ea830f9f46d1 · 2026-02-19T09:24:29.000+01:00
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
@@ -12,6 +12,7 @@
 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;
@@ -218,6 +219,10 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
       builder.asChildOf(parent.context());
     }
     final AgentSpan span = builder.start();
+    final AgentSpan localRootSpan = span.getLocalRootSpan();
+    if (localRootSpan != null) {
+      localRootSpan.setTag(DDTags.MANUAL_KEEP, true);
+    }
     try (final AgentScope scope = tracer.activateSpan(span)) {
       final Message last = messages.get(messages.size() - 1);
       if (isToolCall(last)) {
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
@@ -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.DDTags
 import datadog.trace.api.aiguard.AIGuard
 import datadog.trace.api.telemetry.WafMetricCollector
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
@@ -72,12 +73,15 @@ class AIGuardInternalTests extends DDSpecification {
   protected static final PROMPT = TOOL_OUTPUT + [AIGuard.Message.message('assistant', '2 + 2 is 5'), AIGuard.Message.message('user', '')]
 
   protected AgentSpan span
+  protected AgentSpan localRootSpan
 
   void setup() {
     injectEnvConfig('SERVICE', 'ai_guard_test')
     injectEnvConfig('ENV', 'test')
 
     span = Mock(AgentSpan)
+    localRootSpan = Mock(AgentSpan)
+    span.getLocalRootSpan() >> localRootSpan
     final builder = Mock(AgentTracer.SpanBuilder) {
       start() >> span
     }
@@ -185,6 +189,7 @@ class AIGuardInternalTests extends DDSpecification {
 
     then:
     1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target)
+    1 * localRootSpan.setTag(DDTags.MANUAL_KEEP, true)
     if (suite.target == 'tool') {
       1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc')
     }
diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy
@@ -69,9 +69,15 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest {
     ?.find {
       it.resource == 'ai_guard'
     } as DecodedSpan
+    final rootSpan = traces*.spans
+    ?.flatten()
+    ?.find {
+      it.traceId == span.traceId && it.parentId == 0
+    } as DecodedSpan
     assert span.meta.get('ai_guard.action') == action
     assert span.meta.get('ai_guard.reason') == reason
     assert span.meta.get('ai_guard.target') == 'prompt'
+    assert rootSpan.metrics.get('_sampling_priority_v1') == 2
     final messages = span.metaStruct.get('ai_guard').messages as List<Map<String, Object>>
     assert messages.size() == 2
     with(messages[0]) {
