@@ -16,6 +16,7 @@ import {
1616import type { CloudflareOptions } from './client' ;
1717import { addCloudResourceContext , addCultureContext , addRequest } from './scope-utils' ;
1818import { init } from './sdk' ;
19+ import { endSpanAfterWaitUntil } from './utils/endSpanAfterWaitUntil' ;
1920import { classifyResponseStreaming } from './utils/streaming' ;
2021
2122interface RequestHandlerWrapperOptions {
@@ -107,73 +108,79 @@ export function wrapRequestHandler(
107108 // See: https://developers.cloudflare.com/workers/runtime-apis/performance/
108109
109110 // Use startSpanManual to control when span ends (needed for streaming responses)
110- return startSpanManual ( { name, attributes } , async span => {
111- let res : Response ;
112-
113- try {
114- res = await handler ( ) ;
115- setHttpStatus ( span , res . status ) ;
116-
117- // After the handler runs, the span name might have been updated by nested instrumentation
118- // (e.g., Remix parameterizing routes). The span should already have the correct name
119- // from that instrumentation, so we don't need to do anything here.
120- } catch ( e ) {
121- span . end ( ) ;
122- if ( captureErrors ) {
123- captureException ( e , { mechanism : { handled : false , type : 'auto.http.cloudflare' } } ) ;
124- }
125- waitUntil ?.( flush ( 2000 ) ) ;
126- throw e ;
127- }
111+ return startSpanManual ( { name, attributes } , async rootSpan => {
112+ return startSpanManual ( { name : 'fetch' , attributes } , async fetchSpan => {
113+ const finishSpansAndWaitUntil = ( ) : void => {
114+ fetchSpan . end ( ) ;
115+ waitUntil ?.( flush ( 2000 ) ) ;
116+ waitUntil ?.( endSpanAfterWaitUntil ( rootSpan ) ) ;
117+ } ;
128118
129- // Classify response to detect actual streaming
130- const classification = classifyResponseStreaming ( res ) ;
119+ let res : Response ;
131120
132- if ( classification . isStreaming && res . body ) {
133- // Streaming response detected - monitor consumption to keep span alive
134121 try {
135- const [ clientStream , monitorStream ] = res . body . tee ( ) ;
122+ res = await handler ( ) ;
123+ setHttpStatus ( rootSpan , res . status ) ;
136124
137- // Monitor stream consumption and end span when complete
138- const streamMonitor = ( async ( ) => {
139- const reader = monitorStream . getReader ( ) ;
125+ // After the handler runs, the span name might have been updated by nested instrumentation
126+ // (e.g., Remix parameterizing routes). The span should already have the correct name
127+ // from that instrumentation, so we don't need to do anything here.
128+ } catch ( e ) {
129+ // For errors, we still wait for waitUntil promises before ending the span
130+ // so that any spans created in waitUntil callbacks are captured
131+ if ( captureErrors ) {
132+ captureException ( e , { mechanism : { handled : false , type : 'auto.http.cloudflare' } } ) ;
133+ }
134+ finishSpansAndWaitUntil ( ) ;
135+ throw e ;
136+ }
140137
141- try {
142- let done = false ;
143- while ( ! done ) {
144- const result = await reader . read ( ) ;
145- done = result . done ;
138+ // Classify response to detect actual streaming
139+ const classification = classifyResponseStreaming ( res ) ;
140+
141+ if ( classification . isStreaming && res . body ) {
142+ // Streaming response detected - monitor consumption to keep span alive
143+ try {
144+ const [ clientStream , monitorStream ] = res . body . tee ( ) ;
145+
146+ // Monitor stream consumption and end span when complete
147+ const streamMonitor = ( async ( ) => {
148+ const reader = monitorStream . getReader ( ) ;
149+
150+ try {
151+ let done = false ;
152+ while ( ! done ) {
153+ const result = await reader . read ( ) ;
154+ done = result . done ;
155+ }
156+ } catch {
157+ // Stream error or cancellation - will end span in finally
158+ } finally {
159+ reader . releaseLock ( ) ;
160+ finishSpansAndWaitUntil ( ) ;
146161 }
147- } catch {
148- // Stream error or cancellation - will end span in finally
149- } finally {
150- reader . releaseLock ( ) ;
151- span . end ( ) ;
152- waitUntil ?.( flush ( 2000 ) ) ;
153- }
154- } ) ( ) ;
155-
156- // Keep worker alive until stream monitoring completes (otherwise span won't end)
157- waitUntil ?.( streamMonitor ) ;
158-
159- // Return response with client stream
160- return new Response ( clientStream , {
161- status : res . status ,
162- statusText : res . statusText ,
163- headers : res . headers ,
164- } ) ;
165- } catch ( e ) {
166- // tee() failed (e.g stream already locked) - fall back to non-streaming handling
167- span . end ( ) ;
168- waitUntil ?.( flush ( 2000 ) ) ;
169- return res ;
162+ } ) ( ) ;
163+
164+ // Keep worker alive until stream monitoring completes (otherwise span won't end)
165+ waitUntil ?.( streamMonitor ) ;
166+
167+ // Return response with client stream
168+ return new Response ( clientStream , {
169+ status : res . status ,
170+ statusText : res . statusText ,
171+ headers : res . headers ,
172+ } ) ;
173+ } catch ( e ) {
174+ // tee() failed (e.g stream already locked) - fall back to non-streaming handling
175+ finishSpansAndWaitUntil ( ) ;
176+ return res ;
177+ }
170178 }
171- }
172179
173- // Non-streaming response - end span immediately and return original
174- span . end ( ) ;
175- waitUntil ?. ( flush ( 2000 ) ) ;
176- return res ;
180+ // Non-streaming response - end span after all waitUntil promises complete
181+ finishSpansAndWaitUntil ( ) ;
182+ return res ;
183+ } ) ;
177184 } ) ;
178185 } ,
179186 ) ;
0 commit comments