1  /**
  2   * In-process teammate spawning
  3   *
  4   * Creates and registers an in-process teammate task. Unlike process-based
  5   * teammates (tmux/iTerm2), in-process teammates run in the same Node.js
  6   * process using AsyncLocalStorage for context isolation.
  7   *
  8   * The actual agent execution loop is handled by InProcessTeammateTask
  9   * component (Task #14). This module handles:
 10   * 1. Creating TeammateContext
 11   * 2. Creating linked AbortController
 12   * 3. Registering InProcessTeammateTaskState in AppState
 13   * 4. Returning spawn result for backend
 14   */
 15  
 16  import sample from 'lodash-es/sample.js'
 17  import { getSessionId } from '../../bootstrap/state.js'
 18  import { getSpinnerVerbs } from '../../constants/spinnerVerbs.js'
 19  import { TURN_COMPLETION_VERBS } from '../../constants/turnCompletionVerbs.js'
 20  import type { AppState } from '../../state/AppState.js'
 21  import { createTaskStateBase, generateTaskId } from '../../Task.js'
 22  import type {
 23    InProcessTeammateTaskState,
 24    TeammateIdentity,
 25  } from '../../tasks/InProcessTeammateTask/types.js'
 26  import { createAbortController } from '../abortController.js'
 27  import { formatAgentId } from '../agentId.js'
 28  import { registerCleanup } from '../cleanupRegistry.js'
 29  import { logForDebugging } from '../debug.js'
 30  import { emitTaskTerminatedSdk } from '../sdkEventQueue.js'
 31  import { evictTaskOutput } from '../task/diskOutput.js'
 32  import {
 33    evictTerminalTask,
 34    registerTask,
 35    STOPPED_DISPLAY_MS,
 36  } from '../task/framework.js'
 37  import { createTeammateContext } from '../teammateContext.js'
 38  import {
 39    isPerfettoTracingEnabled,
 40    registerAgent as registerPerfettoAgent,
 41    unregisterAgent as unregisterPerfettoAgent,
 42  } from '../telemetry/perfettoTracing.js'
 43  import { removeMemberByAgentId } from './teamHelpers.js'
 44  
 45  type SetAppStateFn = (updater: (prev: AppState) => AppState) => void
 46  
 47  /**
 48   * Minimal context required for spawning an in-process teammate.
 49   * This is a subset of ToolUseContext - only what spawnInProcessTeammate actually uses.
 50   */
 51  export type SpawnContext = {
 52    setAppState: SetAppStateFn
 53    toolUseId?: string
 54  }
 55  
 56  /**
 57   * Configuration for spawning an in-process teammate.
 58   */
 59  export type InProcessSpawnConfig = {
 60    /** Display name for the teammate, e.g., "researcher" */
 61    name: string
 62    /** Team this teammate belongs to */
 63    teamName: string
 64    /** Initial prompt/task for the teammate */
 65    prompt: string
 66    /** Optional UI color for the teammate */
 67    color?: string
 68    /** Whether teammate must enter plan mode before implementing */
 69    planModeRequired: boolean
 70    /** Optional model override for this teammate */
 71    model?: string
 72  }
 73  
 74  /**
 75   * Result from spawning an in-process teammate.
 76   */
 77  export type InProcessSpawnOutput = {
 78    /** Whether spawn was successful */
 79    success: boolean
 80    /** Full agent ID (format: "name@team") */
 81    agentId: string
 82    /** Task ID for tracking in AppState */
 83    taskId?: string
 84    /** AbortController for this teammate (linked to parent) */
 85    abortController?: AbortController
 86    /** Teammate context for AsyncLocalStorage */
 87    teammateContext?: ReturnType<typeof createTeammateContext>
 88    /** Error message if spawn failed */
 89    error?: string
 90  }
 91  
 92  /**
 93   * Spawns an in-process teammate.
 94   *
 95   * Creates the teammate's context, registers the task in AppState, and returns
 96   * the spawn result. The actual agent execution is driven by the
 97   * InProcessTeammateTask component which uses runWithTeammateContext() to
 98   * execute the agent loop with proper identity isolation.
 99   *
100   * @param config - Spawn configuration
101   * @param context - Context with setAppState for registering task
102   * @returns Spawn result with teammate info
103   */
104  export async function spawnInProcessTeammate(
105    config: InProcessSpawnConfig,
106    context: SpawnContext,
107  ): Promise<InProcessSpawnOutput> {
108    const { name, teamName, prompt, color, planModeRequired, model } = config
109    const { setAppState } = context
110  
111    // Generate deterministic agent ID
112    const agentId = formatAgentId(name, teamName)
113    const taskId = generateTaskId('in_process_teammate')
114  
115    logForDebugging(
116      `[spawnInProcessTeammate] Spawning ${agentId} (taskId: ${taskId})`,
117    )
118  
119    try {
120      // Create independent AbortController for this teammate
121      // Teammates should not be aborted when the leader's query is interrupted
122      const abortController = createAbortController()
123  
124      // Get parent session ID for transcript correlation
125      const parentSessionId = getSessionId()
126  
127      // Create teammate identity (stored as plain data in AppState)
128      const identity: TeammateIdentity = {
129        agentId,
130        agentName: name,
131        teamName,
132        color,
133        planModeRequired,
134        parentSessionId,
135      }
136  
137      // Create teammate context for AsyncLocalStorage
138      // This will be used by runWithTeammateContext() during agent execution
139      const teammateContext = createTeammateContext({
140        agentId,
141        agentName: name,
142        teamName,
143        color,
144        planModeRequired,
145        parentSessionId,
146        abortController,
147      })
148  
149      // Register agent in Perfetto trace for hierarchy visualization
150      if (isPerfettoTracingEnabled()) {
151        registerPerfettoAgent(agentId, name, parentSessionId)
152      }
153  
154      // Create task state
155      const description = `${name}: ${prompt.substring(0, 50)}${prompt.length > 50 ? '...' : ''}`
156  
157      const taskState: InProcessTeammateTaskState = {
158        ...createTaskStateBase(
159          taskId,
160          'in_process_teammate',
161          description,
162          context.toolUseId,
163        ),
164        type: 'in_process_teammate',
165        status: 'running',
166        identity,
167        prompt,
168        model,
169        abortController,
170        awaitingPlanApproval: false,
171        spinnerVerb: sample(getSpinnerVerbs()),
172        pastTenseVerb: sample(TURN_COMPLETION_VERBS),
173        permissionMode: planModeRequired ? 'plan' : 'default',
174        isIdle: false,
175        shutdownRequested: false,
176        lastReportedToolCount: 0,
177        lastReportedTokenCount: 0,
178        pendingUserMessages: [],
179        messages: [], // Initialize to empty array so getDisplayedMessages works immediately
180      }
181  
182      // Register cleanup handler for graceful shutdown
183      const unregisterCleanup = registerCleanup(async () => {
184        logForDebugging(`[spawnInProcessTeammate] Cleanup called for ${agentId}`)
185        abortController.abort()
186        // Task state will be updated by the execution loop when it detects abort
187      })
188      taskState.unregisterCleanup = unregisterCleanup
189  
190      // Register task in AppState
191      registerTask(taskState, setAppState)
192  
193      logForDebugging(
194        `[spawnInProcessTeammate] Registered ${agentId} in AppState`,
195      )
196  
197      return {
198        success: true,
199        agentId,
200        taskId,
201        abortController,
202        teammateContext,
203      }
204    } catch (error) {
205      const errorMessage =
206        error instanceof Error ? error.message : 'Unknown error during spawn'
207      logForDebugging(
208        `[spawnInProcessTeammate] Failed to spawn ${agentId}: ${errorMessage}`,
209      )
210      return {
211        success: false,
212        agentId,
213        error: errorMessage,
214      }
215    }
216  }
217  
218  /**
219   * Kills an in-process teammate by aborting its controller.
220   *
221   * Note: This is the implementation called by InProcessBackend.kill().
222   *
223   * @param taskId - Task ID of the teammate to kill
224   * @param setAppState - AppState setter
225   * @returns true if killed successfully
226   */
227  export function killInProcessTeammate(
228    taskId: string,
229    setAppState: SetAppStateFn,
230  ): boolean {
231    let killed = false
232    let teamName: string | null = null
233    let agentId: string | null = null
234    let toolUseId: string | undefined
235    let description: string | undefined
236  
237    setAppState((prev: AppState) => {
238      const task = prev.tasks[taskId]
239      if (!task || task.type !== 'in_process_teammate') {
240        return prev
241      }
242  
243      const teammateTask = task as InProcessTeammateTaskState
244  
245      if (teammateTask.status !== 'running') {
246        return prev
247      }
248  
249      // Capture identity for cleanup after state update
250      teamName = teammateTask.identity.teamName
251      agentId = teammateTask.identity.agentId
252      toolUseId = teammateTask.toolUseId
253      description = teammateTask.description
254  
255      // Abort the controller to stop execution
256      teammateTask.abortController?.abort()
257  
258      // Call cleanup handler
259      teammateTask.unregisterCleanup?.()
260  
261      // Update task state and remove from teamContext.teammates
262      killed = true
263  
264      // Call pending idle callbacks to unblock any waiters (e.g., engine.waitForIdle)
265      teammateTask.onIdleCallbacks?.forEach(cb => cb())
266  
267      // Remove from teamContext.teammates using the agentId
268      let updatedTeamContext = prev.teamContext
269      if (prev.teamContext && prev.teamContext.teammates && agentId) {
270        const { [agentId]: _, ...remainingTeammates } = prev.teamContext.teammates
271        updatedTeamContext = {
272          ...prev.teamContext,
273          teammates: remainingTeammates,
274        }
275      }
276  
277      return {
278        ...prev,
279        teamContext: updatedTeamContext,
280        tasks: {
281          ...prev.tasks,
282          [taskId]: {
283            ...teammateTask,
284            status: 'killed' as const,
285            notified: true,
286            endTime: Date.now(),
287            onIdleCallbacks: [], // Clear callbacks to prevent stale references
288            messages: teammateTask.messages?.length
289              ? [teammateTask.messages[teammateTask.messages.length - 1]!]
290              : undefined,
291            pendingUserMessages: [],
292            inProgressToolUseIDs: undefined,
293            abortController: undefined,
294            unregisterCleanup: undefined,
295            currentWorkAbortController: undefined,
296          },
297        },
298      }
299    })
300  
301    // Remove from team file (outside state updater to avoid file I/O in callback)
302    if (teamName && agentId) {
303      removeMemberByAgentId(teamName, agentId)
304    }
305  
306    if (killed) {
307      void evictTaskOutput(taskId)
308      // notified:true was pre-set so no XML notification fires; close the SDK
309      // task_started bookend directly. The in-process runner's own
310      // completion/failure emit guards on status==='running' so it won't
311      // double-emit after seeing status:killed.
312      emitTaskTerminatedSdk(taskId, 'stopped', {
313        toolUseId,
314        summary: description,
315      })
316      setTimeout(
317        evictTerminalTask.bind(null, taskId, setAppState),
318        STOPPED_DISPLAY_MS,
319      )
320    }
321  
322    // Release perfetto agent registry entry
323    if (agentId) {
324      unregisterPerfettoAgent(agentId)
325    }
326  
327    return killed
328  }