dd-trace 6.9.0 → 6.11.0

package/packages/dd-trace/src/exporters/agent/writer.js

@@ -3,7 +3,7 @@
 const { inspect } = require('node:util')
 const { channel } = require('dc-polyfill')
 
-const request = require('../common/request')
+const commonRequest = require('../common/request')
 const { logIntegrations, logAgentError } = require('../../startup-log')
 const runtimeMetrics = require('../../runtime_metrics')
 const log = require('../../log')
@@ -15,12 +15,15 @@ const METRIC_PREFIX = 'datadog.tracer.node.exporter.agent'
 const firstFlushChannel = channel('dd-trace:exporter:first-flush')
 
 class AgentWriter extends BaseWriter {
+  #request = commonRequest
+  #requestTracker
+
   constructor (...args) {
     super({
       ...args[0],
       beforeFirstFlush: () => firstFlushChannel.publish(),
     })
-    const { prioritySampler, lookup, protocolVersion, headers } = args[0]
+    const { prioritySampler, lookup, protocolVersion, headers, isTestOptimization } = args[0]
     const AgentEncoder = getEncoder(protocolVersion)
 
     this._prioritySampler = prioritySampler
@@ -28,13 +31,33 @@ class AgentWriter extends BaseWriter {
     this._protocolVersion = protocolVersion
     this._headers = headers
     this._encoder = new AgentEncoder(this)
+    if (isTestOptimization) {
+      this.#request = require('../../ci-visibility/exporters/request')
+      const TestOptimizationRequestTracker = require('../../ci-visibility/exporters/agentless/request-tracker')
+      this.#requestTracker = new TestOptimizationRequestTracker(this)
+    }
+  }
+
+  /**
+   * Flushes payloads, including requests already in flight during Test Optimization finalization.
+   *
+   * @param {(error?: Error) => void} [done]
+   * @param {{ deadline?: number }} [options]
+   * @returns {void}
+   */
+  flush (done, options) {
+    if (this.#requestTracker) {
+      this.#requestTracker.flush(done, options)
+      return
+    }
+    super.flush(done, options)
   }
 
-  _sendPayload (data, count, done) {
+  _sendPayload (data, count, done, flushOptions) {
     runtimeMetrics.increment(`${METRIC_PREFIX}.requests`, true)
 
     const { _headers, _lookup, _protocolVersion, _url } = this
-    makeRequest(_protocolVersion, data, count, _url, _headers, _lookup, (err, res, status, headers) => {
+    const onResponse = (err, res, status, headers) => {
       if (status) {
         runtimeMetrics.increment(`${METRIC_PREFIX}.responses`, true)
         runtimeMetrics.increment(`${METRIC_PREFIX}.responses.by.status`, `status:${status}`, true)
@@ -49,7 +72,7 @@ class AgentWriter extends BaseWriter {
 
       if (err) {
         log.errorWithoutTelemetry('Error sending payload to the agent (status code: %s)', err.status, err)
-        done()
+        done(flushOptions?.deadline === undefined ? undefined : err)
         return
       }
 
@@ -74,7 +97,19 @@ class AgentWriter extends BaseWriter {
         runtimeMetrics.increment(`${METRIC_PREFIX}.errors.by.name`, `name:${e.name}`, true)
       }
       done()
-    })
+    }
+    makeRequest(
+      _protocolVersion,
+      data,
+      count,
+      _url,
+      _headers,
+      _lookup,
+      flushOptions,
+      this.#request,
+      this.#requestTracker,
+      onResponse
+    )
   }
 }
 
@@ -84,7 +119,7 @@ function getEncoder (protocolVersion) {
     : require('../../encode/0.4').AgentEncoder
 }
 
-function makeRequest (version, data, count, url, headers, lookup, cb) {
+function makeRequest (version, data, count, url, headers, lookup, flushOptions, request, requestTracker, cb) {
   const options = {
     path: `/v${version}/traces`,
     method: 'PUT',
@@ -100,16 +135,21 @@ function makeRequest (version, data, count, url, headers, lookup, cb) {
     lookup,
     url,
   }
+  if (flushOptions?.deadline !== undefined) {
+    options.deadline = flushOptions.deadline
+  }
 
   log.debug('Request to the agent: %j', options)
 
-  request(data, options, (err, res, status, headers) => {
+  const onResponse = (err, res, status, headers) => {
     logIntegrations()
     if (status !== 404 && status !== 200 && err) {
       logAgentError({ status, message: err.message ?? inspect(err) })
     }
     cb(err, res, status, headers)
-  })
+  }
+  if (requestTracker) requestTracker.send(request, data, options, onResponse)
+  else request(data, options, onResponse)
 }
 
 module.exports = AgentWriter

package/packages/dd-trace/src/exporters/common/buffering-exporter.js

@@ -24,23 +24,25 @@ class BufferingExporter {
     this._export(trace)
   }
 
-  _export (payload, writer = this._writer, timerKey = '_timer') {
+  _export (payload, writer = this._writer, timerKey = '_timer', deferImmediateFlush = false) {
     if (this._config.isCiVisibility) {
       incrementCountMetric(TELEMETRY_EVENTS_ENQUEUED_FOR_SERIALIZATION, {}, payload.length)
     }
-    writer.append(payload)
+    const appended = writer.append(payload)
 
     const { flushInterval } = this._config
 
-    if (flushInterval === 0) {
+    if (flushInterval === 0 && !deferImmediateFlush) {
       writer.flush()
-    } else if (this[timerKey] === undefined) {
+    } else if (flushInterval !== 0 && this[timerKey] === undefined) {
       this[timerKey] = setTimeout(() => {
         writer.flush()
         this[timerKey] = undefined
       }, flushInterval)
       this[timerKey].unref?.()
     }
+
+    return appended
   }
 
   getUncodedTraces () {

package/packages/dd-trace/src/exporters/common/writer.js

@@ -17,10 +17,10 @@ class Writer {
 
   #isFirstFlush = true
 
-  flush (done = () => {}) {
+  flush (done = () => {}, options) {
     const count = this._encoder.count()
 
-    if (!request.writable) {
+    if (!request.writable && options?.deadline === undefined) {
       this._encoder.reset()
       done()
     } else if (count > 0) {
@@ -44,23 +44,28 @@ class Writer {
         done()
         return
       }
-      this._sendPayload(payload, count, done)
+      if (options === undefined) {
+        this._sendPayload(payload, count, done)
+      } else {
+        this._sendPayload(payload, count, done, options)
+      }
     } else {
       done()
     }
   }
 
-  append (payload) {
-    if (!request.writable) {
+  append (payload, options) {
+    if (!request.writable && options?.deadline === undefined) {
       // eslint-disable-next-line eslint-rules/eslint-log-printf-style
       log.debug(() => `Maximum number of active requests reached. Payload discarded: ${safeJSONStringify(payload)}`)
-      return
+      return false
     }
 
     // eslint-disable-next-line eslint-rules/eslint-log-printf-style
     log.debug(() => `Encoding payload: ${safeJSONStringify(payload)}`)
 
     this._encode(payload)
+    return true
   }
 
   _encode (payload) {

package/packages/dd-trace/src/llmobs/constants/tags.js

@@ -16,6 +16,11 @@ module.exports = {
   PROPAGATED_ML_APP_KEY: '_dd.p.llmobs_ml_app',
   PROPAGATED_SESSION_ID_KEY: '_dd.p.llmobs_sid',
   PARENT_ID_KEY: '_ml_obs.llmobs_parent_id',
+
+  PARENT_AGENT_NAME: '_ml_obs.parent_agent_name',
+  PARENT_AGENT_SPAN_ID: '_ml_obs.parent_agent_span_id',
+  PROPAGATED_PARENT_AGENT_ID_KEY: '_dd.p.llmobs_pagent_span_id',
+  PROPAGATED_PARENT_AGENT_NAME_KEY: '_dd.p.llmobs_pagent_name',
   PROPAGATED_SAMPLE_RATE_KEY: '_dd.p.llmobs_sr',
   PROPAGATED_SAMPLING_DECISION_KEY: '_dd.p.llmobs_sd',
   SAMPLE_RATE: '_ml_obs.sample_rate',

package/packages/dd-trace/src/llmobs/experiments/client.js

@@ -3,6 +3,9 @@
 // Control-plane HTTP client for LLM Obs Experiments. Uses the global `fetch`,
 // so this module adds no new dependency; credentials and site come from config.
 
+const { Dataset, DatasetRecord } = require('./dataset')
+const { ExperimentResult } = require('./result')
+
 const API_BASE_PATH = '/api/v2/llm-obs/v1'
 
 // Control-plane host for a Datadog site, e.g.
@@ -22,12 +25,67 @@ function appHost (site) {
   return site.split('.').length === 2 ? `app.${site}` : site
 }
 
+function datasetRecordFromResource (resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  const id = String(resource?.id ?? attrs.id ?? '')
+  if (id === '') throw new Error('Dataset record response is missing an id')
+  return new DatasetRecord(
+    attrs.input ?? null,
+    attrs.expected_output ?? null,
+    attrs.metadata ?? {},
+    id
+  )
+}
+
+function datasetVersionFromResource (resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  return attrs.valid_from_version ?? attrs.version ?? null
+}
+
+function datasetVersionFromResources (resources) {
+  const versions = resources
+    .map(datasetVersionFromResource)
+    .filter(version => version != null)
+    .map(Number)
+    .filter(Number.isFinite)
+  if (versions.length === 0) return null
+  return Math.max(...versions)
+}
+
+function datasetMutationResultFromResources (resources) {
+  return {
+    records: resources.map(datasetRecordFromResource),
+    version: datasetVersionFromResources(resources),
+  }
+}
+
+function datasetFromResource (client, projectId, resource) {
+  const attrs = resource?.attributes ?? resource ?? {}
+  const version = attrs.current_version ?? null
+  return Dataset.fromExisting(
+    client,
+    String(attrs.name ?? ''),
+    String(attrs.description ?? ''),
+    resource?.id ?? attrs.id ?? null,
+    projectId,
+    [],
+    version,
+    version
+  )
+}
+
+function experimentFromResource (client, resource) {
+  const id = resource?.id
+  return new ExperimentResult(id, [], id == null ? null : `${client.appBase}/llm/experiments/${id}`)
+}
+
 class ExperimentsClient {
   #apiKey
   #appKey
   #site
   #projectName
   #timeout
+  apiBase
   #cachedProjectId
 
   constructor ({ apiKey, appKey, site, projectName, timeout = 30_000 } = {}) {
@@ -36,6 +94,7 @@ class ExperimentsClient {
     this.#site = site
     this.#projectName = projectName
     this.#timeout = timeout
+    this.apiBase = `https://${apiHost(this.#site)}`
     this.#cachedProjectId = null
   }
 
@@ -61,7 +120,7 @@ class ExperimentsClient {
   // Low-level request. Builds https://api.<site><path>, attaches both keys, and
   // returns the parsed JSON body. Throws with status + body on a non-2xx.
   async request (method, path, body) {
-    const url = `https://${apiHost(this.#site)}${path}`
+    const url = `${this.apiBase}${path}`
     const headers = {
       'DD-API-KEY': this.#apiKey,
       'DD-APPLICATION-KEY': this.#appKey,
@@ -92,6 +151,107 @@ class ExperimentsClient {
     return text ? JSON.parse(text) : {}
   }
 
+  jsonApiRequest (method, path, type, attributes) {
+    return this.request(method, path, {
+      data: { type, attributes },
+    })
+  }
+
+  async createProject (name) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/projects`, 'projects', { name })
+    return response?.data ?? null
+  }
+
+  async createDataset (projectId, attributes) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/${projectId}/datasets`, 'datasets', attributes)
+    return datasetFromResource(this, projectId, response?.data ?? null)
+  }
+
+  deleteDataset (projectId, datasetId) {
+    return this.jsonApiRequest('POST', `${API_BASE_PATH}/${projectId}/datasets/delete`, 'datasets', {
+      type: 'soft',
+      dataset_ids: [datasetId],
+    })
+  }
+
+  async listDatasets (projectId, options = {}) {
+    const query = new URLSearchParams()
+    if (options.name !== undefined) query.set('filter[name]', options.name)
+    const response = await this.request('GET', `${API_BASE_PATH}/${projectId}/datasets?${query.toString()}`)
+    const resources = Array.isArray(response?.data) ? response.data : []
+    return resources.map(resource => datasetFromResource(this, projectId, resource))
+  }
+
+  async appendDatasetRecords (projectId, datasetId, records) {
+    const response = await this.jsonApiRequest(
+      'POST',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records`,
+      'datasets',
+      { records }
+    )
+    // The append-records response has used both a top-level `records` array
+    // and JSON:API `data` resources. Accept either so generated/custom record
+    // ids are preserved for experiment row tagging.
+    const resources = Array.isArray(response?.records)
+      ? response.records
+      : (Array.isArray(response?.data) ? response.data : [])
+    return resources.map(datasetRecordFromResource)
+  }
+
+  async batchUpdateDatasetRecords (projectId, datasetId, attributes) {
+    const response = await this.request(
+      'POST',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/batch_update`,
+      {
+        data: {
+          type: 'datasets',
+          id: datasetId,
+          attributes: {
+            insert_records: attributes.insert_records ?? [],
+            update_records: attributes.update_records ?? [],
+            delete_records: attributes.delete_records ?? [],
+            deduplicate: attributes.deduplicate !== false,
+            create_new_version: attributes.create_new_version !== false,
+          },
+        },
+      }
+    )
+    const resources = Array.isArray(response?.records)
+      ? response.records
+      : (Array.isArray(response?.data) ? response.data : [])
+    return datasetMutationResultFromResources(resources)
+  }
+
+  async listDatasetRecords (projectId, datasetId, options = {}) {
+    const query = new URLSearchParams()
+    if (options.cursor) query.set('page[cursor]', options.cursor)
+    if (options.version !== undefined && options.version !== null) query.set('filter[version]', String(options.version))
+    const response = await this.request(
+      'GET',
+      `${API_BASE_PATH}/${projectId}/datasets/${datasetId}/records?${query.toString()}`
+    )
+    const records = Array.isArray(response?.data) ? response.data.map(datasetRecordFromResource) : []
+    return { records, after: response?.meta?.after ?? '' }
+  }
+
+  async createExperiment (attributes) {
+    const response = await this.jsonApiRequest('POST', `${API_BASE_PATH}/experiments`, 'experiments', attributes)
+    return experimentFromResource(this, response?.data ?? null)
+  }
+
+  postExperimentEvents (experimentId, attributes) {
+    return this.jsonApiRequest(
+      'POST',
+      `${API_BASE_PATH}/experiments/${experimentId}/events`,
+      'experiments',
+      attributes
+    )
+  }
+
+  updateExperiment (experimentId, attributes) {
+    return this.jsonApiRequest('PATCH', `${API_BASE_PATH}/experiments/${experimentId}`, 'experiments', attributes)
+  }
+
   // Resolve the project id for `name`, creating it if absent. The create
   // endpoint is get-or-create on name, so repeated calls return the same id.
   // Cached after the first resolution.
@@ -100,14 +260,12 @@ class ExperimentsClient {
 
     let response
     try {
-      response = await this.request('POST', `${API_BASE_PATH}/projects`, {
-        data: { type: 'projects', attributes: { name } },
-      })
+      response = await this.createProject(name)
     } catch (err) {
       throw new Error(`Failed to create or get project '${name}': ${err.message}`)
     }
 
-    this.#cachedProjectId = response?.data?.id ?? null
+    this.#cachedProjectId = response?.id ?? null
     return this.#cachedProjectId
   }
 }