@@ -11,6 +11,7 @@ import { createValidationErrorBlocks } from '../utils/create-validation-error-bl
1111import type { ChatMessage , ContentBlock } from '../chat'
1212import type { AgentDefinition , ToolName } from '@codebuff/sdk'
1313import type { SetStateAction } from 'react'
14+ import type { ElapsedTimeTracker } from './use-elapsed-time'
1415
1516const hiddenToolNames = new Set < ToolName | 'spawn_agent_inline' > ( [
1617 'spawn_agent_inline' ,
@@ -100,7 +101,7 @@ interface UseSendMessageOptions {
100101 abortControllerRef : React . MutableRefObject < AbortController | null >
101102 agentId ?: string
102103 onBeforeMessageSend ?: ( ) => Promise < { success : boolean ; errors : Array < { id : string ; message : string } > } >
103- setMainAgentStreamStartTime : ( time : number | null ) => void
104+ mainAgentTimer : ElapsedTimeTracker
104105 scrollToLatest : ( ) => void
105106 availableWidth ?: number
106107}
@@ -124,7 +125,7 @@ export const useSendMessage = ({
124125 abortControllerRef,
125126 agentId,
126127 onBeforeMessageSend,
127- setMainAgentStreamStartTime ,
128+ mainAgentTimer ,
128129 scrollToLatest,
129130 availableWidth = 80 ,
130131} : UseSendMessageOptions ) => {
@@ -256,6 +257,9 @@ export const useSendMessage = ({
256257 const { agentMode } = params
257258 const timestamp = formatTimestamp ( )
258259
260+ // Start timer immediately when message is sent
261+ mainAgentTimer . start ( )
262+
259263 // Add user message to UI first
260264 const userMessage : ChatMessage = {
261265 id : `user-${ Date . now ( ) } ` ,
@@ -597,8 +601,6 @@ export const useSendMessage = ({
597601 if ( ! hasReceivedContent ) {
598602 hasReceivedContent = true
599603 setIsWaitingForResponse ( false )
600- // Main agent started streaming - set timer
601- setMainAgentStreamStartTime ( Date . now ( ) )
602604 }
603605
604606 const previous = rootStreamBufferRef . current ?? ''
@@ -637,7 +639,6 @@ export const useSendMessage = ({
637639 if ( ! hasReceivedContent && ! event . agentId ) {
638640 hasReceivedContent = true
639641 setIsWaitingForResponse ( false )
640- setMainAgentStreamStartTime ( Date . now ( ) )
641642 } else if ( ! hasReceivedContent ) {
642643 hasReceivedContent = true
643644 setIsWaitingForResponse ( false )
@@ -1309,7 +1310,7 @@ export const useSendMessage = ({
13091310 setCanProcessQueue ( true )
13101311 updateChainInProgress ( false )
13111312 setIsWaitingForResponse ( false )
1312- setMainAgentStreamStartTime ( null )
1313+ mainAgentTimer . stop ( )
13131314
13141315 if ( ( result as any ) ?. credits !== undefined ) {
13151316 actualCredits = ( result as any ) . credits
@@ -1344,7 +1345,7 @@ export const useSendMessage = ({
13441345 setCanProcessQueue ( true )
13451346 updateChainInProgress ( false )
13461347 setIsWaitingForResponse ( false )
1347- setMainAgentStreamStartTime ( null )
1348+ mainAgentTimer . stop ( )
13481349
13491350 if ( isAborted ) {
13501351 applyMessageUpdate ( ( prev ) =>
@@ -1421,7 +1422,7 @@ export const useSendMessage = ({
14211422 addActiveSubagent ,
14221423 removeActiveSubagent ,
14231424 onBeforeMessageSend ,
1424- setMainAgentStreamStartTime ,
1425+ mainAgentTimer ,
14251426 scrollToLatest ,
14261427 availableWidth ,
14271428 ] ,
0 commit comments