yzcode-cli 1.0.1 → 1.0.3

package/tasks/LocalMainSessionTask.ts

@@ -0,0 +1,479 @@
+/**
+ * LocalMainSessionTask - Handles backgrounding the main session query.
+ *
+ * When user presses Ctrl+B twice during a query, the session is "backgrounded":
+ * - The query continues running in the background
+ * - The UI clears to a fresh prompt
+ * - A notification is sent when the query completes
+ *
+ * This reuses the LocalAgentTask state structure since the behavior is similar.
+ */
+
+import type { UUID } from 'crypto'
+import { randomBytes } from 'crypto'
+import {
+  OUTPUT_FILE_TAG,
+  STATUS_TAG,
+  SUMMARY_TAG,
+  TASK_ID_TAG,
+  TASK_NOTIFICATION_TAG,
+  TOOL_USE_ID_TAG,
+} from '../constants/xml.js'
+import { type QueryParams, query } from '../query.js'
+import { roughTokenCountEstimation } from '../services/tokenEstimation.js'
+import type { SetAppState } from '../Task.js'
+import { createTaskStateBase } from '../Task.js'
+import type {
+  AgentDefinition,
+  CustomAgentDefinition,
+} from '../tools/AgentTool/loadAgentsDir.js'
+import { asAgentId } from '../types/ids.js'
+import type { Message } from '../types/message.js'
+import { createAbortController } from '../utils/abortController.js'
+import {
+  runWithAgentContext,
+  type SubagentContext,
+} from '../utils/agentContext.js'
+import { registerCleanup } from '../utils/cleanupRegistry.js'
+import { logForDebugging } from '../utils/debug.js'
+import { logError } from '../utils/log.js'
+import { enqueuePendingNotification } from '../utils/messageQueueManager.js'
+import { emitTaskTerminatedSdk } from '../utils/sdkEventQueue.js'
+import {
+  getAgentTranscriptPath,
+  recordSidechainTranscript,
+} from '../utils/sessionStorage.js'
+import {
+  evictTaskOutput,
+  getTaskOutputPath,
+  initTaskOutputAsSymlink,
+} from '../utils/task/diskOutput.js'
+import { registerTask, updateTaskState } from '../utils/task/framework.js'
+import type { LocalAgentTaskState } from './LocalAgentTask/LocalAgentTask.js'
+
+// Main session tasks use LocalAgentTaskState with agentType='main-session'
+export type LocalMainSessionTaskState = LocalAgentTaskState & {
+  agentType: 'main-session'
+}
+
+/**
+ * Default agent definition for main session tasks when no agent is specified.
+ */
+const DEFAULT_MAIN_SESSION_AGENT: CustomAgentDefinition = {
+  agentType: 'main-session',
+  whenToUse: 'Main session query',
+  source: 'userSettings',
+  getSystemPrompt: () => '',
+}
+
+/**
+ * Generate a unique task ID for main session tasks.
+ * Uses 's' prefix to distinguish from agent tasks ('a' prefix).
+ */
+const TASK_ID_ALPHABET = '0123456789abcdefghijklmnopqrstuvwxyz'
+
+function generateMainSessionTaskId(): string {
+  const bytes = randomBytes(8)
+  let id = 's'
+  for (let i = 0; i < 8; i++) {
+    id += TASK_ID_ALPHABET[bytes[i]! % TASK_ID_ALPHABET.length]
+  }
+  return id
+}
+
+/**
+ * Register a backgrounded main session task.
+ * Called when the user backgrounds the current session query.
+ *
+ * @param description - Description of the task
+ * @param setAppState - State setter function
+ * @param mainThreadAgentDefinition - Optional agent definition if running with --agent
+ * @param existingAbortController - Optional abort controller to reuse (for backgrounding an active query)
+ * @returns Object with task ID and abort signal for stopping the background query
+ */
+export function registerMainSessionTask(
+  description: string,
+  setAppState: SetAppState,
+  mainThreadAgentDefinition?: AgentDefinition,
+  existingAbortController?: AbortController,
+): { taskId: string; abortSignal: AbortSignal } {
+  const taskId = generateMainSessionTaskId()
+
+  // Link output to an isolated per-task transcript file (same layout as
+  // sub-agents). Do NOT use getTranscriptPath() — that's the main session's
+  // file, and writing there from a background query after /clear would corrupt
+  // the post-clear conversation. The isolated path lets this task survive
+  // /clear: the symlink re-link in clearConversation handles session ID changes.
+  void initTaskOutputAsSymlink(
+    taskId,
+    getAgentTranscriptPath(asAgentId(taskId)),
+  )
+
+  // Use the existing abort controller if provided (important for backgrounding an active query)
+  // This ensures that aborting the task will abort the actual query
+  const abortController = existingAbortController ?? createAbortController()
+
+  const unregisterCleanup = registerCleanup(async () => {
+    // Clean up on process exit
+    setAppState(prev => {
+      const { [taskId]: removed, ...rest } = prev.tasks
+      return { ...prev, tasks: rest }
+    })
+  })
+
+  // Use provided agent definition or default
+  const selectedAgent = mainThreadAgentDefinition ?? DEFAULT_MAIN_SESSION_AGENT
+
+  // Create task state - already backgrounded since this is called when user backgrounds
+  const taskState: LocalMainSessionTaskState = {
+    ...createTaskStateBase(taskId, 'local_agent', description),
+    type: 'local_agent',
+    status: 'running',
+    agentId: taskId,
+    prompt: description,
+    selectedAgent,
+    agentType: 'main-session',
+    abortController,
+    unregisterCleanup,
+    retrieved: false,
+    lastReportedToolCount: 0,
+    lastReportedTokenCount: 0,
+    isBackgrounded: true, // Already backgrounded
+    pendingMessages: [],
+    retain: false,
+    diskLoaded: false,
+  }
+
+  logForDebugging(
+    `[LocalMainSessionTask] Registering task ${taskId} with description: ${description}`,
+  )
+  registerTask(taskState, setAppState)
+
+  // Verify task was registered by checking state
+  setAppState(prev => {
+    const hasTask = taskId in prev.tasks
+    logForDebugging(
+      `[LocalMainSessionTask] After registration, task ${taskId} exists in state: ${hasTask}`,
+    )
+    return prev
+  })
+
+  return { taskId, abortSignal: abortController.signal }
+}
+
+/**
+ * Complete the main session task and send notification.
+ * Called when the backgrounded query finishes.
+ */
+export function completeMainSessionTask(
+  taskId: string,
+  success: boolean,
+  setAppState: SetAppState,
+): void {
+  let wasBackgrounded = true
+  let toolUseId: string | undefined
+
+  updateTaskState<LocalMainSessionTaskState>(taskId, setAppState, task => {
+    if (task.status !== 'running') {
+      return task
+    }
+
+    // Track if task was backgrounded (for notification decision)
+    wasBackgrounded = task.isBackgrounded ?? true
+    toolUseId = task.toolUseId
+
+    task.unregisterCleanup?.()
+
+    return {
+      ...task,
+      status: success ? 'completed' : 'failed',
+      endTime: Date.now(),
+      messages: task.messages?.length ? [task.messages.at(-1)!] : undefined,
+    }
+  })
+
+  void evictTaskOutput(taskId)
+
+  // Only send notification if task is still backgrounded (not foregrounded)
+  // If foregrounded, user is watching it directly - no notification needed
+  if (wasBackgrounded) {
+    enqueueMainSessionNotification(
+      taskId,
+      'Background session',
+      success ? 'completed' : 'failed',
+      setAppState,
+      toolUseId,
+    )
+  } else {
+    // Foregrounded: no XML notification (TUI user is watching), but SDK
+    // consumers still need to see the task_started bookend close.
+    // Set notified so evictTerminalTask/generateTaskAttachments eviction
+    // guards pass; the backgrounded path sets this inside
+    // enqueueMainSessionNotification's check-and-set.
+    updateTaskState(taskId, setAppState, task => ({ ...task, notified: true }))
+    emitTaskTerminatedSdk(taskId, success ? 'completed' : 'failed', {
+      toolUseId,
+      summary: 'Background session',
+    })
+  }
+}
+
+/**
+ * Enqueue a notification about the backgrounded session completing.
+ */
+function enqueueMainSessionNotification(
+  taskId: string,
+  description: string,
+  status: 'completed' | 'failed',
+  setAppState: SetAppState,
+  toolUseId?: string,
+): void {
+  // Atomically check and set notified flag to prevent duplicate notifications.
+  let shouldEnqueue = false
+  updateTaskState(taskId, setAppState, task => {
+    if (task.notified) {
+      return task
+    }
+    shouldEnqueue = true
+    return { ...task, notified: true }
+  })
+
+  if (!shouldEnqueue) {
+    return
+  }
+
+  const summary =
+    status === 'completed'
+      ? `Background session "${description}" completed`
+      : `Background session "${description}" failed`
+
+  const toolUseIdLine = toolUseId
+    ? `\n<${TOOL_USE_ID_TAG}>${toolUseId}</${TOOL_USE_ID_TAG}>`
+    : ''
+
+  const outputPath = getTaskOutputPath(taskId)
+  const message = `<${TASK_NOTIFICATION_TAG}>
+<${TASK_ID_TAG}>${taskId}</${TASK_ID_TAG}>${toolUseIdLine}
+<${OUTPUT_FILE_TAG}>${outputPath}</${OUTPUT_FILE_TAG}>
+<${STATUS_TAG}>${status}</${STATUS_TAG}>
+<${SUMMARY_TAG}>${summary}</${SUMMARY_TAG}>
+</${TASK_NOTIFICATION_TAG}>`
+
+  enqueuePendingNotification({ value: message, mode: 'task-notification' })
+}
+
+/**
+ * Foreground a main session task - mark it as foregrounded so its output
+ * appears in the main view. The background query keeps running.
+ * Returns the task's accumulated messages, or undefined if task not found.
+ */
+export function foregroundMainSessionTask(
+  taskId: string,
+  setAppState: SetAppState,
+): Message[] | undefined {
+  let taskMessages: Message[] | undefined
+
+  setAppState(prev => {
+    const task = prev.tasks[taskId]
+    if (!task || task.type !== 'local_agent') {
+      return prev
+    }
+
+    taskMessages = (task as LocalMainSessionTaskState).messages
+
+    // Restore previous foregrounded task to background if it exists
+    const prevId = prev.foregroundedTaskId
+    const prevTask = prevId ? prev.tasks[prevId] : undefined
+    const restorePrev =
+      prevId && prevId !== taskId && prevTask?.type === 'local_agent'
+
+    return {
+      ...prev,
+      foregroundedTaskId: taskId,
+      tasks: {
+        ...prev.tasks,
+        ...(restorePrev && { [prevId]: { ...prevTask, isBackgrounded: true } }),
+        [taskId]: { ...task, isBackgrounded: false },
+      },
+    }
+  })
+
+  return taskMessages
+}
+
+/**
+ * Check if a task is a main session task (vs a regular agent task).
+ */
+export function isMainSessionTask(
+  task: unknown,
+): task is LocalMainSessionTaskState {
+  if (
+    typeof task !== 'object' ||
+    task === null ||
+    !('type' in task) ||
+    !('agentType' in task)
+  ) {
+    return false
+  }
+  return (
+    task.type === 'local_agent' &&
+    (task as LocalMainSessionTaskState).agentType === 'main-session'
+  )
+}
+
+// Max recent activities to keep for display
+const MAX_RECENT_ACTIVITIES = 5
+
+type ToolActivity = {
+  toolName: string
+  input: Record<string, unknown>
+}
+
+/**
+ * Start a fresh background session with the given messages.
+ *
+ * Spawns an independent query() call with the current messages and registers it
+ * as a background task. The caller's foreground query continues running normally.
+ */
+export function startBackgroundSession({
+  messages,
+  queryParams,
+  description,
+  setAppState,
+  agentDefinition,
+}: {
+  messages: Message[]
+  queryParams: Omit<QueryParams, 'messages'>
+  description: string
+  setAppState: SetAppState
+  agentDefinition?: AgentDefinition
+}): string {
+  const { taskId, abortSignal } = registerMainSessionTask(
+    description,
+    setAppState,
+    agentDefinition,
+  )
+
+  // Persist the pre-backgrounding conversation to the task's isolated
+  // transcript so TaskOutput shows context immediately. Subsequent messages
+  // are written incrementally below.
+  void recordSidechainTranscript(messages, taskId).catch(err =>
+    logForDebugging(`bg-session initial transcript write failed: ${err}`),
+  )
+
+  // Wrap in agent context so skill invocations scope to this task's agentId
+  // (not null). This lets clearInvokedSkills(preservedAgentIds) selectively
+  // preserve this task's skills across /clear. AsyncLocalStorage isolates
+  // concurrent async chains — this wrapper doesn't affect the foreground.
+  const agentContext: SubagentContext = {
+    agentId: taskId,
+    agentType: 'subagent',
+    subagentName: 'main-session',
+    isBuiltIn: true,
+  }
+
+  void runWithAgentContext(agentContext, async () => {
+    try {
+      const bgMessages: Message[] = [...messages]
+      const recentActivities: ToolActivity[] = []
+      let toolCount = 0
+      let tokenCount = 0
+      let lastRecordedUuid: UUID | null = messages.at(-1)?.uuid ?? null
+
+      for await (const event of query({
+        messages: bgMessages,
+        ...queryParams,
+      })) {
+        if (abortSignal.aborted) {
+          // Aborted mid-stream — completeMainSessionTask won't be reached.
+          // chat:killAgents path already marked notified + emitted; stopTask path did not.
+          let alreadyNotified = false
+          updateTaskState(taskId, setAppState, task => {
+            alreadyNotified = task.notified === true
+            return alreadyNotified ? task : { ...task, notified: true }
+          })
+          if (!alreadyNotified) {
+            emitTaskTerminatedSdk(taskId, 'stopped', {
+              summary: description,
+            })
+          }
+          return
+        }
+
+        if (
+          event.type !== 'user' &&
+          event.type !== 'assistant' &&
+          event.type !== 'system'
+        ) {
+          continue
+        }
+
+        bgMessages.push(event)
+
+        // Per-message write (matches runAgent.ts pattern) — gives live
+        // TaskOutput progress and keeps the transcript file current even if
+        // /clear re-links the symlink mid-run.
+        void recordSidechainTranscript([event], taskId, lastRecordedUuid).catch(
+          err => logForDebugging(`bg-session transcript write failed: ${err}`),
+        )
+        lastRecordedUuid = event.uuid
+
+        if (event.type === 'assistant') {
+          for (const block of event.message.content) {
+            if (block.type === 'text') {
+              tokenCount += roughTokenCountEstimation(block.text)
+            } else if (block.type === 'tool_use') {
+              toolCount++
+              const activity: ToolActivity = {
+                toolName: block.name,
+                input: block.input as Record<string, unknown>,
+              }
+              recentActivities.push(activity)
+              if (recentActivities.length > MAX_RECENT_ACTIVITIES) {
+                recentActivities.shift()
+              }
+            }
+          }
+        }
+
+        setAppState(prev => {
+          const task = prev.tasks[taskId]
+          if (!task || task.type !== 'local_agent') return prev
+          const prevProgress = task.progress
+          if (
+            prevProgress?.tokenCount === tokenCount &&
+            prevProgress.toolUseCount === toolCount &&
+            task.messages === bgMessages
+          ) {
+            return prev
+          }
+          return {
+            ...prev,
+            tasks: {
+              ...prev.tasks,
+              [taskId]: {
+                ...task,
+                progress: {
+                  tokenCount,
+                  toolUseCount: toolCount,
+                  recentActivities:
+                    prevProgress?.toolUseCount === toolCount
+                      ? prevProgress.recentActivities
+                      : [...recentActivities],
+                },
+                messages: bgMessages,
+              },
+            },
+          }
+        })
+      }
+
+      completeMainSessionTask(taskId, true, setAppState)
+    } catch (error) {
+      logError(error)
+      completeMainSessionTask(taskId, false, setAppState)
+    }
+  })
+
+  return taskId
+}