dd-trace 5.121.0 → 5.122.0

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

@@ -1,90 +1,18 @@
 'use strict'
 
-const rfdc = require('../../../../vendor/dist/rfdc')({ proto: false, circles: false })
-const { HTTP_CLIENT_IP, NETWORK_CLIENT_IP, HTTP_USERAGENT } = require('../../../../ext/tags')
-const { USER_ID, USER_SESSION_ID } = require('../appsec/addresses')
-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')
-const { AI_GUARD } = require('../standalone/product')
+const AIGuardClient = require('./client')
+const { createEvaluationOutcome, EvaluationReporter } = require('./evaluation')
+const { AIGuardAbortError } = require('./errors')
 const NoopAIGuard = require('./noop')
-const executeRequest = require('./client')
 const TAGS = require('./tags')
 
-const aiguardMetrics = telemetryMetrics.manager.namespace('ai_guard')
-
-const ALLOW = 'ALLOW'
-
-// Tags from the service entry span that must be mirrored onto every AI Guard span
-// so anomaly detection pipelines can process each span independently.
-const SERVICE_ENTRY_TAG_MAPPINGS = [
-  [HTTP_USERAGENT, TAGS.HTTP_USERAGENT_TAG_KEY],
-  [HTTP_CLIENT_IP, TAGS.HTTP_CLIENT_IP_TAG_KEY],
-  [NETWORK_CLIENT_IP, TAGS.NETWORK_CLIENT_IP_TAG_KEY],
-  [USER_ID, TAGS.USR_ID_TAG_KEY],
-  [USER_SESSION_ID, TAGS.USR_SESSION_ID_TAG_KEY],
-]
-
-/**
- * Reports a telemetry error
- *
- * @param {string} errorType - The error type constant (client_error, bad_status, bad_response)
- * @param {{ source: string, integration: string }} telemetryTags - Source and integration tags
- */
-function reportTelemetryError (errorType, telemetryTags) {
-  aiguardMetrics.count(TAGS.TELEMETRY_REQUESTS, { error: true, ...telemetryTags }).inc(1)
-  aiguardMetrics.count(TAGS.TELEMETRY_ERROR, { type: errorType, ...telemetryTags }).inc(1)
-}
-
-class AIGuardAbortError extends Error {
-  constructor (reason, tags, tagProbs, sds) {
-    super(reason)
-    this.name = 'AIGuardAbortError'
-    this.reason = reason
-    this.tags = tags
-    this.tagProbabilities = tagProbs
-    this.sds = sds || []
-  }
-}
-
-class AIGuardClientError extends Error {
-  constructor (message, opts = {}) {
-    super(message)
-    this.name = 'AIGuardClientError'
-    if (opts.errors) {
-      this.errors = opts.errors
-    }
-    if (opts.cause) {
-      this.cause = opts.cause
-    }
-  }
-}
-
-/**
- * Resolves the AI Guard host for a given Datadog site. Sites with a single subdomain level
- * (e.g. `datadoghq.com`, `ddog-gov.com`) are served from the `app.` subdomain, while regional
- * sites (e.g. `us3.datadoghq.com`, `ap1.datadoghq.com`) are used as-is.
- *
- * @param {string} site - Datadog site (e.g. `datadoghq.com`, `us3.datadoghq.com`)
- * @returns {string} The host to use for the AI Guard endpoint
- */
-function aiGuardHost (site) {
-  return site.split('.').length === 2 ? `app.${site}` : site
-}
-
 class AIGuard extends NoopAIGuard {
   #initialized
   #tracer
-  #headers
-  #evaluateUrl
-  #timeout
-  #maxMessagesLength
-  #maxContentSize
+  #client
+  #reporter
   #meta
-  #config
 
   /**
    * @param {import('../tracer')} tracer - Tracer instance
@@ -99,223 +27,39 @@ class AIGuard extends NoopAIGuard {
       return
     }
     this.#tracer = tracer
-    this.#headers = {
-      'DD-API-KEY': config.DD_API_KEY,
-      'DD-APPLICATION-KEY': config.DD_APP_KEY,
-      'DD-AI-GUARD-VERSION': tracerVersion,
-      'DD-AI-GUARD-SOURCE': 'SDK',
-      'DD-AI-GUARD-LANGUAGE': 'nodejs',
-    }
-    const endpoint = config.experimental.aiguard.endpoint || `https://${aiGuardHost(config.site)}/api/v2/ai-guard`
-    this.#evaluateUrl = `${endpoint}/evaluate`
-    this.#timeout = config.experimental.aiguard.timeout
-    this.#maxMessagesLength = config.experimental.aiguard.maxMessagesLength
-    this.#maxContentSize = config.experimental.aiguard.maxContentSize
+    this.#client = new AIGuardClient(config)
+    this.#reporter = new EvaluationReporter(config)
     this.#meta = { service: config.service, env: config.env }
-    this.#config = config
     this.#initialized = true
   }
 
-  /**
-   * Returns a safe copy of the messages to be serialized into the meta struct.
-   *
-   * - Clones each message so callers cannot mutate the data set in the meta struct.
-   * - Truncates the list of messages and `content` fields emitting metrics accordingly.
-   *
-   * @param {import('../../../../index').aiguard.Message[]} messages
-   * @param {{ source: string, integration: string }} telemetryTags
-   */
-  #buildMessagesForMetaStruct (messages, telemetryTags) {
-    const size = Math.min(messages.length, this.#maxMessagesLength)
-    if (messages.length > size) {
-      aiguardMetrics.count(TAGS.TELEMETRY_TRUNCATED, { type: 'messages', ...telemetryTags }).inc(1)
-    }
-    const result = []
-    let contentTruncated = false
-    for (let i = messages.length - size; i < messages.length; i++) {
-      const message = rfdc(messages[i])
-      if (message.content?.length > this.#maxContentSize) {
-        contentTruncated = true
-        message.content = message.content.slice(0, this.#maxContentSize)
-      }
-      result.push(message)
-    }
-    if (contentTruncated) {
-      aiguardMetrics.count(TAGS.TELEMETRY_TRUNCATED, { type: 'content', ...telemetryTags }).inc(1)
-    }
-    return result
-  }
-
-  #isToolCall (message) {
-    return message.tool_calls || message.tool_call_id
-  }
-
-  #getToolName (message, history) {
-    // 1. assistant message with tool calls
-    if (message.tool_calls) {
-      const names = message.tool_calls.map((tool) => tool.function.name)
-      return names.length === 0 ? null : names.join(',')
-    }
-    // 2. assistant message with tool output (search the linked tool call in reverse order)
-    const id = message.tool_call_id
-    for (let i = history.length - 2; i >= 0; i--) {
-      const item = history[i]
-      if (item.tool_calls) {
-        for (const toolCall of item.tool_calls) {
-          if (toolCall.id === id) {
-            return toolCall.function.name
-          }
-        }
-      }
-    }
-    return null
-  }
-
-  #setRootSpanClientIpTags (rootSpan) {
-    if (!rootSpan) return
-
-    const currentTags = rootSpan.context().getTags()
-    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)
-    }
-  }
-
-  /**
-   * Copies service entry span tags needed by anomaly detection to the AI Guard span.
-   *
-   * @param {import('../opentracing/span')} guardSpan
-   * @param {import('../opentracing/span')} rootSpan
-   */
-  #copyServiceEntryTagsToGuardSpan (guardSpan, rootSpan) {
-    const rootTags = rootSpan.context().getTags()
-    for (const [sourceTag, destTag] of SERVICE_ENTRY_TAG_MAPPINGS) {
-      const value = rootTags[sourceTag]
-      if (value !== undefined && value !== null) {
-        guardSpan.setTag(destTag, value)
-      }
-    }
-  }
-
   evaluate (messages, opts) {
     if (!this.#initialized) {
       return super.evaluate(messages, opts)
     }
     const { block = true, source = TAGS.SOURCE_SDK, integration = TAGS.INTEGRATION_NONE, childOf } = opts ?? {}
-    const telemetryTags = { source, integration }
     // Only pass `childOf` when truthy so `tracer.trace`'s default (`scope().active()`)
     // still applies for SDK callers that don't supply an explicit parent.
     const traceOpts = childOf ? { childOf } : {}
     return this.#tracer.trace(TAGS.RESOURCE, traceOpts, async (span) => {
-      const last = messages.at(-1)
-      const target = this.#isToolCall(last) ? 'tool' : 'prompt'
-      span.setTag(TAGS.TARGET_TAG_KEY, target)
-      if (target === 'tool') {
-        const name = this.#getToolName(last, messages)
-        if (name) {
-          span.setTag(TAGS.TOOL_NAME_TAG_KEY, name)
-        }
-      }
-      const metaStruct = {
-        messages: this.#buildMessagesForMetaStruct(messages, telemetryTags),
-      }
-      span.meta_struct = {
-        [TAGS.META_STRUCT_KEY]: metaStruct,
-      }
-      const rootSpan = span.context()?._trace?.started?.[0]
-      if (rootSpan) {
-        this.#setRootSpanClientIpTags(rootSpan)
-        this.#copyServiceEntryTagsToGuardSpan(span, 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(TAGS.EVENT_TAG_KEY, 'true')
-      }
-      let response
+      const report = this.#reporter.start(span, messages, { source, integration })
+      let evaluation
       try {
-        const payload = {
-          data: {
-            attributes: {
-              messages,
-              meta: this.#meta,
-            },
-          },
-        }
-        response = await executeRequest(
-          payload,
-          { url: this.#evaluateUrl, headers: this.#headers, timeout: this.#timeout })
-      } catch (e) {
-        reportTelemetryError(TAGS.ERROR_TYPE_CLIENT, telemetryTags)
-        throw new AIGuardClientError(`Unexpected error calling AI Guard service: ${e.message}`, { cause: e })
-      }
-      if (response.status !== 200) {
-        reportTelemetryError(TAGS.ERROR_TYPE_STATUS, telemetryTags)
-        throw new AIGuardClientError(
-          `AI Guard service call failed, status ${response.status}`,
-          { errors: response.body?.errors })
+        evaluation = await this.#client.evaluate(messages, this.#meta)
+      } catch (error) {
+        this.#reporter.fail(report, error.telemetryType ?? TAGS.ERROR_TYPE_CLIENT)
+        throw error
       }
-      const attr = response.body?.data?.attributes
-      if (!attr?.action) {
-        reportTelemetryError(TAGS.ERROR_TYPE_RESPONSE, telemetryTags)
-        throw new AIGuardClientError(`AI Guard service returned unexpected response : ${response.body}`)
-      }
-      const action = attr.action
-      const reason = attr.reason
-      const tags = attr.tags ?? []
-      if (tags.length > 0) {
-        metaStruct.attack_categories = tags
-      }
-      const sdsFindings = attr.sds_findings ?? []
-      if (sdsFindings.length > 0) {
-        metaStruct.sds = sdsFindings
-      }
-      const tagProbabilities = attr.tag_probs ?? {}
-      if (attr.tag_probs) {
-        metaStruct.tag_probs = tagProbabilities
-      }
-      const blockingEnabled = attr.is_blocking_enabled ?? false
-      const shouldBlock = block && blockingEnabled && action !== ALLOW
-      aiguardMetrics.count(TAGS.TELEMETRY_REQUESTS, {
-        action,
-        error: false,
-        block: shouldBlock,
-        ...telemetryTags,
-      }).inc(1)
-      span.setTag(TAGS.ACTION_TAG_KEY, action)
-      if (reason) {
-        span.setTag(TAGS.REASON_TAG_KEY, reason)
-      }
-      if (shouldBlock) {
-        span.setTag(TAGS.BLOCKED_TAG_KEY, 'true')
-        throw new AIGuardAbortError(reason, tags, tagProbabilities, sdsFindings)
+
+      const outcome = createEvaluationOutcome(evaluation, block)
+      this.#reporter.finish(report, outcome)
+
+      if (outcome.shouldBlock) {
+        const { reason, tags, tagProbabilities, sds } = outcome.result
+        throw new AIGuardAbortError(reason, tags, tagProbabilities, sds)
       }
-      return { action, reason, tags, tagProbabilities, sds: sdsFindings }
+
+      return outcome.result
     })
   }
 }

package/packages/dd-trace/src/appsec/iast/taint-tracking/rewriter.js

@@ -97,15 +97,14 @@ function getPrepareStackTraceAccessor () {
 
 function getCompileMethodFn (compileMethod) {
   let delegate = function (content, filename) {
+    if (isDdTrace(filename) || !isPrivateModule(filename) || !config.iast?.enabled) {
+      return compileMethod.apply(this, [content, filename])
+    }
+
+    // TODO when we have CJS support for orchestrion and taint-tracking, add
+    // them here as appropriate
+    let rewrittenContent
     try {
-      if (isDdTrace(filename)) {
-        return compileMethod.apply(this, [content, filename])
-      }
-      if (!isPrivateModule(filename) || !config.iast?.enabled) {
-        return compileMethod.apply(this, [content, filename])
-      }
-      // TODO when we have CJS support for orchestrion and taint-tracking, add
-      // them here as appropriate
       const rewritten = rewriter.rewrite(content, filename, ['iast'])
 
       incrementTelemetryIfNeeded(rewritten.metrics)
@@ -114,13 +113,11 @@ function getCompileMethodFn (compileMethod) {
         hardcodedSecretCh.publish(rewritten.literalsResult)
       }
 
-      if (rewritten?.content) {
-        return compileMethod.apply(this, [rewritten.content, filename])
-      }
+      rewrittenContent = rewritten?.content
     } catch (e) {
       log.error('Error rewriting file %s', filename, e)
     }
-    return compileMethod.apply(this, [content, filename])
+    return compileMethod.apply(this, [rewrittenContent || content, filename])
   }
 
   const shim = function (...args) {

package/packages/dd-trace/src/ci-visibility/dynamic-instrumentation/worker/index.js

@@ -92,6 +92,7 @@ function removeEmptyCaptureProperties (value) {
       if (current.fields === null || (
         current.fields &&
         typeof current.fields === 'object' &&
+        // eslint-disable-next-line no-restricted-syntax -- snapshot field names are user variables; no key to probe
         Object.keys(current.fields).length === 0
       )) {
         delete current.fields

package/packages/dd-trace/src/ci-visibility/exporters/agent-proxy/index.js

@@ -4,6 +4,7 @@ const AgentWriter = require('../../../exporters/agent/writer')
 const AgentlessWriter = require('../agentless/writer')
 const CoverageWriter = require('../agentless/coverage-writer')
 const CiVisibilityExporter = require('../ci-visibility-exporter')
+const request = require('../request')
 const { fetchAgentInfo } = require('../../../agent/info')
 const { DEBUGGER_INPUT_V1 } = require('../../../debugger/constants')
 
@@ -43,9 +44,21 @@ class AgentProxyCiVisibilityExporter extends CiVisibilityExporter {
       testOptimization,
     } = config
 
+    const initializationController = new AbortController()
+    const initializationOptions = { signal: initializationController.signal }
+    this._initializationRequest = {
+      controller: initializationController,
+      options: initializationOptions,
+    }
+
     fetchAgentInfo(this._url, (err, agentInfo) => {
+      const initializationFinalFlush = this._initializationRequest?.finalFlush
+      this._initializationRequest = undefined
+      const initializationAborted = initializationController.signal.aborted
+      const agentInfoError = err || (initializationAborted ? initializationController.signal.reason : undefined)
+
       this._isInitialized = true
-      let latestEvpProxyVersion = getLatestEvpProxyVersion(err, agentInfo)
+      let latestEvpProxyVersion = getLatestEvpProxyVersion(agentInfoError, agentInfo)
       const isEvpCompatible = latestEvpProxyVersion >= 2
       this._isGzipCompatible = latestEvpProxyVersion >= 4
 
@@ -72,7 +85,7 @@ class AgentProxyCiVisibilityExporter extends CiVisibilityExporter {
         // path with evpProxyPrefix and sets X-Datadog-EVP-Subdomain: api (see uploadTestScreenshot).
         this._testScreenshotUploadUrl = this._url
         if (testOptimization.DD_TEST_FAILED_TEST_REPLAY_ENABLED) {
-          const canFowardLogs = getCanForwardDebuggerLogs(err, agentInfo)
+          const canFowardLogs = getCanForwardDebuggerLogs(agentInfoError, agentInfo)
           if (canFowardLogs) {
             const DynamicInstrumentationLogsWriter = require('../agentless/di-logs-writer')
             this._logsWriter = new DynamicInstrumentationLogsWriter({
@@ -89,14 +102,25 @@ class AgentProxyCiVisibilityExporter extends CiVisibilityExporter {
           lookup,
           protocolVersion,
           headers,
+          isTestOptimization: true,
         })
         // coverages will never be used, so we discard them
         this._coverageBuffer = []
       }
       this._resolveCanUseCiVisProtocol(isEvpCompatible)
+      if (initializationAborted) {
+        this.resetUncodedTraces()
+        this.resetDeferredTestSuiteSpans()
+        return
+      }
+      if (isEvpCompatible) {
+        if (initializationFinalFlush) this.exportDeferredTestSuiteSpans()
+      } else {
+        this.resetDeferredTestSuiteSpans()
+      }
       this.exportUncodedTraces()
       this.exportUncodedCoverages()
-    })
+    }, initializationOptions, request)
   }
 
   setUrl (url, coverageUrl) {

package/packages/dd-trace/src/ci-visibility/exporters/agentless/coverage-writer.js

@@ -1,11 +1,8 @@
 'use strict'
 const getConfig = require('../../../config')
-const request = require('../../../exporters/common/request')
 const log = require('../../../log')
 const { safeJSONStringify } = require('../../../exporters/common/util')
 
-const { CoverageCIVisibilityEncoder } = require('../../../encode/coverage-ci-visibility')
-const BaseWriter = require('../../../exporters/common/writer')
 const {
   incrementCountMetric,
   distributionMetric,
@@ -15,16 +12,34 @@ const {
   TELEMETRY_ENDPOINT_PAYLOAD_REQUESTS_ERRORS,
   TELEMETRY_ENDPOINT_PAYLOAD_DROPPED,
 } = require('../../../ci-visibility/telemetry')
+const { CoverageCIVisibilityEncoder } = require('../../../encode/coverage-ci-visibility')
+const BaseWriter = require('../../../exporters/common/writer')
+const request = require('../request')
+const TestOptimizationRequestTracker = require('./request-tracker')
 
 class Writer extends BaseWriter {
+  #requestTracker
+
   constructor ({ url, evpProxyPrefix = '' }) {
     super(...arguments)
+    this.#requestTracker = new TestOptimizationRequestTracker(this)
     this._url = url
     this._encoder = new CoverageCIVisibilityEncoder(this)
     this._evpProxyPrefix = evpProxyPrefix
   }
 
-  _sendPayload (form, _, done) {
+  /**
+   * Flushes buffered coverage, waiting for tracked requests during finalization.
+   *
+   * @param {(error?: Error) => void} [done]
+   * @param {{ deadline?: number }} [options]
+   * @returns {void}
+   */
+  flush (done, options) {
+    this.#requestTracker.flush(done, options)
+  }
+
+  _sendPayload (form, _, done, flushOptions) {
     const options = {
       path: '/api/v2/citestcov',
       method: 'POST',
@@ -34,6 +49,7 @@ class Writer extends BaseWriter {
       },
       timeout: 15_000,
       url: this._url,
+      deadline: flushOptions?.deadline,
     }
 
     if (this._evpProxyPrefix) {
@@ -50,7 +66,7 @@ class Writer extends BaseWriter {
     incrementCountMetric(TELEMETRY_ENDPOINT_PAYLOAD_REQUESTS, { endpoint: 'code_coverage' })
     distributionMetric(TELEMETRY_ENDPOINT_PAYLOAD_BYTES, { endpoint: 'code_coverage' }, form.size())
 
-    request(form, options, (err, res, statusCode) => {
+    this.#requestTracker.send(request, form, options, (err, res, statusCode) => {
       distributionMetric(
         TELEMETRY_ENDPOINT_PAYLOAD_REQUESTS_MS,
         { endpoint: 'code_coverage' },
@@ -66,7 +82,7 @@ class Writer extends BaseWriter {
           { endpoint: 'code_coverage' }
         )
         log.error('Error sending CI coverage payload', err)
-        done()
+        done(err)
         return
       }
       log.debug('Response from the intake:', res)

package/packages/dd-trace/src/ci-visibility/exporters/agentless/di-logs-writer.js

@@ -1,27 +1,42 @@
 'use strict'
 const getConfig = require('../../../config')
-const request = require('../../../exporters/common/request')
 const log = require('../../../log')
 const { safeJSONStringify } = require('../../../exporters/common/util')
 const { JSONEncoder } = require('../../encode/json-encoder')
 const { DEBUGGER_INPUT_V1 } = require('../../../debugger/constants')
-
 const BaseWriter = require('../../../exporters/common/writer')
 
+const request = require('../request')
+const TestOptimizationRequestTracker = require('./request-tracker')
+
 // Writer used by the integration between Dynamic Instrumentation and Test Visibility
 // It is used to encode and send logs to both the logs intake directly and the
 // `/debugger/v1/input` endpoint in the agent, which is a proxy to the logs intake.
 class DynamicInstrumentationLogsWriter extends BaseWriter {
+  #requestTracker
+
   // TODO: what's a good value for timeout for the logs intake?
   constructor ({ url, timeout = 15_000, isAgentProxy = false }) {
     super(...arguments)
+    this.#requestTracker = new TestOptimizationRequestTracker(this)
     this._url = url
     this._encoder = new JSONEncoder()
     this._isAgentProxy = isAgentProxy
     this.timeout = timeout
   }
 
-  _sendPayload (data, _, done) {
+  /**
+   * Flushes buffered logs, waiting for tracked requests during finalization.
+   *
+   * @param {(error?: Error) => void} [done]
+   * @param {{ deadline?: number }} [options]
+   * @returns {void}
+   */
+  flush (done, options) {
+    this.#requestTracker.flush(done, options)
+  }
+
+  _sendPayload (data, _, done, flushOptions) {
     const options = {
       path: '/api/v2/logs',
       method: 'POST',
@@ -31,6 +46,7 @@ class DynamicInstrumentationLogsWriter extends BaseWriter {
       },
       timeout: this.timeout,
       url: this._url,
+      deadline: flushOptions?.deadline,
     }
 
     if (this._isAgentProxy) {
@@ -41,10 +57,10 @@ class DynamicInstrumentationLogsWriter extends BaseWriter {
     // eslint-disable-next-line eslint-rules/eslint-log-printf-style
     log.debug(() => `Request to the logs intake: ${safeJSONStringify(options)}`)
 
-    request(data, options, (err, res) => {
+    this.#requestTracker.send(request, data, options, (err, res) => {
       if (err) {
         log.error('Error sending DI logs payload', err)
-        done()
+        done(err)
         return
       }
       log.debug('Response from the logs intake:', res)