1  /**
  2   * LocalMainSessionTask - Handles backgrounding the main session query.
  3   *
  4   * When user presses Ctrl+B twice during a query, the session is "backgrounded":
  5   * - The query continues running in the background
  6   * - The UI clears to a fresh prompt
  7   * - A notification is sent when the query completes
  8   *
  9   * This reuses the LocalAgentTask state structure since the behavior is similar.
 10   */
 11  
 12  import type { UUID } from 'crypto'
 13  import { randomBytes } from 'crypto'
 14  import {
 15    OUTPUT_FILE_TAG,
 16    STATUS_TAG,
 17    SUMMARY_TAG,
 18    TASK_ID_TAG,
 19    TASK_NOTIFICATION_TAG,
 20    TOOL_USE_ID_TAG,
 21  } from '../constants/xml.js'
 22  import { type QueryParams, query } from '../query.js'
 23  import { roughTokenCountEstimation } from '../services/tokenEstimation.js'
 24  import type { SetAppState } from '../Task.js'
 25  import { createTaskStateBase } from '../Task.js'
 26  import type {
 27    AgentDefinition,
 28    CustomAgentDefinition,
 29  } from '../tools/AgentTool/loadAgentsDir.js'
 30  import { asAgentId } from '../types/ids.js'
 31  import type { Message } from '../types/message.js'
 32  import { createAbortController } from '../utils/abortController.js'
 33  import {
 34    runWithAgentContext,
 35    type SubagentContext,
 36  } from '../utils/agentContext.js'
 37  import { registerCleanup } from '../utils/cleanupRegistry.js'
 38  import { logForDebugging } from '../utils/debug.js'
 39  import { logError } from '../utils/log.js'
 40  import { enqueuePendingNotification } from '../utils/messageQueueManager.js'
 41  import { emitTaskTerminatedSdk } from '../utils/sdkEventQueue.js'
 42  import {
 43    getAgentTranscriptPath,
 44    recordSidechainTranscript,
 45  } from '../utils/sessionStorage.js'
 46  import {
 47    evictTaskOutput,
 48    getTaskOutputPath,
 49    initTaskOutputAsSymlink,
 50  } from '../utils/task/diskOutput.js'
 51  import { registerTask, updateTaskState } from '../utils/task/framework.js'
 52  import type { LocalAgentTaskState } from './LocalAgentTask/LocalAgentTask.js'
 53  
 54  // Main session tasks use LocalAgentTaskState with agentType='main-session'
 55  export type LocalMainSessionTaskState = LocalAgentTaskState & {
 56    agentType: 'main-session'
 57  }
 58  
 59  /**
 60   * Default agent definition for main session tasks when no agent is specified.
 61   */
 62  const DEFAULT_MAIN_SESSION_AGENT: CustomAgentDefinition = {
 63    agentType: 'main-session',
 64    whenToUse: 'Main session query',
 65    source: 'userSettings',
 66    getSystemPrompt: () => '',
 67  }
 68  
 69  /**
 70   * Generate a unique task ID for main session tasks.
 71   * Uses 's' prefix to distinguish from agent tasks ('a' prefix).
 72   */
 73  const TASK_ID_ALPHABET = '0123456789abcdefghijklmnopqrstuvwxyz'
 74  
 75  function generateMainSessionTaskId(): string {
 76    const bytes = randomBytes(8)
 77    let id = 's'
 78    for (let i = 0; i < 8; i++) {
 79      id += TASK_ID_ALPHABET[bytes[i]! % TASK_ID_ALPHABET.length]
 80    }
 81    return id
 82  }
 83  
 84  /**
 85   * Register a backgrounded main session task.
 86   * Called when the user backgrounds the current session query.
 87   *
 88   * @param description - Description of the task
 89   * @param setAppState - State setter function
 90   * @param mainThreadAgentDefinition - Optional agent definition if running with --agent
 91   * @param existingAbortController - Optional abort controller to reuse (for backgrounding an active query)
 92   * @returns Object with task ID and abort signal for stopping the background query
 93   */
 94  export function registerMainSessionTask(
 95    description: string,
 96    setAppState: SetAppState,
 97    mainThreadAgentDefinition?: AgentDefinition,
 98    existingAbortController?: AbortController,
 99  ): { taskId: string; abortSignal: AbortSignal } {
100    const taskId = generateMainSessionTaskId()
101  
102    // Link output to an isolated per-task transcript file (same layout as
103    // sub-agents). Do NOT use getTranscriptPath() — that's the main session's
104    // file, and writing there from a background query after /clear would corrupt
105    // the post-clear conversation. The isolated path lets this task survive
106    // /clear: the symlink re-link in clearConversation handles session ID changes.
107    void initTaskOutputAsSymlink(
108      taskId,
109      getAgentTranscriptPath(asAgentId(taskId)),
110    )
111  
112    // Use the existing abort controller if provided (important for backgrounding an active query)
113    // This ensures that aborting the task will abort the actual query
114    const abortController = existingAbortController ?? createAbortController()
115  
116    const unregisterCleanup = registerCleanup(async () => {
117      // Clean up on process exit
118      setAppState(prev => {
119        const { [taskId]: removed, ...rest } = prev.tasks
120        return { ...prev, tasks: rest }
121      })
122    })
123  
124    // Use provided agent definition or default
125    const selectedAgent = mainThreadAgentDefinition ?? DEFAULT_MAIN_SESSION_AGENT
126  
127    // Create task state - already backgrounded since this is called when user backgrounds
128    const taskState: LocalMainSessionTaskState = {
129      ...createTaskStateBase(taskId, 'local_agent', description),
130      type: 'local_agent',
131      status: 'running',
132      agentId: taskId,
133      prompt: description,
134      selectedAgent,
135      agentType: 'main-session',
136      abortController,
137      unregisterCleanup,
138      retrieved: false,
139      lastReportedToolCount: 0,
140      lastReportedTokenCount: 0,
141      isBackgrounded: true, // Already backgrounded
142      pendingMessages: [],
143      retain: false,
144      diskLoaded: false,
145    }
146  
147    logForDebugging(
148      `[LocalMainSessionTask] Registering task ${taskId} with description: ${description}`,
149    )
150    registerTask(taskState, setAppState)
151  
152    // Verify task was registered by checking state
153    setAppState(prev => {
154      const hasTask = taskId in prev.tasks
155      logForDebugging(
156        `[LocalMainSessionTask] After registration, task ${taskId} exists in state: ${hasTask}`,
157      )
158      return prev
159    })
160  
161    return { taskId, abortSignal: abortController.signal }
162  }
163  
164  /**
165   * Complete the main session task and send notification.
166   * Called when the backgrounded query finishes.
167   */
168  export function completeMainSessionTask(
169    taskId: string,
170    success: boolean,
171    setAppState: SetAppState,
172  ): void {
173    let wasBackgrounded = true
174    let toolUseId: string | undefined
175  
176    updateTaskState<LocalMainSessionTaskState>(taskId, setAppState, task => {
177      if (task.status !== 'running') {
178        return task
179      }
180  
181      // Track if task was backgrounded (for notification decision)
182      wasBackgrounded = task.isBackgrounded ?? true
183      toolUseId = task.toolUseId
184  
185      task.unregisterCleanup?.()
186  
187      return {
188        ...task,
189        status: success ? 'completed' : 'failed',
190        endTime: Date.now(),
191        messages: task.messages?.length ? [task.messages.at(-1)!] : undefined,
192      }
193    })
194  
195    void evictTaskOutput(taskId)
196  
197    // Only send notification if task is still backgrounded (not foregrounded)
198    // If foregrounded, user is watching it directly - no notification needed
199    if (wasBackgrounded) {
200      enqueueMainSessionNotification(
201        taskId,
202        'Background session',
203        success ? 'completed' : 'failed',
204        setAppState,
205        toolUseId,
206      )
207    } else {
208      // Foregrounded: no XML notification (TUI user is watching), but SDK
209      // consumers still need to see the task_started bookend close.
210      // Set notified so evictTerminalTask/generateTaskAttachments eviction
211      // guards pass; the backgrounded path sets this inside
212      // enqueueMainSessionNotification's check-and-set.
213      updateTaskState(taskId, setAppState, task => ({ ...task, notified: true }))
214      emitTaskTerminatedSdk(taskId, success ? 'completed' : 'failed', {
215        toolUseId,
216        summary: 'Background session',
217      })
218    }
219  }
220  
221  /**
222   * Enqueue a notification about the backgrounded session completing.
223   */
224  function enqueueMainSessionNotification(
225    taskId: string,
226    description: string,
227    status: 'completed' | 'failed',
228    setAppState: SetAppState,
229    toolUseId?: string,
230  ): void {
231    // Atomically check and set notified flag to prevent duplicate notifications.
232    let shouldEnqueue = false
233    updateTaskState(taskId, setAppState, task => {
234      if (task.notified) {
235        return task
236      }
237      shouldEnqueue = true
238      return { ...task, notified: true }
239    })
240  
241    if (!shouldEnqueue) {
242      return
243    }
244  
245    const summary =
246      status === 'completed'
247        ? `Background session "${description}" completed`
248        : `Background session "${description}" failed`
249  
250    const toolUseIdLine = toolUseId
251      ? `\n<${TOOL_USE_ID_TAG}>${toolUseId}</${TOOL_USE_ID_TAG}>`
252      : ''
253  
254    const outputPath = getTaskOutputPath(taskId)
255    const message = `<${TASK_NOTIFICATION_TAG}>
256  <${TASK_ID_TAG}>${taskId}</${TASK_ID_TAG}>${toolUseIdLine}
257  <${OUTPUT_FILE_TAG}>${outputPath}</${OUTPUT_FILE_TAG}>
258  <${STATUS_TAG}>${status}</${STATUS_TAG}>
259  <${SUMMARY_TAG}>${summary}</${SUMMARY_TAG}>
260  </${TASK_NOTIFICATION_TAG}>`
261  
262    enqueuePendingNotification({ value: message, mode: 'task-notification' })
263  }
264  
265  /**
266   * Foreground a main session task - mark it as foregrounded so its output
267   * appears in the main view. The background query keeps running.
268   * Returns the task's accumulated messages, or undefined if task not found.
269   */
270  export function foregroundMainSessionTask(
271    taskId: string,
272    setAppState: SetAppState,
273  ): Message[] | undefined {
274    let taskMessages: Message[] | undefined
275  
276    setAppState(prev => {
277      const task = prev.tasks[taskId]
278      if (!task || task.type !== 'local_agent') {
279        return prev
280      }
281  
282      taskMessages = (task as LocalMainSessionTaskState).messages
283  
284      // Restore previous foregrounded task to background if it exists
285      const prevId = prev.foregroundedTaskId
286      const prevTask = prevId ? prev.tasks[prevId] : undefined
287      const restorePrev =
288        prevId && prevId !== taskId && prevTask?.type === 'local_agent'
289  
290      return {
291        ...prev,
292        foregroundedTaskId: taskId,
293        tasks: {
294          ...prev.tasks,
295          ...(restorePrev && { [prevId]: { ...prevTask, isBackgrounded: true } }),
296          [taskId]: { ...task, isBackgrounded: false },
297        },
298      }
299    })
300  
301    return taskMessages
302  }
303  
304  /**
305   * Check if a task is a main session task (vs a regular agent task).
306   */
307  export function isMainSessionTask(
308    task: unknown,
309  ): task is LocalMainSessionTaskState {
310    if (
311      typeof task !== 'object' ||
312      task === null ||
313      !('type' in task) ||
314      !('agentType' in task)
315    ) {
316      return false
317    }
318    return (
319      task.type === 'local_agent' &&
320      (task as LocalMainSessionTaskState).agentType === 'main-session'
321    )
322  }
323  
324  // Max recent activities to keep for display
325  const MAX_RECENT_ACTIVITIES = 5
326  
327  type ToolActivity = {
328    toolName: string
329    input: Record<string, unknown>
330  }
331  
332  /**
333   * Start a fresh background session with the given messages.
334   *
335   * Spawns an independent query() call with the current messages and registers it
336   * as a background task. The caller's foreground query continues running normally.
337   */
338  export function startBackgroundSession({
339    messages,
340    queryParams,
341    description,
342    setAppState,
343    agentDefinition,
344  }: {
345    messages: Message[]
346    queryParams: Omit<QueryParams, 'messages'>
347    description: string
348    setAppState: SetAppState
349    agentDefinition?: AgentDefinition
350  }): string {
351    const { taskId, abortSignal } = registerMainSessionTask(
352      description,
353      setAppState,
354      agentDefinition,
355    )
356  
357    // Persist the pre-backgrounding conversation to the task's isolated
358    // transcript so TaskOutput shows context immediately. Subsequent messages
359    // are written incrementally below.
360    void recordSidechainTranscript(messages, taskId).catch(err =>
361      logForDebugging(`bg-session initial transcript write failed: ${err}`),
362    )
363  
364    // Wrap in agent context so skill invocations scope to this task's agentId
365    // (not null). This lets clearInvokedSkills(preservedAgentIds) selectively
366    // preserve this task's skills across /clear. AsyncLocalStorage isolates
367    // concurrent async chains — this wrapper doesn't affect the foreground.
368    const agentContext: SubagentContext = {
369      agentId: taskId,
370      agentType: 'subagent',
371      subagentName: 'main-session',
372      isBuiltIn: true,
373    }
374  
375    void runWithAgentContext(agentContext, async () => {
376      try {
377        const bgMessages: Message[] = [...messages]
378        const recentActivities: ToolActivity[] = []
379        let toolCount = 0
380        let tokenCount = 0
381        let lastRecordedUuid: UUID | null = messages.at(-1)?.uuid ?? null
382  
383        for await (const event of query({
384          messages: bgMessages,
385          ...queryParams,
386        })) {
387          if (abortSignal.aborted) {
388            // Aborted mid-stream — completeMainSessionTask won't be reached.
389            // chat:killAgents path already marked notified + emitted; stopTask path did not.
390            let alreadyNotified = false
391            updateTaskState(taskId, setAppState, task => {
392              alreadyNotified = task.notified === true
393              return alreadyNotified ? task : { ...task, notified: true }
394            })
395            if (!alreadyNotified) {
396              emitTaskTerminatedSdk(taskId, 'stopped', {
397                summary: description,
398              })
399            }
400            return
401          }
402  
403          if (
404            event.type !== 'user' &&
405            event.type !== 'assistant' &&
406            event.type !== 'system'
407          ) {
408            continue
409          }
410  
411          bgMessages.push(event)
412  
413          // Per-message write (matches runAgent.ts pattern) — gives live
414          // TaskOutput progress and keeps the transcript file current even if
415          // /clear re-links the symlink mid-run.
416          void recordSidechainTranscript([event], taskId, lastRecordedUuid).catch(
417            err => logForDebugging(`bg-session transcript write failed: ${err}`),
418          )
419          lastRecordedUuid = event.uuid
420  
421          if (event.type === 'assistant') {
422            for (const block of event.message.content) {
423              if (block.type === 'text') {
424                tokenCount += roughTokenCountEstimation(block.text)
425              } else if (block.type === 'tool_use') {
426                toolCount++
427                const activity: ToolActivity = {
428                  toolName: block.name,
429                  input: block.input as Record<string, unknown>,
430                }
431                recentActivities.push(activity)
432                if (recentActivities.length > MAX_RECENT_ACTIVITIES) {
433                  recentActivities.shift()
434                }
435              }
436            }
437          }
438  
439          setAppState(prev => {
440            const task = prev.tasks[taskId]
441            if (!task || task.type !== 'local_agent') return prev
442            const prevProgress = task.progress
443            if (
444              prevProgress?.tokenCount === tokenCount &&
445              prevProgress.toolUseCount === toolCount &&
446              task.messages === bgMessages
447            ) {
448              return prev
449            }
450            return {
451              ...prev,
452              tasks: {
453                ...prev.tasks,
454                [taskId]: {
455                  ...task,
456                  progress: {
457                    tokenCount,
458                    toolUseCount: toolCount,
459                    recentActivities:
460                      prevProgress?.toolUseCount === toolCount
461                        ? prevProgress.recentActivities
462                        : [...recentActivities],
463                  },
464                  messages: bgMessages,
465                },
466              },
467            }
468          })
469        }
470  
471        completeMainSessionTask(taskId, true, setAppState)
472      } catch (error) {
473        logError(error)
474        completeMainSessionTask(taskId, false, setAppState)
475      }
476    })
477  
478    return taskId
479  }