fix(ai): fix infinite recursion when using BedrockChatLanguageModel (#9487)

Author: sabrenner

packages/datadog-instrumentations/src/ai.js

@@ -230,7 +230,13 @@ for (const hook of getHooks('ai')) {
     // generateObject, streamObject)
     tracingChannel('orchestrion:ai:resolveLanguageModel').subscribe({
       end (ctx) {
-        wrapModelWithLifecycle(ctx.result)
+        const model = ctx.arguments[0]
+        if (typeof model !== 'string' && model !== ctx.result) {
+          wrapModelWithLifecycle(model)
+          wrappedModels.add(ctx.result)
+        } else {
+          wrapModelWithLifecycle(ctx.result)
+        }
       },
     })
 

packages/datadog-plugin-ai/test/integration-test/bedrock-recursion.mjs

@@ -0,0 +1,38 @@
+import assert from 'node:assert/strict'
+import { webcrypto } from 'node:crypto'
+
+import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'
+import { generateText } from 'ai'
+
+globalThis.crypto ??= webcrypto
+
+const response = {
+  stopReason: 'end_turn',
+  output: {
+    message: {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Reproduction completed.' }],
+    },
+  },
+  usage: {
+    inputTokens: 3,
+    outputTokens: 3,
+    totalTokens: 6,
+  },
+  metrics: { latencyMs: 1 },
+}
+
+const bedrock = createAmazonBedrock({
+  region: 'us-east-1',
+  fetch: () => Promise.resolve(new Response(JSON.stringify(response), {
+    status: 200,
+    headers: { 'Content-Type': 'application/json' },
+  })),
+})
+
+const result = await generateText({
+  model: bedrock('anthropic.claude-3-haiku-20240307-v1:0'),
+  prompt: 'Run the Bedrock recursion reproduction.',
+})
+
+assert.strictEqual(result.text, 'Reproduction completed.')

packages/datadog-plugin-ai/test/integration-test/bedrock-recursion.spec.js

@@ -0,0 +1,62 @@
+'use strict'
+
+const assert = require('node:assert/strict')
+const { inspect } = require('node:util')
+
+const {
+  FakeAgent,
+  sandboxCwd,
+  useSandbox,
+  spawnPluginIntegrationTestProcAndExpectExit,
+  stopProc,
+} = require('../../../../integration-tests/helpers')
+const { withVersions } = require('../../../dd-trace/test/setup/mocha')
+
+describe('Bedrock recursion regression', () => {
+  let agent
+  let proc
+
+  withVersions('ai', 'ai', '>=6.0.0 <7.0.0', aiVersion => {
+    withVersions('ai', '@ai-sdk/amazon-bedrock', '^3.0.0', bedrockVersion => {
+      useSandbox([
+        `ai@${aiVersion}`,
+        `@ai-sdk/amazon-bedrock@${bedrockVersion}`,
+      ], false, [
+        './packages/datadog-plugin-ai/test/integration-test/bedrock-recursion.mjs',
+      ])
+
+      beforeEach(async () => {
+        agent = await new FakeAgent().start()
+      })
+
+      afterEach(async () => {
+        await stopProc(proc)
+        await agent.stop()
+      })
+
+      it('does not recurse when AI SDK adapts a v2 Bedrock model', async () => {
+        const received = agent.assertMessageReceived(({ payload }) => {
+          assert.ok(Array.isArray(payload), `Expected array, got ${inspect(payload)}`)
+          assert.ok(payload.flat().some(span => span.name === 'ai.generateText'))
+        })
+
+        proc = await spawnPluginIntegrationTestProcAndExpectExit(
+          sandboxCwd(),
+          'bedrock-recursion.mjs',
+          agent.port,
+          {
+            AWS_ACCESS_KEY_ID: 'test-access-key',
+            AWS_SECRET_ACCESS_KEY: 'test-secret-key',
+            AWS_REGION: 'us-east-1',
+            DD_LLMOBS_ENABLED: '1',
+            DD_LLMOBS_ML_APP: 'test',
+            NODE_OPTIONS: '--import dd-trace/initialize.mjs',
+          },
+          ['--stack-size=128']
+        )
+
+        await received
+      }).timeout(20000)
+    })
+  })
+})