Something went wrong. Try again.
Pi extension for async subagent delegation with truncation, artifacts, and session sharing
Something went wrong. Try again.
8.6 kB · 264 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265import * 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";
const MAX_TOOL_PAYLOAD_BYTES = 32 * 1024;const TOOL_PAYLOAD_TRUNCATION_MARKER = "\n\n… payload truncated";
function boundedPayload(value: unknown, maxBytes = MAX_TOOL_PAYLOAD_BYTES): string | undefined { let text: string; if (typeof value === "string") text = value; else { try { const serialized = JSON.stringify(value, null, 2); if (serialized === undefined) return undefined; text = serialized; } catch { return undefined; } } if (!text.trim()) return undefined; const payload = Buffer.from(text, "utf-8"); if (payload.length <= maxBytes) return text; const markerBytes = Buffer.byteLength(TOOL_PAYLOAD_TRUNCATION_MARKER, "utf-8"); let end = Math.max(0, maxBytes - markerBytes); while (end > 0 && (payload[end]! & 0xc0) === 0x80) end--; return `${payload.subarray(0, end).toString("utf-8")}${TOOL_PAYLOAD_TRUNCATION_MARKER}`;}
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; toolCallId?: string; toolName?: string; args?: unknown; isError?: boolean;}
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); if (message.role === "toolResult") { const output = boundedPayload(text); writeRecord({ ...baseRecord("message"), sourceEventType, role: message.role, toolCallId: message.toolCallId, toolName: message.toolName, isError: message.isError, ...(output ? { text: output, outputTruncated: output.includes("… payload truncated") } : {}), message: { role: message.role, toolCallId: message.toolCallId, toolName: message.toolName, isError: message.isError, content: output ? [{ type: "text", text: output }] : [], ...(message.timestamp !== undefined ? { timestamp: message.timestamp } : {}), }, }); return; } 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); const argsPayload = boundedPayload(args); writeRecord({ ...baseRecord("tool_start"), sourceEventType: event.type, ...(event.toolCallId ? { toolCallId: event.toolCallId } : {}), toolName: event.toolName, ...(Object.keys(args).length > 0 ? { argsPreview: extractToolArgsPreview(args) } : {}), ...(argsPayload ? { argsPayload } : {}), }); return; } if (event.type === "tool_execution_end") { writeRecord({ ...baseRecord("tool_end"), sourceEventType: event.type, ...(event.toolCallId ? { toolCallId: event.toolCallId } : {}), ...(event.toolName ? { toolName: event.toolName } : {}), ...(typeof event.isError === "boolean" ? { isError: event.isError } : {}), }); } }, 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; }, };}