Something went wrong. Try again.
source dump of claude code forked from oppi.li/claude-code
Something went wrong. Try again.
15 kB · 479 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480/** * 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 displayconst 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}