Add setting of manual.keep on the root span for AI Guard

lib/datadog/ai_guard/evaluation.rb

@@ -10,6 +10,12 @@ def perform(messages, allow_raise: false)
           raise ArgumentError, "Messages must not be empty" if messages&.empty?
 
           Tracing.trace(Ext::SPAN_NAME) do |span, trace|
+            trace.keep!
+            trace.set_tag(
+              Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER,
+              Tracing::Sampling::Ext::Decision::AI_GUARD
+            )
+
             if (last_message = messages.last)
               if last_message.tool_call
                 span.set_tag(Ext::TARGET_TAG, "tool")

lib/datadog/tracing/sampling/ext.rb

@@ -47,6 +47,8 @@ module Decision
           MANUAL = '-4'
           # Formerly AppSec.
           ASM = '-5'
+          # AI Guard.
+          AI_GUARD = '-13'
           # Dynamically configured rule, explicitly created by the user.
           REMOTE_USER_RULE = '-11'
           # Dynamically configured  rule, automatically generated by Datadog.

sig/datadog/tracing/sampling/ext.rbs

@@ -21,6 +21,7 @@ module Datadog
           TRACE_SAMPLING_RULE: ::String
           MANUAL: ::String
           ASM: ::String
+          AI_GUARD: ::String
         end
       end
     end

spec/datadog/ai_guard/evaluation_spec.rb

@@ -51,6 +51,16 @@
       expect(ai_guard_span).not_to be_nil
     end
 
+    it "sets manual.keep on the trace with AI Guard decision maker" do
+      described_class.perform([
+        Datadog::AIGuard.message(role: :user, content: "Some content")
+      ])
+
+      trace = traces.first
+      expect(trace.sampling_priority).to eq(Datadog::Tracing::Sampling::Ext::Priority::USER_KEEP)
+      expect(trace.send(:sampling_decision_maker)).to eq('-13')
+    end
+
     it "sets target tag to 'prompt' when last message is a prompt" do
       described_class.perform([
         Datadog::AIGuard.message(role: :system, content: "Some content"),