Something went wrong. Try again.
Pi extension for async subagent delegation with truncation, artifacts, and session sharing
Something went wrong. Try again.
6.7 kB · 212 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213import * as fs from "node:fs";import * as path from "node:path";import type { Message } from "@earendil-works/pi-ai";import { extractTextFromContent, extractToolArgsPreview } from "./utils.ts";
export const CHILD_TRANSCRIPT_ARTIFACT_VERSION = 1;const DEFAULT_MAX_CHILD_TRANSCRIPT_BYTES = 50 * 1024 * 1024;
type ChildTranscriptSource = "foreground" | "async";type ChildTranscriptRecordType = "message" | "tool_start" | "tool_end" | "stdout" | "stderr" | "truncated";
type ChildTranscriptMessage = Message & { model?: string; errorMessage?: string; stopReason?: string; usage?: unknown;};
interface ChildTranscriptEvent { type?: string; message?: ChildTranscriptMessage; toolName?: string; args?: unknown;}
interface ChildTranscriptWriterInput { transcriptPath: string; source: ChildTranscriptSource; runId: string; agent: string; childIndex?: number; cwd: string; maxBytes?: number;}
export interface ChildTranscriptWriter { path: string; writeInitialUserMessage(prompt: string): void; writeChildEvent(event: ChildTranscriptEvent): void; writeStdoutLine(line: string): void; writeStderrLine(line: string): void; writeStderrText(text: string): void; getError(): string | undefined;}
function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error);}
function finiteNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined;}
function normalizeUsage(value: unknown): { input: number; output: number; cacheRead: number; cacheWrite: number; cost: number } | undefined { if (!value || typeof value !== "object") return undefined; const raw = value as Record<string, unknown>; const rawCost = raw.cost; const cost = rawCost && typeof rawCost === "object" ? finiteNumber((rawCost as { total?: unknown }).total) ?? 0 : finiteNumber(rawCost) ?? 0; return { input: finiteNumber(raw.input) ?? finiteNumber(raw.inputTokens) ?? 0, output: finiteNumber(raw.output) ?? finiteNumber(raw.outputTokens) ?? 0, cacheRead: finiteNumber(raw.cacheRead) ?? 0, cacheWrite: finiteNumber(raw.cacheWrite) ?? 0, cost, };}
function eventArgs(event: ChildTranscriptEvent): Record<string, unknown> { return event.args && typeof event.args === "object" && !Array.isArray(event.args) ? event.args as Record<string, unknown> : {};}
export function createChildTranscriptWriter(input: ChildTranscriptWriterInput): ChildTranscriptWriter { let bytesWritten = 0; let writeError: string | undefined; let truncated = false; const maxBytes = input.maxBytes ?? DEFAULT_MAX_CHILD_TRANSCRIPT_BYTES;
const baseRecord = (recordType: ChildTranscriptRecordType) => { const ts = Date.now(); return { version: CHILD_TRANSCRIPT_ARTIFACT_VERSION, recordType, source: input.source, runId: input.runId, agent: input.agent, ...(input.childIndex !== undefined ? { childIndex: input.childIndex } : {}), cwd: input.cwd, ts, timestamp: new Date(ts).toISOString(), }; };
const writeTruncatedMarker = () => { truncated = true; const marker = `${JSON.stringify({ ...baseRecord("truncated"), maxBytes, message: `Child transcript exceeded ${maxBytes} bytes; further records were omitted.`, })}\n`; const markerBytes = Buffer.byteLength(marker, "utf-8"); if (bytesWritten + markerBytes > maxBytes) return false; try { fs.appendFileSync(input.transcriptPath, marker, "utf-8"); bytesWritten += markerBytes; return true; } catch (error) { writeError = `Failed to write child transcript '${input.transcriptPath}': ${errorMessage(error)}`; return false; } };
const writeRecord = (record: Record<string, unknown>) => { if (writeError || truncated) return; const line = `${JSON.stringify(record)}\n`; const bytes = Buffer.byteLength(line, "utf-8"); if (bytesWritten + bytes > maxBytes) { writeTruncatedMarker(); return; } const markerProbe = `${JSON.stringify({ ...baseRecord("truncated"), maxBytes, message: `Child transcript exceeded ${maxBytes} bytes; further records were omitted.`, })}\n`; if (bytesWritten + bytes + Buffer.byteLength(markerProbe, "utf-8") > maxBytes) { writeTruncatedMarker(); return; } try { fs.appendFileSync(input.transcriptPath, line, "utf-8"); bytesWritten += bytes; } catch (error) { writeError = `Failed to write child transcript '${input.transcriptPath}': ${errorMessage(error)}`; } };
try { fs.mkdirSync(path.dirname(input.transcriptPath), { recursive: true }); fs.writeFileSync(input.transcriptPath, "", "utf-8"); } catch (error) { writeError = `Failed to initialize child transcript '${input.transcriptPath}': ${errorMessage(error)}`; }
const writeMessage = (sourceEventType: string, message: ChildTranscriptMessage) => { const text = extractTextFromContent(message.content); writeRecord({ ...baseRecord("message"), sourceEventType, role: message.role, ...(text ? { text } : {}), ...(message.model ? { model: message.model } : {}), ...(message.stopReason ? { stopReason: message.stopReason } : {}), ...(message.errorMessage ? { errorMessage: message.errorMessage } : {}), ...(message.usage ? { usage: normalizeUsage(message.usage) } : {}), message, }); };
return { path: input.transcriptPath, writeInitialUserMessage(prompt: string) { writeRecord({ ...baseRecord("message"), sourceEventType: "initial_prompt", role: "user", text: prompt, message: { role: "user", content: [{ type: "text", text: prompt }] }, }); }, writeChildEvent(event: ChildTranscriptEvent) { if ((event.type === "message_end" || event.type === "tool_result_end") && event.message) { writeMessage(event.type, event.message); return; } if (event.type === "tool_execution_start" && event.toolName) { const args = eventArgs(event); writeRecord({ ...baseRecord("tool_start"), sourceEventType: event.type, toolName: event.toolName, ...(Object.keys(args).length > 0 ? { argsPreview: extractToolArgsPreview(args) } : {}), }); return; } if (event.type === "tool_execution_end") { writeRecord({ ...baseRecord("tool_end"), sourceEventType: event.type, ...(event.toolName ? { toolName: event.toolName } : {}), }); } }, writeStdoutLine(line: string) { if (!line.trim()) return; writeRecord({ ...baseRecord("stdout"), text: line }); }, writeStderrLine(line: string) { if (!line.trim()) return; writeRecord({ ...baseRecord("stderr"), text: line }); }, writeStderrText(text: string) { for (const line of text.split(/\r?\n/)) this.writeStderrLine(line); }, getError() { return writeError; }, };}