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.
16 kB · 507 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508import type { ChildProcess } from "node:child_process";import { randomUUID } from "node:crypto";import { existsSync } from "node:fs";import { createRequire } from "node:module";import { join, sep } from "node:path";import { fileURLToPath } from "node:url";import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js";import { getPackageDir, isBunBinary } from "../../config.js";import type { DeleteSessionFileResult } from "../../core/session-file-actions.js";import { deleteSessionFile } from "../../core/session-file-actions.js";import { findSessionFileBySelector, getDefaultSessionDir, readSessionInfo, type SessionInfo, SessionManager,} from "../../core/session-manager.js";import { spawnHidden } from "../../utils/child-process.js";
export const DAEMON_CATALOG_ROLE_ENV = "PRIME_AGENT_INTERNAL_DAEMON_CATALOG";const DAEMON_CATALOG_START_TIMEOUT_MS = 30_000;
export function isDaemonCatalogSourcePath(modulePath: string, packageDir: string): boolean { return modulePath.startsWith(`${join(packageDir, "src")}${sep}`);}
function resolveDaemonCatalogEntrypoint(): string { const packageDir = getPackageDir(); const sourceEntrypoint = join(packageDir, "src", "modes", "daemon", "daemon-catalog-entry.ts"); const compiledEntrypoint = join(packageDir, "dist", "modes", "daemon", "daemon-catalog-entry.js"); const runningFromSource = isDaemonCatalogSourcePath(fileURLToPath(import.meta.url), packageDir); const candidates = runningFromSource ? [sourceEntrypoint, compiledEntrypoint] : [compiledEntrypoint, sourceEntrypoint]; const entrypoint = candidates.find((candidate) => existsSync(candidate)); if (entrypoint) return entrypoint; throw new Error("Cannot locate the daemon catalog entrypoint");}
interface SessionInfoWire extends Omit<SessionInfo, "created" | "modified"> { created: string; modified: string;}
type CatalogRequest = | { type: "request"; id: string; command: "list"; cwd?: string; sessionDir?: string } | { type: "request"; id: string; command: "resolve"; selector: string; cwd: string; sessionDir?: string } | { type: "request"; id: string; command: "rename"; sessionPath: string; name: string } | { type: "request"; id: string; command: "delete"; sessionPath: string } | { type: "request"; id: string; command: "archive"; sessionPath: string; sessionId: string } | { type: "request"; id: string; command: "mark_interrupted"; sessionPath: string; activeSessionId: string; operations: string[]; } | { type: "request"; id: string; command: "shutdown" };
type CatalogOutbound = | { type: "ready" } | { type: "progress"; id: string; loaded: number; total: number } | { type: "session"; id: string; session: SessionInfoWire } | { type: "response"; id: string; success: true; data?: unknown } | { type: "response"; id: string; success: false; error: string };
interface CatalogListCallbacks { onProgress?: (loaded: number, total: number) => void; onSession?: (session: SessionInfo) => void;}
function serializeSessionInfo(session: SessionInfo): SessionInfoWire { return { ...session, created: session.created.toISOString(), modified: session.modified.toISOString(), };}
function deserializeSessionInfo(session: SessionInfoWire): SessionInfo { return { ...session, created: new Date(session.created), modified: new Date(session.modified), };}
export function resolveCatalogSessionMatch( sessions: readonly SessionInfo[], selector: string,): SessionInfo | undefined { const matches = sessions.filter((session) => session.id.startsWith(selector) || session.name === selector); if (matches.length > 1) { throw new Error(`Ambiguous session selector "${selector}"`); } return matches[0];}
function isCatalogOutbound(value: unknown): value is CatalogOutbound { if (!value || typeof value !== "object") { return false; } const candidate = value as { type?: unknown; id?: unknown }; return ( candidate.type === "ready" || ((candidate.type === "progress" || candidate.type === "session" || candidate.type === "response") && typeof candidate.id === "string") );}
function isCatalogRequest(value: unknown): value is CatalogRequest { if (!value || typeof value !== "object") { return false; } const candidate = value as { type?: unknown; id?: unknown; command?: unknown }; return ( candidate.type === "request" && typeof candidate.id === "string" && (candidate.command === "list" || candidate.command === "resolve" || candidate.command === "rename" || candidate.command === "delete" || candidate.command === "archive" || candidate.command === "mark_interrupted" || candidate.command === "shutdown") );}
function sendCatalogMessage(message: CatalogOutbound): void { if (process.send) { process.send(message); }}
export function isDaemonCatalogProcess(environment: NodeJS.ProcessEnv = process.env): boolean { return environment[DAEMON_CATALOG_ROLE_ENV] === "1";}
export async function runDaemonCatalogProcess(): Promise<never> { process.on("disconnect", () => process.exit(0)); process.on("message", (value: unknown) => { if (!isCatalogRequest(value)) { return; } void handleCatalogRequest(value); }); sendCatalogMessage({ type: "ready" }); return new Promise(() => {});}
async function handleCatalogRequest(request: CatalogRequest): Promise<void> { try { switch (request.command) { case "list": { const callbacks = { onProgress: (loaded: number, total: number) => sendCatalogMessage({ type: "progress", id: request.id, loaded, total }), onSession: (session: SessionInfo) => sendCatalogMessage({ type: "session", id: request.id, session: serializeSessionInfo(session) }), }; const sessions = request.cwd ? await SessionManager.list(request.cwd, request.sessionDir, callbacks) : await SessionManager.listAll(callbacks, request.sessionDir); sendCatalogMessage({ type: "response", id: request.id, success: true, data: { sessions: sessions.map(serializeSessionInfo) }, }); return; } case "resolve": { // Selectors that name one transcript resolve from the directory listing and its // header cwd; the catalog scans below stay the fallback for names, ambiguous // selectors, and sessions whose header disagrees with the file name. const directMatch = findSessionFileBySelector( request.selector, request.sessionDir ?? getDefaultSessionDir(request.cwd), request.cwd, ); if (directMatch) { sendCatalogMessage({ type: "response", id: request.id, success: true, data: { sessionPath: directMatch.path }, }); return; } const localMatch = resolveCatalogSessionMatch( await SessionManager.list(request.cwd, request.sessionDir), request.selector, ); if (localMatch) { sendCatalogMessage({ type: "response", id: request.id, success: true, data: { sessionPath: localMatch.path }, }); return; } const globalMatch = resolveCatalogSessionMatch( await SessionManager.listAll(undefined, request.sessionDir), request.selector, ); if (globalMatch) { sendCatalogMessage({ type: "response", id: request.id, success: true, data: { sessionPath: globalMatch.path }, }); return; } throw new Error(`No session found matching '${request.selector}'`); } case "rename": SessionManager.open(request.sessionPath).appendSessionInfo(request.name.trim()); sendCatalogMessage({ type: "response", id: request.id, success: true }); return; case "delete": sendCatalogMessage({ type: "response", id: request.id, success: true, data: await deleteSessionFile(request.sessionPath), }); return; case "archive": { const session = await readSessionInfo(request.sessionPath); if (!session || session.id !== request.sessionId) { sendCatalogMessage({ type: "response", id: request.id, success: true, data: { archived: false }, }); return; } if (session.state?.status !== "archived") { SessionManager.open(request.sessionPath).appendSessionState({ status: "archived" }); } sendCatalogMessage({ type: "response", id: request.id, success: true, data: { archived: true }, }); return; } case "mark_interrupted": SessionManager.open(request.sessionPath).appendCustomMessageEntry( "prime-agent.worker_recovery", "<prime_agent_worker_interrupted>\nThe isolated session worker stopped during in-flight work. The saved transcript was recovered, but uncertain model, tool, bash, or child-agent work was not replayed. Inspect external side effects before continuing.\n</prime_agent_worker_interrupted>", false, { activeSessionId: request.activeSessionId, operations: request.operations, }, ); sendCatalogMessage({ type: "response", id: request.id, success: true }); return; case "shutdown": sendCatalogMessage({ type: "response", id: request.id, success: true }); setImmediate(() => process.exit(0)); return; } } catch (error) { sendCatalogMessage({ type: "response", id: request.id, success: false, error: error instanceof Error ? error.message : String(error), }); }}
export class DaemonCatalogClient { private child?: ChildProcess; private starting?: Promise<void>; private readonly pending = new Map< string, { resolve: (data: unknown) => void; reject: (error: Error) => void; callbacks?: CatalogListCallbacks; timeout: ReturnType<typeof setTimeout>; } >();
constructor(private readonly onDiagnostic: (message: string) => void) {}
async start(): Promise<void> { if (this.child?.connected) { return; } if (this.starting) { return this.starting; } this.starting = this.spawnCatalog().finally(() => { this.starting = undefined; }); return this.starting; }
async list(cwd?: string, sessionDir?: string, callbacks?: CatalogListCallbacks): Promise<SessionInfo[]> { const data = await this.request<{ sessions: SessionInfoWire[] }>( { type: "request", id: randomUUID(), command: "list", cwd, sessionDir }, callbacks, ); return data.sessions.map(deserializeSessionInfo); }
async rename(sessionPath: string, name: string): Promise<void> { await this.request({ type: "request", id: randomUUID(), command: "rename", sessionPath, name }); }
async resolve(selector: string, cwd: string, sessionDir?: string): Promise<string> { const data = await this.request<{ sessionPath: string }>({ type: "request", id: randomUUID(), command: "resolve", selector, cwd, sessionDir, }); return data.sessionPath; }
delete(sessionPath: string): Promise<DeleteSessionFileResult> { return this.request({ type: "request", id: randomUUID(), command: "delete", sessionPath }); }
async archive(sessionPath: string, sessionId: string): Promise<boolean> { const data = await this.request<{ archived: boolean }>({ type: "request", id: randomUUID(), command: "archive", sessionPath, sessionId, }); return data.archived; }
async markInterrupted(sessionPath: string, activeSessionId: string, operations: string[]): Promise<void> { await this.request({ type: "request", id: randomUUID(), command: "mark_interrupted", sessionPath, activeSessionId, operations, }); }
async stop(): Promise<void> { const child = this.child; if (!child) { return; } await this.request({ type: "request", id: randomUUID(), command: "shutdown" }).catch(() => undefined); child.disconnect(); this.child = undefined; }
private async spawnCatalog(): Promise<void> { let command: string; let args: string[]; let environment = createCliSubprocessEnv({ ...process.env, [DAEMON_CATALOG_ROLE_ENV]: "1" }); if (isBunBinary) { const launch = createCliSubprocessLaunchSpec(["--version"]); command = launch.command; args = launch.args; } else { const catalogEntry = resolveDaemonCatalogEntrypoint(); const execArgs = catalogEntry.endsWith(".ts") ? [...process.execArgv, "--import", createRequire(import.meta.url).resolve("tsx")] : process.execArgv; const launch = createCliSubprocessLaunchSpec([], undefined, execArgs, catalogEntry); command = launch.command; args = launch.args; environment = createCliSubprocessEnv(environment, catalogEntry, execArgs); } const child = spawnHidden(command, args, { cwd: process.cwd(), env: environment, stdio: ["ignore", "ignore", "ignore", "ipc"], }); this.child = child; child.on("message", (value: unknown) => this.handleMessage(value)); child.on("error", (error) => this.handleClose(child, error)); child.on("exit", (code, signal) => this.handleClose(child, new Error(`Daemon catalog exited (${signal ?? code ?? "unknown"})`)), ); await new Promise<void>((resolveReady, rejectReady) => { const timeout = setTimeout(() => { cleanup(); const error = new Error("Timed out starting daemon catalog"); this.handleClose(child, error); if (child.connected) { child.disconnect(); } child.kill("SIGKILL"); rejectReady(error); }, DAEMON_CATALOG_START_TIMEOUT_MS); const cleanup = () => { clearTimeout(timeout); child.off("message", onMessage); child.off("error", onError); child.off("exit", onExit); }; const onMessage = (value: unknown) => { if (isCatalogOutbound(value) && value.type === "ready") { cleanup(); resolveReady(); } }; const onError = (error: Error) => { cleanup(); rejectReady(error); }; const onExit = (code: number | null, signal: NodeJS.Signals | null) => { onError(new Error(`Daemon catalog exited during startup (${signal ?? code ?? "unknown"})`)); }; child.on("message", onMessage); child.once("error", onError); child.once("exit", onExit); }); }
private async request<T = void>(request: CatalogRequest, callbacks?: CatalogListCallbacks): Promise<T> { await this.start(); const child = this.child; if (!child?.connected) { throw new Error("Daemon catalog is not connected"); } return new Promise<T>((resolveRequest, rejectRequest) => { const timeout = setTimeout( () => { if (!this.pending.delete(request.id)) { return; } child.kill("SIGKILL"); rejectRequest(new Error(`Timed out waiting for daemon catalog ${request.command}`)); }, 5 * 60 * 1000, ); this.pending.set(request.id, { resolve: (data) => resolveRequest(data as T), reject: rejectRequest, callbacks, timeout, }); child.send(request, (error) => { if (!error) { return; } const pending = this.pending.get(request.id); if (pending) { clearTimeout(pending.timeout); this.pending.delete(request.id); } rejectRequest(error); }); }); }
private handleMessage(value: unknown): void { if (!isCatalogOutbound(value) || value.type === "ready") { return; } const pending = this.pending.get(value.id); if (!pending) { return; } if (value.type === "progress") { pending.callbacks?.onProgress?.(value.loaded, value.total); return; } if (value.type === "session") { pending.callbacks?.onSession?.(deserializeSessionInfo(value.session)); return; } this.pending.delete(value.id); clearTimeout(pending.timeout); if (value.success) { pending.resolve(value.data); } else { pending.reject(new Error(value.error)); } }
private handleClose(child: ChildProcess, error: Error): void { if (this.child !== child) { return; } this.child = undefined; this.onDiagnostic(error.message); for (const [id, pending] of this.pending) { clearTimeout(pending.timeout); pending.reject(error); this.pending.delete(id); } }}