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
   })
 }
 
@@ -103,6 +125,7 @@ function wrapCreate (create) {
 
       let apiPromise
       try {
+        // Anthropic starts the request eagerly; the input verdict only gates result delivery.
         apiPromise = create.apply(this, args)
       } catch (error) {
         finish(ctx, null, error)
@@ -143,17 +166,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 +175,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 raw access; output evaluation supports only the common json() and text() readers.
       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()