dd-trace 6.18.0 → 6.19.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (55) hide show
  1. package/LICENSE-3rdparty.csv +1 -0
  2. package/ci/vitest-no-worker-init-setup.mjs +22 -12
  3. package/index.d.ts +3 -2
  4. package/package.json +6 -6
  5. package/packages/datadog-instrumentations/src/anthropic.js +111 -9
  6. package/packages/datadog-instrumentations/src/claude-agent-sdk.js +5 -1
  7. package/packages/datadog-instrumentations/src/cucumber.js +40 -4
  8. package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/playwright.js +8 -0
  9. package/packages/datadog-instrumentations/src/helpers/rewriter/targets.json +2 -0
  10. package/packages/datadog-instrumentations/src/jest/session-error.js +102 -0
  11. package/packages/datadog-instrumentations/src/jest.js +3 -10
  12. package/packages/datadog-instrumentations/src/playwright.js +26 -0
  13. package/packages/datadog-instrumentations/src/vitest-main.js +4 -1
  14. package/packages/datadog-instrumentations/src/vitest-worker.js +34 -3
  15. package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js +83 -29
  16. package/packages/datadog-plugin-cucumber/src/index.js +19 -4
  17. package/packages/datadog-plugin-cypress/src/cypress-plugin.js +7 -0
  18. package/packages/datadog-plugin-fetch/src/index.js +3 -0
  19. package/packages/datadog-plugin-oracledb/src/connection-parser.js +3 -1
  20. package/packages/datadog-plugin-undici/src/index.js +5 -0
  21. package/packages/dd-trace/src/aiguard/client.js +38 -19
  22. package/packages/dd-trace/src/aiguard/integrations/anthropic.js +29 -4
  23. package/packages/dd-trace/src/aiguard/integrations/index.js +1 -1
  24. package/packages/dd-trace/src/aiguard/integrations/openai.js +20 -56
  25. package/packages/dd-trace/src/aiguard/integrations/stream.js +60 -0
  26. package/packages/dd-trace/src/aiguard/messages/anthropic.js +64 -0
  27. package/packages/dd-trace/src/config/generated-config-types.d.ts +4 -4
  28. package/packages/dd-trace/src/config/supported-configurations.json +5 -5
  29. package/packages/dd-trace/src/constants.js +2 -0
  30. package/packages/dd-trace/src/llmobs/constants/tags.js +31 -8
  31. package/packages/dd-trace/src/llmobs/gen-ai-tags.js +127 -0
  32. package/packages/dd-trace/src/llmobs/plugins/ai/ddTelemetry.js +18 -0
  33. package/packages/dd-trace/src/llmobs/plugins/ai/vercelTelemetry.js +33 -8
  34. package/packages/dd-trace/src/llmobs/plugins/anthropic/index.js +77 -41
  35. package/packages/dd-trace/src/llmobs/plugins/base.js +154 -15
  36. package/packages/dd-trace/src/llmobs/plugins/bedrockruntime.js +162 -10
  37. package/packages/dd-trace/src/llmobs/plugins/claude-agent-sdk/index.js +65 -16
  38. package/packages/dd-trace/src/llmobs/plugins/genai/index.js +14 -0
  39. package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js +34 -35
  40. package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js +10 -0
  41. package/packages/dd-trace/src/llmobs/plugins/langchain/index.js +25 -2
  42. package/packages/dd-trace/src/llmobs/plugins/langgraph/index.js +8 -7
  43. package/packages/dd-trace/src/llmobs/plugins/openai/index.js +13 -0
  44. package/packages/dd-trace/src/llmobs/plugins/openai/realtime.js +5 -0
  45. package/packages/dd-trace/src/llmobs/plugins/vertexai.js +7 -0
  46. package/packages/dd-trace/src/llmobs/prompts/prompt.js +2 -1
  47. package/packages/dd-trace/src/llmobs/span_processor.js +10 -65
  48. package/packages/dd-trace/src/llmobs/tagger.js +2 -36
  49. package/packages/dd-trace/src/opentelemetry/trace/index.js +5 -0
  50. package/packages/dd-trace/src/plugins/util/test.js +4 -0
  51. package/packages/dd-trace/src/profiler.js +51 -5
  52. package/packages/dd-trace/src/profiling/profiler.js +24 -20
  53. package/packages/dd-trace/src/ritm.js +12 -9
  54. package/packages/dd-trace/src/span_processor.js +6 -1
  55. package/vendor/dist/@datadog/openfeature-node-server/index.js +1 -1
@@ -4,6 +4,7 @@ const fs = require('node:fs')
4
4
  const path = require('node:path')
5
5
  const { performance } = require('node:perf_hooks')
6
6
  const { fileURLToPath } = require('node:url')
7
+ const { deserialize } = require('node:v8')
7
8
  const { isMainThread, parentPort } = require('node:worker_threads')
8
9
 
9
10
  const { channel } = require('dc-polyfill')
@@ -47,6 +48,7 @@ const {
47
48
 
48
49
  const EFD_SUITE_ADMISSION_TIMEOUT_MS = 5000
49
50
  const logSubmissionFlushCh = channel('ci:log-submission:flush')
51
+ const wrappedTestRunners = new WeakSet()
50
52
  const taskToCtx = new WeakMap()
51
53
  const taskToTestProperties = new WeakMap()
52
54
  const taskToStatuses = new WeakMap()
@@ -412,6 +414,14 @@ function finishEfdSuiteAdmissionRequest (requestId, allowed) {
412
414
  * @param {unknown} message
413
415
  */
414
416
  function handleEfdSuiteAdmissionResponse (message) {
417
+ // Older Vitest 4 forks serialize parent messages explicitly; newer forks and threads send plain values.
418
+ if (Buffer.isBuffer(message) || (message?.type === 'Buffer' && Array.isArray(message.data))) {
419
+ try {
420
+ message = deserialize(Buffer.from(message))
421
+ } catch {
422
+ return
423
+ }
424
+ }
415
425
  if (!Array.isArray(message) || message[0] !== VITEST_WORKER_EFD_SUITE_ADMISSION_RESPONSE_CODE) return
416
426
 
417
427
  const { allowed, requestId } = message[1] || {}
@@ -520,6 +530,11 @@ function wrapSuiteHookFn (hookType, fn, fallbackTask) {
520
530
  }
521
531
 
522
532
  function wrapVitestTestRunner (VitestTestRunner) {
533
+ // ESM hooks can expose the same runner through multiple module identities.
534
+ // Wrapping its lifecycle methods twice would report each retry twice.
535
+ if (wrappedTestRunners.has(VitestTestRunner)) return
536
+ wrappedTestRunners.add(VitestTestRunner)
537
+
523
538
  // `onBeforeRunTask` is run before any repetition or attempt is run
524
539
  // `onBeforeRunTask` is an async function
525
540
  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => async function (task) {
@@ -623,6 +638,13 @@ function wrapVitestTestRunner (VitestTestRunner) {
623
638
  shimmer.wrap(VitestTestRunner.prototype, 'onAfterRunTask', onAfterRunTask => function (task) {
624
639
  const { isTestManagementTestsEnabled } = getProvidedContext()
625
640
 
641
+ // Vitest 5 aggregates failed repetitions into the final state. Preserve the last attempt's actual result.
642
+ if ((attemptToFixTasks.has(task) || efdRetryTasks.has(task)) &&
643
+ task.result.state === 'fail' && taskToCtx.get(task)?.status === 'pass' &&
644
+ (task.result.errors?.length ?? 0) === (taskToReportedErrorCount.get(task) ?? 0)) {
645
+ task.result.state = 'pass'
646
+ }
647
+
626
648
  if (isTestManagementTestsEnabled) {
627
649
  const isAttemptingToFix = attemptToFixTasks.has(task)
628
650
  const isQuarantined = quarantinedTasks.has(task)
@@ -745,7 +767,12 @@ function wrapVitestTestRunner (VitestTestRunner) {
745
767
  }
746
768
  }
747
769
 
748
- const lastExecutionStatus = task.result.state
770
+ let lastExecutionStatus = task.result.state
771
+ if (lastExecutionStatus === 'run' && numRepetition > 0) {
772
+ // Vitest 5 resets the state before each repetition but retains errors from earlier attempts.
773
+ const reportedErrorCount = taskToReportedErrorCount.get(task) ?? 0
774
+ lastExecutionStatus = (task.result.errors?.length ?? 0) > reportedErrorCount ? 'fail' : 'pass'
775
+ }
749
776
  const isAtf = attemptToFixTasks.has(task)
750
777
  const isEfd = efdRetryTasks.has(task)
751
778
  const shouldTrackStatuses = isEfd || isAtf
@@ -1101,14 +1128,18 @@ function getStartTestsWrapper (frameworkVersion) {
1101
1128
  isDisabled: disabledTasks.has(task),
1102
1129
  isQuarantined: quarantinedTasks.has(task),
1103
1130
  })
1104
- } else if (state === 'pass' && !isSwitchedStatus) {
1131
+ } else if ((state === 'pass' && !isSwitchedStatus) || switchedStatus === 'pass') {
1132
+ // Attempt-to-fix can pass its last attempt while an earlier failure still fails the suite.
1133
+ if (switchedStatus === 'pass' && errors?.length) {
1134
+ testSuiteError = testError
1135
+ }
1105
1136
  if (testCtx) {
1106
1137
  const isSkippedByTestManagement =
1107
1138
  !attemptToFixTasks.has(task) && (disabledTasks.has(task) || quarantinedTasks.has(task))
1108
1139
  const promises = {}
1109
1140
  testPassCh.publish({
1110
1141
  task,
1111
- finalStatus: isSkippedByTestManagement ? 'skip' : 'pass',
1142
+ finalStatus: isSkippedByTestManagement ? 'skip' : state,
1112
1143
  earlyFlakeAbortReason: efdSlowAbortedTasks.has(task) ? 'slow' : undefined,
1113
1144
  promises,
1114
1145
  ...testCtx.currentStore,
@@ -23,6 +23,15 @@ const PROVIDER = {
23
23
  MISTRAL: 'MISTRAL',
24
24
  }
25
25
 
26
+ /**
27
+ * @typedef {{
28
+ * inputTokens?: number,
29
+ * outputTokens?: number,
30
+ * cacheReadTokens?: number,
31
+ * cacheWriteTokens?: number,
32
+ * }} StreamedUsage
33
+ */
34
+
26
35
  /**
27
36
  * Coerce the chunks into a single response body.
28
37
  *
@@ -44,10 +53,7 @@ function extractTextAndResponseReasonFromStream (chunks, modelProvider, modelNam
44
53
  }
45
54
 
46
55
  let message = ''
47
- let inputTokens = 0
48
- let outputTokens = 0
49
- let cacheReadTokens = 0
50
- let cacheWriteTokens = 0
56
+ let usage
51
57
 
52
58
  for (const { chunk: { bytes } } of chunks) {
53
59
  const body = JSON.parse(Buffer.from(bytes).toString('utf8'))
@@ -56,8 +62,6 @@ function extractTextAndResponseReasonFromStream (chunks, modelProvider, modelNam
56
62
  case PROVIDER.AMAZON: {
57
63
  if (body?.outputText) {
58
64
  message += body?.outputText
59
- inputTokens = body?.inputTextTokenCount
60
- outputTokens = body?.totalOutputTextTokenCount
61
65
  } else if (body?.contentBlockDelta?.delta?.text) {
62
66
  message += body.contentBlockDelta.delta.text
63
67
  }
@@ -79,9 +83,6 @@ function extractTextAndResponseReasonFromStream (chunks, modelProvider, modelNam
79
83
  message += body.delta.text
80
84
  }
81
85
 
82
- if (body.message?.usage?.input_tokens) inputTokens = body.message.usage.input_tokens
83
- if (body.message?.usage?.output_tokens) outputTokens = body.message.usage.output_tokens
84
-
85
86
  break
86
87
  }
87
88
  case PROVIDER.COHERE: {
@@ -101,24 +102,61 @@ function extractTextAndResponseReasonFromStream (chunks, modelProvider, modelNam
101
102
  }
102
103
  }
103
104
 
104
- // by default, it seems newer versions of the AWS SDK include the input/output token counts in the response body
105
- const invocationMetrics = body['amazon-bedrock-invocationMetrics']
106
- if (invocationMetrics) {
107
- inputTokens = invocationMetrics.inputTokenCount
108
- outputTokens = invocationMetrics.outputTokenCount
109
- cacheReadTokens = invocationMetrics.cacheReadInputTokenCount
110
- cacheWriteTokens = invocationMetrics.cacheWriteInputTokenCount
105
+ usage = mergeStreamedUsage(usage, body, modelProviderUpper)
106
+ }
107
+
108
+ return new Generation({ message, role: 'assistant', ...usage })
109
+ }
110
+
111
+ /**
112
+ * Merge the token counts one streamed `invokeModel` frame reports into the running totals. Each
113
+ * provider spells them differently, and most report none until the trailing invocation metrics,
114
+ * so a frame carrying no counts leaves the totals as they were.
115
+ *
116
+ * @param {StreamedUsage} [usage]
117
+ * @param {object} body parsed chunk body
118
+ * @param {string} modelProviderUpper
119
+ * @returns {StreamedUsage | undefined}
120
+ */
121
+ function mergeStreamedUsage (usage, body, modelProviderUpper) {
122
+ // by default, it seems newer versions of the AWS SDK include the input/output token counts in
123
+ // the response body; any provider can send them, and they supersede what the frames reported
124
+ const invocationMetrics = body['amazon-bedrock-invocationMetrics']
125
+ if (invocationMetrics) {
126
+ return {
127
+ inputTokens: invocationMetrics.inputTokenCount,
128
+ outputTokens: invocationMetrics.outputTokenCount,
129
+ cacheReadTokens: invocationMetrics.cacheReadInputTokenCount,
130
+ cacheWriteTokens: invocationMetrics.cacheWriteInputTokenCount,
111
131
  }
112
132
  }
113
133
 
114
- return new Generation({
115
- message,
116
- role: 'assistant',
117
- inputTokens,
118
- outputTokens,
119
- cacheReadTokens,
120
- cacheWriteTokens,
121
- })
134
+ switch (modelProviderUpper) {
135
+ case PROVIDER.AMAZON: {
136
+ const { inputTextTokenCount, totalOutputTextTokenCount } = body
137
+ if (inputTextTokenCount === undefined && totalOutputTextTokenCount === undefined) return usage
138
+
139
+ return {
140
+ ...usage,
141
+ inputTokens: inputTextTokenCount ?? usage?.inputTokens,
142
+ outputTokens: totalOutputTextTokenCount ?? usage?.outputTokens,
143
+ }
144
+ }
145
+ case PROVIDER.ANTHROPIC: {
146
+ // `message_start` reports the initial counts under `message.usage`; the closing
147
+ // `message_delta` reports the final output count at the top level
148
+ const chunkUsage = body.message?.usage ?? body.usage
149
+ if (!chunkUsage) return usage
150
+
151
+ return {
152
+ ...usage,
153
+ inputTokens: chunkUsage.input_tokens ?? usage?.inputTokens,
154
+ outputTokens: chunkUsage.output_tokens ?? usage?.outputTokens,
155
+ }
156
+ }
157
+ }
158
+
159
+ return usage
122
160
  }
123
161
 
124
162
  class Generation {
@@ -568,13 +606,27 @@ function resolveToolResultItem (item) {
568
606
  return `[Unsupported content type(s): ${getContentBlockType(item)}]`
569
607
  }
570
608
 
609
+ /**
610
+ * Normalize a Converse usage object onto the LLMObs metric names. Returns undefined when the
611
+ * response reported no counts at all, so a caller cannot mistake an empty record for a measurement.
612
+ *
613
+ * @param {object} [usage]
614
+ * @returns {{
615
+ * inputTokens?: number, outputTokens?: number, cacheReadTokens?: number, cacheWriteTokens?: number
616
+ * } | undefined}
617
+ */
571
618
  function buildUsage (usage = {}) {
572
- return {
573
- inputTokens: usage.inputTokens,
574
- outputTokens: usage.outputTokens,
575
- cacheReadTokens: usage.cacheReadInputTokens ?? usage.cacheReadInputTokenCount,
576
- cacheWriteTokens: usage.cacheWriteInputTokens ?? usage.cacheWriteInputTokenCount,
619
+ const inputTokens = usage.inputTokens
620
+ const outputTokens = usage.outputTokens
621
+ const cacheReadTokens = usage.cacheReadInputTokens ?? usage.cacheReadInputTokenCount
622
+ const cacheWriteTokens = usage.cacheWriteInputTokens ?? usage.cacheWriteInputTokenCount
623
+
624
+ if (inputTokens === undefined && outputTokens === undefined &&
625
+ cacheReadTokens === undefined && cacheWriteTokens === undefined) {
626
+ return
577
627
  }
628
+
629
+ return { inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens }
578
630
  }
579
631
 
580
632
  /**
@@ -714,6 +766,8 @@ function extractTextAndResponseReasonConverseFromStream (chunks) {
714
766
 
715
767
  module.exports = {
716
768
  Generation,
769
+ buildUsage,
770
+ mergeStreamedUsage,
717
771
  RequestParams,
718
772
  extractTextAndResponseReasonFromStream,
719
773
  parseModelId,
@@ -449,7 +449,8 @@ class CucumberPlugin extends CiPlugin {
449
449
  })
450
450
 
451
451
  this.addSub('ci:cucumber:is-modified-test', ({
452
- scenarios,
452
+ gherkinNodes,
453
+ gherkinScopeRanges,
453
454
  testFileAbsolutePath,
454
455
  modifiedFiles,
455
456
  stepIds,
@@ -457,11 +458,25 @@ class CucumberPlugin extends CiPlugin {
457
458
  setIsModified,
458
459
  }) => {
459
460
  const testScenarioPath = getTestSuitePath(testFileAbsolutePath, this.repositoryRoot || process.cwd())
460
- for (const scenario of scenarios) {
461
+ for (const [startLine, endLine] of gherkinScopeRanges) {
462
+ if (isModifiedTest(testScenarioPath, startLine, endLine, modifiedFiles, 'cucumber')) {
463
+ setIsModified(true)
464
+ return
465
+ }
466
+ }
467
+ for (const gherkinNode of gherkinNodes) {
468
+ const lastStep = gherkinNode.steps.at(-1)
469
+ let endLine = lastStep?.location.line ?? gherkinNode.location.line
470
+ if (lastStep?.dataTable?.rows.length) {
471
+ endLine = lastStep.dataTable.rows.at(-1).location.line
472
+ } else if (lastStep?.docString) {
473
+ const { content, location } = lastStep.docString
474
+ endLine = location.line + (content ? content.split('\n').length + 1 : 1)
475
+ }
461
476
  const isModified = isModifiedTest(
462
477
  testScenarioPath,
463
- scenario.location.line,
464
- scenario.steps.at(-1).location.line,
478
+ gherkinNode.location.line,
479
+ endLine,
465
480
  modifiedFiles,
466
481
  'cucumber'
467
482
  )
@@ -94,6 +94,7 @@ const {
94
94
  const { ORIGIN_KEY, COMPONENT } = require('../../dd-trace/src/constants')
95
95
  const { RESOURCE_NAME } = require('../../../ext/tags')
96
96
  const getConfig = require('../../dd-trace/src/config')
97
+ const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
97
98
  const {
98
99
  SCREENSHOT_UPLOAD_RESULT_ERROR,
99
100
  SCREENSHOT_UPLOAD_RESULT_UPLOADED,
@@ -857,6 +858,12 @@ class CypressPlugin {
857
858
 
858
859
  const testOptimizationConfig = getConfig().testOptimization
859
860
  this.rumFlushWaitMillis = testOptimizationConfig.DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS
861
+ if (getValueFromEnvSources('DD_TEST_FAILURE_SCREENSHOTS_ENABLED', true) === true) {
862
+ cypressConfig.screenshotOnRunFailure = true
863
+ }
864
+ if (getValueFromEnvSources('DD_TEST_FAILURE_VIDEOS_ENABLED', true) === true) {
865
+ cypressConfig.video = true
866
+ }
860
867
  this.warnIfMisconfiguredTestFailureScreenshots(cypressConfig, tracer, testOptimizationConfig)
861
868
  this.warnIfMisconfiguredTestFailureVideos(cypressConfig, tracer, testOptimizationConfig)
862
869
 
@@ -17,6 +17,9 @@ class FetchPlugin extends HttpClientPlugin {
17
17
 
18
18
  const store = super.bindStart(ctx)
19
19
 
20
+ // Native undici diagnostics also fire for requests already traced by this wrapper.
21
+ store.skipUndici = true
22
+
20
23
  for (const name in options.headers) {
21
24
  if (!req.headers.has(name)) {
22
25
  req.headers.set(name, options.headers[name])
@@ -3,6 +3,8 @@
3
3
  const { URL } = require('url')
4
4
  const log = require('../../dd-trace/src/log')
5
5
 
6
+ const easyConnectProtocol = /^tcps?:\/\//i
7
+
6
8
  function parseOracleDescriptor (descriptor) {
7
9
  const hostnameMatch = descriptor.match(/HOST\s*=\s*([^)]+)/i)
8
10
  const hostname = hostnameMatch?.[1] || 'localhost' // Default Oracle hostname
@@ -24,7 +26,7 @@ module.exports = function getDBInformation (connAttrs) {
24
26
  return parseOracleDescriptor(connectString)
25
27
  }
26
28
  try {
27
- const url = new URL(`oracle://${connectString}`)
29
+ const url = new URL(easyConnectProtocol.test(connectString) ? connectString : `oracle://${connectString}`)
28
30
  return {
29
31
  hostname: url.hostname || 'localhost', // Default Oracle hostname
30
32
  port: url.port || '1521', // Default Oracle port
@@ -47,6 +47,8 @@ class UndiciPlugin extends HttpClientPlugin {
47
47
  if (!request) return
48
48
 
49
49
  const store = storage('legacy').getStore()
50
+ if (store?.skipUndici) return
51
+
50
52
  const { origin = '', path = '/' } = request
51
53
  const method = request.method?.toUpperCase() ?? 'GET'
52
54
 
@@ -207,6 +209,9 @@ class UndiciPlugin extends HttpClientPlugin {
207
209
 
208
210
  const store = super.bindStart(ctx)
209
211
 
212
+ // Undici 5.0 emits native diagnostics even though its fetch needs the legacy wrapper.
213
+ store.skipUndici = true
214
+
210
215
  // Inject trace headers back into the request
211
216
  for (const name of Object.keys(options.headers)) {
212
217
  if (!req.headers.has(name)) {
@@ -1,10 +1,15 @@
1
1
  'use strict'
2
2
 
3
3
  const tracerVersion = require('../../../../package.json').version
4
+ const { createAgents } = require('../exporters/common/agents')
5
+ const request = require('../exporters/common/request')
4
6
  const { AIGuardClientError } = require('./errors')
5
7
  const { parseEvaluationResponse } = require('./evaluation')
6
8
  const TAGS = require('./tags')
7
9
 
10
+ // Evaluations can run concurrently on the application's request path.
11
+ const { httpAgent, httpsAgent } = createAgents(16)
12
+
8
13
  /**
9
14
  * Resolves the AI Guard host for a Datadog site.
10
15
  *
@@ -18,29 +23,43 @@ function aiGuardHost (site) {
18
23
  * Sends a request to the AI Guard service.
19
24
  *
20
25
  * @param {object} body
21
- * @param {{ url: string, headers: Record<string, string>, timeout: number }} opts
26
+ * @param {{ url: string, headers: Record<string, string|undefined>, timeout: number }} opts
22
27
  * @returns {Promise<{ status: number, body: unknown }>}
23
28
  */
24
- async function executeRequest (body, opts) {
25
- const postData = JSON.stringify(body)
26
- const headers = {
27
- 'Content-Type': 'application/json',
28
- 'Content-Length': Buffer.byteLength(postData),
29
- ...opts.headers,
30
- }
29
+ function executeRequest (body, opts) {
30
+ return new Promise((resolve, reject) => {
31
+ const postData = JSON.stringify(body)
32
+ const url = new URL(opts.url)
33
+ request(postData, {
34
+ url,
35
+ method: 'POST',
36
+ headers: {
37
+ 'Content-Type': 'application/json',
38
+ ...opts.headers,
39
+ },
40
+ agent: url.protocol === 'https:' ? httpsAgent : httpAgent,
41
+ timeout: opts.timeout,
42
+ signal: AbortSignal.timeout(opts.timeout),
43
+ retry: false,
44
+ includeErrorResponseBody: true,
45
+ }, (error, result, status) => {
46
+ if (status === undefined) {
47
+ reject(error || new Error('AI Guard request completed without a status'))
48
+ return
49
+ }
31
50
 
32
- const response = await fetch(opts.url, {
33
- method: 'POST',
34
- headers,
35
- body: postData,
36
- signal: AbortSignal.timeout(opts.timeout),
51
+ let responseBody
52
+ try {
53
+ const responseText = error ? /** @type {Error & { responseBody?: string }} */ (error).responseBody : result
54
+ if (typeof responseText !== 'string') throw new TypeError('AI Guard response body is missing')
55
+ responseBody = JSON.parse(responseText)
56
+ } catch (cause) {
57
+ reject(cause)
58
+ return
59
+ }
60
+ resolve({ status, body: responseBody })
61
+ })
37
62
  })
38
-
39
- const responseBody = await response.json()
40
- return {
41
- status: response.status,
42
- body: responseBody,
43
- }
44
63
  }
45
64
 
46
65
  class AIGuardClient {
@@ -2,10 +2,15 @@
2
2
 
3
3
  const { channel } = require('dc-polyfill')
4
4
 
5
- const { getMessagesInputMessages, getMessagesOutputMessages } = require('../messages/anthropic')
5
+ const {
6
+ getMessagesInputMessages,
7
+ getMessagesOutputMessages,
8
+ getStreamedMessagesOutputMessages,
9
+ } = require('../messages/anthropic')
6
10
  const { decode } = require('../messages/utils')
7
11
  const { SOURCE_AUTO } = require('../tags')
8
12
  const { evaluate } = require('./evaluate')
13
+ const { interceptStream } = require('./stream')
9
14
 
10
15
  const messagesPrepareChannel = channel('dd-trace:anthropic:messages:prepare')
11
16
  const messagesInterceptChannel = channel('dd-trace:anthropic:messages:intercept')
@@ -13,18 +18,21 @@ const messagesInterceptChannel = channel('dd-trace:anthropic:messages:intercept'
13
18
  let isEnabled = false
14
19
  let aiguard
15
20
  let opts
21
+ let analyzeStreamResponses
16
22
 
17
23
  /**
18
24
  * Subscribes AI Guard to the Anthropic interception channels.
19
25
  *
20
26
  * @param {object} aiguardInstance
21
27
  * @param {boolean} block
28
+ * @param {boolean} analyzeStreams
22
29
  */
23
- function enable (aiguardInstance, block) {
30
+ function enable (aiguardInstance, block, analyzeStreams) {
24
31
  if (isEnabled) return
25
32
 
26
33
  aiguard = aiguardInstance
27
34
  opts = { block, source: SOURCE_AUTO, integration: 'anthropic' }
35
+ analyzeStreamResponses = analyzeStreams
28
36
 
29
37
  messagesPrepareChannel.subscribe(onMessagesPrepare)
30
38
  messagesInterceptChannel.subscribe(onMessagesIntercept)
@@ -40,6 +48,7 @@ function disable () {
40
48
 
41
49
  aiguard = undefined
42
50
  opts = undefined
51
+ analyzeStreamResponses = undefined
43
52
  isEnabled = false
44
53
  }
45
54
 
@@ -76,19 +85,35 @@ function onMessagesIntercept (ctx) {
76
85
  return inputEvaluation
77
86
  }
78
87
 
88
+ const isStream = ctx.arguments[0].stream
89
+ if (isStream && !analyzeStreamResponses) return
90
+
79
91
  // One model call has one output however many readers observe it: `parse`, `json()`, `text()`
80
92
  // and every `clone()` of the raw response share this callback.
81
93
  let outputEvaluation
82
94
  ctx.onResult = body => {
83
95
  if (!isEnabled) return body
84
96
 
97
+ if (isStream) {
98
+ outputEvaluation ??= interceptStream(body, chunks => {
99
+ if (!isEnabled) return
100
+ const outputMessages = decode(
101
+ () => getStreamedMessagesOutputMessages(chunks),
102
+ null,
103
+ 'AIGuard: unable to decode the streamed Anthropic response: %s'
104
+ )
105
+ if (!outputMessages?.length) return
106
+ return evaluate(ctx, aiguard, [[...inputMessages, ...outputMessages]], opts)
107
+ })
108
+ return outputEvaluation
109
+ }
110
+
85
111
  const outputMessages = decode(
86
112
  () => getMessagesOutputMessages(body),
87
113
  null,
88
114
  'AIGuard: unable to decode Anthropic response body: %s'
89
115
  )
90
- if (outputMessages === null) return body
91
- if (!outputMessages.length) return body
116
+ if (!outputMessages?.length) return body
92
117
 
93
118
  outputEvaluation ??= evaluate(ctx, aiguard, [[...inputMessages, ...outputMessages]], opts)
94
119
  return outputEvaluation.then(() => body)
@@ -16,7 +16,7 @@ let isEnabled = false
16
16
  function enable (aiguard, block, analyzeStreamResponses) {
17
17
  if (isEnabled) return
18
18
 
19
- anthropic.enable(aiguard, block)
19
+ anthropic.enable(aiguard, block, analyzeStreamResponses)
20
20
  openai.enable(aiguard, block, analyzeStreamResponses)
21
21
  vercelAi.enable(aiguard, block, analyzeStreamResponses)
22
22
 
@@ -13,6 +13,7 @@ const {
13
13
  const { decode } = require('../messages/utils')
14
14
  const { SOURCE_AUTO } = require('../tags')
15
15
  const { evaluate } = require('./evaluate')
16
+ const { interceptStream } = require('./stream')
16
17
 
17
18
  const chatCompletionsInterceptChannel = channel('dd-trace:openai:chat.completions:intercept')
18
19
  const responsesInterceptChannel = channel('dd-trace:openai:responses:intercept')
@@ -74,9 +75,16 @@ function onChatCompletions (ctx) {
74
75
 
75
76
  if (isStream) {
76
77
  outputEvaluation ??= interceptStream(body, chunks => {
77
- return getStreamedChatCompletionsOutputMessages(chunks)
78
- .map(message => [...inputMessages, message])
79
- }, ctx)
78
+ if (!isEnabled) return
79
+ const conversations = decode(
80
+ () => getStreamedChatCompletionsOutputMessages(chunks)
81
+ .map(message => [...inputMessages, message]),
82
+ null,
83
+ 'AIGuard: unable to decode the streamed OpenAI response: %s'
84
+ )
85
+ if (!conversations?.length) return
86
+ return evaluate(ctx, aiguard, conversations, opts)
87
+ })
80
88
  return outputEvaluation
81
89
  }
82
90
 
@@ -112,9 +120,15 @@ function onResponses (ctx) {
112
120
 
113
121
  if (isStream) {
114
122
  outputEvaluation ??= interceptStream(body, chunks => {
115
- const outputMessages = getStreamedResponsesOutputMessages(chunks)
116
- return outputMessages.length ? [[...inputMessages, ...outputMessages]] : []
117
- }, ctx)
123
+ if (!isEnabled) return
124
+ const outputMessages = decode(
125
+ () => getStreamedResponsesOutputMessages(chunks),
126
+ null,
127
+ 'AIGuard: unable to decode the streamed OpenAI response: %s'
128
+ )
129
+ if (!outputMessages?.length) return
130
+ return evaluate(ctx, aiguard, [[...inputMessages, ...outputMessages]], opts)
131
+ })
118
132
  return outputEvaluation
119
133
  }
120
134
 
@@ -130,54 +144,4 @@ function onResponses (ctx) {
130
144
  }
131
145
  }
132
146
 
133
- /**
134
- * Uses the OpenAI SDK's stream splitting support to evaluate one branch and return the other.
135
- *
136
- * @param {object} stream
137
- * @param {(chunks: Array<object>) => Array<Array<object>>} getConversations
138
- * @param {object} ctx
139
- * @returns {object|Promise<object>}
140
- */
141
- function interceptStream (stream, getConversations, ctx) {
142
- if (typeof stream?.tee !== 'function') return stream
143
-
144
- let branches
145
- try {
146
- branches = stream.tee()
147
- } catch {
148
- return stream
149
- }
150
-
151
- const [evaluationStream, resultStream] = branches
152
- return drainStream(evaluationStream).then(chunks => {
153
- if (!isEnabled) return resultStream
154
- const conversations = decode(
155
- () => getConversations(chunks),
156
- null,
157
- 'AIGuard: unable to decode the streamed OpenAI response: %s'
158
- )
159
- if (!conversations?.length) return resultStream
160
- return evaluate(ctx, aiguard, conversations, opts).then(() => resultStream)
161
- }, () => resultStream)
162
- }
163
-
164
- /**
165
- * @param {object} stream
166
- * @returns {Promise<Array<object>>}
167
- */
168
- function drainStream (stream) {
169
- const chunks = []
170
- const iterator = stream[Symbol.asyncIterator]()
171
-
172
- function readAll () {
173
- return iterator.next().then(({ done, value }) => {
174
- if (done) return chunks
175
- chunks.push(value)
176
- return readAll()
177
- })
178
- }
179
-
180
- return readAll()
181
- }
182
-
183
147
  module.exports = { enable, disable }