datadog 2.28.0 → 2.41.0

data/lib/datadog/ai_guard/contrib/rails/patcher.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require_relative "../../../core/utils/only_once"
+require_relative "../rack/request_middleware"
+require_relative "../../../tracing/contrib"
+require_relative "../../../tracing/contrib/rack/middlewares"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rails
+        # Patcher for AI Guard on Rails. Inserts the AI Guard Rack middleware
+        # right after the Tracing Rack middleware so the request span is
+        # already active when AI Guard tags the client IP.
+        module Patcher
+          BEFORE_INITIALIZE_ONLY_ONCE_PER_APP = Hash.new { |h, key| h[key] = Datadog::Core::Utils::OnlyOnce.new }
+
+          module_function
+
+          def patched?
+            !!Patcher.instance_variable_get(:@patched)
+          end
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            patch_before_initialize
+            Patcher.instance_variable_set(:@patched, true)
+          end
+
+          def patch_before_initialize
+            ::ActiveSupport.on_load(:before_initialize) do
+              Datadog::AIGuard::Contrib::Rails::Patcher.before_initialize(self)
+            end
+          end
+
+          def before_initialize(app)
+            BEFORE_INITIALIZE_ONLY_ONCE_PER_APP[app].run do
+              # Middleware must be added before the application is initialized.
+              # Otherwise the middleware stack will be frozen.
+              add_middleware(app) if Datadog.configuration.tracing[:rails][:middleware]
+            end
+          end
+
+          def add_middleware(app)
+            if include_middleware?(Datadog::Tracing::Contrib::Rack::TraceMiddleware, app)
+              app.middleware.insert_after(
+                Datadog::Tracing::Contrib::Rack::TraceMiddleware,
+                Datadog::AIGuard::Contrib::Rack::RequestMiddleware
+              )
+            else
+              app.middleware.insert_before(0, Datadog::AIGuard::Contrib::Rack::RequestMiddleware)
+            end
+          end
+
+          def include_middleware?(middleware, app)
+            found = false
+
+            # find tracer middleware reference in Rails::Configuration::MiddlewareStackProxy
+            app.middleware.instance_variable_get(:@operations).each do |operation|
+              args = case operation
+              when Array
+                # rails 5.2
+                _op, args = operation
+                args
+              when Proc
+                if operation.binding.local_variables.include?(:args)
+                  # rails 6.0, 6.1
+                  operation.binding.local_variable_get(:args)
+                else
+                  # rails 7.0 uses ... to pass args
+                  # steep:ignore:start
+                  args_getter = Class.new do
+                    def method_missing(_op, *args) # standard:disable Style/MissingRespondToMissing
+                      args
+                    end
+                  end.new
+                  # steep:ignore:end
+                  operation.call(args_getter)
+                end
+              else
+                # unknown, pass through
+                []
+              end
+
+              found = true if args.include?(middleware)
+            end
+
+            found
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/contrib/ruby_llm/chat_instrumentation.rb

@@ -14,13 +14,51 @@ module Datadog
                     AIGuard.assistant(id: tool_call_id, tool_name: tool_call.name, arguments: tool_call.arguments.to_s)
                   end
                 elsif message.tool_result?
-                  AIGuard.tool(tool_call_id: message.tool_call_id, content: message.content)
+                  build_ai_guard_tool(message)
                 else
-                  AIGuard.message(role: message.role, content: message.content)
+                  build_ai_guard_message(message)
                 end
               end
 
-              AIGuard.evaluate(*ai_guard_messages, allow_raise: true)
+              AIGuard.evaluate(*ai_guard_messages)
+            end
+
+            private
+
+            def build_ai_guard_message(message)
+              content = message.content
+
+              case content
+              when ::RubyLLM::Content
+                AIGuard.message(role: message.role) do |m|
+                  m.text(content.text.to_s) if content.text
+
+                  # Calling attachment.for_llm triggers lazy loading of file contents.
+                  # The result is memoized, so providers won't re-read.
+                  content.attachments.each do |attachment|
+                    case attachment.type
+                    when :image
+                      m.image_url(attachment.for_llm)
+                    when :text
+                      m.text(attachment.for_llm)
+                    end
+                    # Skip :pdf, :audio, :video, :unknown — not supported by AIGuard
+                  end
+                end
+              else
+                AIGuard.message(role: message.role, content: content)
+              end
+            end
+
+            def build_ai_guard_tool(message)
+              content = message.content
+              # Tools can return Content or Content::Raw objects (e.g. with attachments),
+              # but AIGuard.tool expects a String. Extract text when content is a Content object.
+              case content
+              when ::RubyLLM::Content
+                content = content.text.to_s
+              end
+              AIGuard.tool(tool_call_id: message.tool_call_id, content: content)
             end
           end
 

data/lib/datadog/ai_guard/evaluation/content_builder.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Builder for collecting content parts inside a message block.
+      #
+      # Used via the block form of {Datadog::AIGuard.message}:
+      #
+      #   Datadog::AIGuard.message(role: :user) do |m|
+      #     m.text("What's in this image?")
+      #     m.image_url("https://example.com/img.png")
+      #   end
+      class ContentBuilder
+        attr_reader :parts
+
+        def initialize
+          @parts = []
+        end
+
+        def text(text)
+          @parts << ContentPart::Text.new(text)
+        end
+
+        def image_url(url)
+          @parts << ContentPart::ImageURL.new(url)
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/evaluation/content_part.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Evaluation
+      # Namespace for content part types used in multi-modal messages.
+      module ContentPart
+        # A text content part.
+        class Text
+          attr_reader :text
+
+          def initialize(text)
+            @text = text
+          end
+
+          def type
+            :text
+          end
+        end
+
+        # An image URL content part. Accepts an absolute URL or a base64 data URI.
+        class ImageURL
+          attr_reader :url
+
+          def initialize(url)
+            @url = url
+          end
+
+          def type
+            :image_url
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/evaluation/no_op_result.rb

@@ -5,12 +5,14 @@ module Datadog
     module Evaluation
       # Class for emulating AI Guard evaluation result when AI Guard is disabled.
       class NoOpResult
-        attr_reader :action, :reason, :tags
+        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
 
         def initialize
           @action = Result::ALLOW_ACTION
           @reason = "AI Guard is disabled"
           @tags = []
+          @sds_findings = []
+          @tag_probabilities = {}
         end
 
         def allow?

data/lib/datadog/ai_guard/evaluation/request.rb

@@ -36,23 +36,15 @@ module Datadog
                 messages: @serialized_messages,
                 meta: {
                   service: Datadog.configuration.service,
-                  env: Datadog.configuration.env
-                }
-              }
-            }
+                  env: Datadog.configuration.env,
+                },
+              },
+            },
           }
         end
 
         def serialize_messages(messages)
-          serialized_messages = []
-
-          messages.each do |message|
-            serialized_messages << serialize_message(message)
-
-            break if serialized_messages.count == Datadog.configuration.ai_guard.max_messages_length
-          end
-
-          serialized_messages
+          messages.map { |message| serialize_message(message) }
         end
 
         def serialize_message(message)
@@ -64,17 +56,30 @@ module Datadog
                   id: message.tool_call.id,
                   function: {
                     name: message.tool_call.tool_name,
-                    arguments: message.tool_call.arguments
-                  }
-                }
-              ]
+                    arguments: message.tool_call.arguments,
+                  },
+                },
+              ],
             }
+          elsif message.content.is_a?(::Array)
+            {role: message.role, content: serialize_content_parts(message.content)}
           elsif message.tool_call_id
             {role: message.role, tool_call_id: message.tool_call_id, content: message.content}
           else
             {role: message.role, content: message.content}
           end
         end
+
+        def serialize_content_parts(parts)
+          parts.map do |part|
+            case part
+            when ContentPart::Text
+              {type: "text", text: part.text}
+            when ContentPart::ImageURL
+              {type: "image_url", image_url: {url: part.url}}
+            end
+          end
+        end
       end
     end
   end

data/lib/datadog/ai_guard/evaluation/result.rb

@@ -9,7 +9,7 @@ module Datadog
         DENY_ACTION = "DENY"
         ABORT_ACTION = "ABORT"
 
-        attr_reader :action, :reason, :tags
+        attr_reader :action, :reason, :tags, :sds_findings, :tag_probabilities
 
         def initialize(raw_response)
           attributes = raw_response.fetch("data").fetch("attributes")
@@ -17,7 +17,9 @@ module Datadog
           @action = attributes.fetch("action")
           @reason = attributes.fetch("reason")
           @tags = attributes.fetch("tags")
+          @tag_probabilities = attributes.fetch("tag_probs")
           @is_blocking_enabled = attributes.fetch("is_blocking_enabled")
+          @sds_findings = attributes.fetch("sds_findings", [])
         rescue KeyError => e
           raise AIGuardClientError, "Missing key: \"#{e.key}\""
         end

data/lib/datadog/ai_guard/evaluation.rb

@@ -6,10 +6,24 @@ module Datadog
     # and creating `ai_guard` span with required tags
     module Evaluation
       class << self
-        def perform(messages, allow_raise: false)
+        def perform(messages, allow_raise: true)
           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
+            )
+            trace.set_distributed_source(Ext::PRODUCT_BIT)
+            trace.set_tag(Ext::EVENT_TAG, true)
+            trace.set_tag(Ext::TRACE_EXECUTED_TAG, "1")
+
+            Ext::TRACE_ANOMALY_DETECTION_TAGS.each do |tag|
+              value = trace.get_tag(tag)
+              span.set_tag(tag, value) if value
+            end
+
             if (last_message = messages.last)
               if last_message.tool_call
                 span.set_tag(Ext::TARGET_TAG, "tool")
@@ -34,8 +48,10 @@ module Datadog
             span.set_metastruct_tag(
               Ext::METASTRUCT_TAG,
               {
-                messages: truncate_content(request.serialized_messages),
-                attack_categories: result.tags
+                messages: truncate_content(truncate_messages(request.serialized_messages)),
+                attack_categories: result.tags,
+                sds: result.sds_findings,
+                tag_probs: result.tag_probabilities,
               }
             )
 
@@ -56,16 +72,44 @@ module Datadog
 
         private
 
+        def truncate_messages(serialized_messages)
+          max_length = Datadog.configuration.ai_guard.max_messages_length
+          serialized_messages.first(max_length)
+        end
+
+        # Truncates content in serialized messages to stay within the configured byte limit.
+        # For multi-modal messages, only text parts are truncated; image URLs are left intact.
         def truncate_content(serialized_messages)
+          max_bytes = Datadog.configuration.ai_guard.max_content_size_bytes
+
           serialized_messages.map do |message| # steep:ignore
             next message unless message[:content]
 
-            {
-              **message,
-              content: message[:content].byteslice(0, Datadog.configuration.ai_guard.max_content_size_bytes)
-            }
+            if message[:content].is_a?(::Array)
+              serialized_content = message[:content].map do |part|
+                if part[:text]
+                  {**part, text: truncate_content_value(part[:text].to_s, max_bytes)}
+                else
+                  part
+                end
+              end
+
+              {**message, content: serialized_content}
+            else
+              {
+                **message,
+                content: truncate_content_value(message[:content], max_bytes),
+              }
+            end
           end
         end
+
+        def truncate_content_value(content, max_bytes)
+          truncated = content.byteslice(0, max_bytes).to_s
+          return truncated if truncated.valid_encoding?
+
+          truncated.scrub("")
+        end
       end
     end
   end

data/lib/datadog/ai_guard/ext.rb

@@ -4,13 +4,28 @@ module Datadog
   module AIGuard
     # AI Guard specific constants
     module Ext
+      PRODUCT_BIT = 0x20
+
       SPAN_NAME = "ai_guard"
       TARGET_TAG = "ai_guard.target"
       TOOL_NAME_TAG = "ai_guard.tool_name"
       ACTION_TAG = "ai_guard.action"
       REASON_TAG = "ai_guard.reason"
       BLOCKED_TAG = "ai_guard.blocked"
+      EVENT_TAG = "ai_guard.event"
+
       METASTRUCT_TAG = "ai_guard"
+
+      TRACE_EXECUTED_TAG = "_dd.ai_guard.executed"
+      TRACE_HTTP_USERAGENT_TAG = "_dd.ai_guard.http.useragent"
+      TRACE_HTTP_CLIENT_IP_TAG = "_dd.ai_guard.http.client_ip"
+      TRACE_NETWORK_CLIENT_IP_TAG = "_dd.ai_guard.network.client.ip"
+
+      TRACE_ANOMALY_DETECTION_TAGS = [
+        TRACE_HTTP_USERAGENT_TAG,
+        TRACE_HTTP_CLIENT_IP_TAG,
+        TRACE_NETWORK_CLIENT_IP_TAG,
+      ].freeze
     end
   end
 end

data/lib/datadog/ai_guard.rb

@@ -3,6 +3,8 @@
 require_relative "core/configuration"
 require_relative "ai_guard/configuration"
 
+require_relative "ai_guard/contrib/rack/integration"
+require_relative "ai_guard/contrib/rails/integration"
 require_relative "ai_guard/contrib/ruby_llm/integration"
 
 module Datadog
@@ -10,7 +12,7 @@ module Datadog
   module AIGuard
     Core::Configuration::Settings.extend(Configuration::Settings)
 
-    # This error is raised when user passes `allow_raise: true` to Evaluation.perform
+    # This error is raised when `allow_raise` is set to true (the default) in Evaluation.perform
     # and AI Guard considers the messages not safe. Intended to be rescued by the user.
     #
     # WARNING: This name must not change, since front-end is using it.
@@ -50,6 +52,10 @@ module Datadog
         Datadog.send(:components).ai_guard&.logger
       end
 
+      def telemetry
+        Datadog.send(:components).ai_guard&.telemetry
+      end
+
       # Evaluates one or more messages using AI Guard API.
       #
       # Example:
@@ -68,13 +74,14 @@ module Datadog
       #   One or more message objects to be evaluated.
       # @param allow_raise [Boolean]
       #   Whether this method may raise an exception when evaluation result is not ALLOW.
+      #   Defaults to true.
       #
       # @return [Datadog::AIGuard::Evaluation::Result]
       #   The result of AI Guard evaluation.
       # @raise [Datadog::AIGuard::AIGuardAbortError]
       #   If the evaluation results in DENY or ABORT action and `allow_raise` is set to true
       # @public_api
-      def evaluate(*messages, allow_raise: false)
+      def evaluate(*messages, allow_raise: true)
         if enabled?
           Evaluation.perform(messages, allow_raise: allow_raise)
         else
@@ -84,25 +91,42 @@ module Datadog
 
       # Builds a generic evaluation message.
       #
-      # Example:
+      # Accepts either a string content or a block for multi-modal content parts:
       #
       # ```
+      # # String content:
       # Datadog::AIGuard.message(role: :user, content: "Hello, assistant")
+      #
+      # # Multi-modal content with block:
+      # Datadog::AIGuard.message(role: :user) do |m|
+      #   m.text("What's in this image?")
+      #   m.image_url("https://example.com/img.png")
+      # end
       # ```
       #
       # @param role [Symbol]
       #   The role associated with the message.
       #   Must be one of `:assistant`, `:tool`, `:system`, `:developer`, or `:user`.
-      # @param content [String]
-      #   The textual content of the message.
+      # @param content [String, nil]
+      #   The textual content of the message. Cannot be combined with a block.
+      # @yield [builder] A block for building multi-modal content parts.
+      # @yieldparam builder [Datadog::AIGuard::Evaluation::ContentBuilder]
       #
       # @return [Datadog::AIGuard::Evaluation::Message]
       #   A new message instance with the given role and content.
       # @raise [ArgumentError]
-      #   If an invalid role is provided.
+      #   If both content and a block are provided, or if an invalid role is provided.
       # @public_api
-      def message(role:, content:)
-        Evaluation::Message.new(role: role, content: content)
+      def message(role:, content: nil)
+        if block_given?
+          raise ArgumentError, "Cannot pass both content and a block" if content
+
+          builder = Evaluation::ContentBuilder.new
+          yield builder
+          Evaluation::Message.new(role: role, content: builder.parts)
+        else
+          Evaluation::Message.new(role: role, content: content)
+        end
       end
 
       # Builds an assistant message representing a tool call initiated by the model.
@@ -153,3 +177,5 @@ module Datadog
     end
   end
 end
+
+require_relative "ai_guard/autoload"

data/lib/datadog/appsec/actions_handler/serializable_backtrace.rb

@@ -3,7 +3,7 @@
 module Datadog
   module AppSec
     module ActionsHandler
-      # This module serves encapsulates MessagePack serialization for caller locations.
+      # Encapsulates serialisation of caller locations.
       #
       # It serializes part of the stack:
       # up to 32 frames (configurable)
@@ -21,50 +21,21 @@ module Datadog
         def to_msgpack(packer = nil)
           # JRuby doesn't pass the packer
           packer ||= MessagePack::Packer.new
-
-          packer.write_map_header(3)
-
-          packer.write('id')
-          packer.write(@stack_id.encode('UTF-8'))
-
-          packer.write('language')
-          packer.write('ruby'.encode('UTF-8'))
-
-          serializable_locations_map = build_serializable_locations_map
-
-          packer.write('frames')
-          packer.write_array_header(serializable_locations_map.size)
-
-          serializable_locations_map.each do |frame_id, location|
-            packer.write_map_header(6)
-
-            packer.write('id')
-            packer.write(frame_id)
-
-            packer.write('text')
-            packer.write(location.to_s.encode('UTF-8'))
-
-            packer.write('file')
-            packer.write(location.path&.encode('UTF-8'))
-
-            packer.write('line')
-            packer.write(location.lineno)
-
-            class_name, function_name = location.label&.match(CLASS_AND_FUNCTION_NAME_REGEX)&.captures
-
-            packer.write('class_name')
-            packer.write(class_name&.encode('UTF-8'))
-
-            packer.write('function')
-            packer.write(function_name&.encode('UTF-8'))
-          end
-
+          packer.write(to_h)
           packer
         end
 
+        def to_h
+          {
+            "id" => @stack_id.encode("UTF-8"),
+            "language" => "ruby".encode("UTF-8"),
+            "frames" => build_serializable_frames,
+          }
+        end
+
         private
 
-        def build_serializable_locations_map
+        def build_serializable_frames
           max_depth = Datadog.configuration.appsec.stack_trace.max_depth
           top_percent = Datadog.configuration.appsec.stack_trace.top_percentage
 
@@ -72,15 +43,23 @@ module Datadog
           drop_until_idx = @locations.size - (max_depth - drop_from_idx)
 
           frame_idx = -1
-          @locations.each_with_object({}) do |location, map|
+          @locations.each_with_object([]) do |location, frames|
             # we are dropping frames from library code without increasing frame index
-            next if location.path&.include?('lib/datadog')
+            next if location.path&.include?("lib/datadog")
 
             frame_idx += 1
 
             next if max_depth != 0 && frame_idx >= drop_from_idx && frame_idx < drop_until_idx
 
-            map[frame_idx] = location
+            class_name, function_name = location.label&.match(CLASS_AND_FUNCTION_NAME_REGEX)&.captures
+            frames << {
+              "id" => frame_idx,
+              "text" => location.to_s.encode("UTF-8"),
+              "file" => location.path&.encode("UTF-8"),
+              "line" => location.lineno,
+              "class_name" => class_name&.encode("UTF-8"),
+              "function" => function_name&.encode("UTF-8"),
+            }
           end
         end
       end

data/lib/datadog/appsec/actions_handler.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative 'actions_handler/serializable_backtrace'
+require_relative "actions_handler/serializable_backtrace"
 
 module Datadog
   module AppSec
@@ -11,10 +11,10 @@ module Datadog
       def handle(actions_hash)
         # handle actions according their precedence
         # stack and schema generation should be done before we throw an interrupt signal
-        generate_stack(actions_hash['generate_stack']) if actions_hash.key?('generate_stack')
-        generate_schema(actions_hash['generate_schema']) if actions_hash.key?('generate_schema')
-        interrupt_execution(actions_hash['redirect_request']) if actions_hash.key?('redirect_request')
-        interrupt_execution(actions_hash['block_request']) if actions_hash.key?('block_request')
+        generate_stack(actions_hash["generate_stack"]) if actions_hash.key?("generate_stack")
+        generate_schema(actions_hash["generate_schema"]) if actions_hash.key?("generate_schema")
+        interrupt_execution(actions_hash["redirect_request"]) if actions_hash.key?("redirect_request")
+        interrupt_execution(actions_hash["block_request"]) if actions_hash.key?("block_request")
       end
 
       def interrupt_execution(action_params)
@@ -24,7 +24,7 @@ module Datadog
       def generate_stack(action_params)
         return unless Datadog.configuration.appsec.stack_trace.enabled
 
-        stack_id = action_params['stack_id']
+        stack_id = action_params["stack_id"]
         return unless stack_id
 
         active_span = AppSec.active_context&.span

data/lib/datadog/appsec/anonymizer.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'digest/sha2'
+require "digest/sha2"
 
 module Datadog
   module AppSec