Differentiate between handling of stream: true and client.messages.stream

andreiborza · andreiborza · commit 35a1a0235465 · 2025-09-25T13:42:52.000+02:00
diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario.mjs
@@ -29,6 +29,55 @@ function startMockAnthropicServer() {
       return;
     }
 
+    // Check if streaming is requested
+    if (req.body.stream === true) {
+      res.writeHead(200, {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache',
+        Connection: 'keep-alive',
+      });
+
+      // Send streaming events
+      const events = [
+        {
+          type: 'message_start',
+          message: {
+            id: 'msg_stream123',
+            type: 'message',
+            role: 'assistant',
+            model,
+            content: [],
+            usage: { input_tokens: 10 },
+          },
+        },
+        { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } },
+        { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Hello ' } },
+        { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'from ' } },
+        { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'stream!' } },
+        { type: 'content_block_stop', index: 0 },
+        {
+          type: 'message_delta',
+          delta: { stop_reason: 'end_turn', stop_sequence: null },
+          usage: { output_tokens: 15 },
+        },
+        { type: 'message_stop' },
+      ];
+
+      events.forEach((event, index) => {
+        setTimeout(() => {
+          res.write(`event: ${event.type}\n`);
+          res.write(`data: ${JSON.stringify(event)}\n\n`);
+
+          if (index === events.length - 1) {
+            res.end();
+          }
+        }, index * 10); // Small delay between events
+      });
+
+      return;
+    }
+
+    // Non-streaming response
     res.send({
       id: 'msg_mock123',
       type: 'message',
@@ -92,8 +141,32 @@ async function run() {
 
     // Fourth test: models.retrieve
     await client.models.retrieve('claude-3-haiku-20240307');
+
+    // Fifth test: streaming via messages.create
+    const stream = await client.messages.create({
+      model: 'claude-3-haiku-20240307',
+      messages: [{ role: 'user', content: 'What is the capital of France?' }],
+      stream: true,
+    });
+
+    for await (const _ of stream) {
+      void _;
+    }
+
+    // Sixth test: streaming via messages.stream
+    await client.messages
+      .stream({
+        model: 'claude-3-haiku-20240307',
+        messages: [{ role: 'user', content: 'What is the capital of France?' }],
+      })
+      .on('streamEvent', () => {
+        Sentry.captureMessage('stream event from user-added event listener captured');
+      });
   });
 
+  // Wait for the stream event handler to finish
+  await Sentry.flush(2000);
+
   server.close();
 }
 
diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts b/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts
@@ -152,6 +152,30 @@ describe('Anthropic integration', () => {
         origin: 'auto.ai.anthropic',
         status: 'ok',
       }),
+      // Fifth span - messages.create with stream: true
+      expect.objectContaining({
+        data: expect.objectContaining({
+          'gen_ai.operation.name': 'messages',
+          'gen_ai.request.model': 'claude-3-haiku-20240307',
+          'gen_ai.request.stream': true,
+        }),
+        description: 'messages claude-3-haiku-20240307 stream-response',
+        op: 'gen_ai.messages',
+        origin: 'auto.ai.anthropic',
+        status: 'ok',
+      }),
+      // Sixth span - messages.stream
+      expect.objectContaining({
+        data: expect.objectContaining({
+          'gen_ai.operation.name': 'messages',
+          'gen_ai.request.model': 'claude-3-haiku-20240307',
+          'gen_ai.request.stream': true,
+        }),
+        description: 'messages claude-3-haiku-20240307 stream-response',
+        op: 'gen_ai.messages',
+        origin: 'auto.ai.anthropic',
+        status: 'ok',
+      }),
     ]),
   };
 
@@ -189,6 +213,21 @@ describe('Anthropic integration', () => {
     ]),
   };
 
+  const EXPECTED_MODEL_ERROR = {
+    exception: {
+      values: [
+        {
+          type: 'Error',
+          value: '404 Model not found',
+        },
+      ],
+    },
+  };
+
+  const EXPECTED_STREAM_EVENT_HANDLER_MESSAGE = {
+    message: 'stream event from user-added event listener captured',
+  };
+
   createEsmAndCjsTests(__dirname, 'scenario-manual-client.mjs', 'instrument.mjs', (createRunner, test) => {
     test('creates anthropic related spans when manually insturmenting client', async () => {
       await createRunner()
@@ -202,8 +241,9 @@ describe('Anthropic integration', () => {
   createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createRunner, test) => {
     test('creates anthropic related spans with sendDefaultPii: false', async () => {
       await createRunner()
-        .ignore('event')
+        .expect({ event: EXPECTED_MODEL_ERROR })
         .expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE })
+        .expect({ event: EXPECTED_STREAM_EVENT_HANDLER_MESSAGE })
         .start()
         .completed();
     });
@@ -212,8 +252,9 @@ describe('Anthropic integration', () => {
   createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
     test('creates anthropic related spans with sendDefaultPii: true', async () => {
       await createRunner()
-        .ignore('event')
+        .expect({ event: EXPECTED_MODEL_ERROR })
         .expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE })
+        .expect({ event: EXPECTED_STREAM_EVENT_HANDLER_MESSAGE })
         .start()
         .completed();
     });
@@ -222,8 +263,9 @@ describe('Anthropic integration', () => {
   createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-with-options.mjs', (createRunner, test) => {
     test('creates anthropic related spans with custom options', async () => {
       await createRunner()
-        .ignore('event')
+        .expect({ event: EXPECTED_MODEL_ERROR })
         .expect({ transaction: EXPECTED_TRANSACTION_WITH_OPTIONS })
+        .expect({ event: EXPECTED_STREAM_EVENT_HANDLER_MESSAGE })
         .start()
         .completed();
     });
diff --git a/packages/core/src/utils/anthropic-ai/index.ts b/packages/core/src/utils/anthropic-ai/index.ts
@@ -25,8 +25,14 @@ import {
 } from '../ai/gen-ai-attributes';
 import { buildMethodPath, getFinalOperationName, getSpanOperation, setTokenUsageAttributes } from '../ai/utils';
 import { handleCallbackErrors } from '../handleCallbackErrors';
-import { instrumentStream } from './streaming';
-import type { AnthropicAiInstrumentedMethod, AnthropicAiOptions, AnthropicAiResponse, ContentBlock } from './types';
+import { instrumentAsyncIterableStream, instrumentMessageStream } from './streaming';
+import type {
+  AnthropicAiInstrumentedMethod,
+  AnthropicAiOptions,
+  AnthropicAiResponse,
+  AnthropicAiStreamingEvent,
+  ContentBlock,
+} from './types';
 import { shouldInstrument } from './utils';
 
 /**
@@ -209,7 +215,45 @@ function instrumentMethod<T extends unknown[], R>(
       const isStreamRequested = Boolean(params?.stream);
       const isStreamingMethod = methodPath === 'messages.stream';
 
-      if (isStreamRequested || isStreamingMethod) {
+      if (isStreamRequested) {
+        return startSpanManual(
+          {
+            name: `${operationName} ${model} stream-response`,
+            op: getSpanOperation(methodPath),
+            attributes: requestAttributes as Record<string, SpanAttributeValue>,
+          },
+          async span => {
+            try {
+              if (options.recordInputs && params) {
+                addPrivateRequestAttributes(span, params);
+              }
+
+              const result = await originalMethod.apply(context, args);
+              return instrumentAsyncIterableStream(
+                result as AsyncIterable<AnthropicAiStreamingEvent>,
+                span,
+                options.recordOutputs ?? false,
+              ) as unknown as R;
+            } catch (error) {
+              captureException(error, {
+                mechanism: {
+                  handled: false,
+                  type: 'auto.ai.anthropic',
+                  data: {
+                    function: methodPath,
+                  },
+                },
+              });
+
+              if (span.isRecording()) {
+                span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+                span.end();
+              }
+              throw error;
+            }
+          },
+        );
+      } else if (isStreamingMethod) {
         // Create span for instrumentation using startSpanManual
         return startSpanManual(
           {
@@ -224,7 +268,7 @@ function instrumentMethod<T extends unknown[], R>(
               }
 
               const messageStream = target.apply(context, args);
-              return instrumentStream(messageStream, span, options.recordOutputs ?? false);
+              return instrumentMessageStream(messageStream, span, options.recordOutputs ?? false);
             } catch (error) {
               captureException(error, {
                 mechanism: {
diff --git a/packages/core/src/utils/anthropic-ai/streaming.ts b/packages/core/src/utils/anthropic-ai/streaming.ts
@@ -259,10 +259,86 @@ function finalizeStreamSpan(state: StreamingState, span: Span, recordOutputs: bo
   }
 }
 
+/**
+ * Instruments an async iterable stream of Anthropic events, updates the span with
+ * streaming attributes and (optionally) the aggregated output text, and yields
+ * each event from the input stream unchanged.
+ */
+export async function* instrumentAsyncIterableStream(
+  stream: AsyncIterable<AnthropicAiStreamingEvent>,
+  span: Span,
+  recordOutputs: boolean,
+): AsyncGenerator<AnthropicAiStreamingEvent, void, unknown> {
+  const state: StreamingState = {
+    responseTexts: [],
+    finishReasons: [],
+    responseId: '',
+    responseModel: '',
+    promptTokens: undefined,
+    completionTokens: undefined,
+    cacheCreationInputTokens: undefined,
+    cacheReadInputTokens: undefined,
+    toolCalls: [],
+    activeToolBlocks: {},
+  };
+
+  try {
+    for await (const event of stream) {
+      processEvent(event, state, recordOutputs, span);
+      yield event;
+    }
+  } finally {
+    // Set common response attributes if available
+    if (state.responseId) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_ID_ATTRIBUTE]: state.responseId,
+      });
+    }
+    if (state.responseModel) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_MODEL_ATTRIBUTE]: state.responseModel,
+      });
+    }
+
+    setTokenUsageAttributes(
+      span,
+      state.promptTokens,
+      state.completionTokens,
+      state.cacheCreationInputTokens,
+      state.cacheReadInputTokens,
+    );
+
+    span.setAttributes({
+      [GEN_AI_RESPONSE_STREAMING_ATTRIBUTE]: true,
+    });
+
+    if (state.finishReasons.length > 0) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE]: JSON.stringify(state.finishReasons),
+      });
+    }
+
+    if (recordOutputs && state.responseTexts.length > 0) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: state.responseTexts.join(''),
+      });
+    }
+
+    // Set tool calls if any were captured
+    if (recordOutputs && state.toolCalls.length > 0) {
+      span.setAttributes({
+        [GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE]: JSON.stringify(state.toolCalls),
+      });
+    }
+
+    span.end();
+  }
+}
+
 /**
  * Instruments a MessageStream by registering event handlers and preserving the original stream API.
  */
-export function instrumentStream<R extends { on: (...args: unknown[]) => void }>(
+export function instrumentMessageStream<R extends { on: (...args: unknown[]) => void }>(
   stream: R,
   span: Span,
   recordOutputs: boolean,
