datadog 2.31.0 → 2.33.0

data/lib/datadog/ai_guard/configuration.rb

@@ -1,11 +1,114 @@
 # frozen_string_literal: true
 
+require "uri"
+require_relative "configuration/ext"
+
 module Datadog
   module AIGuard
     # Configuration module for AI Guard
     module Configuration
+      # AI Guard specific settings
+      module Settings
+        def self.extended(base)
+          base = base.singleton_class unless base.is_a?(Class)
+          add_settings!(base)
+        end
+
+        def self.add_settings!(base)
+          base.class_eval do
+            # AI Guard specific configurations.
+            # @public_api
+            #
+            # Steep does not update `self` for this `class_eval` block.
+            # @type self: Datadog::Core::Configuration::Base::_DslContext
+            settings :ai_guard do
+              # Enable AI Guard.
+              #
+              # You can use this option to skip calls to AI Guard API without having to remove library as a whole.
+              #
+              # @default `DD_AI_GUARD_ENABLED`, otherwise `false`
+              # @return [Boolean]
+              option :enabled do |o|
+                o.type :bool
+                o.env Ext::ENV_AI_GUARD_ENABLED
+                o.default false
+              end
+
+              define_method(:instrument) do |integration_name|
+                return unless enabled # steep:ignore
+
+                if (registered_integration = Datadog::AIGuard::Contrib::Integration.registry[integration_name])
+                  klass = registered_integration.klass
+                  if klass.loaded? && klass.compatible?
+                    instance = klass.new
+                    instance.patcher.patch unless instance.patcher.patched?
+                  end
+                end
+              end
+
+              # AI Guard API endpoint path.
+              #
+              # @default `DD_AI_GUARD_ENDPOINT`, otherwise `nil`
+              # @return [String, nil]
+              option :endpoint do |o|
+                o.type :string, nilable: true
+                o.env Ext::ENV_AI_GUARD_ENDPOINT
+
+                o.setter do |value|
+                  next unless value
+
+                  uri = URI(value.to_s)
+                  raise ArgumentError, "Please provide an absolute URI that includes a protocol" unless uri.absolute?
+
+                  uri.to_s.delete_suffix("/")
+                end
+              end
+
+              # Datadog Application key.
+              #
+              # @default `DD_APP_KEY` environment variable, otherwise `nil`
+              # @return [String, nil]
+              option :app_key do |o|
+                o.type :string, nilable: true
+                o.env Ext::ENV_APP_KEY
+              end
+
+              # Request timeout in milliseconds.
+              #
+              # @default `DD_AI_GUARD_TIMEOUT`, otherwise 10 000 ms
+              # @return [Integer]
+              option :timeout_ms do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_TIMEOUT
+                o.default 10_000
+              end
+
+              # Maximum content size in bytes.
+              # Content that exceeds the maximum allowed size is truncated before
+              # being stored in the current span context.
+              #
+              # @default `DD_AI_GUARD_MAX_CONTENT_SIZE`, otherwise 524 228 bytes
+              # @return [Integer]
+              option :max_content_size_bytes do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_MAX_CONTENT_SIZE
+                o.default 512 * 1024
+              end
+
+              # Maximum number of messages.
+              # Older messages are omitted once the message limit is reached.
+              #
+              # @default `DD_AI_GUARD_MAX_MESSAGES_LENGTH`, otherwise 16 messages
+              # @return [Integer]
+              option :max_messages_length do |o|
+                o.type :int
+                o.env Ext::ENV_AI_GUARD_MAX_MESSAGES_LENGTH
+                o.default 16
+              end
+            end
+          end
+        end
+      end
     end
   end
 end
-
-require_relative "configuration/settings"

data/lib/datadog/ai_guard/contrib/auto_instrument.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      # Auto-instrumentation for AI Guard integrations
+      module AutoInstrument
+        def self.patch_all
+          integrations = []
+
+          Datadog::AIGuard::Contrib::Integration.registry.each_value do |integration|
+            next unless integration.klass.auto_instrument?
+
+            integrations << integration.name
+          end
+
+          integrations.each do |integration_name|
+            Datadog.configuration.ai_guard.instrument(integration_name)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/contrib/rack/integration.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "patcher"
+require_relative "request_middleware"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # Rack integration for AI Guard
+        class Integration
+          include Datadog::AIGuard::Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new("1.1.0")
+
+          register_as :rack, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs["rack"]&.version
+          end
+
+          def self.loaded?
+            !defined?(::Rack).nil?
+          end
+
+          def self.compatible?
+            super && !!(version&.>= MINIMUM_VERSION)
+          end
+
+          def self.auto_instrument?
+            false
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # Patcher for Rack integration
+        module Patcher
+          module_function
+
+          def patched?
+            !!Patcher.instance_variable_get(:@patched)
+          end
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            Patcher.instance_variable_set(:@patched, true)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/ai_guard/contrib/rack/request_middleware.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require_relative "../../../tracing/client_ip"
+require_relative "../../../tracing/contrib/rack/header_collection"
+require_relative "../../../tracing/metadata/ext"
+require_relative "../../ext"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rack
+        # AI Guard Rack middleware.
+        #
+        # Inserted into the middleware stack right after
+        # Datadog::Tracing::Contrib::Rack::TraceMiddleware (i.e. nested inside
+        # it). This ordering matters: on the way out of the request, AI Guard's
+        # `ensure` block unwinds *before* Tracing's ensure, while Tracing's
+        # request span is still live. We need that, because Tracing's ensure
+        # calls `request_span.finish`, which builds a frozen `Span` snapshot of
+        # the meta hash — any `set_tag` call after that point mutates the
+        # `SpanOperation` but never reaches the exported `Span`.
+        #
+        # So while the span is still active, we tag `http.client_ip` and
+        # `network.client.ip` on it — but only when an AI Guard span was
+        # actually recorded during the request.
+        class RequestMiddleware
+          NETWORK_CLIENT_IP_TAG = "network.client.ip"
+
+          def initialize(app, opt = {})
+            @app = app
+          end
+
+          def call(env)
+            @app.call(env)
+          ensure
+            tag_client_ip(env) if consume_ai_guard_executed_flag
+          end
+
+          private
+
+          # AI Guard's evaluation flow sets `ai_guard.executed` on the trace
+          # whenever an AI Guard span is created during the request. We read
+          # it here to know whether to tag client IP, then clear it so the
+          # internal flag does not propagate to the exported trace.
+          #
+          # `Tracing.active_trace` is publicly typed as `TraceSegment?` but at
+          # runtime returns a `TraceOperation`, which exposes `get_tag` and
+          # `clear_tag`. Pre-existing sig mismatch — hence the steep:ignore.
+          # steep:ignore:start
+          def consume_ai_guard_executed_flag
+            trace = Datadog::Tracing.active_trace
+            return false unless trace
+            return false unless trace.get_tag(Datadog::AIGuard::Ext::SERVICE_ENTRY_EXECUTED_TAG) == "1"
+
+            trace.clear_tag(Datadog::AIGuard::Ext::SERVICE_ENTRY_EXECUTED_TAG)
+            true
+          end
+          # steep:ignore:end
+
+          def tag_client_ip(env)
+            span = Datadog::Tracing.active_span
+            return unless span
+
+            if span.get_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP).nil?
+              headers = Datadog::Tracing::Contrib::Rack::Header::RequestHeaderCollection.new(env)
+              Datadog::Tracing::ClientIp.set_client_ip_tag!(
+                span,
+                headers: headers,
+                remote_ip: env["REMOTE_ADDR"]
+              )
+            end
+
+            if env["REMOTE_ADDR"] && span.get_tag(NETWORK_CLIENT_IP_TAG).nil?
+              span.set_tag(NETWORK_CLIENT_IP_TAG, env["REMOTE_ADDR"])
+            end
+          rescue => e
+            Datadog::AIGuard.telemetry&.report(e, description: "AI Guard: failed to tag client IP on root span")
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "patcher"
+
+module Datadog
+  module AIGuard
+    module Contrib
+      module Rails
+        # Rails integration for AI Guard
+        class Integration
+          include Datadog::AIGuard::Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new("4")
+
+          register_as :rails, auto_patch: false
+
+          def self.version
+            Gem.loaded_specs["railties"]&.version
+          end
+
+          def self.loaded?
+            !defined?(::Rails).nil?
+          end
+
+          def self.compatible?
+            super && !!(version&.>= MINIMUM_VERSION)
+          end
+
+          def self.auto_instrument?
+            true
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

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/evaluation.rb

@@ -15,6 +15,8 @@ module Datadog
               Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER,
               Tracing::Sampling::Ext::Decision::AI_GUARD
             )
+            trace.set_tag(Ext::EVENT_TAG, true)
+            trace.set_tag(Ext::SERVICE_ENTRY_EXECUTED_TAG, "1")
 
             if (last_message = messages.last)
               if last_message.tool_call

data/lib/datadog/ai_guard/ext.rb

@@ -10,6 +10,8 @@ module Datadog
       ACTION_TAG = "ai_guard.action"
       REASON_TAG = "ai_guard.reason"
       BLOCKED_TAG = "ai_guard.blocked"
+      EVENT_TAG = "ai_guard.event"
+      SERVICE_ENTRY_EXECUTED_TAG = "_dd.ai_guard.executed"
       METASTRUCT_TAG = "ai_guard"
     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
@@ -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:
@@ -171,3 +177,5 @@ module Datadog
     end
   end
 end
+
+require_relative "ai_guard/autoload"

data/lib/datadog/appsec/autoload.rb

@@ -7,7 +7,7 @@ if %w[1 true].include?((Datadog::DATADOG_ENV['DD_APPSEC_ENABLED'] || '').downcas
   rescue => e
     Kernel.warn(
       '[datadog] AppSec failed to instrument. No security check will be performed. error: ' \
-      " #{e.class}: #{e}"
+      " #{e.class}: #{e.message}"
     )
   end
 end

data/lib/datadog/appsec/component.rb

@@ -47,7 +47,7 @@ module Datadog
           security_engine = SecurityEngine::Engine.new(appsec_settings: settings.appsec, telemetry: telemetry)
           new(security_engine: security_engine)
         rescue => e
-          Datadog.logger.warn("AppSec is disabled: #{e.class}: #{e}; there may be additional logged errors above")
+          Datadog.logger.warn("AppSec is disabled: #{e.class}: #{e.message}; there may be additional logged errors above")
 
           # Not reporting to telemetry here because some of the rescued exceptions
           # have already been reported by the code that raised them