quick refactor

Author: RulaKhaled

packages/cloudflare/src/request.ts

@@ -108,7 +108,7 @@ export function wrapRequestHandler(
             res = await handler();
             setHttpStatus(span, res.status);
           } catch (e) {
-            span.end(); // End span on error
+            span.end();
             if (captureErrors) {
               captureException(e, { mechanism: { handled: false, type: 'auto.http.cloudflare' } });
             }
@@ -119,40 +119,29 @@ export function wrapRequestHandler(
           // Classify response to detect actual streaming
           const classification = await classifyResponseStreaming(res);
 
-          if (classification.isStreaming) {
+          if (classification.isStreaming && classification.response.body) {
             // 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
+              // Safety timeout - abort reading and end span after 5s even if stream hasn't finished
+              const timeout = setTimeout(() => reader.cancel(), 5000);
 
               try {
                 let done = false;
                 while (!done) {
                   const result = await reader.read();
                   done = result.done;
                 }
-                clearTimeout(timeout);
-                span.end();
-              } catch (err) {
-                clearTimeout(timeout);
-                span.end();
+              } catch {
+                // Stream error or cancellation - will end span in finally
               } finally {
+                clearTimeout(timeout);
                 reader.releaseLock();
+                span.end();
               }
             })();
 

packages/cloudflare/src/utils/streaming.ts

@@ -4,7 +4,18 @@ export type StreamingGuess = {
 };
 
 /**
+ * Classifies a Response as streaming or non-streaming.
  *
+ * Uses multiple heuristics:
+ * - Content-Type: text/event-stream → streaming
+ * - Content-Length header present → not streaming
+ * - Otherwise: probes stream with timeout to detect behavior
+ *
+ * Note: Probing will tee() the stream and return a new Response object.
+ *
+ * @param res - The Response to classify
+ * @param opts.timeoutMs - Probe timeout in ms (default: 25)
+ * @returns Classification result with safe-to-return Response
  */
 export async function classifyResponseStreaming(
   res: Response,
@@ -16,57 +27,61 @@ export async function classifyResponseStreaming(
     return { response: res, isStreaming: false };
   }
 
-  const ct = res.headers.get('content-type') ?? '';
-  const cl = res.headers.get('content-length');
+  const contentType = res.headers.get('content-type') ?? '';
+  const contentLength = res.headers.get('content-length');
 
-  // Definitive streaming indicators
-  if (/^text\/event-stream\b/i.test(ct)) {
+  // Fast path: Server-Sent Events
+  if (/^text\/event-stream\b/i.test(contentType)) {
     return { response: res, isStreaming: true };
   }
 
-  // Definitive non-streaming indicators
-  if (cl && /^\d+$/.test(cl)) {
+  // Fast path: Content-Length indicates buffered response
+  if (contentLength && /^\d+$/.test(contentLength)) {
     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();
+  // Uncertain - probe the stream to determine behavior
+  // After tee(), must use the teed stream (original is locked)
+  const [probeStream, passStream] = res.body.tee();
+  const reader = probeStream.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 probeResult = await Promise.race([
+    // Try to read first chunk
+    (async () => {
+      try {
+        const { value, done } = await reader.read();
+        reader.releaseLock();
 
-  const timeout = new Promise<{ arrivedBytes: number; done: boolean }>(r =>
-    setTimeout(() => r({ arrivedBytes: 0, done: false }), timeoutMs),
-  );
+        if (done) {
+          return { arrivedBytes: 0, done: true };
+        }
 
-  const peek = await Promise.race([firstChunkPromise, timeout]);
+        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 };
+      }
+    })(),
+    // Timeout if first chunk takes too long
+    new Promise<{ arrivedBytes: number; done: boolean }>(resolve =>
+      setTimeout(() => resolve({ arrivedBytes: 0, done: false }), timeoutMs),
+    ),
+  ]);
 
-  // We must return the teed response since original is now locked
-  const preserved = new Response(pass, res);
+  const teededResponse = new Response(passStream, 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;
+  // Determine if streaming based on probe result
+  if (probeResult.done) {
+    // Stream completed immediately - buffered
+    return { response: teededResponse, isStreaming: false };
+  } else if (probeResult.arrivedBytes === 0) {
+    // Timeout waiting - definitely streaming
+    return { response: teededResponse, isStreaming: true };
   } else {
-    // Got first chunk - streaming if no Content-Length
-    isStreaming = cl == null;
+    // Got chunk quickly - streaming if no Content-Length
+    return { response: teededResponse, isStreaming: contentLength == null };
   }
-
-  return { response: preserved, isStreaming };
 }

packages/cloudflare/test/request.test.ts

@@ -90,7 +90,7 @@ describe('withSentry', () => {
   });
 
   test('flush must be called when all waitUntil are done', async () => {
-    const flush = vi.spyOn(SentryCore.Client.prototype, 'flush');
+    const flushSpy = vi.spyOn(SentryCore.Client.prototype, 'flush');
     vi.useFakeTimers();
     onTestFinished(() => {
       vi.useRealTimers();
@@ -104,13 +104,17 @@ describe('withSentry', () => {
 
     await wrapRequestHandler({ options: MOCK_OPTIONS, request: new Request('https://example.com'), context }, () => {
       addDelayedWaitUntil(context);
-      return new Response('test');
+      const response = new Response('test');
+      // Add Content-Length to skip probing
+      response.headers.set('content-length', '4');
+      return response;
     });
-    expect(flush).not.toBeCalled();
+    expect(flushSpy).not.toBeCalled();
     expect(waitUntil).toBeCalled();
-    vi.advanceTimersToNextTimerAsync().then(() => vi.runAllTimers());
+    await vi.advanceTimersToNextTimerAsync();
+    vi.runAllTimers();
     await Promise.all(waits);
-    expect(flush).toHaveBeenCalledOnce();
+    expect(flushSpy).toHaveBeenCalledOnce();
   });
 
   describe('scope instrumentation', () => {
@@ -305,7 +309,7 @@ describe('withSentry', () => {
       mockRequest.headers.set('content-length', '10');
 
       let sentryEvent: Event = {};
-      const result = await wrapRequestHandler(
+      await wrapRequestHandler(
         {
           options: {
             ...MOCK_OPTIONS,
@@ -320,14 +324,10 @@ describe('withSentry', () => {
         },
         () => {
           SentryCore.captureMessage('sentry-trace');
-          const response = new Response('test');
-          return response;
+          return new Response('test');
         },
       );
 
-      // 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));