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.
20 kB · 543 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544import { spawn } from "node:child_process";import { existsSync, mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";import { createServer, type Server, type Socket } from "node:net";import { tmpdir } from "node:os";import { dirname, join } from "node:path";import { afterEach, describe, expect, it, vi } from "vitest";import { ensureInteractiveDaemonRunning, probeDaemonVersion, probeRunningDaemonSessions, shouldStartDaemonEarly, shutdownDaemonAndWait,} from "../src/cli/daemon-launch.js";import { ENV_AGENT_DIR, getDaemonLogPath, VERSION } from "../src/config.js";import { DAEMON_PROTOCOL_VERSION, DAEMON_SCHEMA_ID } from "../src/modes/daemon/daemon-protocol.js";
interface FakeDaemonOptions { /** Sessions returned for a `list` command. */ sessions?: Array<Record<string, unknown>>; busyClientOwnedSessionCount?: number; /** When true, the `list` command responds with a failure. */ failList?: boolean; /** When false, the server ignores `shutdown` and stays up. */ respondToShutdown?: boolean; protocolVersion?: number; appVersion?: string; schemaId?: string; firstSchemaId?: string; serverCapabilities?: string[]; shouldSendHello?: (connectionIndex: number) => boolean; onConnection?: (connectionIndex: number) => void; onCommand?: (command: { type: string }) => void;}
interface FakeDaemon { socketPath: string; close: () => Promise<void>;}
function send(socket: Socket, message: unknown): void { socket.write(`${JSON.stringify(message)}\n`);}
async function startFakeDaemon(options: FakeDaemonOptions = {}): Promise<FakeDaemon> { const dir = mkdtempSync(join(tmpdir(), "pa-launch-")); const socketPath = join(dir, "d.sock"); let connectionIndex = 0; const server: Server = createServer((socket) => { const currentConnectionIndex = connectionIndex++; options.onConnection?.(currentConnectionIndex); socket.on("error", () => undefined); if (options.shouldSendHello?.(currentConnectionIndex) ?? true) { send(socket, { type: "daemon_hello", socketPath, protocol: { name: "prime-agent.daemon", version: options.protocolVersion ?? DAEMON_PROTOCOL_VERSION }, appVersion: options.appVersion, schemaId: currentConnectionIndex === 0 && options.firstSchemaId ? options.firstSchemaId : (options.schemaId ?? DAEMON_SCHEMA_ID), clientId: "fake-client", serverCapabilities: options.serverCapabilities ?? [], }); } let buffer = ""; socket.on("data", (chunk) => { buffer += chunk.toString(); let newline = buffer.indexOf("\n"); while (newline !== -1) { const line = buffer.slice(0, newline); buffer = buffer.slice(newline + 1); newline = buffer.indexOf("\n"); if (!line.trim()) { continue; } const wire = JSON.parse(line) as { type: string; id: string; command?: { type: string; id: string }; }; const command = wire.type === "command" && wire.command ? wire.command : wire; options.onCommand?.(command); if (command.type === "list") { send(socket, { type: "response", command: "list", id: wire.id, ...(options.failList ? { success: false, error: "list failed" } : { success: true, data: { sessions: options.sessions ?? [], busyClientOwnedSessionCount: options.busyClientOwnedSessionCount ?? 0, }, }), }); } else if (command.type === "shutdown") { if (options.respondToShutdown === false) { continue; } send(socket, { type: "response", command: "shutdown", id: wire.id, success: true }); server.close(); socket.end(); } } }); }); await new Promise<void>((resolve) => server.listen(socketPath, resolve)); return { socketPath, close: () => new Promise<void>((resolve) => { server.close(() => resolve()); rmSync(dir, { recursive: true, force: true }); }), };}
async function startCrashingDaemon(): Promise<FakeDaemon> { const dir = mkdtempSync(join(tmpdir(), "pa-launch-crash-")); const socketPath = join(dir, "d.sock"); const child = spawn( process.execPath, [ "-e", `const { createServer } = require("node:net");const socketPath = process.argv[1];const send = (socket, message) => socket.write(JSON.stringify(message) + "\\n");const server = createServer((socket) => { send(socket, { type: "daemon_hello", socketPath, protocol: { name: "prime-agent.daemon", version: ${DAEMON_PROTOCOL_VERSION} }, schemaId: ${JSON.stringify(DAEMON_SCHEMA_ID)}, }); let buffer = ""; socket.on("data", (chunk) => { buffer += chunk.toString(); const newline = buffer.indexOf("\\n"); if (newline === -1) return; const wire = JSON.parse(buffer.slice(0, newline)); const command = wire.command ?? wire; if (command.type !== "shutdown") return; send(socket, { type: "response", command: "shutdown", id: wire.id ?? command.id, success: true }); socket.end(() => process.kill(process.pid, "SIGKILL")); });});server.listen(socketPath, () => process.stdout.write("ready\\n"));`, socketPath, ], { stdio: ["ignore", "pipe", "ignore"] }, ); await new Promise<void>((resolve, reject) => { child.stdout?.once("data", () => resolve()); child.once("error", reject); child.once("exit", (code, signal) => reject(new Error(`Daemon exited before listening: ${code ?? signal}`))); }); return { socketPath, close: async () => { if (child.exitCode === null && child.signalCode === null) { child.kill("SIGKILL"); await new Promise<void>((resolve) => child.once("close", () => resolve())); } rmSync(dir, { recursive: true, force: true }); }, };}
describe("probeRunningDaemonSessions", () => { const cleanups: Array<() => Promise<void>> = []; afterEach(async () => { await Promise.all(cleanups.splice(0).map((fn) => fn())); });
it("reports unreachable when no daemon is running", async () => { const result = await probeRunningDaemonSessions(join(tmpdir(), "pa-launch-missing.sock")); expect(result).toEqual({ reachable: false }); });
it("returns only sessions that have an activeSessionId", async () => { const daemon = await startFakeDaemon({ sessions: [ { id: "a", activeSessionId: "a", isStreaming: true }, { id: "b" }, // saved-only, no activeSessionId { id: "c", activeSessionId: "c", isStreaming: false }, ], }); cleanups.push(daemon.close); const result = await probeRunningDaemonSessions(daemon.socketPath); expect(result.reachable).toBe(true); expect(result.reachable ? result.activeSessions?.map((session) => session.activeSessionId) : undefined).toEqual([ "a", "c", ]); });
it("returns an empty list when the daemon is reachable but idle", async () => { const daemon = await startFakeDaemon({ sessions: [] }); cleanups.push(daemon.close); expect(await probeRunningDaemonSessions(daemon.socketPath)).toEqual({ reachable: true, activeSessions: [] }); });
it("leaves activeSessions undefined when reachable but the list fails", async () => { const daemon = await startFakeDaemon({ failList: true }); cleanups.push(daemon.close); expect(await probeRunningDaemonSessions(daemon.socketPath)).toEqual({ reachable: true }); });});
describe("shouldStartDaemonEarly", () => { it.each([ ["interactive", []], ["print", ["--print", "hello"]], ["json", ["--mode", "json", "hello"]], ["rpc", ["--mode", "rpc"]], ["no-session", ["--no-session"]], ["escaped command word", ["--", "status"]], ["print with a command word", ["--print", "status"]], ["prompt that only starts like a command", ["--offline", "statuses", "of", "my", "agents"]], ["extension flag with a command-like value", ["--extension-option", "status"]], ["value flag before a -- separator", ["--cwd", "--", "status"]], ["resume selector that is an @file reference", ["--resume", "@prompt.md", "status"]], ["prompt-value flag before a command word", ["--system-prompt", "be terse", "status"]], ["appended prompt-value flag before a command word", ["--append-system-prompt", "be terse", "status"]], ])("starts early for the %s client", (_label, args) => { expect(shouldStartDaemonEarly(args, false)).toBe(true); });
it.each([ ["daemon process", ["--mode", "daemon"]], ["help", ["--help"]], ["version", ["--version"]], ["model listing", ["--list-models"]], ["management command after global flags", ["--daemon-socket", "/tmp/prime.sock", "status"]], ["nested command after global flags", ["--offline", "model", "list"]], ["help after global flags", ["--offline", "help"]], ["nested help after global flags", ["--offline", "help", "status"]], ["help before global flags", ["help", "--verbose"]], ["boolean run flag before the command", ["--verbose", "model", "list"]], ["unknown short option before the command", ["-x", "model", "list"]], ["inline extension flag before the command", ["--extension-option=status", "model", "list"]], ["startup benchmark", []], ])("does not start early for %s", (label, args) => { expect(shouldStartDaemonEarly(args, label === "startup benchmark")).toBe(false); });});
describe("ensureInteractiveDaemonRunning", () => { const cleanups: Array<() => Promise<void>> = []; afterEach(async () => { await Promise.all(cleanups.splice(0).map((fn) => fn())); });
it("rejects a busy pre-session-action daemon before attach", async () => { const commands: string[] = []; const daemon = await startFakeDaemon({ protocolVersion: 3, appVersion: VERSION, sessions: [{ id: "active-1", activeSessionId: "active-1", isStreaming: true }], onCommand: (command) => commands.push(command.type), }); cleanups.push(daemon.close);
await expect(ensureInteractiveDaemonRunning(daemon.socketPath)).rejects.toThrow("incompatible"); expect(commands).not.toContain("list"); expect(commands).not.toContain("shutdown"); });
it("does not replace a stale daemon with busy private client sessions", async () => { const commands: string[] = []; const daemon = await startFakeDaemon({ protocolVersion: DAEMON_PROTOCOL_VERSION, appVersion: VERSION, schemaId: "stale-schema", busyClientOwnedSessionCount: 1, onCommand: (command) => commands.push(command.type), }); cleanups.push(daemon.close);
await expect(ensureInteractiveDaemonRunning(daemon.socketPath)).rejects.toThrow("stale"); expect(commands).toContain("list"); expect(commands).not.toContain("shutdown"); });
it("does not treat a live daemon as absent when cold startup delays the first connection", async () => { const daemon = await startFakeDaemon({ protocolVersion: DAEMON_PROTOCOL_VERSION, appVersion: VERSION, schemaId: DAEMON_SCHEMA_ID, }); cleanups.push(daemon.close);
const probe = probeDaemonVersion(daemon.socketPath); Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 300);
await expect(probe).resolves.toMatchObject({ status: "current" }); });
it("classifies a connected daemon without hello as unresponsive", async () => { const daemon = await startFakeDaemon({ shouldSendHello: () => false }); cleanups.push(daemon.close);
await expect(probeDaemonVersion(daemon.socketPath, 10)).resolves.toEqual({ status: "unresponsive" }); });
it("reuses a daemon that becomes current inside the startup window", async () => { vi.useFakeTimers(); let resolveFirstConnection = () => {}; let resolveSecondConnection = () => {}; const firstConnection = new Promise<void>((resolve) => { resolveFirstConnection = resolve; }); const secondConnection = new Promise<void>((resolve) => { resolveSecondConnection = resolve; }); const daemon = await startFakeDaemon({ appVersion: VERSION, shouldSendHello: (connectionIndex) => connectionIndex > 0, onConnection: (connectionIndex) => { if (connectionIndex === 0) resolveFirstConnection(); if (connectionIndex === 1) resolveSecondConnection(); }, }); cleanups.push(daemon.close);
try { const ensuring = ensureInteractiveDaemonRunning(daemon.socketPath); await firstConnection; await vi.advanceTimersByTimeAsync(2000); await secondConnection; await expect(ensuring).resolves.toBeUndefined(); } finally { vi.useRealTimers(); } });
it("leaves a connected unresponsive daemon running after the startup window", async () => { vi.useFakeTimers(); const connections: Array<() => void> = []; const connected = [0, 1].map( (index) => new Promise<void>((resolve) => { connections[index] = resolve; }), ); const commands: string[] = []; const daemon = await startFakeDaemon({ shouldSendHello: () => false, onConnection: (connectionIndex) => connections[connectionIndex]?.(), onCommand: (command) => commands.push(command.type), }); cleanups.push(daemon.close);
try { const ensuring = ensureInteractiveDaemonRunning(daemon.socketPath); const rejected = expect(ensuring).rejects.toThrow(/accepted connections but did not finish startup/); await connected[0]; await vi.advanceTimersByTimeAsync(2000); await connected[1]; await vi.advanceTimersByTimeAsync(30_000); await rejected; expect(commands).not.toContain("list"); expect(commands).not.toContain("shutdown"); } finally { vi.useRealTimers(); } });
it("cancels replacement when a stale-looking daemon becomes current", async () => { const commands: string[] = []; const daemon = await startFakeDaemon({ appVersion: VERSION, firstSchemaId: "stale-schema", schemaId: DAEMON_SCHEMA_ID, onCommand: (command) => commands.push(command.type), }); cleanups.push(daemon.close);
await expect(ensureInteractiveDaemonRunning(daemon.socketPath)).resolves.toBeUndefined(); expect(commands).toContain("list"); expect(commands).not.toContain("shutdown"); });
it("fails fast with the daemon log tail when the spawned daemon exits during startup", async () => { const dir = mkdtempSync(join(tmpdir(), "pa-launch-startup-crash-")); const entrypoint = join(dir, "crash.mjs"); const socketPath = join(dir, "d.sock"); const originalAgentDir = process.env[ENV_AGENT_DIR]; process.env[ENV_AGENT_DIR] = join(dir, "agent"); const logPath = getDaemonLogPath(socketPath); mkdirSync(dirname(logPath), { recursive: true }); const script = `import { appendFileSync } from "node:fs"; appendFileSync(${JSON.stringify(logPath)}, "supervisor: fatal startup failure\\n"); process.exit(23);`; writeFileSync(entrypoint, script); const originalEntrypoint = process.argv[1]!; process.argv[1] = entrypoint; const startedAt = Date.now();
try { await expect(ensureInteractiveDaemonRunning(socketPath)).rejects.toThrow( /Prime Agent daemon exited during startup \(code 23\)\.[\s\S]*fatal startup failure/, ); expect(Date.now() - startedAt).toBeLessThan(10_000); } finally { process.argv[1] = originalEntrypoint; if (originalAgentDir === undefined) delete process.env[ENV_AGENT_DIR]; else process.env[ENV_AGENT_DIR] = originalAgentDir; rmSync(dir, { recursive: true, force: true }); } });
it("names the missing daemon log when the daemon crashes before logging", async () => { const dir = mkdtempSync(join(tmpdir(), "pa-launch-startup-silent-")); const entrypoint = join(dir, "crash.mjs"); const socketPath = join(dir, "d.sock"); const originalAgentDir = process.env[ENV_AGENT_DIR]; process.env[ENV_AGENT_DIR] = join(dir, "agent"); writeFileSync(entrypoint, "process.exit(7);"); const originalEntrypoint = process.argv[1]!; process.argv[1] = entrypoint;
try { await expect(ensureInteractiveDaemonRunning(socketPath)).rejects.toThrow( /exited during startup \(code 7\)\. The daemon wrote nothing to its log/, ); } finally { process.argv[1] = originalEntrypoint; if (originalAgentDir === undefined) delete process.env[ENV_AGENT_DIR]; else process.env[ENV_AGENT_DIR] = originalAgentDir; rmSync(dir, { recursive: true, force: true }); } });
it("surfaces a spawn error when the child never emits exit", async () => { const dir = mkdtempSync(join(tmpdir(), "pa-launch-spawn-error-")); const socketPath = join(dir, "d.sock"); const originalAgentDir = process.env[ENV_AGENT_DIR]; process.env[ENV_AGENT_DIR] = join(dir, "agent");
try { await expect(ensureInteractiveDaemonRunning(socketPath, join(dir, "missing-cwd"))).rejects.toThrow( /Failed to spawn Prime Agent daemon: .*ENOENT/, ); } finally { if (originalAgentDir === undefined) delete process.env[ENV_AGENT_DIR]; else process.env[ENV_AGENT_DIR] = originalAgentDir; rmSync(dir, { recursive: true, force: true }); } });
it("succeeds when the spawned child loses the race to an already-serving daemon", async () => { const dir = mkdtempSync(join(tmpdir(), "pa-launch-startup-race-")); const entrypoint = join(dir, "loser.mjs"); const socketPath = join(dir, "d.sock"); const originalAgentDir = process.env[ENV_AGENT_DIR]; process.env[ENV_AGENT_DIR] = join(dir, "agent"); writeFileSync(entrypoint, "process.exit(7);"); const originalEntrypoint = process.argv[1]!; process.argv[1] = entrypoint; const server: Server = createServer((socket) => { socket.on("error", () => undefined); send(socket, { type: "daemon_hello", socketPath, protocol: { name: "prime-agent.daemon", version: DAEMON_PROTOCOL_VERSION }, appVersion: VERSION, schemaId: DAEMON_SCHEMA_ID, clientId: "fake-client", serverCapabilities: [], }); });
try { const ensurePromise = ensureInteractiveDaemonRunning(socketPath); // Let the child exit first, then bring up the winning daemon inside // the exit grace window. await new Promise((resolve) => setTimeout(resolve, 300)); await new Promise<void>((resolve) => server.listen(socketPath, resolve)); await expect(ensurePromise).resolves.toBeUndefined(); } finally { process.argv[1] = originalEntrypoint; if (originalAgentDir === undefined) delete process.env[ENV_AGENT_DIR]; else process.env[ENV_AGENT_DIR] = originalAgentDir; await new Promise<void>((resolve) => server.close(() => resolve())); rmSync(dir, { recursive: true, force: true }); } });
it("does not attribute a previous run's log to a daemon that crashed before logging", async () => { const dir = mkdtempSync(join(tmpdir(), "pa-launch-startup-stale-")); const entrypoint = join(dir, "crash.mjs"); const socketPath = join(dir, "d.sock"); const originalAgentDir = process.env[ENV_AGENT_DIR]; process.env[ENV_AGENT_DIR] = join(dir, "agent"); const logPath = getDaemonLogPath(socketPath); mkdirSync(dirname(logPath), { recursive: true }); writeFileSync(logPath, "supervisor: error from an earlier run in /tmp/\u00fcml\u00e4ut-p\u00e4th\n"); writeFileSync(entrypoint, "process.exit(7);"); const originalEntrypoint = process.argv[1]!; process.argv[1] = entrypoint;
try { await expect(ensureInteractiveDaemonRunning(socketPath)).rejects.toThrow( /exited during startup \(code 7\)\. The daemon wrote nothing to its log/, ); } finally { process.argv[1] = originalEntrypoint; if (originalAgentDir === undefined) delete process.env[ENV_AGENT_DIR]; else process.env[ENV_AGENT_DIR] = originalAgentDir; rmSync(dir, { recursive: true, force: true }); } });});
describe("shutdownDaemonAndWait", () => { const cleanups: Array<() => Promise<void>> = []; afterEach(async () => { await Promise.all(cleanups.splice(0).map((fn) => fn())); });
it("returns true immediately when no daemon is running", async () => { expect(await shutdownDaemonAndWait(join(tmpdir(), "pa-launch-missing2.sock"))).toBe(true); });
it("stops a running daemon and returns true", async () => { const daemon = await startFakeDaemon({ sessions: [{ id: "a", activeSessionId: "a", isStreaming: false }] }); cleanups.push(daemon.close); expect(await shutdownDaemonAndWait(daemon.socketPath)).toBe(true); });
it("accepts an exited daemon that leaves a stale socket behind", async () => { if (process.platform === "win32") { return; }
const daemon = await startCrashingDaemon(); cleanups.push(daemon.close); expect(await shutdownDaemonAndWait(daemon.socketPath, 100)).toBe(true); expect(existsSync(daemon.socketPath)).toBe(true); });});