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