Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
17 kB · 533 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534import type { ChildProcess, StdioOptions } from "node:child_process";import { randomUUID } from "node:crypto";import { chmodSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import type { AgentSession } from "../core/agent-session.js";import type { AgentSessionRuntime } from "../core/agent-session-runtime.js";import { clearOrphanProcessJournal, killOrphanProcess, ORPHAN_PROCESS_JOURNAL_ENV, readActiveOrphanProcesses, shouldReapOrphanProcess,} from "../core/orphan-process-journal.js";import { SESSION_LEASE_OWNER_ID_ENV, SESSION_LEASES_ENABLED_ENV } from "../core/session-lease.js";import { attachJsonlLineReader, serializeJsonLine } from "../modes/rpc/jsonl.js";import { spawnHidden } from "../utils/child-process.js";import { isHelpCommandRequest, PUBLIC_COMMAND_NAMES, REMOVED_COMMAND_NAMES } from "./command-registry.js";import { type CliSubprocessLaunchSpec, createCliSubprocessLaunchSpec } from "./subprocess-launch.js";
const OWNED_WORKER_ENV = "PRIME_AGENT_INTERNAL_OWNED_WORKER";const OWNED_RECOVERY_DESCRIPTOR_ENV = "PRIME_AGENT_INTERNAL_OWNED_RECOVERY_DESCRIPTOR";const OWNED_PROFILE_ENV = "PRIME_AGENT_INTERNAL_OWNED_PROFILE";
let closeOwnerWatch: (() => void) | undefined;
export type OwnedSessionWorkerProfile = "print" | "json" | "rpc" | "interactive-ephemeral";
export function isOwnedSessionWorkerProcess(environment: NodeJS.ProcessEnv = process.env): boolean { return environment[OWNED_WORKER_ENV] === "1";}
interface OwnedSessionRecoveryDescriptor { version: 1; profile: OwnedSessionWorkerProfile; sessionId: string; sessionFile?: string; cwd: string; updatedAt: string;}
const NON_SESSION_FLAGS = new Set(["--help", "-h", "--version", "-v", "--list-models", "--export"]);
const NON_SESSION_COMMANDS = new Set([...PUBLIC_COMMAND_NAMES, ...REMOVED_COMMAND_NAMES]);
function valueAfter(args: readonly string[], flag: string): string | undefined { const index = args.indexOf(flag); return index === -1 ? undefined : args[index + 1];}
function hasNonSessionOperation(args: readonly string[]): boolean { if (args.some((arg) => NON_SESSION_FLAGS.has(arg) || arg.startsWith("--export="))) { return true; } const first = args[0]; if (first === "help") { return isHelpCommandRequest(args.slice(1)); } return first !== undefined && NON_SESSION_COMMANDS.has(first);}
function isStartupBenchmark(environment: NodeJS.ProcessEnv): boolean { const value = environment.PI_STARTUP_BENCHMARK?.toLowerCase(); return value === "1" || value === "true" || value === "yes";}
export function classifyOwnedSessionWorkerInvocation( args: readonly string[], stdinIsTTY: boolean | undefined, environment: NodeJS.ProcessEnv = process.env,): OwnedSessionWorkerProfile | undefined { if (isOwnedSessionWorkerProcess(environment) || hasNonSessionOperation(args)) { return undefined; }
const mode = valueAfter(args, "--mode"); if (mode === "daemon") { return undefined; } if (mode === "rpc") { return "rpc"; } if (mode === "json") { return "json"; } if (args.includes("--print") || args.includes("-p") || stdinIsTTY === false) { return "print"; } if (args.includes("--no-session") || isStartupBenchmark(environment)) { return "interactive-ephemeral"; } return undefined;}
export type OwnedWorkerLaunchSpec = CliSubprocessLaunchSpec;
export function createOwnedWorkerLaunchSpec( args: readonly string[], executable = process.execPath, execArgs: readonly string[] = process.execArgv, entrypoint = process.argv[1],): OwnedWorkerLaunchSpec { return createCliSubprocessLaunchSpec(args, executable, execArgs, entrypoint);}
function readOwnedRecoveryDescriptor(path: string): OwnedSessionRecoveryDescriptor | undefined { try { const value = JSON.parse(readFileSync(path, "utf8")) as Partial<OwnedSessionRecoveryDescriptor>; if ( value.version === 1 && typeof value.sessionId === "string" && typeof value.cwd === "string" && (value.sessionFile === undefined || typeof value.sessionFile === "string") ) { return value as OwnedSessionRecoveryDescriptor; } } catch { // The worker may have stopped before creating a recoverable session. } return undefined;}
function writeOwnedRecoveryDescriptor(path: string, profile: OwnedSessionWorkerProfile, session: AgentSession): void { const descriptor: OwnedSessionRecoveryDescriptor = { version: 1, profile, sessionId: session.sessionId, ...(session.sessionFile ? { sessionFile: session.sessionFile } : {}), cwd: session.sessionManager.getCwd(), updatedAt: new Date().toISOString(), }; const tempPath = `${path}.${process.pid}.tmp`; writeFileSync(tempPath, `${JSON.stringify(descriptor)}\n`, { mode: 0o600 }); chmodSync(tempPath, 0o600); renameSync(tempPath, path);}
export function installOwnedSessionRecoveryTracking(runtime: AgentSessionRuntime): void { const path = process.env[OWNED_RECOVERY_DESCRIPTOR_ENV]; const profile = process.env[OWNED_PROFILE_ENV] as OwnedSessionWorkerProfile | undefined; if (!path || !profile) { return; } let unsubscribeSession: (() => void) | undefined; const bind = (session: AgentSession) => { unsubscribeSession?.(); writeOwnedRecoveryDescriptor(path, profile, session); let lastSessionFile = session.sessionFile; unsubscribeSession = session.subscribe((event) => { if (event.type !== "message_start" && event.type !== "session_info_changed") { return; } if (session.sessionFile !== lastSessionFile) { lastSessionFile = session.sessionFile; writeOwnedRecoveryDescriptor(path, profile, session); } }); }; bind(runtime.session); runtime.onSessionReplaced(bind);}
export function createRpcRecoveryArgs(args: readonly string[], sessionPath: string): string[] { const recovered: string[] = []; for (let index = 0; index < args.length; index++) { const arg = args[index]!; if (arg === "--resume" || arg === "-r" || arg === "--fork") { index++; continue; } if (arg.startsWith("--resume=")) { continue; } if (arg === "--continue" || arg === "-c") { continue; } recovered.push(arg); } return [...recovered, "--resume", sessionPath];}
function exitCodeForSignal(signal: NodeJS.Signals | null): number { if (signal === "SIGHUP") { return 129; } if (signal === "SIGINT") { return 130; } if (signal === "SIGTERM") { return 143; } return 1;}
function forwardSignal(child: ChildProcess, signal: NodeJS.Signals): void { if (child.exitCode === null && child.signalCode === null) { child.kill(signal); }}
export async function runOwnedSessionWorkerFrontend( args: readonly string[], profile: OwnedSessionWorkerProfile,): Promise<number> { const interactive = profile === "interactive-ephemeral"; const recoveryDescriptorPath = join(tmpdir(), `prime-agent-owned-${process.pid}-${randomUUID().slice(0, 12)}.json`); const orphanProcessJournalPath = `${recoveryDescriptorPath}.orphans.jsonl`; let currentChild: ChildProcess | undefined; let terminating = false; let terminationSignal: NodeJS.Signals | undefined; let stdinEnded = false; let currentRpcInput: NodeJS.WritableStream | undefined; let currentRpcOutput: NodeJS.ReadableStream | undefined; let rpcStdoutPaused = false; let detachRpcInput: (() => void) | undefined; let detachRpcOutput: (() => void) | undefined; const bufferedRpcInput: string[] = []; const pendingRpcCommands = new Map<string, { publicId?: string; command: string }>(); const anonymousRpcIdPrefix = `prime-agent-owned-${randomUUID()}`; let anonymousRpcCommandId = 0;
const prepareRpcInput = (line: string): string => { try { const command = JSON.parse(line) as { id?: unknown; type?: unknown } | null; if ( !command || Array.isArray(command) || typeof command.type !== "string" || command.type === "extension_ui_response" || command.type === "ack_result" ) { return `${line}\n`; } const publicId = typeof command.id === "string" ? command.id : undefined; const internalId = publicId ?? `${anonymousRpcIdPrefix}-${++anonymousRpcCommandId}`; pendingRpcCommands.set(internalId, { publicId, command: command.type }); return publicId !== undefined ? `${line}\n` : serializeJsonLine({ ...command, id: internalId }); } catch { // The worker preserves the existing parse-error response contract. return `${line}\n`; } }; const observeRpcOutput = (line: string) => { let parsed: unknown; try { parsed = JSON.parse(line) as unknown; } catch { return; } if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { return; } const response = parsed as { id?: unknown; type?: unknown; command?: unknown }; const id = typeof response.id === "string" ? response.id : undefined; const pending = id !== undefined ? pendingRpcCommands.get(id) : undefined; let framed = `${line}\n`; if (response.type === "response" && pending?.publicId === undefined) { const { id: _internalId, ...publicResponse } = response; framed = serializeJsonLine(publicResponse); } if (!process.stdout.write(framed) && !rpcStdoutPaused) { rpcStdoutPaused = true; currentRpcOutput?.pause(); process.stdout.once("drain", () => { rpcStdoutPaused = false; currentRpcOutput?.resume(); }); } if (response.type !== "response" || typeof response.command !== "string") { return; } if (id !== undefined && pending?.command === response.command) { pendingRpcCommands.delete(id); } }; const failPendingRpcCommands = () => { for (const pending of pendingRpcCommands.values()) { process.stdout.write( serializeJsonLine({ ...(pending.publicId !== undefined ? { id: pending.publicId } : {}), type: "response", command: pending.command, success: false, error: "The isolated session worker stopped during this command; its result is uncertain and was not replayed", }), ); } pendingRpcCommands.clear(); }; const reapWorkerResources = (workerPid: number | undefined) => { if (!workerPid) { return; } if (process.platform !== "win32") { try { process.kill(-workerPid, "SIGKILL"); } catch { // The worker process group may already be fully reaped. } } for (const orphan of readActiveOrphanProcesses(orphanProcessJournalPath, workerPid)) { if (!shouldReapOrphanProcess(orphan)) { continue; } killOrphanProcess(orphan.pid); } clearOrphanProcessJournal(orphanProcessJournalPath); };
if (profile === "rpc") { detachRpcInput = attachJsonlLineReader(process.stdin, (line) => { const framed = prepareRpcInput(line); const input = currentRpcInput; if (input?.writable) { if (!input.write(framed)) { process.stdin.pause(); input.once("drain", () => { if (currentRpcInput === input && !stdinEnded) { process.stdin.resume(); } }); } } else { bufferedRpcInput.push(framed); } }); process.stdin.once("end", () => { stdinEnded = true; currentRpcInput?.end(); }); }
const spawnWorker = (workerArgs: readonly string[]): ChildProcess => { const launch = createOwnedWorkerLaunchSpec(workerArgs); const bridgeStdin = profile === "rpc" || process.stdin.isTTY !== true; const stdio: StdioOptions = interactive ? ["inherit", "inherit", "inherit", "ipc"] : [bridgeStdin ? "pipe" : "inherit", "pipe", "pipe", "ipc"]; const child = spawnHidden(launch.command, launch.args, { cwd: process.cwd(), detached: process.platform !== "win32", env: { ...process.env, [OWNED_WORKER_ENV]: "1", [OWNED_RECOVERY_DESCRIPTOR_ENV]: recoveryDescriptorPath, [OWNED_PROFILE_ENV]: profile, [ORPHAN_PROCESS_JOURNAL_ENV]: orphanProcessJournalPath, [SESSION_LEASES_ENABLED_ENV]: "1", [SESSION_LEASE_OWNER_ID_ENV]: `owned-${randomUUID()}`, }, stdio, }); currentChild = child; if (!interactive) { const childInput = child.stdin ?? undefined; const childOutput = child.stdout ?? undefined; const childError = child.stderr ?? undefined; if ((bridgeStdin && !childInput) || !childOutput || !childError) { child.kill("SIGTERM"); throw new Error("Owned session worker did not expose bridged stdio"); } childError.pipe(process.stderr, { end: false }); if (profile === "rpc") { if (!childInput) { throw new Error("Owned RPC worker did not expose stdin"); } currentRpcInput = childInput; currentRpcOutput = childOutput; if (rpcStdoutPaused) { childOutput.pause(); } detachRpcOutput = attachJsonlLineReader(childOutput, observeRpcOutput); for (const buffered of bufferedRpcInput.splice(0)) { childInput.write(buffered); } if (stdinEnded) { childInput.end(); } } else { if (childInput) { process.stdin.pipe(childInput); } childOutput.pipe(process.stdout, { end: false }); } } return child; };
const signals: NodeJS.Signals[] = ["SIGINT", "SIGTERM"]; if (process.platform !== "win32") { signals.push("SIGHUP"); } const signalHandlers = signals.map((signal) => { const handler = () => { terminating = true; terminationSignal ??= signal; if (currentChild) { forwardSignal(currentChild, signal); } }; process.on(signal, handler); return { signal, handler }; });
try { let workerArgs = [...args]; let recoveryAttempt = 0; while (true) { if (terminating) { return exitCodeForSignal(terminationSignal ?? null); } const workerStartedAt = Date.now(); const child = spawnWorker(workerArgs); const workerPid = child.pid; const exit = await new Promise<{ code: number; signal: NodeJS.Signals | null }>((resolveExit, reject) => { child.once("error", reject); child.once("close", (code, signal) => resolveExit({ code: code ?? exitCodeForSignal(signal), signal })); }); currentChild = undefined; currentRpcInput = undefined; currentRpcOutput = undefined; if (profile === "rpc" && !stdinEnded) { process.stdin.resume(); } detachRpcOutput?.(); detachRpcOutput = undefined; if (!interactive && profile !== "rpc" && child.stdin) { process.stdin.unpipe(child.stdin); } if (child.connected) { child.disconnect(); } reapWorkerResources(workerPid); const rpcCrashed = profile === "rpc" && !terminating && (exit.code !== 0 || exit.signal !== null || pendingRpcCommands.size > 0); const workerExitCode = rpcCrashed && exit.code === 0 ? 1 : exit.code; if (Date.now() - workerStartedAt >= 60_000) { recoveryAttempt = 0; } if (rpcCrashed) { failPendingRpcCommands(); } const shouldRecover = rpcCrashed && !stdinEnded && recoveryAttempt < 3; if (!shouldRecover) { return terminationSignal ? exitCodeForSignal(terminationSignal) : workerExitCode; } const descriptor = readOwnedRecoveryDescriptor(recoveryDescriptorPath); if (!descriptor?.sessionFile) { return terminationSignal ? exitCodeForSignal(terminationSignal) : workerExitCode; } workerArgs = createRpcRecoveryArgs(args, descriptor.sessionFile); const retryDelay = [250, 1000, 5000][recoveryAttempt] ?? 5000; recoveryAttempt++; await new Promise((resolveDelay) => setTimeout(resolveDelay, retryDelay)); if (terminating) { return exitCodeForSignal(terminationSignal ?? null); } } } finally { for (const { signal, handler } of signalHandlers) { process.off(signal, handler); } detachRpcInput?.(); detachRpcOutput?.(); rmSync(recoveryDescriptorPath, { force: true }); clearOrphanProcessJournal(orphanProcessJournalPath); }}
export async function maybeRunOwnedSessionWorkerFrontend( args: readonly string[], forceLegacyFrontend = false,): Promise<boolean> { if (!forceLegacyFrontend && process.env.PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND !== "1") { return false; } const profile = classifyOwnedSessionWorkerInvocation(args, process.stdin.isTTY); if (!profile) { return false; } process.exitCode = await runOwnedSessionWorkerFrontend(args, profile); return true;}
export function installOwnedSessionWorkerOwnerWatch(): void { if (!isOwnedSessionWorkerProcess()) { return; } if (!process.channel) { throw new Error("Owned session worker is missing its owner channel"); }
let ownerGone = false; const terminate = () => { if (ownerGone) { return; } ownerGone = true; closeOwnerWatch = undefined; const forceTimer = setTimeout(() => { if (process.platform !== "win32") { try { process.kill(-process.pid, "SIGKILL"); return; } catch { // Fall through to terminating only this process. } } process.exit(143); }, 5000); forceTimer.unref(); process.kill(process.pid, "SIGTERM"); }; process.once("disconnect", terminate); process.channel.unref(); closeOwnerWatch = () => { if (ownerGone) { return; } ownerGone = true; process.off("disconnect", terminate); if (process.connected) { process.disconnect(); } closeOwnerWatch = undefined; };}
export function closeOwnedSessionWorkerOwnerWatch(): void { closeOwnerWatch?.();}