feat(aiguard): set manual.keep on root span after AI Guard evaluation (#7758)

smola · web-flow · commit c9ba1450ed41 · 2026-03-13T14:35:46.000Z
feat(aiguard): set manual.keep on root span after AI Guard evaluation

Ensure AI Guard traces are always retained by setting USER_KEEP sampling
priority with AI_GUARD mechanism (13) on the root span after successful
evaluations, bypassing sampling rules.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;

Co-authored-by: santiago.mola &lt;santiago.mola@datadoghq.com&gt;
diff --git a/packages/dd-trace/src/aiguard/sdk.js b/packages/dd-trace/src/aiguard/sdk.js
@@ -4,6 +4,8 @@ const rfdc = require('../../../../vendor/dist/rfdc')({ proto: false, circles: fa
 const log = require('../log')
 const telemetryMetrics = require('../telemetry/metrics')
 const tracerVersion = require('../../../../package.json').version
+const { keepTrace } = require('../priority_sampler')
+const { AI_GUARD } = require('../standalone/product')
 const NoopAIGuard = require('./noop')
 const executeRequest = require('./client')
 const {
@@ -153,6 +155,12 @@ class AIGuard extends NoopAIGuard {
       span.meta_struct = {
         [AI_GUARD_META_STRUCT_KEY]: metaStruct,
       }
+      const rootSpan = span.context()?._trace?.started?.[0]
+      if (rootSpan) {
+        // keepTrace must be called before executeRequest so the sampling decision
+        // is propagated correctly to outgoing HTTP client calls.
+        keepTrace(rootSpan, AI_GUARD)
+      }
       let response
       try {
         const payload = {
diff --git a/packages/dd-trace/src/constants.js b/packages/dd-trace/src/constants.js
@@ -17,6 +17,7 @@ module.exports = {
   SAMPLING_MECHANISM_SPAN: 8,
   SAMPLING_MECHANISM_REMOTE_USER: 11,
   SAMPLING_MECHANISM_REMOTE_DYNAMIC: 12,
+  SAMPLING_MECHANISM_AI_GUARD: 13,
   SPAN_SAMPLING_MECHANISM: '_dd.span_sampling.mechanism',
   SPAN_SAMPLING_RULE_RATE: '_dd.span_sampling.rule_rate',
   SPAN_SAMPLING_MAX_PER_SECOND: '_dd.span_sampling.max_per_second',
diff --git a/packages/dd-trace/src/standalone/product.js b/packages/dd-trace/src/standalone/product.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { SAMPLING_MECHANISM_APPSEC } = require('../constants')
+const { SAMPLING_MECHANISM_APPSEC, SAMPLING_MECHANISM_AI_GUARD } = require('../constants')
 const RateLimiter = require('../rate_limiter')
 
 /**
@@ -26,6 +26,7 @@ const PRODUCTS = {
   DSM: { id: 1 << 2 },
   DJM: { id: 1 << 3 },
   DBM: { id: 1 << 4 },
+  AI_GUARD: { id: 1 << 5, mechanism: SAMPLING_MECHANISM_AI_GUARD },
 }
 
 module.exports = {
diff --git a/packages/dd-trace/test/aiguard/index.spec.js b/packages/dd-trace/test/aiguard/index.spec.js
@@ -15,6 +15,8 @@ const { assertObjectContains } = require('../../../../integration-tests/helpers'
 const tracerVersion = require('../../../../package.json').version
 const telemetryMetrics = require('../../src/telemetry/metrics')
 const appsecNamespace = telemetryMetrics.manager.namespace('appsec')
+const { USER_KEEP } = require('../../../../ext/priority')
+const { SAMPLING_MECHANISM_AI_GUARD, DECISION_MAKER_KEY } = require('../../src/constants')
 
 describe('AIGuard SDK', () => {
   const config = {
@@ -477,4 +479,77 @@ describe('AIGuard SDK', () => {
       assertFetch(toolCall, `${endpoint}/evaluate`)
     })
   }
+
+  describe('manual keep on root span', () => {
+    const assertRootSpanKept = async () => {
+      await agent.assertSomeTraces(traces => {
+        const rootSpan = traces[0][0]
+        assert.strictEqual(rootSpan.metrics._sampling_priority_v1, USER_KEEP)
+        assert.strictEqual(rootSpan.meta[DECISION_MAKER_KEY], `-${SAMPLING_MECHANISM_AI_GUARD}`)
+      })
+    }
+
+    it('sets USER_KEEP on root span after ALLOW evaluation', async () => {
+      mockFetch({
+        body: { data: { attributes: { action: 'ALLOW', reason: 'OK', tags: [], is_blocking_enabled: false } } },
+      })
+
+      await tracer.trace('root', async () => {
+        await aiguard.evaluate(prompt)
+      })
+
+      await assertRootSpanKept()
+    })
+
+    it('sets USER_KEEP on root span after DENY evaluation (non-blocking)', async () => {
+      mockFetch({
+        body: {
+          data: { attributes: { action: 'DENY', reason: 'denied', tags: ['deny_tag'], is_blocking_enabled: false } },
+        },
+      })
+
+      await tracer.trace('root', async () => {
+        await aiguard.evaluate(prompt, { block: false })
+      })
+
+      await assertRootSpanKept()
+    })
+
+    it('keeps trace even when auto-sampling would drop it', async () => {
+      // Configure sampler to drop all traces (0% sample rate)
+      tracer._tracer._prioritySampler.configure('test', { sampleRate: 0 })
+
+      try {
+        mockFetch({
+          body: { data: { attributes: { action: 'ALLOW', reason: 'OK', tags: [], is_blocking_enabled: false } } },
+        })
+
+        await tracer.trace('root', async () => {
+          await aiguard.evaluate(prompt)
+        })
+
+        await assertRootSpanKept()
+      } finally {
+        tracer._tracer._prioritySampler.configure('test', {})
+      }
+    })
+
+    it('sets USER_KEEP on root span after ABORT evaluation (blocking)', async () => {
+      mockFetch({
+        body: {
+          data: { attributes: { action: 'ABORT', reason: 'blocked', tags: ['tag'], is_blocking_enabled: true } },
+        },
+      })
+
+      await tracer.trace('root', async () => {
+        try {
+          await aiguard.evaluate(prompt, { block: true })
+        } catch {
+          // expected AIGuardAbortError
+        }
+      })
+
+      await assertRootSpanKept()
+    })
+  })
 })
