fix(server-utils): Mark AI streaming spans as errored when the stream fails midway

dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts

@@ -453,18 +453,23 @@ describe('Anthropic integration', () => {
             expect(streamInitErrorSpan!.status).toBe('error');
 
             const createMidwayErrorSpan = spans.find(
-              span => span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'error-stream-midway' && span.status === 'ok',
+              span =>
+                span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'error-stream-midway' &&
+                span.attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE]?.value === true,
             );
             expect(createMidwayErrorSpan!.name).toBe('chat error-stream-midway');
-            expect(createMidwayErrorSpan!.attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE].value).toBe(true);
+            expect(createMidwayErrorSpan!.status).toBe('error');
             expect(createMidwayErrorSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true);
             expect(createMidwayErrorSpan!.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('This stream will ');
 
             const streamMidwayErrorSpan = spans.find(
-              span => span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'error-stream-midway' && span.status === 'error',
+              span =>
+                span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'error-stream-midway' &&
+                span.attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE] === undefined,
             );
             expect(streamMidwayErrorSpan!.name).toBe('chat error-stream-midway');
-            expect(streamMidwayErrorSpan!.attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE]).toBeUndefined();
+            expect(streamMidwayErrorSpan!.status).toBe('error');
+            expect(streamMidwayErrorSpan!.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('This stream will ');
           },
         })
         .start()

dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-stream-midway-error.mjs

@@ -0,0 +1,103 @@
+import { GoogleGenAI } from '@google/genai';
+import * as Sentry from '@sentry/node';
+import express from 'express';
+
+function startMockGoogleGenAIServer() {
+  const app = express();
+  app.use(express.json());
+
+  app.post('/v1beta/models/:model\\:streamGenerateContent', (req, res) => {
+    res.setHeader('Content-Type', 'application/json');
+    res.setHeader('Transfer-Encoding', 'chunked');
+
+    const chunks = [
+      {
+        candidates: [{ content: { parts: [{ text: 'This stream will ' }], role: 'model' }, index: 0 }],
+        responseId: 'mock-midway-response-id',
+        modelVersion: 'gemini-1.5-pro',
+        usageMetadata: { promptTokenCount: 8, candidatesTokenCount: 4, totalTokenCount: 12 },
+      },
+      { candidates: [{ content: { parts: [{ text: 'be cut off' }], role: 'model' }, index: 0 }] },
+    ];
+    chunks.forEach((chunk, index) => {
+      setTimeout(() => res.write(`data: ${JSON.stringify(chunk)}\n\n`), index * 10);
+    });
+
+    // The midway case drops the connection after the chunks are parsed, without a final chunk.
+    // The other cases need the body to stay open long enough for their abort to hit a pending read.
+    setTimeout(
+      () => (req.params.model === 'midway-error' ? res.destroy() : res.end()),
+      req.params.model === 'midway-error' ? 70 : 1000,
+    );
+  });
+
+  return new Promise(resolve => {
+    const server = app.listen(0, () => {
+      resolve(server);
+    });
+  });
+}
+
+async function run() {
+  const server = await startMockGoogleGenAIServer();
+
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const client = new GoogleGenAI({
+      apiKey: 'mock-api-key',
+      httpOptions: { baseUrl: `http://localhost:${server.address().port}` },
+    });
+
+    try {
+      const stream = await client.models.generateContentStream({
+        model: 'midway-error',
+        contents: [{ role: 'user', parts: [{ text: 'This will fail midway' }] }],
+      });
+      for await (const _ of stream) {
+        void _;
+      }
+    } catch {
+      // Error expected
+    }
+
+    const stream = await client.models.generateContentStream({
+      model: 'early-break',
+      contents: [{ role: 'user', parts: [{ text: 'The caller stops reading after the first chunk' }] }],
+    });
+    for await (const _ of stream) {
+      break;
+    }
+
+    const controller = new AbortController();
+    try {
+      const abortedStream = await client.models.generateContentStream({
+        model: 'caller-abort',
+        contents: [{ role: 'user', parts: [{ text: 'The caller aborts after the first chunk' }] }],
+        config: { abortSignal: controller.signal },
+      });
+      for await (const _ of abortedStream) {
+        controller.abort();
+      }
+    } catch {
+      // Error expected
+    }
+
+    try {
+      const timedOutStream = await client.models.generateContentStream({
+        model: 'timeout',
+        contents: [{ role: 'user', parts: [{ text: 'The SDK timeout fires mid-stream' }] }],
+        config: { httpOptions: { timeout: 100 } },
+      });
+      for await (const _ of timedOutStream) {
+        void _;
+      }
+    } catch {
+      // Error expected
+    }
+  });
+
+  await Sentry.flush(2000);
+
+  server.close();
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts

@@ -585,4 +585,39 @@ describe('Google GenAI integration', () => {
         .completed();
     });
   });
+
+  createEsmAndCjsTests(
+    __dirname,
+    'scenario-stream-midway-error.mjs',
+    'instrument-with-pii.mjs',
+    (createRunner, test) => {
+      test('marks the span as errored when the stream fails midway, but not when the caller breaks early', async () => {
+        await createRunner()
+          .unordered()
+          .expect({
+            span: container => {
+              const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === EXPECTED_ORIGIN);
+              expect(spans).toHaveLength(4);
+
+              const midwaySpan = spans.find(span => span.name === 'generate_content midway-error');
+              expect(midwaySpan!.status).toBe('error');
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true);
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_ID]?.value).toBe('mock-midway-response-id');
+              expect(midwaySpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS]?.value).toBe(8);
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('This stream will be cut off');
+
+              const earlyBreakSpan = spans.find(span => span.name === 'generate_content early-break');
+              expect(earlyBreakSpan!.status).toBe('ok');
+              expect(earlyBreakSpan!.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('This stream will ');
+
+              // `cancelled` is sent as `ok` under span streaming, so a timeout must not be recorded as one.
+              expect(spans.find(span => span.name === 'generate_content caller-abort')!.status).toBe('ok');
+              expect(spans.find(span => span.name === 'generate_content timeout')!.status).toBe('error');
+            },
+          })
+          .start()
+          .completed();
+      });
+    },
+  );
 });

dev-packages/node-integration-tests/suites/tracing/openai/scenario-stream-midway-error.mjs

@@ -0,0 +1,81 @@
+import * as Sentry from '@sentry/node';
+import express from 'express';
+import OpenAI from 'openai';
+
+function startMockServer() {
+  const app = express();
+  app.use(express.json());
+
+  app.post('/openai/chat/completions', (req, res) => {
+    res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', Connection: 'keep-alive' });
+
+    const chunks = [
+      {
+        id: 'chatcmpl-midway',
+        object: 'chat.completion.chunk',
+        model: req.body.model,
+        choices: [{ index: 0, delta: { role: 'assistant', content: 'This stream will ' } }],
+      },
+      {
+        id: 'chatcmpl-midway',
+        object: 'chat.completion.chunk',
+        model: req.body.model,
+        choices: [{ index: 0, delta: { content: 'be cut off' } }],
+      },
+    ];
+    chunks.forEach((chunk, index) => {
+      setTimeout(() => res.write(`data: ${JSON.stringify(chunk)}\n\n`), index * 10);
+    });
+
+    // The midway case drops the connection after the chunks are parsed, without a final chunk or `[DONE]`.
+    setTimeout(
+      () => (req.body.model === 'midway-error' ? res.destroy() : res.end('data: [DONE]\n\n')),
+      chunks.length * 10 + 50,
+    );
+  });
+
+  return new Promise(resolve => {
+    const server = app.listen(0, () => {
+      resolve(server);
+    });
+  });
+}
+
+async function run() {
+  const server = await startMockServer();
+
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const client = new OpenAI({
+      baseURL: `http://localhost:${server.address().port}/openai`,
+      apiKey: 'mock-api-key',
+    });
+
+    try {
+      const stream = await client.chat.completions.create({
+        model: 'midway-error',
+        messages: [{ role: 'user', content: 'This will fail midway' }],
+        stream: true,
+      });
+      for await (const _ of stream) {
+        void _;
+      }
+    } catch {
+      // Error expected
+    }
+
+    const stream = await client.chat.completions.create({
+      model: 'early-break',
+      messages: [{ role: 'user', content: 'The caller stops reading after the first chunk' }],
+      stream: true,
+    });
+    for await (const _ of stream) {
+      break;
+    }
+  });
+
+  await Sentry.flush(2000);
+
+  server.close();
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/openai/test.ts

@@ -1692,4 +1692,34 @@ describe('OpenAI integration', () => {
       expect(rawSpans[0]!.attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBeUndefined();
     });
   });
+
+  createEsmAndCjsTests(
+    __dirname,
+    'scenario-stream-midway-error.mjs',
+    'instrument-with-pii.mjs',
+    (createRunner, test) => {
+      test('marks the span as errored when the stream fails midway, but not when the caller breaks early', async () => {
+        await createRunner()
+          .unordered()
+          .expect({
+            span: container => {
+              const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.openai');
+              expect(spans).toHaveLength(2);
+
+              const midwaySpan = spans.find(span => span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'midway-error');
+              expect(midwaySpan!.status).toBe('error');
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true);
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_ID]?.value).toBe('chatcmpl-midway');
+              expect(midwaySpan!.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('This stream will be cut off');
+
+              const earlyBreakSpan = spans.find(span => span.attributes[GEN_AI_REQUEST_MODEL]?.value === 'early-break');
+              expect(earlyBreakSpan!.status).toBe('ok');
+              expect(earlyBreakSpan!.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('This stream will ');
+            },
+          })
+          .start()
+          .completed();
+      });
+    },
+  );
 });

packages/server-utils/src/ai/anthropic-ai/streaming.ts

@@ -247,6 +247,9 @@ export async function* instrumentAsyncIterableStream(
       processEvent(event, state, recordOutputs, span);
       yield event;
     }
+  } catch (error) {
+    span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+    throw error;
   } finally {
     endStreamSpan(span, state, recordOutputs);
   }
@@ -296,7 +299,7 @@ export function instrumentMessageStream<R extends { on: (...args: unknown[]) =>
 
     if (span.isRecording()) {
       span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
-      span.end();
+      endStreamSpan(span, state, recordOutputs);
     }
   });
 

packages/server-utils/src/ai/google-genai/index.ts

@@ -333,7 +333,12 @@ function instrumentMethod<T extends unknown[], R>(
                 addPrivateRequestAttributes(span, attributeParams, operationName);
               }
               const stream = await target.apply(context, args);
-              return instrumentStream(stream, span, Boolean(options.recordOutputs)) as R;
+              return instrumentStream(
+                stream,
+                span,
+                Boolean(options.recordOutputs),
+                (attributeParams?.config as { abortSignal?: AbortSignal } | undefined)?.abortSignal,
+              ) as R;
             } catch (error) {
               span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
               span.end();

packages/server-utils/src/ai/google-genai/streaming.ts

@@ -122,6 +122,7 @@ export async function* instrumentStream(
   stream: AsyncIterable<GoogleGenAIResponse>,
   span: Span,
   recordOutputs: boolean,
+  abortSignal?: AbortSignal,
 ): AsyncGenerator<GoogleGenAIResponse, void, unknown> {
   const state: StreamingState = {
     responseTexts: [],
@@ -135,6 +136,9 @@ export async function* instrumentStream(
       processChunk(chunk, state, recordOutputs, span);
       yield chunk;
     }
+  } catch (error) {
+    span.setStatus({ code: SPAN_STATUS_ERROR, message: abortSignal?.aborted ? 'cancelled' : 'internal_error' });
+    throw error;
   } finally {
     if (recordOutputs) {
       setOutputMessagesAttribute(span, state.outputParts);

packages/server-utils/src/ai/openai/streaming.ts

@@ -282,6 +282,9 @@ export async function* instrumentStream<T>(
       }
       yield event;
     }
+  } catch (error) {
+    span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
+    throw error;
   } finally {
     const allToolCalls = [...state.chatCompletionToolCalls.values(), ...state.responsesApiToolCalls];
     endStreamSpan(span, { ...state, toolCalls: allToolCalls }, recordOutputs);

packages/server-utils/src/integrations/google-genai.ts

@@ -63,7 +63,7 @@ function instrumentGoogleGenai(options: GoogleGenAIOptions): void {
             );
           }
         },
-        deferSpanEnd: ({ span, data }) => wrapStreamResult(span, data, options),
+        deferSpanEnd: ({ span, data }) => wrapStreamResult(span, data, operation, options),
       },
     );
   }
@@ -136,15 +136,26 @@ function isAsyncIterable(value: unknown): value is AsyncIterableStream {
  * For a stream we patch `result[Symbol.asyncIterator]` in place so `instrumentStream` ends the
  * span when iteration finishes.
  */
-function wrapStreamResult(span: Span, data: OrchestrionChannelContext, options: GoogleGenAIOptions): boolean {
+function wrapStreamResult(
+  span: Span,
+  data: OrchestrionChannelContext,
+  operation: string,
+  options: GoogleGenAIOptions,
+): boolean {
   const result = data.result;
   if (!isAsyncIterable(result)) {
     return false;
   }
 
   const { recordOutputs } = resolveAIRecordingOptions(options);
   const iterate = result[Symbol.asyncIterator].bind(result);
-  const instrumented = instrumentStream({ [Symbol.asyncIterator]: iterate }, span, recordOutputs ?? false);
+  const params = resolveChatParams(operation, data.arguments?.[0] as Record<string, unknown> | undefined, data.self);
+  const instrumented = instrumentStream(
+    { [Symbol.asyncIterator]: iterate },
+    span,
+    recordOutputs ?? false,
+    (params?.config as { abortSignal?: AbortSignal } | undefined)?.abortSignal,
+  );
   result[Symbol.asyncIterator] = () => instrumented;
 
   return true;