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.
11 kB · 359 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360import { createHash, randomUUID } from "node:crypto";import { existsSync, mkdirSync, readFileSync, realpathSync, renameSync, rmSync, writeFileSync } from "node:fs";import { basename, dirname, join, resolve } from "node:path";import { lockSync } from "proper-lockfile";import { execFileSyncHidden, isProcessAlive } from "../utils/child-process.js";
export const SESSION_LEASES_ENABLED_ENV = "PRIME_AGENT_INTERNAL_SESSION_LEASES";export const SESSION_LEASE_OWNER_ID_ENV = "PRIME_AGENT_INTERNAL_SESSION_LEASE_OWNER_ID";
interface SessionLeaseOwner { version: 1; token: string; pid: number; processStartId?: string; activeSessionId?: string; sessionPath: string; createdAt: string;}
export class SessionAlreadyActiveError extends Error { readonly code = "session_already_active" as const;
constructor( readonly sessionPath: string, readonly activeSessionId?: string, ) { super( activeSessionId ? `Session is already active in ${activeSessionId}: ${sessionPath}` : `Session is already active in another process: ${sessionPath}`, ); this.name = "SessionAlreadyActiveError"; }}
export class SessionLease { private released = false;
constructor( readonly sessionPath: string, private readonly directory: string, private readonly token: string, ) {}
release(): void { if (this.released) { return; } this.released = true; try { withLeaseGuard(this.directory, () => { const owner = readLeaseOwner(this.directory); if (typeof owner === "object" && owner.token === this.token) { reclaimStaleLease(this.directory); } }); } catch { // Lease cleanup is best-effort. A stale owner is reclaimed by the next process. } }}
function leasesEnabled(environment: NodeJS.ProcessEnv): boolean { const value = environment[SESSION_LEASES_ENABLED_ENV]?.toLowerCase(); return value === "1" || value === "true" || value === "yes";}
function leaseDirectory(agentDir: string, sessionPath: string): string { const key = createHash("sha256").update(sessionPath).digest("hex"); return join(agentDir, "session-leases", `${key}.lock`);}
export function canonicalSessionPath(sessionPath: string): string { const resolvedPath = resolve(sessionPath); try { return realpathSync(resolvedPath); } catch { try { return join(realpathSync(dirname(resolvedPath)), basename(resolvedPath)); } catch { return resolvedPath; } }}
// An unreadable owner may hold a live lease. Only a missing owner is safely absent.function readLeaseOwner(directory: string): SessionLeaseOwner | "absent" | "unreadable" { const ownerPath = join(directory, "owner.json"); let raw: string; try { raw = readFileSync(ownerPath, "utf8"); } catch (error) { return (error as NodeJS.ErrnoException).code === "ENOENT" ? "absent" : "unreadable"; } try { const parsed = JSON.parse(raw) as Partial<SessionLeaseOwner>; if ( parsed.version !== 1 || typeof parsed.token !== "string" || typeof parsed.pid !== "number" || typeof parsed.sessionPath !== "string" || typeof parsed.createdAt !== "string" ) { throw new TypeError(`Corrupt session lease owner file: ${ownerPath} (missing or invalid required fields)`); } return parsed as SessionLeaseOwner; } catch (error) { if (error instanceof SyntaxError || error instanceof TypeError) { throw new Error(`Corrupt session lease owner file: ${ownerPath} - ${error.message}`); } throw error; }}
interface ProcessQueryOptions { env?: NodeJS.ProcessEnv;}
type ProcessQuery = (command: string, args: string[], options?: ProcessQueryOptions) => string;
function runProcessQuery(command: string, args: string[], options?: ProcessQueryOptions): string { return execFileSyncHidden(command, args, { encoding: "utf8", stdio: ["ignore", "pipe", "ignore"], env: options?.env, });}
export function getWindowsProcessStartId(pid: number, query: ProcessQuery = runProcessQuery): string | undefined { if (!Number.isInteger(pid) || pid <= 0) { return undefined; } try { const startTicks = query("powershell.exe", [ "-NoLogo", "-NoProfile", "-NonInteractive", "-Command", `([System.Diagnostics.Process]::GetProcessById(${pid})).StartTime.ToUniversalTime().Ticks`, ]).trim(); return /^\d+$/.test(startTicks) ? `win:${startTicks}` : undefined; } catch { return undefined; }}
export function getPsProcessStartId(pid: number, query: ProcessQuery = runProcessQuery): string | undefined { if (!Number.isInteger(pid) || pid <= 0) { return undefined; } try { // `lstart` is rendered in the subprocess timezone and locale, so pin both for a durable identity. const startTime = query("ps", ["-p", String(pid), "-o", "lstart="], { env: { ...process.env, LC_ALL: "C", LC_TIME: "C", LANG: "C", TZ: "UTC" }, }).trim(); return startTime ? `ps:${startTime}` : undefined; } catch { return undefined; }}
export function getProcessStartId(pid: number): string | undefined { if (!Number.isInteger(pid) || pid <= 0) { return undefined; } if (process.platform === "win32") { return getWindowsProcessStartId(pid); } try { const stat = readFileSync(`/proc/${pid}/stat`, "utf8"); const commandEnd = stat.lastIndexOf(")"); const fields = stat.slice(commandEnd + 2).split(" "); const startTime = fields[19]; if (startTime) { return `proc:${startTime}`; } } catch { // Fall through to the portable process listing used on macOS and BSD. } return getPsProcessStartId(pid);}
let currentProcessStartId: string | undefined;let currentProcessStartIdRead = false;
function getCurrentProcessStartId(): string | undefined { if (!currentProcessStartIdRead) { currentProcessStartId = getProcessStartId(process.pid); currentProcessStartIdRead = true; } return currentProcessStartId;}
function isLeaseOwnerAlive(owner: SessionLeaseOwner): boolean { if (!isProcessAlive(owner.pid)) { return false; } if (!owner.processStartId) { return true; } const currentStartId = getProcessStartId(owner.pid); return currentStartId === undefined || currentStartId === owner.processStartId;}
function withLeaseGuard<T>(directory: string, action: () => T): T { let release: (() => void) | undefined; let guardCompromised = false; for (let attempt = 0; attempt < 100; attempt++) { try { release = lockSync(directory, { realpath: false, lockfilePath: `${directory}.guard`, stale: 5000, onCompromised: () => { guardCompromised = true; }, }); break; } catch (error) { if ((error as NodeJS.ErrnoException).code !== "ELOCKED") { throw error; } if (attempt === 99) { throw new Error(`Could not coordinate session lease: ${directory}`); } Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 10); } } if (!release) { throw new Error(`Could not coordinate session lease: ${directory}`); } const assertGuardHeld = () => { if (guardCompromised) throw new Error(`Session lease guard was compromised: ${directory}`); }; try { assertGuardHeld(); const result = action(); assertGuardHeld(); return result; } finally { if (guardCompromised) { try { release(); } catch { // The compromised guard no longer owns a lock that can be safely released. } } else { release(); } }}
export function isRenameTargetContention( directory: string, code: string | undefined, platform: string = process.platform,): boolean { // POSIX: renameSync into an existing directory raises EEXIST or ENOTEMPTY. if (code === "EEXIST" || code === "ENOTEMPTY") { return true; } // Windows: renameSync into an existing directory raises EPERM or EACCES // instead of EEXIST. Only treat them as contention when the target // actually exists so real permission errors still propagate. if ((code === "EPERM" || code === "EACCES") && platform === "win32") { try { return existsSync(directory); } catch { return false; } } return false;}
function reclaimStaleLease(directory: string): boolean { const stalePath = `${directory}.stale-${process.pid}-${randomUUID()}`; const sleepBuffer = new Int32Array(new SharedArrayBuffer(4)); for (let attempt = 1; ; attempt++) { try { renameSync(directory, stalePath); break; } catch (error) { const code = (error as NodeJS.ErrnoException).code; if (code === "ENOENT") return true; const transient = process.platform === "win32" && (code === "EBUSY" || code === "EPERM" || code === "EACCES"); if (!transient || attempt >= 8) return false; Atomics.wait(sleepBuffer, 0, 0, 10 * attempt); } } try { rmSync(stalePath, { recursive: true, force: true, maxRetries: 8, retryDelay: 10 }); } catch { // The quarantined directory no longer owns the lease path. } return true;}
export function acquireSessionLease( sessionPath: string | undefined, agentDir: string, environment: NodeJS.ProcessEnv = process.env,): SessionLease | undefined { if (!sessionPath || !leasesEnabled(environment)) { return undefined; } const canonicalPath = canonicalSessionPath(sessionPath); const root = join(agentDir, "session-leases"); mkdirSync(root, { recursive: true, mode: 0o700 }); const directory = leaseDirectory(agentDir, canonicalPath);
return withLeaseGuard(directory, () => { for (let attempt = 0; attempt < 3; attempt++) { const token = randomUUID(); const candidateDirectory = `${directory}.candidate-${process.pid}-${token}`; const owner: SessionLeaseOwner = { version: 1, token, pid: process.pid, processStartId: getCurrentProcessStartId(), activeSessionId: environment[SESSION_LEASE_OWNER_ID_ENV], sessionPath: canonicalPath, createdAt: new Date().toISOString(), }; mkdirSync(candidateDirectory, { mode: 0o700 }); writeFileSync(join(candidateDirectory, "owner.json"), `${JSON.stringify(owner, null, 2)}\n`, { mode: 0o600, }); try { renameSync(candidateDirectory, directory); return new SessionLease(canonicalPath, directory, token); } catch (error) { const err = error as NodeJS.ErrnoException; rmSync(candidateDirectory, { recursive: true, force: true }); if (err.code === "ENOENT") { // Candidate vanished - treat as retryable race. continue; } if (isRenameTargetContention(directory, err.code)) { const existingOwner = readLeaseOwner(directory); if (existingOwner === "unreadable") { continue; } if (existingOwner !== "absent" && isLeaseOwnerAlive(existingOwner)) { throw new SessionAlreadyActiveError(canonicalPath, existingOwner.activeSessionId); } reclaimStaleLease(directory); continue; } throw error; } }
const owner = existsSync(directory) ? readLeaseOwner(directory) : undefined; if (typeof owner === "object" && isLeaseOwnerAlive(owner)) { throw new SessionAlreadyActiveError(canonicalPath, owner.activeSessionId); } throw new Error(`Could not acquire session lease: ${canonicalPath}`); });}