fixup

Author: IlyasShabi

packages/datadog-instrumentations/src/anthropic.js

@@ -47,16 +47,38 @@ function waitForVerdict (promise, verdict) {
 /**
  * @param {object} response
  * @param {'json'|'text'} method
+ * @param {object} ctx
  * @param {(body: object|string) => Promise<void>|undefined} getVerdict
  */
-function wrapResponseReader (response, method, getVerdict) {
+function wrapResponseReader (response, method, ctx, getVerdict) {
   if (typeof response[method] !== 'function') return
 
   shimmer.wrap(response, method, original => function (...args) {
-    return original.apply(this, args).then(body => {
-      const verdict = getVerdict(body)
-      return verdict ? verdict.then(() => body) : body
-    })
+    return original.apply(this, args)
+      .then(body => finishResult(ctx, body, getVerdict))
+      .catch(error => {
+        if (!ctx.finished) finish(ctx, null, error)
+        throw error
+      })
+  })
+}
+
+/**
+ * @param {object} ctx
+ * @param {object|string} result
+ * @param {(body: object|string) => Promise<void>|undefined} getVerdict
+ * @returns {object|string|Promise<object|string>}
+ */
+function finishResult (ctx, result, getVerdict) {
+  const verdict = getVerdict(result)
+  if (!verdict) {
+    finish(ctx, result, null)
+    return result
+  }
+
+  return verdict.then(() => {
+    finish(ctx, result, null)
+    return result
   })
 }
 
@@ -143,17 +165,7 @@ function wrapCreate (create) {
               shimmer.wrap(response, Symbol.asyncIterator, iterator => wrapStreamIterator(iterator, ctx))
               return response
             }
-            const verdict = getAfterVerdict(response)
-            if (!verdict) {
-              finish(ctx, response, null)
-              return response
-            }
-            // Finish after evaluation so a block propagates the error to anthropic.request
-            // and the span wraps its child instead of closing before it.
-            return verdict.then(() => {
-              finish(ctx, response, null)
-              return response
-            })
+            return finishResult(ctx, response, getAfterVerdict)
           }).catch(error => {
             if (!ctx.finished) finish(ctx, null, error)
             throw error
@@ -162,19 +174,17 @@ function wrapCreate (create) {
         return parseResult
       })
 
-      // Gate `.asResponse()` callers on the before verdict so raw-response paths still block,
-      // then evaluate output only if the caller consumes the JSON body.
+      // Gate `.asResponse()` callers on the before verdict, then finish with a supported body reader's result.
       shimmer.wrap(apiPromise, 'asResponse', origAsResponse => function (...asResponseArgs) {
         return waitForVerdict(origAsResponse.apply(this, asResponseArgs), getBeforeVerdict())
           .then(response => {
-            if (!stream && hasLifecycle && wrappedResponse !== response) {
+            if (!stream && wrappedResponse !== response) {
               wrappedResponse = response
-              wrapResponseReader(response, 'json', getAfterVerdict)
-              wrapResponseReader(response, 'text', getAfterVerdict)
+              wrapResponseReader(response, 'json', ctx, getAfterVerdict)
+              wrapResponseReader(response, 'text', ctx, getAfterVerdict)
             }
 
             if (afterVerdict) return afterVerdict.then(() => response)
-            if (!stream && !ctx.finished && !parseResult) finish(ctx, null, null)
             return response
           })
           .catch(error => {

packages/datadog-instrumentations/test/anthropic-lifecycle.spec.js

@@ -260,7 +260,7 @@ describe('anthropic lifecycle instrumentation', () => {
     }
   })
 
-  it('finishes after a resolving before subscriber leaves', async () => {
+  it('finishes after reading a raw response when the before subscriber leaves', async () => {
     const apmChannel = tracingChannel('apm:anthropic:request')
     const apmHandlers = { start () {} }
     let asyncEndCount = 0
@@ -282,18 +282,19 @@ describe('anthropic lifecycle instrumentation', () => {
     messages._nextApiPromise = apiPromise
 
     try {
-      assert.strictEqual(
-        await messages.create({ messages: [{ role: 'user', content: 'Hi' }] }).asResponse(),
-        apiPromise._rawResponse
-      )
+      const response = await messages.create({ messages: [{ role: 'user', content: 'Hi' }] }).asResponse()
+
+      assert.strictEqual(response, apiPromise._rawResponse)
+      assert.strictEqual(asyncEndCount, 0)
+      await response.json()
       assert.strictEqual(asyncEndCount, 1)
     } finally {
       apmChannel.unsubscribe(apmHandlers)
       unsubscribe()
     }
   })
 
-  it('finishes an unconsumed raw response and evaluates it when read', async () => {
+  it('finishes a raw response with its result when read', async () => {
     const apmChannel = tracingChannel('apm:anthropic:request')
     const apmHandlers = { start () {} }
     let asyncEndCtx
@@ -312,18 +313,42 @@ describe('anthropic lifecycle instrumentation', () => {
       const response = await messages.create({ messages: [{ role: 'user', content: 'Hi' }] }).asResponse()
 
       assert.strictEqual(calls.length, 1)
-      assert.ok(asyncEndCtx, 'asyncEnd was not published')
-      assert.strictEqual(asyncEndCtx.finished, true)
+      assert.strictEqual(asyncEndCtx, undefined)
       assert.strictEqual(response.bodyUsed, false)
       assert.deepStrictEqual(await response.json(), body)
       assert.strictEqual(calls.length, 2)
       assert.deepStrictEqual(calls[1].body, body)
+      assert.ok(asyncEndCtx, 'asyncEnd was not published')
+      assert.strictEqual(asyncEndCtx.finished, true)
+      assert.deepStrictEqual(asyncEndCtx.result, body)
     } finally {
       apmChannel.unsubscribe(apmHandlers)
       unsubscribe()
     }
   })
 
+  it('finishes a raw response reader when only tracing is active', async () => {
+    const apmChannel = tracingChannel('apm:anthropic:request')
+    const apmHandlers = { start () {} }
+    let asyncEndCtx
+    apmHandlers.asyncEnd = ctx => { asyncEndCtx = ctx }
+    apmChannel.subscribe(apmHandlers)
+
+    const body = { role: 'assistant', content: [{ type: 'text', text: 'Hi' }] }
+    const messages = new Messages()
+    messages._nextApiPromise = new FakeAPIPromise(body)
+
+    try {
+      const response = await messages.create({ messages: [{ role: 'user', content: 'Hi' }] }).asResponse()
+
+      assert.strictEqual(asyncEndCtx, undefined)
+      assert.deepStrictEqual(await response.json(), body)
+      assert.deepStrictEqual(asyncEndCtx.result, body)
+    } finally {
+      apmChannel.unsubscribe(apmHandlers)
+    }
+  })
+
   it('evaluates repeated asResponse() calls once', async () => {
     const { calls, unsubscribe } = subscribeAutoResolve([
       messagesBeforeChannel,
@@ -377,11 +402,12 @@ describe('anthropic lifecycle instrumentation', () => {
       )
       sinon.assert.notCalled(response.clone)
       assert.strictEqual(calls.length, 0)
-      assert.strictEqual(asyncEndCount, 1)
+      assert.strictEqual(asyncEndCount, 0)
       assert.deepStrictEqual(await response.json(), body)
       sinon.assert.calledOnce(readJson)
       assert.strictEqual(calls.length, 1)
       assert.strictEqual(calls[0].body, body)
+      assert.strictEqual(asyncEndCount, 1)
     } finally {
       apmChannel.unsubscribe(apmHandlers)
       unsubscribe()
@@ -508,22 +534,33 @@ describe('anthropic lifecycle instrumentation', () => {
     }
   })
 
-  it('keeps a raw-terminal span finished when parsing starts later', async () => {
+  it('retains the parsed result when parsing starts after raw response access', async () => {
     const apmChannel = tracingChannel('apm:anthropic:request')
     const apmHandlers = { start () {} }
     let asyncEndCount = 0
-    apmHandlers.asyncEnd = () => { asyncEndCount++ }
+    let asyncEndCtx
+    apmHandlers.asyncEnd = ctx => {
+      asyncEndCount++
+      asyncEndCtx = ctx
+    }
     apmChannel.subscribe(apmHandlers)
 
     const { unsubscribe } = subscribeAutoResolve([messagesAfterChannel])
+    const body = {
+      role: 'assistant',
+      content: [{ type: 'text', text: 'Hi' }],
+      usage: { input_tokens: 1, output_tokens: 2 },
+    }
     const messages = new Messages()
-    messages._nextApiPromise = new FakeAPIPromise({ role: 'assistant', content: [] })
+    messages._nextApiPromise = new FakeAPIPromise(body)
     const apiPromise = messages.create({ messages: [{ role: 'user', content: 'Hi' }] })
 
     try {
       await apiPromise.asResponse()
-      await apiPromise.parse()
+      assert.strictEqual(asyncEndCount, 0)
+      assert.strictEqual(await apiPromise.parse(), body)
       assert.strictEqual(asyncEndCount, 1)
+      assert.strictEqual(asyncEndCtx.result, body)
     } finally {
       apmChannel.unsubscribe(apmHandlers)
       unsubscribe()

packages/dd-trace/src/aiguard/messages/anthropic.js

@@ -228,9 +228,10 @@ function convertServerToolResultContent (content) {
 /**
  * Converts a single Anthropic message to zero or more normalized messages.
  * Assistant `tool_use` blocks become an assistant `tool_calls` message.
- * User `tool_result` blocks become one `tool` message per block, emitted
- * before any accompanying text so the chat-style timeline is preserved.
- * Text/image blocks are merged into a single message per role.
+ * Tool result blocks become one `tool` message per block. When the message also carries a tool call
+ * (built-in server tools return the call and its result together), the assistant tool-call message
+ * precedes the results; otherwise (user turns) the results precede any accompanying text so the
+ * chat-style timeline is preserved. Text/image blocks are merged into a single message per role.
  *
  * @param {{role: string, content: string|Array<AnthropicContentBlock>}} message
  * @returns {Array<object>}
@@ -245,16 +246,21 @@ function convertAnthropicMessage (message) {
   if (!Array.isArray(content)) return []
 
   const { parts, toolCalls, toolResults, hasImages } = walkContentBlocks(content)
-  const messages = [...toolResults]
   const messageContent = partsToContent(parts, hasImages)
 
+  let assistantMessage
   if (messageContent != null || toolCalls.length) {
-    const converted = { role }
-    if (messageContent != null) converted.content = messageContent
-    if (toolCalls.length) converted.tool_calls = toolCalls
-    messages.push(converted)
+    assistantMessage = { role }
+    if (messageContent != null) assistantMessage.content = messageContent
+    if (toolCalls.length) assistantMessage.tool_calls = toolCalls
+  }
+
+  if (toolCalls.length) {
+    return [assistantMessage, ...toolResults]
   }
 
+  const messages = [...toolResults]
+  if (assistantMessage) messages.push(assistantMessage)
   return messages
 }
 

packages/dd-trace/test/aiguard/messages/anthropic.spec.js

@@ -207,6 +207,33 @@ describe('aiguard/messages/anthropic', () => {
   })
 
   describe('built-in server-tool blocks', () => {
+    it('emits the server tool call before its matching tool result', () => {
+      const message = {
+        role: 'assistant',
+        content: [
+          { type: 'server_tool_use', id: 'srv_1', name: 'web_search', input: { query: 'datadog' } },
+          {
+            type: 'web_search_tool_result',
+            tool_use_id: 'srv_1',
+            content: [{ type: 'web_search_result', title: 'Datadog', url: 'https://datadoghq.com' }],
+          },
+          { type: 'text', text: 'Based on the search, Datadog is an observability platform.' },
+        ],
+      }
+      assert.deepStrictEqual(convertAnthropicMessage(message), [
+        {
+          role: 'assistant',
+          content: 'Based on the search, Datadog is an observability platform.',
+          tool_calls: [{ id: 'srv_1', function: { name: 'web_search', arguments: '{"query":"datadog"}' } }],
+        },
+        {
+          role: 'tool',
+          tool_call_id: 'srv_1',
+          content: 'Datadog\nhttps://datadoghq.com',
+        },
+      ])
+    })
+
     it('converts server_tool_use blocks to tool_calls', () => {
       const message = {
         role: 'assistant',