Skip to content

Commit 8ac91a9

Browse files
brandonkachenbrandonkachen
committed
Improve abort controller handling and retry logic
- Simplified abort controller: SDK and CLI share single controller - Removed duplicate CLI timeout (SDK timeout is sole source) - Fixed retry trigger to work without connection status dependency - Added queue guard to prevent concurrent message submissions - Silent retries for recoverable errors (no false interruptions) - Improved UX: retries happen automatically in background
1 parent 5ff0a02 commit 8ac91a9

File tree

3 files changed

+105
-81
lines changed

3 files changed

+105
-81
lines changed

‎cli/src/chat.tsx‎

Lines changed: 57 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
import { startTransition, useCallback, useEffect, useMemo, useRef, useState } from 'react'
1+
import {
2+
startTransition,
3+
useCallback,
4+
useEffect,
5+
useMemo,
6+
useRef,
7+
useState,
8+
} from 'react'
29
import { useKeyboard } from '@opentui/react'
310
import { useShallow } from 'zustand/react/shallow'
411

@@ -203,44 +210,52 @@ export const Chat = ({
203210
const { statusMessage } = useClipboard()
204211
const [showReconnectionMessage, setShowReconnectionMessage] = useState(false)
205212
const [connectionEstablished, setConnectionEstablished] = useState(0) // Increment to trigger retry check
206-
const { setTimeout: setReconnectionTimeout, clearTimeout: clearReconnectionTimeout } = useSafeTimeout()
213+
const {
214+
setTimeout: setReconnectionTimeout,
215+
clearTimeout: clearReconnectionTimeout,
216+
} = useSafeTimeout()
207217
const retryPendingMessagesRef = useRef<(() => Promise<void>) | null>(null)
208218
const processFailedMessagesRef = useRef<(() => void) | null>(null)
209219
const retryScheduledRef = useRef(false)
210220

211-
const handleReconnection = useCallback((isInitialConnection: boolean) => {
212-
logger.debug(
213-
{ isInitialConnection },
214-
`[Connection] ${isInitialConnection ? 'Initial connection' : 'Reconnection'} callback triggered`
215-
)
216-
217-
// Invalidate auth queries to allow network status to clear after reconnection
218-
queryClient.invalidateQueries({ queryKey: authQueryKeys.all })
219-
logger.debug('[Connection] Invalidated auth queries to refresh network status')
221+
const handleReconnection = useCallback(
222+
(isInitialConnection: boolean) => {
223+
logger.debug(
224+
{ isInitialConnection },
225+
`[Connection] ${isInitialConnection ? 'Initial connection' : 'Reconnection'} callback triggered`,
226+
)
220227

221-
// Process any failed messages and schedule them for retry (batched)
222-
if (processFailedMessagesRef.current) {
223-
processFailedMessagesRef.current()
224-
}
228+
// Invalidate auth queries to allow network status to clear after reconnection
229+
queryClient.invalidateQueries({ queryKey: authQueryKeys.all })
230+
logger.debug(
231+
'[Connection] Invalidated auth queries to refresh network status',
232+
)
225233

226-
// Batch state updates using startTransition to prevent Bun crashes from cascading updates
227-
startTransition(() => {
228-
// Only show reconnection message if it's not the initial connection
229-
if (!isInitialConnection) {
230-
setShowReconnectionMessage(true)
231-
232-
// Hide the message after the configured duration using safe timeout
233-
setReconnectionTimeout(() => {
234-
startTransition(() => {
235-
setShowReconnectionMessage(false)
236-
})
237-
}, RECONNECTION_MESSAGE_DURATION_MS)
234+
// Process any failed messages and schedule them for retry (batched)
235+
if (processFailedMessagesRef.current) {
236+
processFailedMessagesRef.current()
238237
}
239238

240-
// Always trigger retry check (for both initial connection and reconnection)
241-
setConnectionEstablished(prev => prev + 1)
242-
})
243-
}, [queryClient, setReconnectionTimeout])
239+
// Batch state updates using startTransition to prevent Bun crashes from cascading updates
240+
startTransition(() => {
241+
// Only show reconnection message if it's not the initial connection
242+
if (!isInitialConnection) {
243+
setShowReconnectionMessage(true)
244+
245+
// Hide the message after the configured duration using safe timeout
246+
setReconnectionTimeout(() => {
247+
startTransition(() => {
248+
setShowReconnectionMessage(false)
249+
})
250+
}, RECONNECTION_MESSAGE_DURATION_MS)
251+
}
252+
253+
// Always trigger retry check (for both initial connection and reconnection)
254+
setConnectionEstablished((prev) => prev + 1)
255+
})
256+
},
257+
[queryClient, setReconnectionTimeout],
258+
)
244259

245260
const isConnected = useConnectionStatus(handleReconnection)
246261
const isConnectedRef = useRef(isConnected)
@@ -350,7 +365,7 @@ export const Chat = ({
350365

351366
return block
352367
})
353-
368+
354369
// Return original array reference if nothing changed
355370
return foundTarget ? result : blocks
356371
}
@@ -529,9 +544,9 @@ export const Chat = ({
529544
retryPendingMessages,
530545
processFailedMessages,
531546
} = useSendMessage({
532-
messages,
533-
allToggleIds,
534-
setMessages,
547+
messages,
548+
allToggleIds,
549+
setMessages,
535550
setFocusedAgentId,
536551
setInputFocused,
537552
inputRef,
@@ -582,19 +597,20 @@ export const Chat = ({
582597
retryPendingMessagesRef.current = retryPendingMessages
583598
processFailedMessagesRef.current = processFailedMessages
584599

585-
// Trigger retry when connection is established and we have pending messages
600+
// Trigger retry when we have pending messages (either connection restored or error occurred)
586601
useEffect(() => {
587-
// Prevent race condition: only schedule retry if one isn't already scheduled
588-
if (connectionEstablished > 0 && pendingRetryCount > 0 && !retryScheduledRef.current) {
602+
const shouldTriggerRetry =
603+
pendingRetryCount > 0 && !retryScheduledRef.current
604+
605+
if (shouldTriggerRetry) {
589606
logger.debug(
590607
{ pendingRetryCount, connectionEstablished },
591-
`[RETRY-EFFECT] Scheduling retry for ${pendingRetryCount} pending message(s) after ${RECONNECTION_RETRY_DELAY_MS}ms delay`
608+
`[RETRY-EFFECT] Scheduling retry for ${pendingRetryCount} pending message(s) after ${RECONNECTION_RETRY_DELAY_MS}ms delay`,
592609
)
593610
retryScheduledRef.current = true
594611

595-
// Small delay to ensure the connection is fully established
612+
// Small delay before retrying (gives connection time to stabilize)
596613
const timer = setTimeout(() => {
597-
logger.debug('[RETRY-EFFECT] Executing scheduled retry')
598614
retryPendingMessagesRef.current?.()
599615
retryScheduledRef.current = false
600616
}, RECONNECTION_RETRY_DELAY_MS)

‎cli/src/commands/router.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,11 +215,25 @@ export async function routeUserPrompt(params: {
215215
return
216216
}
217217

218+
// Prevent submitting new messages while another is processing
219+
if (isChainInProgressRef.current) {
220+
// Queue the message instead of sending immediately
221+
addToQueue(trimmed)
222+
setInputValue({ text: '', cursorPosition: 0, lastEditDueToNav: false })
223+
setTimeout(() => {
224+
scrollToLatest()
225+
}, 0)
226+
return
227+
}
228+
229+
// Set flag immediately to block concurrent submissions
230+
isChainInProgressRef.current = true
231+
218232
sendMessage({ content: trimmed, agentMode, postUserMessage })
219233

220234
setTimeout(() => {
221235
scrollToLatest()
222236
}, 0)
223-
237+
224238
return
225239
}

‎cli/src/hooks/use-send-message.ts‎

Lines changed: 33 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1173,8 +1173,7 @@ export const useSendMessage = ({
11731173
) {
11741174
const { agentId, chunk } = event
11751175

1176-
const previous =
1177-
agentStreamAccumulatorsRef.current[agentId] ?? ''
1176+
const previous = agentStreamAccumulatorsRef.current[agentId] ?? ''
11781177
if (!chunk) {
11791178
return
11801179
}
@@ -1347,10 +1346,9 @@ export const useSendMessage = ({
13471346
addActiveSubagent(event.agentId)
13481347

13491348
let foundExistingBlock = false
1350-
for (const [
1351-
tempId,
1352-
info,
1353-
] of Object.entries(spawnAgentsMapRef.current)) {
1349+
for (const [tempId, info] of Object.entries(
1350+
spawnAgentsMapRef.current,
1351+
)) {
13541352
const eventType = event.agentType || ''
13551353
const storedType = info.agentType || ''
13561354
// Match if exact match, or if eventType ends with storedType (e.g., 'codebuff/file-picker@0.0.2' matches 'file-picker')
@@ -1998,43 +1996,42 @@ export const useSendMessage = ({
19981996
const errorMessage =
19991997
error instanceof Error ? error.message : 'Unknown error occurred'
20001998

2001-
// Mark message as interrupted when an error occurs
2002-
markAiMessageInterrupted(aiMessageId)
2003-
2004-
applyMessageUpdate((prev) =>
2005-
prev.map((msg) => {
2006-
if (msg.id !== aiMessageId) {
2007-
return msg
2008-
}
2009-
const updatedContent =
2010-
msg.content + `\n\n**Error:** ${errorMessage}`
2011-
return {
2012-
...msg,
2013-
content: updatedContent,
2014-
}
2015-
}),
2016-
)
2017-
2018-
applyMessageUpdate((prev) =>
2019-
prev.map((msg) => {
2020-
if (msg.id !== aiMessageId) {
2021-
return msg
2022-
}
2023-
return { ...msg, isComplete: true }
2024-
}),
2025-
)
2026-
20271999
const pendingAlreadyScheduled =
20282000
userMessageId in pendingRetriesRef.current
20292001
const timedOutDueToSdk =
20302002
isNetworkError(error) && Boolean(error.streamTimedOut)
20312003

20322004
const shouldRetryError =
2033-
timedOutDueToSdk ||
2034-
!isConnectedRef.current ||
2035-
isRetryableError(error)
2005+
timedOutDueToSdk || !isConnectedRef.current || isRetryableError(error)
20362006

2007+
// Only mark as interrupted and show error if NOT retrying
2008+
// (retryable errors will be retried silently)
20372009
if (!shouldRetryError) {
2010+
markAiMessageInterrupted(aiMessageId)
2011+
2012+
applyMessageUpdate((prev) =>
2013+
prev.map((msg) => {
2014+
if (msg.id !== aiMessageId) {
2015+
return msg
2016+
}
2017+
const updatedContent =
2018+
msg.content + `\n\n**Error:** ${errorMessage}`
2019+
return {
2020+
...msg,
2021+
content: updatedContent,
2022+
}
2023+
}),
2024+
)
2025+
2026+
applyMessageUpdate((prev) =>
2027+
prev.map((msg) => {
2028+
if (msg.id !== aiMessageId) {
2029+
return msg
2030+
}
2031+
return { ...msg, isComplete: true }
2032+
}),
2033+
)
2034+
20382035
logger.error(
20392036
{ userMessageId, error: errorMessage },
20402037
'Non-retryable error encountered; not scheduling retry',
@@ -2100,9 +2097,7 @@ export const useSendMessage = ({
21002097

21012098
// Process connection failures and schedule them for retry (batched to avoid state cascades)
21022099
const processFailedMessages = useCallback(() => {
2103-
const failedMessages = Object.entries(
2104-
failedDueToConnectionRef.current,
2105-
)
2100+
const failedMessages = Object.entries(failedDueToConnectionRef.current)
21062101

21072102
if (failedMessages.length === 0) {
21082103
return
@@ -2210,7 +2205,6 @@ export const useSendMessage = ({
22102205
setCanProcessQueue,
22112206
])
22122207

2213-
22142208
return {
22152209
sendMessage,
22162210
clearMessages,

0 commit comments

Comments
 (0)