feat: improve connection handling with automatic retry and reconnection detection

- Add reconnection callback to detect and handle connection restoration
- Implement batch retry processing for messages that failed due to connection
- Add stream stall detection with 4-second timeout after initial output
- Refactor StatusBar to use unified statusIndicatorState pattern
- Display temporary reconnection message on connection restoration
- Track connection state transitions for proper retry triggering
- Add extensive logging for connection and retry flow debugging

Author: brandonkachen

cli/src/chat.tsx

@@ -34,6 +34,7 @@ import { useQueueControls } from './hooks/use-queue-controls'
 import { useChatStore } from './state/chat-store'
 import { createChatScrollAcceleration } from './utils/chat-scroll-accel'
 import { loadLocalAgents } from './utils/local-agent-registry'
+import { logger } from './utils/logger'
 import { buildMessageTree } from './utils/message-tree-utils'
 import { computeInputLayoutMetrics } from './utils/text-layout'
 import { createMarkdownPalette } from './utils/theme-system'
@@ -192,11 +193,57 @@ export const Chat = ({
   const sendMessageRef = useRef<SendMessageFn>()
 
   const { clipboardMessage } = useClipboard()
-  const isConnected = useConnectionStatus()
+  const [showReconnectionMessage, setShowReconnectionMessage] = useState(false)
+  const [connectionEstablished, setConnectionEstablished] = useState(0) // Increment to trigger retry check
+  const reconnectionTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+  const retryPendingMessagesRef = useRef<(() => Promise<void>) | null>(null)
+  const processFailedMessagesRef = useRef<(() => void) | null>(null)
+
+  const handleReconnection = useCallback((isInitialConnection: boolean) => {
+    logger.info(
+      { isInitialConnection },
+      `[Connection] ${isInitialConnection ? 'Initial connection' : 'Reconnection'} callback triggered`
+    )
+
+    // Process any failed messages and schedule them for retry (batched)
+    if (processFailedMessagesRef.current) {
+      processFailedMessagesRef.current()
+    }
+
+    // Only show reconnection message if it's not the initial connection
+    if (!isInitialConnection) {
+      setShowReconnectionMessage(true)
+
+      // Clear any existing timeout
+      if (reconnectionTimeoutRef.current) {
+        clearTimeout(reconnectionTimeoutRef.current)
+      }
+
+      // Hide the message after 2 seconds
+      reconnectionTimeoutRef.current = setTimeout(() => {
+        setShowReconnectionMessage(false)
+        reconnectionTimeoutRef.current = null
+      }, 2000)
+    }
+
+    // Always trigger retry check (for both initial connection and reconnection)
+    setConnectionEstablished(prev => prev + 1)
+  }, [])
+
+  const isConnected = useConnectionStatus(handleReconnection)
   const isConnectedRef = useRef(isConnected)
   useEffect(() => {
     isConnectedRef.current = isConnected
   }, [isConnected])
+
+  useEffect(() => {
+    return () => {
+      if (reconnectionTimeoutRef.current) {
+        clearTimeout(reconnectionTimeoutRef.current)
+      }
+    }
+  }, [])
+
   const mainAgentTimer = useElapsedTime()
   const timerStartTime = mainAgentTimer.startTime
 
@@ -397,8 +444,13 @@ export const Chat = ({
   // Timer events are currently tracked but not used for UI updates
   // Future: Could be used for analytics or debugging
 
-  const { sendMessage, clearMessages, pendingRetryCount, retryPendingMessages } =
-    useSendMessage({
+  const {
+    sendMessage,
+    clearMessages,
+    pendingRetryCount,
+    retryPendingMessages,
+    processFailedMessages,
+  } = useSendMessage({
       messages,
       allToggleIds,
       setMessages,
@@ -433,6 +485,27 @@ export const Chat = ({
     })
 
   sendMessageRef.current = sendMessage
+  retryPendingMessagesRef.current = retryPendingMessages
+  processFailedMessagesRef.current = processFailedMessages
+
+  // Trigger retry when connection is established and we have pending messages
+  useEffect(() => {
+    if (connectionEstablished > 0 && pendingRetryCount > 0) {
+      logger.info(
+        { pendingRetryCount, connectionEstablished },
+        `[RETRY-EFFECT] CONDITIONS MET! Will retry ${pendingRetryCount} pending message(s) after 500ms delay...`
+      )
+      // Small delay to ensure the connection is fully established
+      const timer = setTimeout(() => {
+        logger.info('[RETRY-EFFECT] Calling retryPendingMessages()')
+        retryPendingMessagesRef.current?.()
+      }, 500)
+      return () => {
+        clearTimeout(timer)
+      }
+    }
+    return undefined
+  }, [connectionEstablished, pendingRetryCount])
 
   const { inputWidth, handleBuildFast, handleBuildMax } = useChatInput({
     inputValue,
@@ -593,6 +666,7 @@ export const Chat = ({
     streamStatus,
     nextCtrlCWillExit,
     isConnected,
+    showReconnectionMessage,
   })
   const hasStatusIndicatorContent = statusIndicatorState.kind !== 'idle'
   const inputBoxTitle = useMemo(() => {
@@ -705,13 +779,12 @@ export const Chat = ({
           <StatusBar
             clipboardMessage={clipboardMessage}
             streamStatus={streamStatus}
+            statusIndicatorState={statusIndicatorState}
             timerStartTime={timerStartTime}
             nextCtrlCWillExit={nextCtrlCWillExit}
             isConnected={isConnected}
             isAtBottom={isAtBottom}
             scrollToLatest={scrollToLatest}
-            pendingRetryCount={pendingRetryCount}
-            retryPendingMessages={retryPendingMessages}
           />
         )}
 

cli/src/components/status-bar.tsx

@@ -2,36 +2,34 @@ import React, { useEffect, useState } from 'react'
 
 import { ShimmerText } from './shimmer-text'
 import { ScrollToBottomButton } from './scroll-to-bottom-button'
-import { Button } from './button'
 import { useTheme } from '../hooks/use-theme'
 import { formatElapsedTime } from '../utils/format-elapsed-time'
 
 import type { StreamStatus } from '../hooks/use-message-queue'
+import type { StatusIndicatorState } from '../utils/status-indicator-state'
 
 const SHIMMER_INTERVAL_MS = 160
 
 interface StatusBarProps {
   clipboardMessage: string | null
   streamStatus: StreamStatus
+  statusIndicatorState: StatusIndicatorState
   timerStartTime: number | null
   nextCtrlCWillExit: boolean
   isConnected: boolean
   isAtBottom: boolean
   scrollToLatest: () => void
-  pendingRetryCount: number
-  retryPendingMessages: () => Promise<void>
 }
 
 export const StatusBar = ({
   clipboardMessage,
   streamStatus,
+  statusIndicatorState,
   timerStartTime,
   nextCtrlCWillExit,
   isConnected,
   isAtBottom,
   scrollToLatest,
-  pendingRetryCount,
-  retryPendingMessages,
 }: StatusBarProps) => {
   const theme = useTheme()
   const [elapsedSeconds, setElapsedSeconds] = useState(0)
@@ -65,43 +63,49 @@ export const StatusBar = ({
   )
 
   const renderStatusIndicator = () => {
-    if (nextCtrlCWillExit) {
-      return (
-        <text fg={theme.secondary} style={{ wrapMode: 'none' }}>
-          Press Ctrl-C again to exit
-        </text>
-      )
+    switch (statusIndicatorState.kind) {
+      case 'ctrlC':
+        return (
+          <text fg={theme.secondary} style={{ wrapMode: 'none' }}>
+            Press Ctrl-C again to exit
+          </text>
+        )
+
+      case 'reconnected':
+        return (
+          <text fg={theme.success || theme.primary} style={{ wrapMode: 'none' }}>
+            ✓ Reconnected
+          </text>
+        )
+
+      case 'clipboard':
+        return (
+          <text fg={theme.primary} style={{ wrapMode: 'none' }}>
+            {statusIndicatorState.message}
+          </text>
+        )
+
+      case 'connecting':
+        return renderShimmer({ text: 'connecting...' })
+
+      case 'waiting':
+        return renderShimmer({
+          text: 'thinking...',
+          interval: SHIMMER_INTERVAL_MS,
+          primaryColor: theme.secondary,
+        })
+
+      case 'streaming':
+        return renderShimmer({
+          text: 'working...',
+          interval: SHIMMER_INTERVAL_MS,
+          primaryColor: theme.secondary,
+        })
+
+      case 'idle':
+      default:
+        return null
     }
-
-    if (clipboardMessage) {
-      return (
-        <text fg={theme.primary} style={{ wrapMode: 'none' }}>
-          {clipboardMessage}
-        </text>
-      )
-    }
-
-    if (!isConnected) {
-      return renderShimmer({ text: 'connecting...' })
-    }
-
-    if (streamStatus === 'waiting') {
-      return renderShimmer({
-        text: 'thinking...',
-        interval: SHIMMER_INTERVAL_MS,
-        primaryColor: theme.secondary,
-      })
-    }
-
-    if (streamStatus === 'streaming') {
-      return renderShimmer({
-        text: 'working...',
-        interval: SHIMMER_INTERVAL_MS,
-        primaryColor: theme.secondary,
-      })
-    }
-
-    return null
   }
 
   const renderElapsedTime = () => {
@@ -118,44 +122,9 @@ export const StatusBar = ({
 
   const statusIndicatorContent = renderStatusIndicator()
   const elapsedTimeContent = renderElapsedTime()
-  const hasPendingRetries = pendingRetryCount > 0
-  const pendingRetryMessage =
-    pendingRetryCount === 1
-      ? 'Message send interrupted'
-      : `${pendingRetryCount} messages interrupted`
-  const handleRetryClick = () => {
-    void retryPendingMessages()
-  }
-  const pendingRetryContent = hasPendingRetries ? (
-    <box
-      style={{
-        flexDirection: 'row',
-        alignItems: 'center',
-        gap: 1,
-      }}
-    >
-      <text fg={theme.primary} style={{ wrapMode: 'none' }}>
-        {pendingRetryMessage}
-      </text>
-      <Button
-        style={{
-          borderStyle: 'round',
-          borderColor: theme.primary,
-          paddingLeft: 1,
-          paddingRight: 1,
-        }}
-        onClick={handleRetryClick}
-      >
-        <text fg={theme.primary} style={{ wrapMode: 'none' }}>
-          Retry now
-        </text>
-      </Button>
-    </box>
-  ) : null
 
   // Only show gray background when there's status indicator or timer content
-  const hasContent =
-    hasPendingRetries || statusIndicatorContent || elapsedTimeContent
+  const hasContent = statusIndicatorContent || elapsedTimeContent
 
   return (
     <box
@@ -176,7 +145,7 @@ export const StatusBar = ({
           flexBasis: 0,
         }}
       >
-        {pendingRetryContent ?? statusIndicatorContent}
+        {statusIndicatorContent}
       </box>
 
       <box style={{ flexShrink: 0 }}>

cli/src/hooks/use-connection-status.ts

@@ -37,15 +37,20 @@ export function getNextInterval(consecutiveSuccesses: number): number {
  * Hook to monitor connection status to the Codebuff backend.
  * Uses adaptive exponential backoff to reduce polling frequency when connection is stable.
  */
-export const useConnectionStatus = () => {
+export const useConnectionStatus = (
+  onReconnected?: (isInitialConnection: boolean) => void,
+) => {
 
-  const [isConnected, setIsConnected] = useState(true)
+  // Start with null to indicate unknown state
+  const [isConnected, setIsConnected] = useState<boolean | null>(null)
 
   useEffect(() => {
     let isMounted = true
     let timeoutId: NodeJS.Timeout | null = null
     let consecutiveSuccesses = 0
     let currentInterval: number = HEALTH_CHECK_CONFIG.INITIAL_INTERVAL
+    // Start with null to properly detect first connection
+    let previousConnected: boolean | null = null
 
     const scheduleNextCheck = (interval: number) => {
       if (!isMounted) return
@@ -56,6 +61,7 @@ export const useConnectionStatus = () => {
       const client = getCodebuffClient()
       if (!client) {
         if (isMounted) {
+          previousConnected = false
           setIsConnected(false)
           consecutiveSuccesses = 0
           currentInterval = HEALTH_CHECK_CONFIG.INITIAL_INTERVAL
@@ -72,6 +78,33 @@ export const useConnectionStatus = () => {
         const connected = await client.checkConnection()
         if (!isMounted) return
 
+        // Detect reconnection (was disconnected/unknown, now connected)
+        // Also handle first connection (for pending messages from previous session)
+        if (connected && (previousConnected === false || previousConnected === null)) {
+          const isInitialConnection = previousConnected === null
+          logger.info(
+            {
+              previousConnected,
+              connected,
+              isInitialConnection,
+              hasCallback: !!onReconnected
+            },
+            isInitialConnection
+              ? '[CONNECTION-CHECK] Initial connection established - checking for pending messages'
+              : '[CONNECTION-CHECK] Connection restored to backend - triggering reconnection handler'
+          )
+          if (onReconnected) {
+            logger.info(
+              { isInitialConnection },
+              '[CONNECTION-CHECK] Calling onReconnected callback'
+            )
+            onReconnected(isInitialConnection)
+          }
+        } else if (!connected && previousConnected === true) {
+          logger.info('[CONNECTION-CHECK] Connection lost!')
+        }
+
+        previousConnected = connected
         setIsConnected(connected)
 
         if (connected) {
@@ -105,6 +138,7 @@ export const useConnectionStatus = () => {
       } catch (error) {
         logger.debug({ error }, 'Connection check failed')
         if (isMounted) {
+          previousConnected = false
           setIsConnected(false)
           consecutiveSuccesses = 0
           currentInterval = HEALTH_CHECK_CONFIG.INITIAL_INTERVAL
@@ -122,7 +156,8 @@ export const useConnectionStatus = () => {
         clearTimeout(timeoutId)
       }
     }
-  }, [])
+  }, [onReconnected])
 
-  return isConnected
+  // Return false while checking initial connection status
+  return isConnected ?? false
 }