@@ -313,17 +313,29 @@ Use this context to calculate relative dates like "yesterday", "last week", "beg
313313 return
314314 }
315315
316+ let finalUsage : any = null
317+ let usageRecorded = false
318+
319+ const recordUsage = async ( ) => {
320+ if ( usageRecorded || ! finalUsage ) {
321+ return
322+ }
323+
324+ usageRecorded = true
325+ await updateUserStatsForWand ( session . user . id , finalUsage , requestId , isBYOK )
326+ }
327+
316328 try {
317329 let buffer = ''
318330 let chunkCount = 0
319- let finalUsage : any = null
320331 let activeEventType : string | undefined
321332
322333 while ( true ) {
323334 const { done, value } = await reader . read ( )
324335
325336 if ( done ) {
326337 logger . info ( `[${ requestId } ] Stream completed. Total chunks: ${ chunkCount } ` )
338+ await recordUsage ( )
327339 controller . enqueue ( encoder . encode ( `data: ${ JSON . stringify ( { done : true } ) } \n\n` ) )
328340 controller . close ( )
329341 break
@@ -353,9 +365,7 @@ Use this context to calculate relative dates like "yesterday", "last week", "beg
353365 if ( data === '[DONE]' ) {
354366 logger . info ( `[${ requestId } ] Received [DONE] signal` )
355367
356- if ( finalUsage ) {
357- await updateUserStatsForWand ( session . user . id , finalUsage , requestId , isBYOK )
358- }
368+ await recordUsage ( )
359369
360370 controller . enqueue (
361371 encoder . encode ( `data: ${ JSON . stringify ( { done : true } ) } \n\n` )
@@ -423,6 +433,12 @@ Use this context to calculate relative dates like "yesterday", "last week", "beg
423433 stack : streamError ?. stack ,
424434 } )
425435
436+ try {
437+ await recordUsage ( )
438+ } catch ( usageError ) {
439+ logger . warn ( `[${ requestId } ] Failed to record usage after stream error` , usageError )
440+ }
441+
426442 const errorData = `data: ${ JSON . stringify ( { error : 'Streaming failed' , done : true } ) } \n\n`
427443 controller . enqueue ( encoder . encode ( errorData ) )
428444 controller . close ( )
0 commit comments