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.
22 kB · 652 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653import { createHash, randomUUID } from "node:crypto";import { mkdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs";import { resolve } from "node:path";import chalk from "chalk";import lockfile from "proper-lockfile";import { ENV_AGENT_DIR, getAgentDir, SELF_UPDATE_INTERACTIVE_CHILD_ENV } from "../config.js";import { ORPHAN_PROCESS_JOURNAL_ENV } from "../core/orphan-process-journal.js";import { getProcessStartId, SESSION_LEASE_OWNER_ID_ENV, SESSION_LEASES_ENABLED_ENV } from "../core/session-lease.js";import { defaultDaemonSocketDir, defaultDaemonSocketPath, normalizeSocketPath } from "../modes/daemon/daemon-socket.js";import { DAEMON_WORKER_ACTIVE_SESSION_ID_ENV, DAEMON_WORKER_RECOVERY_JOURNAL_ENV, DAEMON_WORKER_ROLE_ENV, DAEMON_WORKER_SUPERVISOR_SOCKET_ENV, DAEMON_WORKER_TOKEN_ENV,} from "../modes/daemon/daemon-worker-protocol.js";import { isProcessAlive, spawnHidden } from "../utils/child-process.js";import { createUpdatedCliSubprocessLaunchSpec } from "./subprocess-launch.js";
export const DAEMON_UPDATE_RESTART_COORDINATOR_FLAG = "--internal-update-restart-coordinator";export const DAEMON_UPDATE_RESTART_STATUS_FLAG = "--internal-update-restart-status";export const DAEMON_UPDATE_RESTART_ORIGIN_FLAG = "--internal-update-restart-origin";export const DAEMON_UPDATE_RESTART_ALWAYS_START_FLAG = "--internal-update-restart-always-start";
export type DaemonUpdateRestartPhase = | "starting" | "preparing" | "stopping" | "starting_daemon" | "restoring" | "complete" | "skipped" | "failed";
export interface DaemonUpdateRestartCounts { total: number; restored: number; resumed: number; failed: number;}
export interface DaemonUpdateRestartFailure { sessionFile: string; message: string;}
export interface DaemonUpdateRestartProcessIdentity { pid: number; processStartId?: string; supervisorGeneration?: string; supervisorOwnerToken?: string;}
export interface DaemonUpdateRestartStatus { version: 1; requestId: string; socketPath: string; phase: DaemonUpdateRestartPhase; coordinator: DaemonUpdateRestartProcessIdentity; predecessor?: DaemonUpdateRestartProcessIdentity; successor?: DaemonUpdateRestartProcessIdentity; counts: DaemonUpdateRestartCounts; failures?: DaemonUpdateRestartFailure[]; message?: string; startedAt: string; updatedAt: string; heartbeatAt?: string;}
export interface DaemonUpdateRestartReport { info: string[]; warnings: string[];}
export interface DaemonUpdateRestartCoordinatorRecord extends DaemonUpdateRestartProcessIdentity { version: 1; token: string; requestId: string; socketPath: string; statusPath: string; createdAt: string;}
export interface LaunchDaemonUpdateRestartCoordinatorOptions { socketPath: string; agentDir: string; cwd?: string; originActiveSessionId?: string; timeoutMs?: number; alwaysStart?: boolean;}
export interface AcquireDaemonUpdateRestartCoordinatorOptions { requestId: string; socketPath: string; statusPath: string; registryDir?: string;}
export function resolveDaemonUpdateRestartSocketPath(socketPath?: string): string { return normalizeSocketPath(socketPath ?? defaultDaemonSocketPath());}
const TERMINAL_PHASES: ReadonlySet<DaemonUpdateRestartPhase> = new Set(["complete", "skipped", "failed"]);const ALL_PHASES: ReadonlySet<DaemonUpdateRestartPhase> = new Set([ "starting", "preparing", "stopping", "starting_daemon", "restoring", "complete", "skipped", "failed",]);const DEFAULT_COORDINATOR_PROGRESS_TIMEOUT_MS = 30 * 60_000;const COORDINATOR_LIVENESS_TIMEOUT_MS = 180_000;const COORDINATOR_STATUS_HEARTBEAT_MS = 5000;const COORDINATOR_REGISTRY_LOCK_STALE_MS = 5000;const COORDINATOR_REGISTRY_LOCK_UPDATE_MS = 1000;const COORDINATOR_REGISTRY_LOCK_RETRIES = 500;const COORDINATOR_REGISTRY_LOCK_RETRY_MS = 10;const PROCESS_START_ID_RECHECK_MS = 1000;
export function buildDaemonUpdateRestartReport(status: DaemonUpdateRestartStatus): DaemonUpdateRestartReport { const report: DaemonUpdateRestartReport = { info: [], warnings: [] }; if (status.phase === "failed") { report.warnings.push(`Updated, but could not restart the daemon (${status.message ?? "unknown error"}).`); } if (status.phase !== "complete" && status.phase !== "failed") { return report; } if (status.counts.total > 0) { report.info.push(`Restored ${status.counts.restored} daemon session${status.counts.restored === 1 ? "" : "s"}`); } if (status.counts.resumed > 0) { report.info.push(`Resumed ${status.counts.resumed} interrupted session${status.counts.resumed === 1 ? "" : "s"}`); } if (status.counts.failed > 0) { report.warnings.push( `${status.counts.failed} daemon session${status.counts.failed === 1 ? "" : "s"} could not be restored.`, ); } for (const failure of status.failures ?? []) { report.warnings.push(`Could not restore ${failure.sessionFile}: ${failure.message}`); } if (report.info.length === 0 && status.counts.total === 0 && status.phase === "complete") { report.info.push(status.message ?? "Restarted the daemon"); } return report;}
export function reportDaemonUpdateRestartStatus(status: DaemonUpdateRestartStatus): void { const report = buildDaemonUpdateRestartReport(status); for (const message of report.info) { console.log(chalk.green(message)); } for (const warning of report.warnings) { console.error(chalk.yellow(`Warning: ${warning}`)); }}
export interface RunDaemonRestartOptions { socketPath?: string; agentDir?: string; cwd?: string; originActiveSessionId?: string; timeoutMs?: number; json?: boolean;}
export async function runDaemonRestart(options: RunDaemonRestartOptions = {}): Promise<DaemonUpdateRestartStatus> { const socketPath = resolveDaemonUpdateRestartSocketPath(options.socketPath); const agentDir = options.agentDir ?? getAgentDir(); const cwd = options.cwd ?? process.cwd(); const originActiveSessionId = options.originActiveSessionId ?? process.env[DAEMON_WORKER_ACTIVE_SESSION_ID_ENV]; const status = await launchDaemonUpdateRestartCoordinator({ socketPath, agentDir, cwd, originActiveSessionId, timeoutMs: options.timeoutMs, alwaysStart: true, }); if (options.json) { console.log(JSON.stringify(status, null, 2)); } else { reportDaemonUpdateRestartStatus(status); } if (status.phase === "failed") { process.exitCode = 1; } return status;}
function delay(ms: number): Promise<void> { return new Promise((resolveDelay) => setTimeout(resolveDelay, ms));}
function statusLivenessId(status: DaemonUpdateRestartStatus): string { return status.heartbeatAt ?? status.updatedAt;}
function socketKey(socketPath: string): string { return createHash("sha256").update(normalizeSocketPath(socketPath)).digest("hex");}
function writeJsonAtomically(path: string, value: unknown): void { const tempPath = `${path}.${process.pid}.${randomUUID()}.tmp`; try { writeFileSync(tempPath, `${JSON.stringify(value, null, 2)}\n`, { mode: 0o600 }); renameSync(tempPath, path); } catch (error) { rmSync(tempPath, { force: true }); throw error; }}
function isProcessIdentity(value: unknown): value is DaemonUpdateRestartProcessIdentity { if (!value || typeof value !== "object") { return false; } const identity = value as Partial<DaemonUpdateRestartProcessIdentity>; return ( Number.isInteger(identity.pid) && (identity.pid ?? 0) > 0 && (identity.processStartId === undefined || typeof identity.processStartId === "string") && (identity.supervisorGeneration === undefined || typeof identity.supervisorGeneration === "string") && (identity.supervisorOwnerToken === undefined || typeof identity.supervisorOwnerToken === "string") );}
function isCounts(value: unknown): value is DaemonUpdateRestartCounts { if (!value || typeof value !== "object") { return false; } const counts = value as Partial<DaemonUpdateRestartCounts>; return [counts.total, counts.restored, counts.resumed, counts.failed].every( (entry) => Number.isInteger(entry) && (entry ?? -1) >= 0, );}
function isFailures(value: unknown): value is DaemonUpdateRestartFailure[] { return ( Array.isArray(value) && value.every( (failure) => failure !== null && typeof failure === "object" && typeof (failure as Partial<DaemonUpdateRestartFailure>).sessionFile === "string" && typeof (failure as Partial<DaemonUpdateRestartFailure>).message === "string", ) );}
function isDaemonUpdateRestartStatus(value: unknown): value is DaemonUpdateRestartStatus { if (!value || typeof value !== "object") { return false; } const status = value as Partial<DaemonUpdateRestartStatus>; return ( status.version === 1 && typeof status.requestId === "string" && typeof status.socketPath === "string" && typeof status.phase === "string" && ALL_PHASES.has(status.phase as DaemonUpdateRestartPhase) && isProcessIdentity(status.coordinator) && (status.predecessor === undefined || isProcessIdentity(status.predecessor)) && (status.successor === undefined || isProcessIdentity(status.successor)) && isCounts(status.counts) && (status.failures === undefined || isFailures(status.failures)) && (status.message === undefined || typeof status.message === "string") && typeof status.startedAt === "string" && typeof status.updatedAt === "string" && (status.heartbeatAt === undefined || typeof status.heartbeatAt === "string") );}
export function readDaemonUpdateRestartStatus(path: string): DaemonUpdateRestartStatus | undefined { try { const value = JSON.parse(readFileSync(path, "utf8")) as unknown; return isDaemonUpdateRestartStatus(value) ? value : undefined; } catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") { return undefined; } throw error; }}
function readTerminalDaemonUpdateRestartStatus(path: string): DaemonUpdateRestartStatus | undefined { const status = readDaemonUpdateRestartStatus(path); return status && TERMINAL_PHASES.has(status.phase) ? status : undefined;}
export class DaemonUpdateRestartStatusWriter { private status: DaemonUpdateRestartStatus;
constructor( private readonly path: string, requestId: string, socketPath: string, ) { const now = new Date().toISOString(); const processStartId = getProcessStartId(process.pid); this.status = { version: 1, requestId, socketPath, phase: "starting", coordinator: { pid: process.pid, ...(processStartId ? { processStartId } : {}) }, counts: { total: 0, restored: 0, resumed: 0, failed: 0 }, startedAt: now, updatedAt: now, heartbeatAt: now, }; this.persist(); }
update( update: Partial<Omit<DaemonUpdateRestartStatus, "version" | "requestId" | "socketPath" | "coordinator">>, ): void { const now = new Date().toISOString(); this.status = { ...this.status, ...update, counts: update.counts ? { ...update.counts } : this.status.counts, failures: update.failures ? update.failures.map((failure) => ({ ...failure })) : this.status.failures?.map((failure) => ({ ...failure })), updatedAt: now, heartbeatAt: now, }; this.persist(); }
current(): DaemonUpdateRestartStatus { return { ...this.status, counts: { ...this.status.counts }, failures: this.status.failures?.map((failure) => ({ ...failure })), }; }
touch(): void { this.status = { ...this.status, heartbeatAt: new Date().toISOString() }; this.persist(); }
startHeartbeat(): () => void { const heartbeat = setInterval(() => { try { this.touch(); } catch { // A later phase write will retry; otherwise the parent detects the stale heartbeat. } }, COORDINATOR_STATUS_HEARTBEAT_MS); heartbeat.unref(); return () => clearInterval(heartbeat); }
private persist(): void { writeJsonAtomically(this.path, this.status); }}
function defaultCoordinatorRegistryDir(): string { return resolve(defaultDaemonSocketDir(), "update-restart-coordinators");}
function coordinatorRecordPath(registryDir: string, socketPath: string): string { return resolve(registryDir, `${socketKey(socketPath)}.json`);}
async function withCoordinatorRegistryGuard<T>(registryDir: string, action: () => T | Promise<T>): Promise<T> { mkdirSync(registryDir, { recursive: true, mode: 0o700 }); let compromisedError: Error | undefined; const assertGuardHeld = () => { if (compromisedError) throw new Error(`Coordinator registry guard was compromised: ${compromisedError.message}`); }; const release = await lockfile.lock(registryDir, { realpath: false, lockfilePath: resolve(registryDir, ".guard"), stale: COORDINATOR_REGISTRY_LOCK_STALE_MS, update: COORDINATOR_REGISTRY_LOCK_UPDATE_MS, onCompromised: (error) => { compromisedError ??= error; }, retries: { retries: COORDINATOR_REGISTRY_LOCK_RETRIES, factor: 1, minTimeout: COORDINATOR_REGISTRY_LOCK_RETRY_MS, maxTimeout: COORDINATOR_REGISTRY_LOCK_RETRY_MS, }, }); try { assertGuardHeld(); const result = await action(); assertGuardHeld(); return result; } finally { if (compromisedError) await release().catch(() => undefined); else await release(); }}
function matchesProcessStartId(identity: DaemonUpdateRestartProcessIdentity): boolean { if (!identity.processStartId) { return true; } const observed = getProcessStartId(identity.pid); return observed === undefined || observed === identity.processStartId;}
function isProcessIdentityAlive(identity: DaemonUpdateRestartProcessIdentity): boolean { return isProcessAlive(identity.pid) && matchesProcessStartId(identity);}
function createProcessIdentityLivenessCheck(identity: DaemonUpdateRestartProcessIdentity): () => boolean { let lastStartIdCheckAt: number | undefined; return () => { if (!isProcessAlive(identity.pid)) { return false; } if (!identity.processStartId) { return true; } const now = Date.now(); if (lastStartIdCheckAt !== undefined && now - lastStartIdCheckAt < PROCESS_START_ID_RECHECK_MS) { return true; } lastStartIdCheckAt = now; return matchesProcessStartId(identity); };}
function readCoordinatorRecord(path: string): DaemonUpdateRestartCoordinatorRecord | undefined { try { const value = JSON.parse(readFileSync(path, "utf8")) as unknown; if (!value || typeof value !== "object") { return undefined; } const record = value as Partial<DaemonUpdateRestartCoordinatorRecord>; if ( record.version !== 1 || typeof record.token !== "string" || typeof record.requestId !== "string" || typeof record.socketPath !== "string" || typeof record.statusPath !== "string" || typeof record.createdAt !== "string" || !isProcessIdentity(record) ) { return undefined; } return record as DaemonUpdateRestartCoordinatorRecord; } catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") { return undefined; } throw error; }}
export class DaemonUpdateRestartCoordinatorLease { private released = false;
constructor( readonly record: DaemonUpdateRestartCoordinatorRecord, private readonly registryDir: string, private readonly path: string, ) {}
async release(): Promise<void> { if (this.released) { return; } await withCoordinatorRegistryGuard(this.registryDir, () => { const current = readCoordinatorRecord(this.path); if (current?.token === this.record.token) { rmSync(this.path, { force: true }); } }); this.released = true; }}
export class DaemonUpdateRestartCoordinatorAlreadyRunningError extends Error { constructor(readonly record: DaemonUpdateRestartCoordinatorRecord) { super(`Another daemon update restart is already running for ${record.socketPath}`); this.name = "DaemonUpdateRestartCoordinatorAlreadyRunningError"; }}
export async function acquireDaemonUpdateRestartCoordinator( options: AcquireDaemonUpdateRestartCoordinatorOptions,): Promise<DaemonUpdateRestartCoordinatorLease> { const registryDir = options.registryDir ?? defaultCoordinatorRegistryDir(); const path = coordinatorRecordPath(registryDir, options.socketPath); const token = randomUUID(); const processStartId = getProcessStartId(process.pid); const record: DaemonUpdateRestartCoordinatorRecord = { version: 1, token, requestId: options.requestId, pid: process.pid, ...(processStartId ? { processStartId } : {}), socketPath: options.socketPath, statusPath: options.statusPath, createdAt: new Date().toISOString(), }; await withCoordinatorRegistryGuard(registryDir, () => { const current = readCoordinatorRecord(path); if (current && isProcessIdentityAlive(current)) { throw new DaemonUpdateRestartCoordinatorAlreadyRunningError(current); } rmSync(path, { force: true }); writeJsonAtomically(path, record); }); return new DaemonUpdateRestartCoordinatorLease(record, registryDir, path);}
export async function waitForActiveDaemonUpdateRestartCoordinator( record: DaemonUpdateRestartCoordinatorRecord, progressTimeoutMs = DEFAULT_COORDINATOR_PROGRESS_TIMEOUT_MS,): Promise<DaemonUpdateRestartStatus> { const coordinatorIsAlive = createProcessIdentityLivenessCheck(record); let observedUpdatedAt: string | undefined; let observedLivenessId: string | undefined; let lastProgressAt = Date.now(); let lastLivenessAt = Date.now(); while (true) { const status = readDaemonUpdateRestartStatus(record.statusPath); if (status && status.updatedAt !== observedUpdatedAt) { observedUpdatedAt = status.updatedAt; lastProgressAt = Date.now(); } if (status && statusLivenessId(status) !== observedLivenessId) { observedLivenessId = statusLivenessId(status); lastLivenessAt = Date.now(); } if (status && TERMINAL_PHASES.has(status.phase)) { return status; } if (!coordinatorIsAlive()) { const terminalStatus = readTerminalDaemonUpdateRestartStatus(record.statusPath); if (terminalStatus) { return terminalStatus; } throw new Error(`Active daemon update restart coordinator exited for ${record.socketPath}`); } if (Date.now() - lastLivenessAt >= COORDINATOR_LIVENESS_TIMEOUT_MS) { throw new Error(`Active daemon update restart coordinator stopped reporting liveness on ${record.socketPath}`); } if (Date.now() - lastProgressAt >= progressTimeoutMs) { throw new Error(`Timed out waiting for active daemon update restart progress on ${record.socketPath}`); } await delay(50); }}
function createStatusPath(agentDir: string, socketPath: string, requestId: string): string { const directory = resolve(agentDir, "update-restarts"); mkdirSync(directory, { recursive: true, mode: 0o700 }); return resolve(directory, `${socketKey(socketPath).slice(0, 16)}-${requestId}.json`);}
function coordinatorEnvironment(agentDir: string): NodeJS.ProcessEnv { const environment = { ...process.env }; environment[ENV_AGENT_DIR] = agentDir; delete environment[SELF_UPDATE_INTERACTIVE_CHILD_ENV]; delete environment[DAEMON_WORKER_ROLE_ENV]; delete environment[DAEMON_WORKER_TOKEN_ENV]; delete environment[DAEMON_WORKER_ACTIVE_SESSION_ID_ENV]; delete environment[DAEMON_WORKER_RECOVERY_JOURNAL_ENV]; delete environment[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV]; delete environment[ORPHAN_PROCESS_JOURNAL_ENV]; delete environment[SESSION_LEASES_ENABLED_ENV]; delete environment[SESSION_LEASE_OWNER_ID_ENV]; return environment;}
export async function launchDaemonUpdateRestartCoordinator( options: LaunchDaemonUpdateRestartCoordinatorOptions,): Promise<DaemonUpdateRestartStatus> { const requestId = randomUUID(); const agentDir = resolve(options.agentDir); const socketPath = resolveDaemonUpdateRestartSocketPath(options.socketPath); const statusPath = createStatusPath(agentDir, socketPath, requestId); const inheritedOrigin = process.env[DAEMON_WORKER_ACTIVE_SESSION_ID_ENV]; const originActiveSessionId = options.originActiveSessionId ?? inheritedOrigin; const launch = createUpdatedCliSubprocessLaunchSpec([ "update", DAEMON_UPDATE_RESTART_COORDINATOR_FLAG, "--daemon-socket", socketPath, DAEMON_UPDATE_RESTART_STATUS_FLAG, statusPath, ...(originActiveSessionId ? [DAEMON_UPDATE_RESTART_ORIGIN_FLAG, originActiveSessionId] : []), ...(options.alwaysStart ? [DAEMON_UPDATE_RESTART_ALWAYS_START_FLAG] : []), ]); const child = spawnHidden(launch.command, launch.args, { cwd: options.cwd ?? process.cwd(), detached: true, env: coordinatorEnvironment(agentDir), stdio: "ignore", }); let launchError: Error | undefined; let exitDescription: string | undefined; child.once("error", (error) => { launchError = error; }); child.once("exit", (code, signal) => { exitDescription = signal ? `signal ${signal}` : `code ${code ?? "unknown"}`; }); child.unref();
const progressTimeoutMs = options.timeoutMs ?? DEFAULT_COORDINATOR_PROGRESS_TIMEOUT_MS; let observedUpdatedAt: string | undefined; let observedLivenessId: string | undefined; let lastProgressAt = Date.now(); let lastLivenessAt = Date.now(); while (true) { const status = readDaemonUpdateRestartStatus(statusPath); if (status && status.updatedAt !== observedUpdatedAt) { observedUpdatedAt = status.updatedAt; lastProgressAt = Date.now(); } if (status && statusLivenessId(status) !== observedLivenessId) { observedLivenessId = statusLivenessId(status); lastLivenessAt = Date.now(); } if (status && TERMINAL_PHASES.has(status.phase)) { return status; } if (launchError) { throw launchError; } if (exitDescription) { const terminalStatus = readTerminalDaemonUpdateRestartStatus(statusPath); if (terminalStatus) { return terminalStatus; } throw new Error(`Daemon update restart coordinator exited with ${exitDescription}`); } if (Date.now() - lastLivenessAt >= COORDINATOR_LIVENESS_TIMEOUT_MS) { throw new Error(`Daemon update restart coordinator stopped reporting liveness on ${socketPath}`); } if (Date.now() - lastProgressAt >= progressTimeoutMs) { throw new Error(`Timed out waiting for daemon update restart progress on ${socketPath}`); } await delay(50); }}