1  /**
  2   * Hook for managing session backgrounding (Ctrl+B to background/foreground sessions).
  3   *
  4   * Handles:
  5   * - Calling onBackgroundQuery to spawn a background task for the current query
  6   * - Re-backgrounding foregrounded tasks
  7   * - Syncing foregrounded task messages/state to main view
  8   */
  9  
 10  import { useCallback, useEffect, useRef } from 'react'
 11  import { useAppState, useSetAppState } from '../state/AppState.js'
 12  import type { Message } from '../types/message.js'
 13  
 14  type UseSessionBackgroundingProps = {
 15    setMessages: (messages: Message[] | ((prev: Message[]) => Message[])) => void
 16    setIsLoading: (loading: boolean) => void
 17    resetLoadingState: () => void
 18    setAbortController: (controller: AbortController | null) => void
 19    onBackgroundQuery: () => void
 20  }
 21  
 22  type UseSessionBackgroundingResult = {
 23    /** Call when user wants to background (Ctrl+B) */
 24    handleBackgroundSession: () => void
 25  }
 26  
 27  export function useSessionBackgrounding({
 28    setMessages,
 29    setIsLoading,
 30    resetLoadingState,
 31    setAbortController,
 32    onBackgroundQuery,
 33  }: UseSessionBackgroundingProps): UseSessionBackgroundingResult {
 34    const foregroundedTaskId = useAppState(s => s.foregroundedTaskId)
 35    const foregroundedTask = useAppState(s =>
 36      s.foregroundedTaskId ? s.tasks[s.foregroundedTaskId] : undefined,
 37    )
 38    const setAppState = useSetAppState()
 39    const lastSyncedMessagesLengthRef = useRef<number>(0)
 40  
 41    const handleBackgroundSession = useCallback(() => {
 42      if (foregroundedTaskId) {
 43        // Re-background the foregrounded task
 44        setAppState(prev => {
 45          const taskId = prev.foregroundedTaskId
 46          if (!taskId) return prev
 47          const task = prev.tasks[taskId]
 48          if (!task) {
 49            return { ...prev, foregroundedTaskId: undefined }
 50          }
 51          return {
 52            ...prev,
 53            foregroundedTaskId: undefined,
 54            tasks: {
 55              ...prev.tasks,
 56              [taskId]: { ...task, isBackgrounded: true },
 57            },
 58          }
 59        })
 60        setMessages([])
 61        resetLoadingState()
 62        setAbortController(null)
 63        return
 64      }
 65  
 66      onBackgroundQuery()
 67    }, [
 68      foregroundedTaskId,
 69      setAppState,
 70      setMessages,
 71      resetLoadingState,
 72      setAbortController,
 73      onBackgroundQuery,
 74    ])
 75  
 76    // Sync foregrounded task's messages and loading state to the main view
 77    useEffect(() => {
 78      if (!foregroundedTaskId) {
 79        // Reset when no foregrounded task
 80        lastSyncedMessagesLengthRef.current = 0
 81        return
 82      }
 83  
 84      if (!foregroundedTask || foregroundedTask.type !== 'local_agent') {
 85        setAppState(prev => ({ ...prev, foregroundedTaskId: undefined }))
 86        resetLoadingState()
 87        lastSyncedMessagesLengthRef.current = 0
 88        return
 89      }
 90  
 91      // Sync messages from background task to main view
 92      // Only update if messages have actually changed to avoid redundant renders
 93      const taskMessages = foregroundedTask.messages ?? []
 94      if (taskMessages.length !== lastSyncedMessagesLengthRef.current) {
 95        lastSyncedMessagesLengthRef.current = taskMessages.length
 96        setMessages([...taskMessages])
 97      }
 98  
 99      if (foregroundedTask.status === 'running') {
100        // Check if the task was aborted (user pressed Escape)
101        const taskAbortController = foregroundedTask.abortController
102        if (taskAbortController?.signal.aborted) {
103          // Task was aborted - clear foregrounded state immediately
104          setAppState(prev => {
105            if (!prev.foregroundedTaskId) return prev
106            const task = prev.tasks[prev.foregroundedTaskId]
107            if (!task) return { ...prev, foregroundedTaskId: undefined }
108            return {
109              ...prev,
110              foregroundedTaskId: undefined,
111              tasks: {
112                ...prev.tasks,
113                [prev.foregroundedTaskId]: { ...task, isBackgrounded: true },
114              },
115            }
116          })
117          resetLoadingState()
118          setAbortController(null)
119          lastSyncedMessagesLengthRef.current = 0
120          return
121        }
122  
123        setIsLoading(true)
124        // Set abort controller to the foregrounded task's controller for Escape handling
125        if (taskAbortController) {
126          setAbortController(taskAbortController)
127        }
128      } else {
129        // Task completed - restore to background and clear foregrounded view
130        setAppState(prev => {
131          const taskId = prev.foregroundedTaskId
132          if (!taskId) return prev
133          const task = prev.tasks[taskId]
134          if (!task) return { ...prev, foregroundedTaskId: undefined }
135          return {
136            ...prev,
137            foregroundedTaskId: undefined,
138            tasks: { ...prev.tasks, [taskId]: { ...task, isBackgrounded: true } },
139          }
140        })
141        resetLoadingState()
142        setAbortController(null)
143        lastSyncedMessagesLengthRef.current = 0
144      }
145    }, [
146      foregroundedTaskId,
147      foregroundedTask,
148      setAppState,
149      setMessages,
150      setIsLoading,
151      resetLoadingState,
152      setAbortController,
153    ])
154  
155    return {
156      handleBackgroundSession,
157    }
158  }