1  /**
  2   * Agent context for analytics attribution using AsyncLocalStorage.
  3   *
  4   * This module provides a way to track agent identity across async operations
  5   * without parameter drilling. Supports two agent types:
  6   *
  7   * 1. Subagents (Agent tool): Run in-process for quick, delegated tasks.
  8   *    Context: SubagentContext with agentType: 'subagent'
  9   *
 10   * 2. In-process teammates: Part of a swarm with team coordination.
 11   *    Context: TeammateAgentContext with agentType: 'teammate'
 12   *
 13   * For swarm teammates in separate processes (tmux/iTerm2), use environment
 14   * variables instead: CLAUDE_CODE_AGENT_ID, CLAUDE_CODE_PARENT_SESSION_ID
 15   *
 16   * WHY AsyncLocalStorage (not AppState):
 17   * When agents are backgrounded (ctrl+b), multiple agents can run concurrently
 18   * in the same process. AppState is a single shared state that would be
 19   * overwritten, causing Agent A's events to incorrectly use Agent B's context.
 20   * AsyncLocalStorage isolates each async execution chain, so concurrent agents
 21   * don't interfere with each other.
 22   */
 23  
 24  import { AsyncLocalStorage } from 'async_hooks'
 25  import type { AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS } from '../services/analytics/index.js'
 26  import { isAgentSwarmsEnabled } from './agentSwarmsEnabled.js'
 27  
 28  /**
 29   * Context for subagents (Agent tool agents).
 30   * Subagents run in-process for quick, delegated tasks.
 31   */
 32  export type SubagentContext = {
 33    /** The subagent's UUID (from createAgentId()) */
 34    agentId: string
 35    /** The team lead's session ID (from CLAUDE_CODE_PARENT_SESSION_ID env var), undefined for main REPL subagents */
 36    parentSessionId?: string
 37    /** Agent type - 'subagent' for Agent tool agents */
 38    agentType: 'subagent'
 39    /** The subagent's type name (e.g., "Explore", "Bash", "code-reviewer") */
 40    subagentName?: string
 41    /** Whether this is a built-in agent (vs user-defined custom agent) */
 42    isBuiltIn?: boolean
 43    /** The request_id in the invoking agent that spawned or resumed this agent.
 44     *  For nested subagents this is the immediate invoker, not the root —
 45     *  session_id already bundles the whole tree. Updated on each resume. */
 46    invokingRequestId?: string
 47    /** Whether this invocation is the initial spawn or a subsequent resume
 48     *  via SendMessage. Undefined when invokingRequestId is absent. */
 49    invocationKind?: 'spawn' | 'resume'
 50    /** Mutable flag: has this invocation's edge been emitted to telemetry yet?
 51     *  Reset to false on each spawn/resume; flipped true by
 52     *  consumeInvokingRequestId() on the first terminal API event. */
 53    invocationEmitted?: boolean
 54  }
 55  
 56  /**
 57   * Context for in-process teammates.
 58   * Teammates are part of a swarm and have team coordination.
 59   */
 60  export type TeammateAgentContext = {
 61    /** Full agent ID, e.g., "researcher@my-team" */
 62    agentId: string
 63    /** Display name, e.g., "researcher" */
 64    agentName: string
 65    /** Team name this teammate belongs to */
 66    teamName: string
 67    /** UI color assigned to this teammate */
 68    agentColor?: string
 69    /** Whether teammate must enter plan mode before implementing */
 70    planModeRequired: boolean
 71    /** The team lead's session ID for transcript correlation */
 72    parentSessionId: string
 73    /** Whether this agent is the team lead */
 74    isTeamLead: boolean
 75    /** Agent type - 'teammate' for swarm teammates */
 76    agentType: 'teammate'
 77    /** The request_id in the invoking agent that spawned or resumed this
 78     *  teammate. Undefined for teammates started outside a tool call
 79     *  (e.g. session start). Updated on each resume. */
 80    invokingRequestId?: string
 81    /** See SubagentContext.invocationKind. */
 82    invocationKind?: 'spawn' | 'resume'
 83    /** Mutable flag: see SubagentContext.invocationEmitted. */
 84    invocationEmitted?: boolean
 85  }
 86  
 87  /**
 88   * Discriminated union for agent context.
 89   * Use agentType to distinguish between subagent and teammate contexts.
 90   */
 91  export type AgentContext = SubagentContext | TeammateAgentContext
 92  
 93  const agentContextStorage = new AsyncLocalStorage<AgentContext>()
 94  
 95  /**
 96   * Get the current agent context, if any.
 97   * Returns undefined if not running within an agent context (subagent or teammate).
 98   * Use type guards isSubagentContext() or isTeammateAgentContext() to narrow the type.
 99   */
100  export function getAgentContext(): AgentContext | undefined {
101    return agentContextStorage.getStore()
102  }
103  
104  /**
105   * Run an async function with the given agent context.
106   * All async operations within the function will have access to this context.
107   */
108  export function runWithAgentContext<T>(context: AgentContext, fn: () => T): T {
109    return agentContextStorage.run(context, fn)
110  }
111  
112  /**
113   * Type guard to check if context is a SubagentContext.
114   */
115  export function isSubagentContext(
116    context: AgentContext | undefined,
117  ): context is SubagentContext {
118    return context?.agentType === 'subagent'
119  }
120  
121  /**
122   * Type guard to check if context is a TeammateAgentContext.
123   */
124  export function isTeammateAgentContext(
125    context: AgentContext | undefined,
126  ): context is TeammateAgentContext {
127    if (isAgentSwarmsEnabled()) {
128      return context?.agentType === 'teammate'
129    }
130    return false
131  }
132  
133  /**
134   * Get the subagent name suitable for analytics logging.
135   * Returns the agent type name for built-in agents, "user-defined" for custom agents,
136   * or undefined if not running within a subagent context.
137   *
138   * Safe for analytics metadata: built-in agent names are code constants,
139   * and custom agents are always mapped to the literal "user-defined".
140   */
141  export function getSubagentLogName():
142    | AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
143    | undefined {
144    const context = getAgentContext()
145    if (!isSubagentContext(context) || !context.subagentName) {
146      return undefined
147    }
148    return (
149      context.isBuiltIn ? context.subagentName : 'user-defined'
150    ) as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS
151  }
152  
153  /**
154   * Get the invoking request_id for the current agent context — once per
155   * invocation. Returns the id on the first call after a spawn/resume, then
156   * undefined until the next boundary. Also undefined on the main thread or
157   * when the spawn path had no request_id.
158   *
159   * Sparse edge semantics: invokingRequestId appears on exactly one
160   * tengu_api_success/error per invocation, so a non-NULL value downstream
161   * marks a spawn/resume boundary.
162   */
163  export function consumeInvokingRequestId():
164    | {
165        invokingRequestId: string
166        invocationKind: 'spawn' | 'resume' | undefined
167      }
168    | undefined {
169    const context = getAgentContext()
170    if (!context?.invokingRequestId || context.invocationEmitted) {
171      return undefined
172    }
173    context.invocationEmitted = true
174    return {
175      invokingRequestId: context.invokingRequestId,
176      invocationKind: context.invocationKind,
177    }
178  }