@townco/agent 0.1.112 → 0.1.114

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,46 +1,17 @@
1
- import { spawn } from "node:child_process";
1
+ import * as crypto from "node:crypto";
2
2
  import * as fs from "node:fs/promises";
3
+ import { mkdir } from "node:fs/promises";
3
4
  import * as path from "node:path";
4
- import { PROTOCOL_VERSION } from "@agentclientprotocol/sdk";
5
5
  import { context, propagation, trace } from "@opentelemetry/api";
6
- import { createLogger as coreCreateLogger } from "@townco/core";
7
6
  import { z } from "zod";
8
7
  import { SUBAGENT_MODE_KEY, } from "../../../acp-server/adapter.js";
9
- import { getAbortSignal } from "../../session-context.js";
10
- import { findAvailablePort } from "./port-utils.js";
11
- import { emitSubagentConnection, emitSubagentMessages, hashQuery, } from "./subagent-connections.js";
8
+ import { makeRunnerFromDefinition } from "../../index.js";
9
+ import { bindGeneratorToSessionContext, getAbortSignal, } from "../../session-context.js";
10
+ import { emitSubagentMessages, hashQuery, } from "./subagent-connections.js";
12
11
  /**
13
12
  * Name of the Task tool created by makeSubagentsTool
14
13
  */
15
14
  export const SUBAGENT_TOOL_NAME = "subagent";
16
- /**
17
- * Base port for subagent HTTP servers (avoid conflict with main agents at 3100+)
18
- */
19
- const SUBAGENT_BASE_PORT = 4000;
20
- /**
21
- * Wait for HTTP server to be ready by polling health endpoint
22
- */
23
- async function waitForServerReady(port, timeoutMs = 30000) {
24
- const startTime = Date.now();
25
- const baseDelay = 50;
26
- let attempt = 0;
27
- while (Date.now() - startTime < timeoutMs) {
28
- try {
29
- const response = await fetch(`http://localhost:${port}/health`, {
30
- signal: AbortSignal.timeout(1000),
31
- });
32
- if (response.ok) {
33
- return;
34
- }
35
- }
36
- catch {
37
- // Server not ready yet
38
- }
39
- await new Promise((r) => setTimeout(r, baseDelay * 1.5 ** attempt));
40
- attempt++;
41
- }
42
- throw new Error(`Subagent server at port ${port} did not become ready within ${timeoutMs}ms`);
43
- }
44
15
  /**
45
16
  * Creates a DirectTool that delegates work to one of multiple configured subagents.
46
17
  *
@@ -74,8 +45,10 @@ export function makeSubagentsTool(configs) {
74
45
  let agentPath;
75
46
  let agentDir;
76
47
  if ("path" in config) {
77
- // Direct path variant
78
- agentPath = config.path;
48
+ // Direct path variant - resolve to absolute path
49
+ agentPath = path.isAbsolute(config.path)
50
+ ? config.path
51
+ : path.resolve(process.cwd(), config.path);
79
52
  agentDir = path.dirname(agentPath);
80
53
  }
81
54
  else {
@@ -181,379 +154,165 @@ assistant: "I'm going to use the Task tool to launch the greeting-responder agen
181
154
  };
182
155
  }
183
156
  /**
184
- * Internal function that spawns a subagent HTTP server and queries it.
157
+ * Internal function that runs a subagent in-process and queries it.
185
158
  */
186
159
  async function querySubagent(agentName, agentPath, agentWorkingDirectory, query) {
187
160
  // Get the abort signal from context (set by parent agent's cancellation)
188
161
  const parentAbortSignal = getAbortSignal();
189
162
  // Check if already cancelled before starting
190
163
  if (parentAbortSignal?.aborted) {
191
- throw new Error("Subagent query cancelled before starting");
164
+ throw new Error(`Subagent query cancelled before starting (agent: ${agentName})`);
192
165
  }
193
166
  // Validate that the agent exists
194
167
  try {
195
168
  await fs.access(agentPath);
196
169
  }
197
170
  catch (_error) {
198
- throw new Error(`Agent not found at ${agentPath}. Make sure the agent exists and has an index.ts file.`);
171
+ throw new Error(`Subagent '${agentName}' not found at ${agentPath}. Make sure the agent exists and has an index.ts file.`);
199
172
  }
200
- // Find an available port for this subagent
201
- const port = await findAvailablePort(SUBAGENT_BASE_PORT, 100);
202
- // Create a logger for this subagent instance with port prefix
203
- // Use core logger directly since service name already identifies the subagent
204
- const logger = coreCreateLogger(`subagent:${port}:${agentName}`);
205
- let agentProcess = null;
206
- let sseAbortController = null;
207
- let cleanupTriggered = false;
208
- // Cleanup function to kill the process and abort SSE
209
- const cleanup = () => {
210
- if (cleanupTriggered)
211
- return;
212
- cleanupTriggered = true;
213
- logger.info(`Cleaning up subagent on port ${port} (cancelled by parent)`);
214
- if (sseAbortController) {
215
- sseAbortController.abort();
216
- }
217
- if (agentProcess) {
218
- agentProcess.kill("SIGTERM");
219
- }
220
- };
221
- // Listen for parent abort signal
222
- if (parentAbortSignal) {
223
- parentAbortSignal.addEventListener("abort", cleanup, { once: true });
173
+ // Load agent definition dynamically
174
+ const agentModule = await import(agentPath);
175
+ const agentDefinition = agentModule.default || agentModule.agent;
176
+ // Create runner instance
177
+ const runner = makeRunnerFromDefinition(agentDefinition);
178
+ // Generate unique session ID for isolation
179
+ const subagentSessionId = crypto.randomUUID();
180
+ // Setup session paths
181
+ const sessionDir = path.join(agentWorkingDirectory, ".sessions", subagentSessionId);
182
+ const artifactsDir = path.join(sessionDir, "artifacts");
183
+ await mkdir(artifactsDir, { recursive: true });
184
+ // Prepare OTEL context for distributed tracing
185
+ const activeCtx = context.active();
186
+ const activeSpan = trace.getSpan(activeCtx);
187
+ const otelCarrier = {};
188
+ if (activeSpan) {
189
+ propagation.inject(trace.setSpan(activeCtx, activeSpan), otelCarrier);
224
190
  }
191
+ // Create invoke request
192
+ const invokeRequest = {
193
+ sessionId: subagentSessionId,
194
+ messageId: crypto.randomUUID(),
195
+ prompt: [{ type: "text", text: query }],
196
+ agentDir: agentWorkingDirectory,
197
+ contextMessages: [],
198
+ ...(parentAbortSignal ? { abortSignal: parentAbortSignal } : {}),
199
+ sessionMeta: {
200
+ [SUBAGENT_MODE_KEY]: true,
201
+ ...(Object.keys(otelCarrier).length > 0
202
+ ? { otelTraceContext: otelCarrier }
203
+ : {}),
204
+ },
205
+ };
206
+ // Bind session context and invoke
207
+ let generator = runner.invoke(invokeRequest);
208
+ generator = bindGeneratorToSessionContext({ sessionId: subagentSessionId, sessionDir, artifactsDir }, generator);
209
+ // Consume stream, accumulate results, and emit incremental updates
210
+ let responseText = "";
211
+ const collectedSources = [];
212
+ const currentMessage = {
213
+ id: `subagent-${Date.now()}`,
214
+ content: "",
215
+ contentBlocks: [],
216
+ toolCalls: [],
217
+ };
218
+ const toolCallMap = new Map();
219
+ const queryHash = hashQuery(query);
225
220
  try {
226
- // Get the parent's logs directory to pass to the subagent
227
- const parentLogsDir = process.env.TOWN_LOGS_DIR || path.join(process.cwd(), ".logs");
228
- // Spawn the agent process in HTTP mode
229
- agentProcess = spawn("bun", [agentPath, "http"], {
230
- cwd: agentWorkingDirectory,
231
- env: {
232
- ...process.env,
233
- PORT: String(port),
234
- TOWN_LOGS_DIR: parentLogsDir,
235
- TOWN_SUBAGENT_NAME: agentName,
236
- },
237
- stdio: ["pipe", "pipe", "pipe"],
238
- });
239
- if (!agentProcess.stderr) {
240
- throw new Error("Failed to create stderr pipe for agent process");
241
- }
242
- // Capture stdout and forward to logger
243
- if (agentProcess.stdout) {
244
- agentProcess.stdout.on("data", (data) => {
245
- const lines = data
246
- .toString()
247
- .split("\n")
248
- .filter((line) => line.trim());
249
- for (const line of lines) {
250
- logger.info(line);
221
+ for await (const update of generator) {
222
+ let shouldEmit = false;
223
+ // Handle agent_message_chunk
224
+ if (update.sessionUpdate === "agent_message_chunk") {
225
+ const content = update.content;
226
+ if (content?.type === "text" && typeof content.text === "string") {
227
+ responseText += content.text;
228
+ currentMessage.content += content.text;
229
+ const lastBlock = currentMessage.contentBlocks[currentMessage.contentBlocks.length - 1];
230
+ if (lastBlock && lastBlock.type === "text") {
231
+ lastBlock.text += content.text;
232
+ }
233
+ else {
234
+ currentMessage.contentBlocks.push({
235
+ type: "text",
236
+ text: content.text,
237
+ });
238
+ }
239
+ shouldEmit = true; // Emit after each text chunk
251
240
  }
252
- });
253
- }
254
- // Capture stderr and forward to logger as errors
255
- agentProcess.stderr.on("data", (data) => {
256
- const lines = data
257
- .toString()
258
- .split("\n")
259
- .filter((line) => line.trim());
260
- for (const line of lines) {
261
- logger.error(line);
262
241
  }
263
- });
264
- // Handle process errors
265
- const processErrorPromise = new Promise((_, reject) => {
266
- agentProcess?.on("error", (error) => {
267
- logger.error(`Process error: ${error.message}`);
268
- reject(new Error(`Agent process error: ${error.message}`));
269
- });
270
- agentProcess?.on("exit", (code, signal) => {
271
- if (code !== 0 && code !== null) {
272
- logger.error(`Process exited with code ${code}, signal ${signal}`);
273
- reject(new Error(`Agent process exited unexpectedly with code ${code}, signal ${signal}`));
274
- }
275
- });
276
- });
277
- logger.info(`Starting subagent HTTP server on port ${port}`);
278
- // Wait for server to be ready
279
- await Promise.race([waitForServerReady(port), processErrorPromise]);
280
- logger.info(`Subagent server ready on port ${port}`);
281
- const baseUrl = `http://localhost:${port}`;
282
- // Step 1: Initialize ACP connection
283
- const initResponse = await fetch(`${baseUrl}/rpc`, {
284
- method: "POST",
285
- headers: { "Content-Type": "application/json" },
286
- body: JSON.stringify({
287
- jsonrpc: "2.0",
288
- id: "init-1",
289
- method: "initialize",
290
- params: {
291
- protocolVersion: PROTOCOL_VERSION,
292
- clientCapabilities: {
293
- fs: { readTextFile: false, writeTextFile: false },
294
- },
295
- },
296
- }),
297
- });
298
- if (!initResponse.ok) {
299
- throw new Error(`Initialize failed: HTTP ${initResponse.status}`);
300
- }
301
- // Step 2: Create new session with subagent mode and OTEL context
302
- // Prepare OpenTelemetry trace context to propagate to the subagent
303
- const otelCarrier = {};
304
- const activeCtx = context.active();
305
- const activeSpan = trace.getSpan(activeCtx);
306
- if (process.env.DEBUG_TELEMETRY === "true") {
307
- console.log(`[querySubagent] Active span when tool executes:`, activeSpan?.spanContext());
308
- }
309
- if (activeSpan) {
310
- const ctxForInjection = trace.setSpan(activeCtx, activeSpan);
311
- propagation.inject(ctxForInjection, otelCarrier);
312
- }
313
- const hasOtelContext = Object.keys(otelCarrier).length > 0;
314
- const sessionResponse = await fetch(`${baseUrl}/rpc`, {
315
- method: "POST",
316
- headers: { "Content-Type": "application/json" },
317
- body: JSON.stringify({
318
- jsonrpc: "2.0",
319
- id: "session-1",
320
- method: "session/new",
321
- params: {
322
- cwd: agentWorkingDirectory,
323
- mcpServers: [],
324
- _meta: {
325
- [SUBAGENT_MODE_KEY]: true,
326
- ...(hasOtelContext ? { otelTraceContext: otelCarrier } : {}),
327
- },
328
- },
329
- }),
330
- });
331
- if (!sessionResponse.ok) {
332
- throw new Error(`Session creation failed: HTTP ${sessionResponse.status}`);
333
- }
334
- const sessionResult = (await sessionResponse.json());
335
- const sessionId = sessionResult.result?.sessionId;
336
- if (!sessionId) {
337
- throw new Error("No sessionId in session/new response");
338
- }
339
- // Emit connection info so the GUI can connect directly to this subagent's SSE
340
- const queryHash = hashQuery(query);
341
- emitSubagentConnection(queryHash, { port, sessionId });
342
- // Step 3: Connect to SSE for receiving streaming responses
343
- sseAbortController = new AbortController();
344
- let responseText = "";
345
- // Track citation sources from subagent's web searches/fetches
346
- const collectedSources = [];
347
- // Track full message structure for session storage
348
- const currentMessage = {
349
- id: `subagent-${Date.now()}`,
350
- content: "",
351
- contentBlocks: [],
352
- toolCalls: [],
353
- };
354
- // Map of tool call IDs to their indices in toolCalls array
355
- const toolCallMap = new Map();
356
- const ssePromise = (async () => {
357
- try {
358
- const sseResponse = await fetch(`${baseUrl}/events`, {
359
- headers: { "X-Session-ID": sessionId },
360
- signal: sseAbortController?.signal,
361
- });
362
- if (!sseResponse.ok || !sseResponse.body) {
363
- throw new Error(`SSE connection failed: HTTP ${sseResponse.status}`);
364
- }
365
- const reader = sseResponse.body.getReader();
366
- const decoder = new TextDecoder();
367
- let buffer = "";
368
- while (true) {
369
- const { done, value } = await reader.read();
370
- if (done)
371
- break;
372
- buffer += decoder.decode(value, { stream: true });
373
- const lines = buffer.split("\n");
374
- buffer = lines.pop() || "";
375
- for (const line of lines) {
376
- if (line.startsWith("data:")) {
377
- const data = line.substring(5).trim();
378
- if (!data)
379
- continue;
380
- try {
381
- const message = JSON.parse(data);
382
- const update = message.params?.update;
383
- if (message.method !== "session/update" || !update)
384
- continue;
385
- // Handle agent_message_chunk - accumulate text
386
- if (update.sessionUpdate === "agent_message_chunk") {
387
- const content = update.content;
388
- if (content?.type === "text" &&
389
- typeof content.text === "string") {
390
- responseText += content.text;
391
- currentMessage.content += content.text;
392
- // Add to contentBlocks - append to last text block or create new one
393
- const lastBlock = currentMessage.contentBlocks[currentMessage.contentBlocks.length - 1];
394
- if (lastBlock && lastBlock.type === "text") {
395
- lastBlock.text += content.text;
396
- }
397
- else {
398
- currentMessage.contentBlocks.push({
399
- type: "text",
400
- text: content.text,
401
- });
402
- }
403
- }
404
- }
405
- // Handle tool_call - track new tool calls
406
- if (update.sessionUpdate === "tool_call" && update.toolCallId) {
407
- const toolCall = {
408
- id: update.toolCallId,
409
- title: update.title || "Tool call",
410
- prettyName: update._meta?.prettyName,
411
- icon: update._meta?.icon,
412
- status: update.status ||
413
- "pending",
414
- };
415
- currentMessage.toolCalls.push(toolCall);
416
- toolCallMap.set(update.toolCallId, currentMessage.toolCalls.length - 1);
417
- // Add to contentBlocks for interleaved display
418
- currentMessage.contentBlocks.push({
419
- type: "tool_call",
420
- toolCall,
421
- });
422
- }
423
- // Handle tool_call_update - update existing tool call status
424
- if (update.sessionUpdate === "tool_call_update" &&
425
- update.toolCallId) {
426
- const idx = toolCallMap.get(update.toolCallId);
427
- if (idx !== undefined && currentMessage.toolCalls[idx]) {
428
- if (update.status) {
429
- currentMessage.toolCalls[idx].status =
430
- update.status;
431
- }
432
- // Also update in contentBlocks
433
- const block = currentMessage.contentBlocks.find((b) => b.type === "tool_call" &&
434
- b.toolCall.id === update.toolCallId);
435
- if (block && update.status) {
436
- block.toolCall.status =
437
- update.status;
438
- }
439
- }
440
- }
441
- // Handle sources - collect citation sources from subagent's web searches
442
- if (update.sessionUpdate === "sources" &&
443
- Array.isArray(update.sources)) {
444
- for (const source of update.sources) {
445
- const citationSource = {
446
- id: source.id,
447
- url: source.url,
448
- title: source.title,
449
- toolCallId: source.toolCallId,
450
- };
451
- if (source.snippet)
452
- citationSource.snippet = source.snippet;
453
- if (source.favicon)
454
- citationSource.favicon = source.favicon;
455
- if (source.sourceName)
456
- citationSource.sourceName = source.sourceName;
457
- collectedSources.push(citationSource);
458
- }
459
- logger.info(`Collected ${update.sources.length} sources from subagent`);
460
- }
461
- }
462
- catch {
463
- // Ignore malformed SSE data
464
- }
465
- }
242
+ // Handle tool_call
243
+ if (update.sessionUpdate === "tool_call" && update.toolCallId) {
244
+ const meta = update._meta;
245
+ const toolCall = {
246
+ id: update.toolCallId,
247
+ title: update.title || "Tool call",
248
+ prettyName: meta?.prettyName,
249
+ icon: meta?.icon,
250
+ status: update.status || "pending",
251
+ };
252
+ currentMessage.toolCalls.push(toolCall);
253
+ toolCallMap.set(update.toolCallId, currentMessage.toolCalls.length - 1);
254
+ currentMessage.contentBlocks.push({ type: "tool_call", toolCall });
255
+ shouldEmit = true; // Emit when new tool call appears
256
+ }
257
+ // Handle tool_call_update
258
+ if (update.sessionUpdate === "tool_call_update" && update.toolCallId) {
259
+ const idx = toolCallMap.get(update.toolCallId);
260
+ if (idx !== undefined && currentMessage.toolCalls[idx]) {
261
+ if (update.status) {
262
+ currentMessage.toolCalls[idx].status =
263
+ update.status;
466
264
  }
265
+ const block = currentMessage.contentBlocks.find((b) => b.type === "tool_call" && b.toolCall.id === update.toolCallId);
266
+ if (block && update.status) {
267
+ block.toolCall.status = update.status;
268
+ }
269
+ shouldEmit = true; // Emit when tool status changes
467
270
  }
468
271
  }
469
- catch (error) {
470
- // Ignore AbortError - this is expected when we abort the SSE connection
471
- if (error instanceof DOMException && error.name === "AbortError") {
472
- return;
272
+ // Handle sources (from ACP protocol)
273
+ if ("sources" in update && Array.isArray(update.sources)) {
274
+ const sources = update.sources;
275
+ for (const source of sources) {
276
+ const citationSource = {
277
+ id: source.id,
278
+ url: source.url,
279
+ title: source.title,
280
+ toolCallId: source.toolCallId,
281
+ };
282
+ if (source.snippet)
283
+ citationSource.snippet = source.snippet;
284
+ if (source.favicon)
285
+ citationSource.favicon = source.favicon;
286
+ if (source.sourceName)
287
+ citationSource.sourceName = source.sourceName;
288
+ collectedSources.push(citationSource);
473
289
  }
474
- throw error;
475
- }
476
- })();
477
- // Step 4: Send the prompt with timeout
478
- const timeoutMs = 5 * 60 * 1000; // 5 minutes
479
- const promptPromise = (async () => {
480
- const promptResponse = await fetch(`${baseUrl}/rpc`, {
481
- method: "POST",
482
- headers: { "Content-Type": "application/json" },
483
- body: JSON.stringify({
484
- jsonrpc: "2.0",
485
- id: "prompt-1",
486
- method: "session/prompt",
487
- params: {
488
- sessionId,
489
- prompt: [{ type: "text", text: query }],
490
- },
491
- }),
492
- });
493
- if (!promptResponse.ok) {
494
- throw new Error(`Prompt failed: HTTP ${promptResponse.status}`);
290
+ shouldEmit = true; // Emit when sources are added
495
291
  }
496
- // Wait for prompt to complete (this blocks until agent finishes processing)
497
- const promptResult = (await promptResponse.json());
498
- if (promptResult.error) {
499
- throw new Error(`Prompt error: ${promptResult.error.message || JSON.stringify(promptResult.error)}`);
292
+ // Emit incremental update to parent (for live streaming)
293
+ if (shouldEmit) {
294
+ emitSubagentMessages(queryHash, [{ ...currentMessage }]);
500
295
  }
501
- })();
502
- const timeoutPromise = new Promise((_, reject) => {
503
- setTimeout(() => {
504
- reject(new Error(`Subagent query timed out after ${timeoutMs / 1000} seconds`));
505
- }, timeoutMs);
506
- });
507
- // Create cancellation promise that rejects when parent aborts
508
- const cancellationPromise = parentAbortSignal
509
- ? new Promise((_, reject) => {
510
- if (parentAbortSignal.aborted) {
511
- reject(new Error("Subagent query cancelled"));
512
- }
513
- parentAbortSignal.addEventListener("abort", () => reject(new Error("Subagent query cancelled")), { once: true });
514
- })
515
- : new Promise(() => { }); // Never resolves if no signal
516
- // Wait for prompt to complete with timeout or cancellation
517
- await Promise.race([
518
- promptPromise,
519
- timeoutPromise,
520
- processErrorPromise,
521
- cancellationPromise,
522
- ]);
523
- // Check if cancelled before processing results
524
- if (parentAbortSignal?.aborted) {
525
- throw new Error("Subagent query cancelled");
526
296
  }
527
- // Give SSE a moment to flush remaining messages
528
- await new Promise((r) => setTimeout(r, 100));
529
- // Abort SSE connection
530
- sseAbortController.abort();
531
- // Wait for SSE to finish (with timeout)
532
- await Promise.race([
533
- ssePromise.catch(() => { }), // Ignore abort errors
534
- new Promise((r) => setTimeout(r, 1000)),
535
- ]);
536
- // Emit accumulated messages for session storage
297
+ // Final emit to ensure everything is captured, with completion flag
537
298
  if (currentMessage.content || currentMessage.toolCalls.length > 0) {
538
- emitSubagentMessages(queryHash, [currentMessage]);
299
+ emitSubagentMessages(queryHash, [currentMessage], true);
300
+ }
301
+ else {
302
+ // Even if no messages, emit completion sentinel
303
+ emitSubagentMessages(queryHash, [], true);
539
304
  }
540
305
  return {
541
306
  text: responseText,
542
307
  sources: collectedSources,
543
308
  };
544
309
  }
545
- finally {
546
- // Remove the abort listener to prevent memory leaks
547
- if (parentAbortSignal) {
548
- parentAbortSignal.removeEventListener("abort", cleanup);
549
- }
550
- // Cleanup: abort SSE and kill process
551
- logger.info(`Shutting down subagent on port ${port}`);
552
- if (sseAbortController) {
553
- sseAbortController.abort();
554
- }
555
- if (agentProcess) {
556
- agentProcess.kill();
310
+ catch (error) {
311
+ // Emit completion sentinel even on error to prevent parent from hanging
312
+ emitSubagentMessages(queryHash, [], true);
313
+ if (parentAbortSignal?.aborted) {
314
+ throw new Error("Subagent query cancelled");
557
315
  }
316
+ throw error;
558
317
  }
559
318
  }
@@ -16,6 +16,14 @@ export interface SessionContext {
16
16
  * Tools can access this to emit events like sandbox file changes.
17
17
  */
18
18
  type EmitUpdateCallback = (update: SessionUpdateNotification) => void;
19
+ /**
20
+ * Invocation context for scoping subagent events to a specific invocation.
21
+ * Each invoke() creates its own EventEmitter so subagents emit to the correct parent.
22
+ */
23
+ export interface InvocationContext {
24
+ invocationId: string;
25
+ subagentEventEmitter: import("node:events").EventEmitter;
26
+ }
19
27
  /**
20
28
  * Run a function with an abort signal available via AsyncLocalStorage.
21
29
  * Tools can access the signal using getAbortSignal().
@@ -74,4 +82,14 @@ export declare function hasSessionContext(): boolean;
74
82
  * @returns Absolute path to the tool's output directory
75
83
  */
76
84
  export declare function getToolOutputDir(toolName: string): string;
85
+ /**
86
+ * Get the current invocation context.
87
+ * Returns undefined if called outside of an invocation context.
88
+ */
89
+ export declare function getInvocationContext(): InvocationContext | undefined;
90
+ /**
91
+ * Bind an async generator to an invocation context so that every iteration
92
+ * runs with the invocation context available.
93
+ */
94
+ export declare function bindGeneratorToInvocationContext<T, R, N = unknown>(ctx: InvocationContext, generator: AsyncGenerator<T, R, N>): AsyncGenerator<T, R, N>;
77
95
  export {};
@@ -7,6 +7,7 @@ const sessionStorage = new AsyncLocalStorage();
7
7
  */
8
8
  const abortSignalStorage = new AsyncLocalStorage();
9
9
  const emitUpdateStorage = new AsyncLocalStorage();
10
+ const invocationContextStorage = new AsyncLocalStorage();
10
11
  /**
11
12
  * Run a function with an abort signal available via AsyncLocalStorage.
12
13
  * Tools can access the signal using getAbortSignal().
@@ -156,3 +157,37 @@ export function getToolOutputDir(toolName) {
156
157
  const ctx = getSessionContext();
157
158
  return path.join(ctx.artifactsDir, `tool-${toolName}`);
158
159
  }
160
+ /**
161
+ * Get the current invocation context.
162
+ * Returns undefined if called outside of an invocation context.
163
+ */
164
+ export function getInvocationContext() {
165
+ return invocationContextStorage.getStore();
166
+ }
167
+ /**
168
+ * Bind an async generator to an invocation context so that every iteration
169
+ * runs with the invocation context available.
170
+ */
171
+ export function bindGeneratorToInvocationContext(ctx, generator) {
172
+ const boundNext = (value) => invocationContextStorage.run(ctx, () => generator.next(value));
173
+ const boundReturn = generator.return
174
+ ? (value) => invocationContextStorage.run(ctx, () => generator.return?.(value))
175
+ : undefined;
176
+ const boundThrow = generator.throw
177
+ ? (e) => invocationContextStorage.run(ctx, () => generator.throw?.(e))
178
+ : undefined;
179
+ const boundGenerator = {
180
+ next: boundNext,
181
+ return: boundReturn,
182
+ throw: boundThrow,
183
+ [Symbol.asyncIterator]() {
184
+ return this;
185
+ },
186
+ [Symbol.asyncDispose]: async () => {
187
+ if (Symbol.asyncDispose in generator) {
188
+ await generator[Symbol.asyncDispose]();
189
+ }
190
+ },
191
+ };
192
+ return boundGenerator;
193
+ }