datadog 2.32.0 → 2.33.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 04ff51661b0a7d494e4beaf4d142c723bce97e2c2a14ad708e1bfd12e780129a
-  data.tar.gz: 0a4a3f4ccc112c3c805a035072270781e360f68ae1f3464356abde84b97944ea
+  metadata.gz: ff1c01535f382bb6e3de5b529bf8ec30a6deda1d1c4aba3450bb21dbd80407bf
+  data.tar.gz: 966a01b956c9df2cd8d451df9175b66e0dbae3b7b3da7732ea998c7adca70442
 SHA512:
-  metadata.gz: c664b9d447c8645e4a202f04e5b2f690b85f15b62689f3312a8d5d3ccb2a08ec289a6d07fa414eee8e5dd92a3bea9f040b758fbfed97ab602bdcf4a989e47d6b
-  data.tar.gz: 92ce248e3e351d774b043921c18cd73f3ba064eed1203246a81ef1de9500cb4c6220fc7df81187b5b49bd4062fa87228bf8be8817c599315da8853752740ab96
+  metadata.gz: c2d4111c1b0526122be94729577ba54d1e85fb08f2cd1ae74495d9de5bd805cb883b39c872187a5a10f2d4f67a3153d6bce2dda8192e75f4b0e808c9ab9b6a79
+  data.tar.gz: 413cb6dc25ca7b8a79e9e85b152bfd4eed649f72ca5275a46203e30c12f893914337cbd8e5136767588b67f81eadefdc839a5a4aafebbdf5cc17e7d089798886

data/ext/datadog_profiling_native_extension/clock_id.h

@@ -4,10 +4,18 @@
 #include <time.h>
 #include <ruby.h>
 
+#ifdef __APPLE__
+  #include <mach/mach.h>
+  // On macOS, we use a mach_port_t to identify threads for CPU time queries
+  typedef mach_port_t cpu_time_id_t;
+#else
+  typedef clockid_t cpu_time_id_t;
+#endif
+
 // Contains the operating-system specific identifier needed to fetch CPU-time, and a flag to indicate if we failed to fetch it
 typedef struct {
   bool valid;
-  clockid_t clock_id;
+  cpu_time_id_t clock_id;
 } thread_cpu_time_id;
 
 // Contains the current cpu time, and a flag to indicate if we failed to fetch it

data/ext/datadog_profiling_native_extension/clock_id_from_mach.c

@@ -0,0 +1,73 @@
+#include "extconf.h"
+
+// This file is only compiled on macOS where Mach thread APIs are available;
+// Otherwise we compile clock_id_from_pthread.c (Linux)
+#ifdef HAVE_MACH_THREAD_INFO
+
+#include <pthread.h>
+#include <time.h>
+#include <mach/mach.h>
+#include <mach/thread_info.h>
+
+#include "clock_id.h"
+#include "helpers.h"
+#include "private_vm_api_access.h"
+#include "ruby_helpers.h"
+#include "time_helpers.h"
+
+// Validate that our home-cooked pthread_id_for() matches pthread_self() for the current thread
+void self_test_clock_id(void) {
+  rb_nativethread_id_t expected_pthread_id = pthread_self();
+  rb_nativethread_id_t actual_pthread_id = pthread_id_for(rb_thread_current());
+
+  if (expected_pthread_id != actual_pthread_id) raise_error(rb_eRuntimeError, "pthread_id_for() self-test failed");
+
+  // Also validate that we can get a valid mach thread port for the current thread
+  mach_port_t mach_thread = pthread_mach_thread_np(expected_pthread_id);
+  if (mach_thread == MACH_PORT_NULL) raise_error(rb_eRuntimeError, "pthread_mach_thread_np() self-test failed");
+}
+
+// Safety: This function is assumed never to raise exceptions by callers
+thread_cpu_time_id thread_cpu_time_id_for(VALUE thread) {
+  rb_nativethread_id_t thread_id = pthread_id_for(thread);
+
+  if (thread_id == 0) return (thread_cpu_time_id) {.valid = false};
+
+  mach_port_t mach_thread = pthread_mach_thread_np(thread_id);
+
+  if (mach_thread == MACH_PORT_NULL) {
+    return (thread_cpu_time_id) {.valid = false};
+  }
+
+  return (thread_cpu_time_id) {.valid = true, .clock_id = mach_thread};
+}
+
+thread_cpu_time thread_cpu_time_for(thread_cpu_time_id time_id) {
+  thread_cpu_time error = (thread_cpu_time) {.valid = false};
+
+  if (!time_id.valid) return error;
+
+  // Fast path: clock_gettime(CLOCK_THREAD_CPUTIME_ID) is ~5x cheaper than thread_info()
+  // and gives sub-microsecond precision (vs microsecond), but only measures the calling
+  // thread on macOS (there is no pthread_getcpuclockid() equivalent).
+  if (time_id.clock_id == pthread_mach_thread_np(pthread_self())) {
+    struct timespec ts;
+    if (clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts) == 0) {
+      return (thread_cpu_time) {.valid = true, .result_ns = SECONDS_AS_NS(ts.tv_sec) + ts.tv_nsec};
+    }
+    // Fall through to thread_info on the unlikely failure case.
+  }
+
+  struct thread_basic_info info;
+  mach_msg_type_number_t count = THREAD_BASIC_INFO_COUNT;
+  kern_return_t kr = thread_info(time_id.clock_id, THREAD_BASIC_INFO, (thread_info_t)&info, &count);
+
+  if (kr != KERN_SUCCESS) return error;
+
+  long user_ns = SECONDS_AS_NS(info.user_time.seconds) + MICROS_AS_NS(info.user_time.microseconds);
+  long system_ns = SECONDS_AS_NS(info.system_time.seconds) + MICROS_AS_NS(info.system_time.microseconds);
+
+  return (thread_cpu_time) {.valid = true, .result_ns = user_ns + system_ns};
+}
+
+#endif

data/ext/datadog_profiling_native_extension/clock_id_from_pthread.c

@@ -1,7 +1,7 @@
 #include "extconf.h"
 
 // This file is only compiled on systems where pthread_getcpuclockid() is available;
-// Otherwise we compile clock_id_noop.c
+// Otherwise we compile clock_id_from_mach.c (macOS)
 #ifdef HAVE_PTHREAD_GETCPUCLOCKID
 
 #include <pthread.h>

data/ext/datadog_profiling_native_extension/collectors_thread_context.c

@@ -1201,7 +1201,11 @@ static int per_thread_context_as_ruby_hash(st_data_t key_thread, st_data_t value
     ID2SYM(rb_intern("thread_id")),                       /* => */ rb_str_new2(thread_context->thread_id),
     ID2SYM(rb_intern("thread_invoke_location")),          /* => */ rb_str_new2(thread_context->thread_invoke_location),
     ID2SYM(rb_intern("thread_cpu_time_id_valid?")),       /* => */ thread_context->thread_cpu_time_id.valid ? Qtrue : Qfalse,
-    ID2SYM(rb_intern("thread_cpu_time_id")),              /* => */ CLOCKID2NUM(thread_context->thread_cpu_time_id.clock_id),
+    #ifdef __APPLE__
+      ID2SYM(rb_intern("thread_cpu_time_id")),              /* => */ ULL2NUM(thread_context->thread_cpu_time_id.clock_id),
+    #else
+      ID2SYM(rb_intern("thread_cpu_time_id")),              /* => */ CLOCKID2NUM(thread_context->thread_cpu_time_id.clock_id),
+    #endif
     ID2SYM(rb_intern("cpu_time_at_previous_sample_ns")),  /* => */ LONG2NUM(thread_context->cpu_time_at_previous_sample_ns),
     ID2SYM(rb_intern("wall_time_at_previous_sample_ns")), /* => */ LONG2NUM(thread_context->wall_time_at_previous_sample_ns),
 

data/ext/datadog_profiling_native_extension/extconf.rb

@@ -129,6 +129,9 @@ if RUBY_PLATFORM.include?("linux")
 
   # Not available on macOS
   $defs << "-DHAVE_CLOCK_MONOTONIC_COARSE"
+elsif RUBY_PLATFORM.include?("darwin")
+  # On macOS, we use Mach thread APIs to get per-thread CPU time
+  $defs << "-DHAVE_MACH_THREAD_INFO"
 end
 
 have_func "malloc_stats"

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -416,7 +416,6 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   if (options == Qnil) options = rb_hash_new();
 
   VALUE recorder_instance = rb_hash_fetch(options, ID2SYM(rb_intern("self_instance")));
-  VALUE cpu_time_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("cpu_time_enabled")));
   VALUE alloc_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("alloc_samples_enabled")));
   VALUE heap_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_samples_enabled")));
   VALUE heap_size_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_size_enabled")));
@@ -424,7 +423,6 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   VALUE timeline_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("timeline_enabled")));
   VALUE heap_clean_after_gc_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_clean_after_gc_enabled")));
 
-  ENFORCE_BOOLEAN(cpu_time_enabled);
   ENFORCE_BOOLEAN(alloc_samples_enabled);
   ENFORCE_BOOLEAN(heap_samples_enabled);
   ENFORCE_BOOLEAN(heap_size_enabled);
@@ -440,7 +438,6 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   heap_recorder_set_sample_rate(state->heap_recorder, NUM2INT(heap_sample_every));
 
   uint8_t requested_values_count = ALL_VALUE_TYPES_COUNT -
-    (cpu_time_enabled == Qtrue ? 0 : 1) -
     (alloc_samples_enabled == Qtrue? 0 : 2) -
     (heap_samples_enabled == Qtrue ? 0 : 1) -
     (heap_size_enabled == Qtrue ? 0 : 1) -
@@ -466,12 +463,9 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_WALL_TIME;
   state->position_for[WALL_TIME_VALUE_ID] = next_enabled_pos++;
 
-  if (cpu_time_enabled == Qtrue) {
-    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
-    state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[CPU_TIME_VALUE_ID] = next_disabled_pos++;
-  }
+  // CPU_TIME is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
+  state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
 
   if (alloc_samples_enabled == Qtrue) {
     enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES;

data/ext/datadog_profiling_native_extension/time_helpers.h

@@ -9,6 +9,7 @@
 
 #define SECONDS_AS_NS(value) (value * 1000 * 1000 * 1000L)
 #define MILLIS_AS_NS(value) (value * 1000 * 1000L)
+#define MICROS_AS_NS(value) (value * 1000L)
 
 typedef enum { RAISE_ON_FAILURE, DO_NOT_RAISE_ON_FAILURE } raise_on_failure_setting;
 

data/ext/libdatadog_api/crashtracker.c

@@ -72,6 +72,8 @@ static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUS
     .endpoint = {.url = char_slice_from_ruby_string(agent_base_url)},
     .resolve_frames = DDOG_CRASHT_STACKTRACE_COLLECTION_ENABLED_WITH_SYMBOLS_IN_RECEIVER,
     .timeout_ms = FIX2INT(upload_timeout_seconds) * 1000,
+    .collect_all_threads = true,
+    .max_threads = 128,
   };
 
   ddog_crasht_Metadata metadata = {

data/ext/libdatadog_extconf_helpers.rb

@@ -10,7 +10,7 @@ module Datadog
   module LibdatadogExtconfHelpers
     # Used to make sure the correct gem version gets loaded, as extconf.rb does not get run with "bundle exec" and thus
     # may see multiple libdatadog versions. See https://github.com/DataDog/dd-trace-rb/pull/2531 for the horror story.
-    LIBDATADOG_VERSION = '~> 30.0.0.1.0'
+    LIBDATADOG_VERSION = '~> 33.0.0.1.0'
 
     # Used as an workaround for a limitation with how dynamic linking works in environments where the datadog gem and
     # libdatadog are moved after the extension gets compiled.

data/lib/datadog/ai_guard/autoload.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+if %w[1 true].include?((Datadog::DATADOG_ENV["DD_AI_GUARD_ENABLED"] || "").downcase)
+  begin
+    require_relative "contrib/auto_instrument"
+    Datadog::AIGuard::Contrib::AutoInstrument.patch_all
+  rescue => e
+    Kernel.warn("[datadog] AI Guard failed to auto-instrument. error: #{e.class}: #{e.message}")
+  end
+end

data/lib/datadog/ai_guard/component.rb

@@ -15,7 +15,7 @@ module Datadog
   module AIGuard
     # Component for API Guard product
     class Component
-      attr_reader :api_client, :logger
+      attr_reader :api_client, :logger, :telemetry
 
       def self.build(settings, logger:, telemetry:)
         return unless settings.respond_to?(:ai_guard) && settings.ai_guard.enabled

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

@@ -16,6 +16,7 @@ module Datadog
               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

@@ -11,6 +11,7 @@ module Datadog
       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/contrib/aws_lambda/gateway/watcher.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require_relative '../waf_addresses'
+require_relative '../../../event'
+require_relative '../../../trace_keeper'
+require_relative '../../../security_event'
+require_relative '../../../instrumentation/gateway'
+
+module Datadog
+  module AppSec
+    module Contrib
+      module AwsLambda
+        module Gateway
+          module Watcher
+            class << self
+              def watch
+                gateway = Instrumentation.gateway
+
+                watch_request(gateway)
+                watch_response(gateway)
+              end
+
+              def watch_request(gateway = Instrumentation.gateway)
+                gateway.watch('aws_lambda.request.start') do |stack, payload|
+                  context = payload.context
+                  next stack.call(payload) unless context
+
+                  persistent_data = WAFAddresses.from_request(payload.data)
+                  result = context.run_waf(persistent_data, {}, Datadog.configuration.appsec.waf_timeout)
+
+                  if result.match? || !result.attributes.empty?
+                    context.events.push(
+                      AppSec::SecurityEvent.new(result, trace: context.trace, span: context.span)
+                    )
+                  end
+
+                  if result.match?
+                    AppSec::Event.tag(context, result)
+                    TraceKeeper.keep!(context.trace) if result.keep?
+                    AppSec::ActionsHandler.handle(result.actions)
+                  end
+
+                  stack.call(payload)
+                end
+              end
+
+              def watch_response(gateway = Instrumentation.gateway)
+                gateway.watch('aws_lambda.response.start') do |stack, payload|
+                  context = payload.context
+                  next stack.call(payload) unless context
+
+                  persistent_data = WAFAddresses.from_response(payload.data)
+                  result = context.run_waf(persistent_data, {}, Datadog.configuration.appsec.waf_timeout)
+
+                  if result.match?
+                    AppSec::Event.tag(context, result)
+                    TraceKeeper.keep!(context.trace) if result.keep?
+
+                    context.events.push(
+                      AppSec::SecurityEvent.new(result, trace: context.trace, span: context.span)
+                    )
+
+                    AppSec::ActionsHandler.handle(result.actions)
+                  end
+
+                  stack.call(payload)
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end