fix(cloudflare): Keep root span alive until streaming responses are consumed

Author: RulaKhaled

packages/cloudflare/src/request.ts

@@ -8,13 +8,14 @@ import {
   parseStringToURLObject,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
   setHttpStatus,
-  startSpan,
+  startSpanManual,
   winterCGHeadersToDict,
   withIsolationScope,
 } from '@sentry/core';
 import type { CloudflareOptions } from './client';
 import { addCloudResourceContext, addCultureContext, addRequest } from './scope-utils';
 import { init } from './sdk';
+import { classifyResponseStreaming } from './utils/streaming';
 
 interface RequestHandlerWrapperOptions {
   options: CloudflareOptions;
@@ -98,26 +99,79 @@ export function wrapRequestHandler(
         // Note: This span will not have a duration unless I/O happens in the handler. This is
         // because of how the cloudflare workers runtime works.
         // See: https://developers.cloudflare.com/workers/runtime-apis/performance/
-        return startSpan(
-          {
-            name,
-            attributes,
-          },
-          async span => {
-            try {
-              const res = await handler();
-              setHttpStatus(span, res.status);
-              return res;
-            } catch (e) {
-              if (captureErrors) {
-                captureException(e, { mechanism: { handled: false, type: 'auto.http.cloudflare' } });
-              }
-              throw e;
-            } finally {
+
+        // Use startSpanManual to control when span ends (needed for streaming responses)
+        return startSpanManual({ name, attributes }, async span => {
+          let res: Response;
+
+          try {
+            res = await handler();
+            setHttpStatus(span, res.status);
+          } catch (e) {
+            span.end(); // End span on error
+            if (captureErrors) {
+              captureException(e, { mechanism: { handled: false, type: 'auto.http.cloudflare' } });
+            }
+            waitUntil?.(flush(2000));
+            throw e;
+          }
+
+          // Classify response to detect actual streaming
+          const classification = await classifyResponseStreaming(res);
+
+          if (classification.isStreaming) {
+            // Streaming response detected - monitor consumption to keep span alive
+            if (!classification.response.body) {
+              // Shouldn't happen since isStreaming requires body, but handle gracefully
+              span.end();
               waitUntil?.(flush(2000));
+              return classification.response;
             }
-          },
-        );
+            
+            const [clientStream, monitorStream] = classification.response.body.tee();
+
+            // Monitor stream consumption and end span when complete
+            const streamMonitor = (async () => {
+              const reader = monitorStream.getReader();
+
+              // Safety timeout to prevent infinite loops if stream hangs
+              const timeout = setTimeout(() => {
+                span.end();
+                reader.cancel().catch(() => {});
+              }, 30000); // 30 second max
+
+              try {
+                let done = false;
+                while (!done) {
+                  const result = await reader.read();
+                  done = result.done;
+                }
+                clearTimeout(timeout);
+                span.end();
+              } catch (err) {
+                clearTimeout(timeout);
+                span.end();
+              } finally {
+                reader.releaseLock();
+              }
+            })();
+
+            // Use waitUntil to keep context alive and flush after span ends
+            waitUntil?.(streamMonitor.then(() => flush(2000)));
+
+            // Return response with client stream
+            return new Response(clientStream, {
+              status: classification.response.status,
+              statusText: classification.response.statusText,
+              headers: classification.response.headers,
+            });
+          }
+
+          // Non-streaming response - end span immediately and return original
+          span.end();
+          waitUntil?.(flush(2000));
+          return classification.response;
+        });
       },
     );
   });

packages/cloudflare/src/utils/streaming.ts

@@ -0,0 +1,72 @@
+export type StreamingGuess = {
+  response: Response;
+  isStreaming: boolean;
+};
+
+/**
+ *
+ */
+export async function classifyResponseStreaming(
+  res: Response,
+  opts: { timeoutMs?: number } = {},
+): Promise<StreamingGuess> {
+  const timeoutMs = opts.timeoutMs ?? 25;
+
+  if (!res.body) {
+    return { response: res, isStreaming: false };
+  }
+
+  const ct = res.headers.get('content-type') ?? '';
+  const cl = res.headers.get('content-length');
+
+  // Definitive streaming indicators
+  if (/^text\/event-stream\b/i.test(ct)) {
+    return { response: res, isStreaming: true };
+  }
+
+  // Definitive non-streaming indicators
+  if (cl && /^\d+$/.test(cl)) {
+    return { response: res, isStreaming: false };
+  }
+
+  // Probe the stream to detect streaming behavior
+  // NOTE: This tees the stream and returns a new Response object
+  const [probe, pass] = res.body.tee();
+  const reader = probe.getReader();
+
+  const firstChunkPromise = (async () => {
+    try {
+      const { value, done } = await reader.read();
+      reader.releaseLock();
+      if (done) return { arrivedBytes: 0, done: true };
+      const bytes =
+        value && typeof value === 'object' && 'byteLength' in value ? (value as { byteLength: number }).byteLength : 0;
+      return { arrivedBytes: bytes, done: false };
+    } catch {
+      return { arrivedBytes: 0, done: false };
+    }
+  })();
+
+  const timeout = new Promise<{ arrivedBytes: number; done: boolean }>(r =>
+    setTimeout(() => r({ arrivedBytes: 0, done: false }), timeoutMs),
+  );
+
+  const peek = await Promise.race([firstChunkPromise, timeout]);
+
+  // We must return the teed response since original is now locked
+  const preserved = new Response(pass, res);
+
+  let isStreaming = false;
+  if (peek.done) {
+    // Stream completed immediately
+    isStreaming = false;
+  } else if (peek.arrivedBytes === 0) {
+    // Timeout waiting for first chunk - definitely streaming
+    isStreaming = true;
+  } else {
+    // Got first chunk - streaming if no Content-Length
+    isStreaming = cl == null;
+  }
+
+  return { response: preserved, isStreaming };
+}

packages/cloudflare/test/durableobject.test.ts

@@ -133,11 +133,25 @@ describe('instrumentDurableObjectWithSentry', () => {
       waitUntil,
     } as unknown as ExecutionContext;
     const dObject: any = Reflect.construct(instrumented, [context, {} as any]);
-    expect(() => dObject.fetch(new Request('https://example.com'))).not.toThrow();
-    expect(flush).not.toBeCalled();
-    expect(waitUntil).toHaveBeenCalledOnce();
+
+    // Call fetch (don't await yet)
+    const responsePromise = dObject.fetch(new Request('https://example.com'));
+
+    // Advance past classification timeout and get response
+    vi.advanceTimersByTime(30);
+    const response = await responsePromise;
+
+    // Consume response (triggers span end for buffered responses)
+    await response.text();
+
+    // The flush should now be queued in waitUntil
+    expect(waitUntil).toHaveBeenCalled();
+
+    // Advance to trigger the setTimeout in the handler's waitUntil
     vi.advanceTimersToNextTimer();
     await Promise.all(waitUntil.mock.calls.map(([p]) => p));
+
+    // Now flush should have been called
     expect(flush).toBeCalled();
   });
 

packages/cloudflare/test/handler.test.ts

@@ -72,7 +72,11 @@ describe('withSentry', () => {
         createMockExecutionContext(),
       );
 
-      expect(result).toBe(response);
+      // Response may be wrapped for streaming detection, verify content
+      expect(result?.status).toBe(response.status);
+      if (result) {
+        expect(await result.text()).toBe('test');
+      }
     });
 
     test('merges options from env and callback', async () => {

packages/cloudflare/test/pages-plugin.test.ts

@@ -52,6 +52,8 @@ describe('sentryPagesPlugin', () => {
       pluginArgs: MOCK_OPTIONS,
     });
 
-    expect(result).toBe(response);
+    // Response may be wrapped for streaming detection, verify content
+    expect(result.status).toBe(response.status);
+    expect(await result.text()).toBe('test');
   });
 });

packages/cloudflare/test/request.test.ts

@@ -33,7 +33,9 @@ describe('withSentry', () => {
       { options: MOCK_OPTIONS, request: new Request('https://example.com'), context: createMockExecutionContext() },
       () => response,
     );
-    expect(result).toBe(response);
+    // Response may be wrapped for streaming detection, verify content matches
+    expect(result.status).toBe(response.status);
+    expect(await result.text()).toBe('test');
   });
 
   test('flushes the event after the handler is done using the cloudflare context.waitUntil', async () => {
@@ -48,6 +50,25 @@ describe('withSentry', () => {
     expect(waitUntilSpy).toHaveBeenLastCalledWith(expect.any(Promise));
   });
 
+  test('handles streaming responses correctly', async () => {
+    const stream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode('chunk1'));
+        controller.enqueue(new TextEncoder().encode('chunk2'));
+        controller.close();
+      },
+    });
+    const streamingResponse = new Response(stream);
+
+    const result = await wrapRequestHandler(
+      { options: MOCK_OPTIONS, request: new Request('https://example.com'), context: createMockExecutionContext() },
+      () => streamingResponse,
+    );
+
+    const text = await result.text();
+    expect(text).toBe('chunk1chunk2');
+  });
+
   test("doesn't error if context is undefined", () => {
     expect(() =>
       wrapRequestHandler(
@@ -284,7 +305,7 @@ describe('withSentry', () => {
       mockRequest.headers.set('content-length', '10');
 
       let sentryEvent: Event = {};
-      await wrapRequestHandler(
+      const result = await wrapRequestHandler(
         {
           options: {
             ...MOCK_OPTIONS,
@@ -299,10 +320,17 @@ describe('withSentry', () => {
         },
         () => {
           SentryCore.captureMessage('sentry-trace');
-          return new Response('test');
+          const response = new Response('test');
+          return response;
         },
       );
 
+      // Consume response to trigger span end for non-streaming responses
+      await result.text();
+
+      // Wait for async span end and transaction capture
+      await new Promise(resolve => setTimeout(resolve, 50));
+
       expect(sentryEvent.transaction).toEqual('GET /');
       expect(sentryEvent.spans).toHaveLength(0);
       expect(sentryEvent.contexts?.trace).toEqual({