1  /**
  2   * Teammate utilities for agent swarm coordination
  3   *
  4   * These helpers identify whether this Claude Code instance is running as a
  5   * spawned teammate in a swarm. Teammates receive their identity via CLI
  6   * arguments (--agent-id, --team-name, etc.) which are stored in dynamicTeamContext.
  7   *
  8   * For in-process teammates (running in the same process), AsyncLocalStorage
  9   * provides isolated context per teammate, preventing concurrent overwrites.
 10   *
 11   * Priority order for identity resolution:
 12   * 1. AsyncLocalStorage (in-process teammates) - via teammateContext.ts
 13   * 2. dynamicTeamContext (tmux teammates via CLI args)
 14   */
 15  
 16  // Re-export in-process teammate utilities from teammateContext.ts
 17  export {
 18    createTeammateContext,
 19    getTeammateContext,
 20    isInProcessTeammate,
 21    runWithTeammateContext,
 22    type TeammateContext,
 23  } from './teammateContext.js'
 24  
 25  import type { AppState } from '../state/AppState.js'
 26  import { isEnvTruthy } from './envUtils.js'
 27  import { getTeammateContext } from './teammateContext.js'
 28  
 29  /**
 30   * Returns the parent session ID for this teammate.
 31   * For in-process teammates, this is the team lead's session ID.
 32   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux teammates).
 33   */
 34  export function getParentSessionId(): string | undefined {
 35    const inProcessCtx = getTeammateContext()
 36    if (inProcessCtx) return inProcessCtx.parentSessionId
 37    return dynamicTeamContext?.parentSessionId
 38  }
 39  
 40  /**
 41   * Dynamic team context for runtime team joining.
 42   * When set, these values take precedence over environment variables.
 43   */
 44  let dynamicTeamContext: {
 45    agentId: string
 46    agentName: string
 47    teamName: string
 48    color?: string
 49    planModeRequired: boolean
 50    parentSessionId?: string
 51  } | null = null
 52  
 53  /**
 54   * Set the dynamic team context (called when joining a team at runtime)
 55   */
 56  export function setDynamicTeamContext(
 57    context: {
 58      agentId: string
 59      agentName: string
 60      teamName: string
 61      color?: string
 62      planModeRequired: boolean
 63      parentSessionId?: string
 64    } | null,
 65  ): void {
 66    dynamicTeamContext = context
 67  }
 68  
 69  /**
 70   * Clear the dynamic team context (called when leaving a team)
 71   */
 72  export function clearDynamicTeamContext(): void {
 73    dynamicTeamContext = null
 74  }
 75  
 76  /**
 77   * Get the current dynamic team context (for inspection/debugging)
 78   */
 79  export function getDynamicTeamContext(): typeof dynamicTeamContext {
 80    return dynamicTeamContext
 81  }
 82  
 83  /**
 84   * Returns the agent ID if this session is running as a teammate in a swarm,
 85   * or undefined if running as a standalone session.
 86   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux via CLI args).
 87   */
 88  export function getAgentId(): string | undefined {
 89    const inProcessCtx = getTeammateContext()
 90    if (inProcessCtx) return inProcessCtx.agentId
 91    return dynamicTeamContext?.agentId
 92  }
 93  
 94  /**
 95   * Returns the agent name if this session is running as a teammate in a swarm.
 96   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux via CLI args).
 97   */
 98  export function getAgentName(): string | undefined {
 99    const inProcessCtx = getTeammateContext()
100    if (inProcessCtx) return inProcessCtx.agentName
101    return dynamicTeamContext?.agentName
102  }
103  
104  /**
105   * Returns the team name if this session is part of a team.
106   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux via CLI args) > passed teamContext.
107   * Pass teamContext from AppState to support leaders who don't have dynamicTeamContext set.
108   *
109   * @param teamContext - Optional team context from AppState (for leaders)
110   */
111  export function getTeamName(teamContext?: {
112    teamName: string
113  }): string | undefined {
114    const inProcessCtx = getTeammateContext()
115    if (inProcessCtx) return inProcessCtx.teamName
116    if (dynamicTeamContext?.teamName) return dynamicTeamContext.teamName
117    return teamContext?.teamName
118  }
119  
120  /**
121   * Returns true if this session is running as a teammate in a swarm.
122   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux via CLI args).
123   * For tmux teammates, requires BOTH an agent ID AND a team name.
124   */
125  export function isTeammate(): boolean {
126    // In-process teammates run within the same process
127    const inProcessCtx = getTeammateContext()
128    if (inProcessCtx) return true
129    // Tmux teammates require both agent ID and team name
130    return !!(dynamicTeamContext?.agentId && dynamicTeamContext?.teamName)
131  }
132  
133  /**
134   * Returns the teammate's assigned color,
135   * or undefined if not running as a teammate or no color assigned.
136   * Priority: AsyncLocalStorage (in-process) > dynamicTeamContext (tmux teammates).
137   */
138  export function getTeammateColor(): string | undefined {
139    const inProcessCtx = getTeammateContext()
140    if (inProcessCtx) return inProcessCtx.color
141    return dynamicTeamContext?.color
142  }
143  
144  /**
145   * Returns true if this teammate session requires plan mode before implementation.
146   * When enabled, the teammate must enter plan mode and get approval before writing code.
147   * Priority: AsyncLocalStorage > dynamicTeamContext > env var.
148   */
149  export function isPlanModeRequired(): boolean {
150    const inProcessCtx = getTeammateContext()
151    if (inProcessCtx) return inProcessCtx.planModeRequired
152    if (dynamicTeamContext !== null) {
153      return dynamicTeamContext.planModeRequired
154    }
155    return isEnvTruthy(process.env.CLAUDE_CODE_PLAN_MODE_REQUIRED)
156  }
157  
158  /**
159   * Check if this session is a team lead.
160   *
161   * A session is considered a team lead if:
162   * 1. A team context exists with a leadAgentId, AND
163   * 2. Either:
164   *    - Our CLAUDE_CODE_AGENT_ID matches the leadAgentId, OR
165   *    - We have no CLAUDE_CODE_AGENT_ID set (backwards compat: the original
166   *      session that created the team before agent IDs were standardized)
167   *
168   * @param teamContext - The team context from AppState, if any
169   * @returns true if this session is the team lead
170   */
171  export function isTeamLead(
172    teamContext:
173      | {
174          leadAgentId: string
175        }
176      | undefined,
177  ): boolean {
178    if (!teamContext?.leadAgentId) {
179      return false
180    }
181  
182    // Use getAgentId() for AsyncLocalStorage support (in-process teammates)
183    const myAgentId = getAgentId()
184    const leadAgentId = teamContext.leadAgentId
185  
186    // If my agent ID matches the lead agent ID, I'm the lead
187    if (myAgentId === leadAgentId) {
188      return true
189    }
190  
191    // Backwards compat: if no agent ID is set and we have a team context,
192    // this is the original session that created the team (the lead)
193    if (!myAgentId) {
194      return true
195    }
196  
197    return false
198  }
199  
200  /**
201   * Checks if there are any active in-process teammates running.
202   * Used by headless/print mode to determine if we should wait for teammates
203   * before exiting.
204   */
205  export function hasActiveInProcessTeammates(appState: AppState): boolean {
206    // Check for running in-process teammate tasks
207    for (const task of Object.values(appState.tasks)) {
208      if (task.type === 'in_process_teammate' && task.status === 'running') {
209        return true
210      }
211    }
212    return false
213  }
214  
215  /**
216   * Checks if there are in-process teammates still actively working on tasks.
217   * Returns true if any teammate is running but NOT idle (still processing).
218   * Used to determine if we should wait before sending shutdown prompts.
219   */
220  export function hasWorkingInProcessTeammates(appState: AppState): boolean {
221    for (const task of Object.values(appState.tasks)) {
222      if (
223        task.type === 'in_process_teammate' &&
224        task.status === 'running' &&
225        !task.isIdle
226      ) {
227        return true
228      }
229    }
230    return false
231  }
232  
233  /**
234   * Returns a promise that resolves when all working in-process teammates become idle.
235   * Registers callbacks on each working teammate's task - they call these when idle.
236   * Returns immediately if no teammates are working.
237   */
238  export function waitForTeammatesToBecomeIdle(
239    setAppState: (f: (prev: AppState) => AppState) => void,
240    appState: AppState,
241  ): Promise<void> {
242    const workingTaskIds: string[] = []
243  
244    for (const [taskId, task] of Object.entries(appState.tasks)) {
245      if (
246        task.type === 'in_process_teammate' &&
247        task.status === 'running' &&
248        !task.isIdle
249      ) {
250        workingTaskIds.push(taskId)
251      }
252    }
253  
254    if (workingTaskIds.length === 0) {
255      return Promise.resolve()
256    }
257  
258    // Create a promise that resolves when all working teammates become idle
259    return new Promise<void>(resolve => {
260      let remaining = workingTaskIds.length
261  
262      const onIdle = (): void => {
263        remaining--
264        if (remaining === 0) {
265          // biome-ignore lint/nursery/noFloatingPromises: resolve is a callback, not a Promise
266          resolve()
267        }
268      }
269  
270      // Register callback on each working teammate
271      // Check current isIdle state to handle race where teammate became idle
272      // between our initial snapshot and this callback registration
273      setAppState(prev => {
274        const newTasks = { ...prev.tasks }
275        for (const taskId of workingTaskIds) {
276          const task = newTasks[taskId]
277          if (task && task.type === 'in_process_teammate') {
278            // If task is already idle, call onIdle immediately
279            if (task.isIdle) {
280              onIdle()
281            } else {
282              newTasks[taskId] = {
283                ...task,
284                onIdleCallbacks: [...(task.onIdleCallbacks ?? []), onIdle],
285              }
286            }
287          }
288        }
289        return { ...prev, tasks: newTasks }
290      })
291    })
292  }