dd-trace 5.98.0 → 5.99.0

package/packages/datadog-plugin-cucumber/src/index.js

@@ -6,7 +6,7 @@ const realSetTimeout = setTimeout
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
-const { getEnvironmentVariable, getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
+const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 
 const {
   addIntelligentTestRunnerSpanTags,
@@ -117,7 +117,7 @@ class CucumberPlugin extends CiPlugin {
       finishAllTraceSpans(this.testSessionSpan)
       this.telemetry.count(TELEMETRY_TEST_SESSION, {
         provider: this.ciProviderName,
-        autoInjected: !!getValueFromEnvSources('DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER'),
+        autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
 
       this.libraryConfig = null

package/packages/datadog-plugin-cypress/src/cypress-plugin.js

@@ -100,7 +100,7 @@ const {
 } = require('../../dd-trace/src/plugins/util/env')
 const { DD_MAJOR } = require('../../../version')
 const {
-  resolveOriginalSourcePosition,
+  resolveOriginalSourceFile,
   resolveSourceLineForTest,
   shouldTrustInvocationDetailsLine,
 } = require('./source-map-utils')
@@ -518,8 +518,7 @@ class CypressPlugin {
     this.ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
 
     if (testSuiteAbsolutePath) {
-      const resolvedSuitePosition = resolveOriginalSourcePosition(testSuiteAbsolutePath, 1)
-      const resolvedSuiteAbsolutePath = resolvedSuitePosition ? resolvedSuitePosition.sourceFile : testSuiteAbsolutePath
+      const resolvedSuiteAbsolutePath = resolveOriginalSourceFile(testSuiteAbsolutePath) || testSuiteAbsolutePath
       const testSourceFile = getTestSuitePath(resolvedSuiteAbsolutePath, this.repositoryRoot)
       testSuiteSpanMetadata[TEST_SOURCE_FILE] = testSourceFile
       testSuiteSpanMetadata[TEST_SOURCE_START] = 1
@@ -950,8 +949,9 @@ class CypressPlugin {
         if (this.itrCorrelationId) {
           finishedTest.testSpan.setTag(ITR_CORRELATION_ID, this.itrCorrelationId)
         }
-        const resolvedSpecPosition = spec.absolute ? resolveOriginalSourcePosition(spec.absolute, 1) : null
-        const resolvedSpecAbsolutePath = resolvedSpecPosition ? resolvedSpecPosition.sourceFile : spec.absolute
+        const resolvedSpecAbsolutePath = spec.absolute
+          ? resolveOriginalSourceFile(spec.absolute) || spec.absolute
+          : spec.absolute
         const testSourceFile = resolvedSpecAbsolutePath && this.repositoryRoot
           ? getTestSuitePath(resolvedSpecAbsolutePath, this.repositoryRoot)
           : spec.relative

package/packages/datadog-plugin-cypress/src/source-map-utils.js

@@ -165,6 +165,48 @@ function resolveOriginalSourcePosition (absoluteFilePath, generatedLine) {
   return null
 }
 
+/**
+ * Given a generated file's absolute path, returns the first original source file
+ * referenced by its source map. This is useful for file-level metadata when the
+ * first generated line has no source mapping, such as TypeScript-emitted prologue
+ * statements.
+ * @param {string} absoluteFilePath - Absolute path to the generated file
+ * @returns {string | null}
+ */
+function resolveOriginalSourceFile (absoluteFilePath) {
+  const sourceMap = getCachedSourceMap(absoluteFilePath)
+  if (!sourceMap) return null
+  const { mappings, sources, sourceRoot } = sourceMap
+  if (!mappings || !sources?.length) return null
+
+  const mapDir = path.dirname(absoluteFilePath)
+  const cursor = { pos: 0 }
+  let srcFile = 0
+
+  const lines = mappings.split(';')
+  for (const line of lines) {
+    if (!line) continue
+    cursor.pos = 0
+    while (cursor.pos < line.length) {
+      decodeVLQ(line, cursor) // genCol - not needed
+      if (cursor.pos < line.length && line[cursor.pos] !== ',') {
+        srcFile += decodeVLQ(line, cursor)
+        decodeVLQ(line, cursor) // srcLine - not needed
+        decodeVLQ(line, cursor) // srcCol - not needed
+        if (cursor.pos < line.length && line[cursor.pos] !== ',') {
+          decodeVLQ(line, cursor) // namesIndex - not needed
+        }
+
+        const sourcePath = sources[srcFile]
+        if (!sourcePath) return null
+        return resolveSourcePath(mapDir, sourceRoot, sourcePath)
+      }
+      if (cursor.pos < line.length && line[cursor.pos] === ',') cursor.pos++
+    }
+  }
+  return null
+}
+
 /**
  * Convert a template literal body (the text between backticks, with `${…}` interpolations)
  * into a regex that matches the runtime-evaluated string. Each `${…}` expression is replaced
@@ -294,4 +336,9 @@ function resolveSourceLineForTest (absoluteFilePath, testName, invocationStack)
   return null
 }
 
-module.exports = { resolveOriginalSourcePosition, resolveSourceLineForTest, shouldTrustInvocationDetailsLine }
+module.exports = {
+  resolveOriginalSourceFile,
+  resolveOriginalSourcePosition,
+  resolveSourceLineForTest,
+  shouldTrustInvocationDetailsLine,
+}

package/packages/datadog-plugin-http/src/server.js

@@ -2,10 +2,13 @@
 
 const ServerPlugin = require('../../dd-trace/src/plugins/server')
 const { storage } = require('../../datadog-core')
+const { withRequest } = require('../../dd-trace/src/appsec/store')
 const web = require('../../dd-trace/src/plugins/util/web')
 const { incomingHttpRequestStart, incomingHttpRequestEnd } = require('../../dd-trace/src/appsec/channels')
 const { COMPONENT, SVC_SRC_KEY } = require('../../dd-trace/src/constants')
 
+const legacyStorage = storage('legacy')
+
 class HttpServerPlugin extends ServerPlugin {
   static id = 'http'
 
@@ -17,7 +20,7 @@ class HttpServerPlugin extends ServerPlugin {
   }
 
   start ({ req, res, abortController }) {
-    let store = storage('legacy').getStore()
+    let store = legacyStorage.getStore()
     const { name: schemaServiceName, source: schemaServiceSource } = this.serviceName()
     const service = this.config.service || schemaServiceName
     const serviceSource = (this.config.service && service !== this.tracer._service)
@@ -40,16 +43,13 @@ class HttpServerPlugin extends ServerPlugin {
     span._integrationName = this.constructor.id
 
     const context = web.getContext(req)
+    context.parentStore = store
 
-    if (context) {
-      context.parentStore = store
-    }
-
-    // Only AppSec needs the request scope to be active for any async work that
-    // may be scheduled after the synchronous `request` event returns (e.g.
-    // Fastify).
-    if (incomingHttpRequestStart.hasSubscribers) {
-      store = { ...store, req, res }
+    const appsecActive = incomingHttpRequestStart.hasSubscribers
+    if (appsecActive) {
+      // AppSec, IAST, and AI Guard need req on the store so downstream
+      // subscribers can access them from the async context.
+      store = withRequest(store, req)
     }
 
     this.enter(span, store)
@@ -59,7 +59,7 @@ class HttpServerPlugin extends ServerPlugin {
       context.instrumented = true
     }
 
-    if (incomingHttpRequestStart.hasSubscribers) {
+    if (appsecActive) {
       incomingHttpRequestStart.publish({ req, res, abortController }) // TODO: no need to make a new object here
     }
   }

package/packages/datadog-plugin-jest/src/index.js

@@ -5,7 +5,7 @@ const realSetTimeout = setTimeout
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
-const { getEnvironmentVariable, getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
+const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 const { appClosing: appClosingTelemetry } = require('../../dd-trace/src/telemetry')
 
 const {
@@ -166,7 +166,7 @@ class JestPlugin extends CiPlugin {
 
       this.telemetry.count(TELEMETRY_TEST_SESSION, {
         provider: this.ciProviderName,
-        autoInjected: !!getValueFromEnvSources('DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER'),
+        autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
 
       appClosingTelemetry()

package/packages/datadog-plugin-mocha/src/index.js

@@ -5,7 +5,6 @@ const realDateNow = Date.now.bind(Date)
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
-const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 
 const {
   TEST_STATUS,
@@ -406,7 +405,7 @@ class MochaPlugin extends CiPlugin {
         finishAllTraceSpans(this.testSessionSpan)
         this.telemetry.count(TELEMETRY_TEST_SESSION, {
           provider: this.ciProviderName,
-          autoInjected: !!getValueFromEnvSources('DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER'),
+          autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
         })
       }
       this.libraryConfig = null

package/packages/datadog-plugin-modelcontextprotocol-sdk/src/index.js

@@ -0,0 +1,24 @@
+'use strict'
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+const mcpLLMObsPlugins = require('../../dd-trace/src/llmobs/plugins/modelcontextprotocol-sdk')
+const tracingPlugins = require('./tracing')
+
+const plugins = {}
+
+// CRITICAL: LLMObs plugins MUST come first
+for (const Plugin of mcpLLMObsPlugins) {
+  plugins[Plugin.id] = Plugin
+}
+
+// Tracing plugins second
+for (const Plugin of tracingPlugins) {
+  plugins[Plugin.id] = Plugin
+}
+
+class ModelcontextprotocolSdkPlugin extends CompositePlugin {
+  static id = 'modelcontextprotocol-sdk'
+  static plugins = plugins
+}
+
+module.exports = ModelcontextprotocolSdkPlugin

package/packages/datadog-plugin-modelcontextprotocol-sdk/src/tracing.js

@@ -0,0 +1,55 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+
+class McpToolCallPlugin extends TracingPlugin {
+  static id = 'modelcontextprotocol_client'
+  static prefix = 'tracing:orchestrion:@modelcontextprotocol/sdk:Client_callTool'
+
+  bindStart (ctx) {
+    const params = ctx.arguments?.[0]
+    const toolName = params?.name
+
+    this.startSpan('mcp.client.tool.call', {
+      resource: toolName,
+      type: 'mcp',
+      kind: 'client',
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    const result = ctx.result
+    if (result?.isError) {
+      const span = ctx.currentStore?.span
+      const errorText = result.content?.find?.(c => c.type === 'text')?.text || 'Tool call returned isError: true'
+      span?.setTag('error', new Error(errorText))
+    }
+    super.finish(ctx)
+  }
+}
+
+class McpListToolsPlugin extends TracingPlugin {
+  static id = 'modelcontextprotocol_list_tools'
+  static prefix = 'tracing:orchestrion:@modelcontextprotocol/sdk:Client_listTools'
+
+  bindStart (ctx) {
+    this.startSpan('mcp.tools.list', {
+      resource: 'tools/list',
+      type: 'mcp',
+      kind: 'client',
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    super.finish(ctx)
+  }
+}
+
+module.exports = [
+  McpToolCallPlugin,
+  McpListToolsPlugin,
+]

package/packages/datadog-plugin-mongodb-core/src/index.js

@@ -1,8 +1,6 @@
 'use strict'
 
-const { isTrue } = require('../../dd-trace/src/util')
 const DatabasePlugin = require('../../dd-trace/src/plugins/database')
-const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 
 class MongodbCorePlugin extends DatabasePlugin {
   static id = 'mongodb-core'
@@ -20,11 +18,8 @@ class MongodbCorePlugin extends DatabasePlugin {
   configure (config) {
     super.configure(config)
 
-    const heartbeatFromEnv = getValueFromEnvSources('DD_TRACE_MONGODB_HEARTBEAT_ENABLED')
-
     this.config.heartbeatEnabled = config.heartbeatEnabled ??
-      (heartbeatFromEnv && isTrue(heartbeatFromEnv)) ??
-      true
+      this._tracerConfig.DD_TRACE_MONGODB_HEARTBEAT_ENABLED
   }
 
   bindStart (ctx) {

package/packages/datadog-plugin-playwright/src/index.js

@@ -3,7 +3,6 @@
 const { storage } = require('../../datadog-core')
 const id = require('../../dd-trace/src/id')
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
-const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 
 const {
   finishAllTraceSpans,
@@ -108,7 +107,7 @@ class PlaywrightPlugin extends CiPlugin {
       finishAllTraceSpans(this.testSessionSpan)
       this.telemetry.count(TELEMETRY_TEST_SESSION, {
         provider: this.ciProviderName,
-        autoInjected: !!getValueFromEnvSources('DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER'),
+        autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
       appClosingTelemetry()
       this.tracer._exporter.flush(onDone)
@@ -420,7 +419,7 @@ class PlaywrightPlugin extends CiPlugin {
       span.finish()
 
       finishAllTraceSpans(span)
-      if (getValueFromEnvSources('DD_PLAYWRIGHT_WORKER')) {
+      if (this._tracerConfig.DD_PLAYWRIGHT_WORKER) {
         this.tracer._exporter.flush(onDone)
       }
     })

package/packages/datadog-plugin-vitest/src/index.js

@@ -2,7 +2,6 @@
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
-const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 
 const {
   TEST_STATUS,
@@ -34,6 +33,7 @@ const {
   isModifiedTest,
   TEST_IS_MODIFIED,
   TEST_HAS_DYNAMIC_NAME,
+  TEST_FINAL_STATUS,
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -213,10 +213,13 @@ class VitestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
-    this.addSub('ci:vitest:test:pass', ({ span, task }) => {
+    this.addSub('ci:vitest:test:pass', ({ span, task, finalStatus }) => {
       if (span) {
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', this.getTestTelemetryTags(span))
         span.setTag(TEST_STATUS, 'pass')
+        if (finalStatus) {
+          span.setTag(TEST_FINAL_STATUS, finalStatus)
+        }
         span.finish(this.taskToFinishTime.get(task))
         finishAllTraceSpans(span)
       }
@@ -230,6 +233,7 @@ class VitestPlugin extends CiPlugin {
       promises,
       hasFailedAllRetries,
       attemptToFixFailed,
+      finalStatus,
     }) => {
       if (!span) {
         return
@@ -255,6 +259,9 @@ class VitestPlugin extends CiPlugin {
       if (attemptToFixFailed) {
         span.setTag(TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED, 'false')
       }
+      if (finalStatus) {
+        span.setTag(TEST_FINAL_STATUS, finalStatus)
+      }
       if (duration) {
         span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
       } else {
@@ -273,6 +280,7 @@ class VitestPlugin extends CiPlugin {
           [TEST_SOURCE_FILE]: testSuite,
           [TEST_SOURCE_START]: 1, // we can't get the proper start line in vitest
           [TEST_STATUS]: 'skip',
+          [TEST_FINAL_STATUS]: 'skip',
           ...(isDisabled ? { [TEST_MANAGEMENT_IS_DISABLED]: 'true' } : {}),
           ...(isNew ? { [TEST_IS_NEW]: 'true' } : {}),
         }
@@ -285,12 +293,12 @@ class VitestPlugin extends CiPlugin {
       const { testSuiteAbsolutePath, frameworkVersion } = ctx
 
       // TODO: Handle case where the command is not set
-      this.command = getValueFromEnvSources('DD_CIVISIBILITY_TEST_COMMAND')
+      this.command = this._tracerConfig.DD_CIVISIBILITY_TEST_COMMAND
       this.frameworkVersion = frameworkVersion
       const testSessionSpanContext = this.tracer.extract('text_map', {
         // TODO: Handle case where the session ID or module ID is not set
-        'x-datadog-trace-id': getValueFromEnvSources('DD_CIVISIBILITY_TEST_SESSION_ID'),
-        'x-datadog-parent-id': getValueFromEnvSources('DD_CIVISIBILITY_TEST_MODULE_ID'),
+        'x-datadog-trace-id': this._tracerConfig.DD_CIVISIBILITY_TEST_SESSION_ID,
+        'x-datadog-parent-id': this._tracerConfig.DD_CIVISIBILITY_TEST_MODULE_ID,
       })
 
       const trimmedCommand = DD_MAJOR < 6 ? this.command : 'vitest run'
@@ -415,7 +423,7 @@ class VitestPlugin extends CiPlugin {
       finishAllTraceSpans(this.testSessionSpan)
       this.telemetry.count(TELEMETRY_TEST_SESSION, {
         provider: this.ciProviderName,
-        autoInjected: !!getValueFromEnvSources('DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER'),
+        autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
       this.tracer._exporter.flush(onFinish)
     })

package/packages/datadog-plugin-ws/src/close.js

@@ -57,11 +57,13 @@ class WSClosePlugin extends TracingPlugin {
   }
 
   bindAsyncStart (ctx) {
+    if (!ctx.span) return ctx.parentStore
     if (!ctx.isPeerClose) ctx.span.finish()
     return ctx.parentStore
   }
 
   asyncStart (ctx) {
+    if (!ctx.span) return
     ctx.span.finish()
   }
 

package/packages/datadog-plugin-ws/src/producer.js

@@ -46,11 +46,13 @@ class WSProducerPlugin extends TracingPlugin {
   }
 
   bindAsyncStart (ctx) {
+    if (!ctx.span) return ctx.parentStore
     ctx.span.finish()
     return ctx.parentStore
   }
 
   asyncStart (ctx) {
+    if (!ctx.span) return
     ctx.span.finish()
   }
 

package/packages/datadog-plugin-ws/src/receiver.js

@@ -62,6 +62,7 @@ class WSReceiverPlugin extends TracingPlugin {
   }
 
   asyncStart (ctx) {
+    if (!ctx.span) return
     ctx.span.finish()
   }
 

package/packages/dd-trace/src/aiguard/channels.js

@@ -0,0 +1,8 @@
+'use strict'
+
+const dc = require('dc-polyfill')
+
+module.exports = {
+  aiguardChannel: dc.channel('dd-trace:ai:aiguard'),
+  incomingHttpRequestStart: dc.channel('dd-trace:incomingHttpRequestStart'),
+}

package/packages/dd-trace/src/aiguard/index.js

@@ -1,15 +1,17 @@
 'use strict'
 
-const { channel } = require('dc-polyfill')
 const log = require('../log')
+const { incomingHttpRequestStart, aiguardChannel } = require('./channels')
 const AIGuard = require('./sdk')
 
-const aiguardChannel = channel('dd-trace:ai:aiguard')
-
 let isEnabled = false
 let aiguard
 let block
 
+function onIncomingHttpRequestStart () {
+  // No-op: subscribing ensures the HTTP plugin spreads req onto the store
+}
+
 function enable (tracer, config) {
   if (isEnabled) return
 
@@ -17,6 +19,7 @@ function enable (tracer, config) {
     aiguard = new AIGuard(tracer, config)
     block = config.experimental?.aiguard?.block !== false
 
+    incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
     aiguardChannel.subscribe(onEvaluate)
 
     isEnabled = true
@@ -29,6 +32,7 @@ function enable (tracer, config) {
 function disable () {
   if (!isEnabled) return
 
+  incomingHttpRequestStart.unsubscribe(onIncomingHttpRequestStart)
   aiguardChannel.unsubscribe(onEvaluate)
 
   aiguard = undefined

package/packages/dd-trace/src/aiguard/sdk.js

@@ -1,7 +1,10 @@
 'use strict'
 
 const rfdc = require('../../../../vendor/dist/rfdc')({ proto: false, circles: false })
+const { HTTP_CLIENT_IP, NETWORK_CLIENT_IP } = require('../../../../ext/tags')
+const { getActiveRequest } = require('../appsec/store')
 const log = require('../log')
+const { extractIp } = require('../plugins/util/ip_extractor')
 const telemetryMetrics = require('../telemetry/metrics')
 const tracerVersion = require('../../../../package.json').version
 const { keepTrace } = require('../priority_sampler')
@@ -13,6 +16,7 @@ const {
   AI_GUARD_TARGET_TAG_KEY,
   AI_GUARD_REASON_TAG_KEY,
   AI_GUARD_ACTION_TAG_KEY,
+  AI_GUARD_EVENT_TAG_KEY,
   AI_GUARD_BLOCKED_TAG_KEY,
   AI_GUARD_META_STRUCT_KEY,
   AI_GUARD_TOOL_NAME_TAG_KEY,
@@ -57,6 +61,7 @@ class AIGuard extends NoopAIGuard {
   #maxMessagesLength
   #maxContentSize
   #meta
+  #config
 
   /**
    * @param {import('../tracer')} tracer - Tracer instance
@@ -84,6 +89,7 @@ class AIGuard extends NoopAIGuard {
     this.#maxMessagesLength = config.experimental.aiguard.maxMessagesLength
     this.#maxContentSize = config.experimental.aiguard.maxContentSize
     this.#meta = { service: config.service, env: config.env }
+    this.#config = config
     this.#initialized = true
   }
 
@@ -139,6 +145,42 @@ class AIGuard extends NoopAIGuard {
     return null
   }
 
+  #setRootSpanClientIpTags (rootSpan) {
+    if (!rootSpan) return
+
+    const currentTags = rootSpan.context()._tags
+    const needsHttpClientIp = !Object.hasOwn(currentTags, HTTP_CLIENT_IP)
+    const needsNetworkClientIp = !Object.hasOwn(currentTags, NETWORK_CLIENT_IP)
+
+    if (!needsHttpClientIp && !needsNetworkClientIp) return
+
+    const req = getActiveRequest()
+
+    if (!req) return
+
+    const newTags = {}
+
+    if (needsHttpClientIp) {
+      const clientIp = extractIp(this.#config, req)
+
+      if (clientIp) {
+        newTags[HTTP_CLIENT_IP] = clientIp
+      }
+    }
+
+    if (needsNetworkClientIp) {
+      const networkClientIp = req.socket?.remoteAddress
+
+      if (networkClientIp) {
+        newTags[NETWORK_CLIENT_IP] = networkClientIp
+      }
+    }
+
+    if (Object.keys(newTags).length > 0) {
+      rootSpan.addTags(newTags)
+    }
+  }
+
   evaluate (messages, opts) {
     if (!this.#initialized) {
       return super.evaluate(messages, opts)
@@ -162,9 +204,11 @@ class AIGuard extends NoopAIGuard {
       }
       const rootSpan = span.context()?._trace?.started?.[0]
       if (rootSpan) {
+        this.#setRootSpanClientIpTags(rootSpan)
         // keepTrace must be called before executeRequest so the sampling decision
         // is propagated correctly to outgoing HTTP client calls.
         keepTrace(rootSpan, AI_GUARD)
+        rootSpan.setTag(AI_GUARD_EVENT_TAG_KEY, 'true')
       }
       let response
       try {

package/packages/dd-trace/src/aiguard/tags.js

@@ -7,6 +7,7 @@ module.exports = {
   AI_GUARD_ACTION_TAG_KEY: 'ai_guard.action',
   AI_GUARD_REASON_TAG_KEY: 'ai_guard.reason',
   AI_GUARD_BLOCKED_TAG_KEY: 'ai_guard.blocked',
+  AI_GUARD_EVENT_TAG_KEY: 'ai_guard.event',
   AI_GUARD_META_STRUCT_KEY: 'ai_guard',
 
   AI_GUARD_TELEMETRY_REQUESTS: 'ai_guard.requests',

package/packages/dd-trace/src/appsec/graphql.js

@@ -1,8 +1,8 @@
 'use strict'
 
-const { storage } = require('../../../datadog-core')
 const log = require('../log')
 const web = require('../plugins/util/web')
+const { getActiveRequest } = require('./store')
 const {
   addSpecificEndpoint,
   specificBlockingTypes,
@@ -33,7 +33,7 @@ function disable () {
 }
 
 function onGraphqlStartResolve ({ context, resolverInfo }) {
-  const req = storage('legacy').getStore()?.req
+  const req = getActiveRequest()
 
   if (!req) return
 
@@ -52,7 +52,7 @@ function onGraphqlStartResolve ({ context, resolverInfo }) {
 }
 
 function enterInApolloMiddleware (data) {
-  const req = data?.req || storage('legacy').getStore()?.req
+  const req = data?.req || getActiveRequest()
   if (!req) return
 
   graphqlRequestData.set(req, {
@@ -61,7 +61,7 @@ function enterInApolloMiddleware (data) {
 }
 
 function enterInApolloServerCoreRequest () {
-  const req = storage('legacy').getStore()?.req
+  const req = getActiveRequest()
   if (!req) return
 
   graphqlRequestData.set(req, {
@@ -71,7 +71,7 @@ function enterInApolloServerCoreRequest () {
 }
 
 function enterInApolloRequest () {
-  const req = storage('legacy').getStore()?.req
+  const req = getActiveRequest()
 
   const requestData = graphqlRequestData.get(req)
   if (requestData) {
@@ -83,7 +83,7 @@ function enterInApolloRequest () {
 }
 
 function beforeWriteApolloGraphqlResponse ({ abortController, abortData }) {
-  const req = storage('legacy').getStore()?.req
+  const req = getActiveRequest()
   if (!req) return
 
   const requestData = graphqlRequestData.get(req)