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.
12 kB · 328 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329import { EventEmitter } from "node:events";import { readFileSync } from "node:fs";import { resolve } from "node:path";import { describe, expect, it, vi } from "vitest";import { type HostRequestHandlers, ReplKernelManager } from "../src/core/kernel/index.js";
type ShutdownInternals = { state: "running"; writeLine: (request: Record<string, unknown>) => Promise<void>; handleEvent: (event: Record<string, unknown>) => void; wireChild: (child: ShutdownInternals["child"]) => void; pendingDoneWaiters: Map<string, () => void>; inFlightHostRequests: Set<Promise<void>>; kernelStderr: string; child: EventEmitter & { exitCode: number | null; signalCode: NodeJS.Signals | null; kill: (signal?: NodeJS.Signals | number) => boolean; pid?: number; stdin: { destroyed: boolean; destroy: () => void; on: (event: string, listener: (...args: unknown[]) => void) => void; }; stdout?: { destroy: () => void; on: (event: string, listener: (...args: unknown[]) => void) => void }; stderr?: { destroy: () => void; on: (event: string, listener: (...args: unknown[]) => void) => void; once: (event: string, listener: (...args: unknown[]) => void) => void; }; };};
function configuredManager( onSend: (request: Record<string, unknown>, internals: ShutdownInternals) => void | Promise<void>, hostHandlers?: HostRequestHandlers, withSnapshot = false,): { manager: ReplKernelManager; internals: ShutdownInternals;} { const manager = new ReplKernelManager({ cwd: process.cwd(), hostHandlers, snapshot: withSnapshot ? { path: "/tmp/shutdown-test.dill", manifestPath: "/tmp/shutdown-test.json" } : undefined, }); const internals = manager as unknown as ShutdownInternals; const child = Object.assign(new EventEmitter(), { exitCode: null, signalCode: null, kill: vi.fn(() => true), pid: undefined, stdin: { destroyed: false, destroy: vi.fn(), on: vi.fn() }, stdout: { destroy: vi.fn(), on: vi.fn() }, stderr: { destroy: vi.fn(), on: vi.fn(), once: vi.fn() }, }); Object.assign(internals, { state: "running", writeLine: vi.fn(async (request: Record<string, unknown>) => onSend(request, internals)), child, }); // Attach the real exit/error handlers so teardown-ownership races are exercised. internals.wireChild(child); return { manager, internals };}
describe("ReplKernelManager graceful shutdown", () => { it.each([ ["optional policy", {}], ["snapshot-and-drain policy", { snapshot: true, drainHostRequests: true }], ] as const)("uses the same diagnostic deadline for the %s", async (_label, options) => { vi.useFakeTimers(); try { const { manager, internals } = configuredManager(() => new Promise<void>(() => {}), undefined, true); // The final flush snapshots through the private captureSnapshot (bounded, prune-free). vi.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot").mockResolvedValue( null, ); const shutdown = manager.shutdown(options); await vi.advanceTimersByTimeAsync(5_000); await expect(shutdown).resolves.toBe(true); expect(internals.kernelStderr).toContain("Kernel did not shut down within 5000ms"); expect(internals.child).toBeUndefined(); } finally { vi.useRealTimers(); } });
it("does not finish shutdown before the stdin write settles", async () => { let finishSend: (() => void) | undefined; const sendBlocked = new Promise<void>((resolve) => { finishSend = resolve; }); const { manager, internals } = configuredManager(async (request, state) => { state.handleEvent({ event: "done", id: request.id, status: "ok" }); await sendBlocked; state.child.exitCode = 0; state.child.emit("exit", 0, null); });
let finished = false; const shutdown = manager.shutdown().then(() => { finished = true; }); await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); expect(finished).toBe(false); finishSend?.(); await shutdown; expect(internals.pendingDoneWaiters.size).toBe(0); });
it("finishes promptly when the runtime exits without a shutdown done", async () => { const { manager, internals } = configuredManager((_request, state) => { state.child.exitCode = 0; state.child.emit("exit", 0, null); }); vi.useFakeTimers(); try { const shutdown = manager.shutdown(); await vi.advanceTimersByTimeAsync(100); // True = this call performed the cleanup: startup-failure recovery relies on it to resurrect to idle. await expect(shutdown).resolves.toBe(true); expect(internals.child).toBeUndefined(); } finally { vi.useRealTimers(); } });
it("flushes a snapshot and drains host requests when session teardown selects both policies", async () => { let releaseHandler: (() => void) | undefined; const handlerBlocked = new Promise<void>((resolve) => { releaseHandler = resolve; }); const sentReplies: Record<string, unknown>[] = []; const { manager, internals } = configuredManager( async (request, state) => { if (request.type === "host_reply") sentReplies.push(request); if (request.type === "shutdown") { state.handleEvent({ event: "done", id: request.id, status: "ok" }); state.child.exitCode = 0; state.child.emit("exit", 0, null); } }, { "test.slow": async () => { await handlerBlocked; return { answer: 42, status: "weird" }; }, }, true, ); const snapshot = vi .spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot") .mockResolvedValue(null);
internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } }); const shutdown = manager.shutdown({ snapshot: true, drainHostRequests: true }); await new Promise((resolve) => globalThis.setTimeout(resolve, 0));
expect(snapshot).toHaveBeenCalledOnce(); expect(internals.writeLine).not.toHaveBeenCalled(); releaseHandler?.(); await shutdown;
expect(internals.inFlightHostRequests.size).toBe(0); expect(sentReplies).toEqual([ { type: "host_reply", id: "hr-1", data: { status: "ok", result: { answer: 42, status: "weird" } }, }, ]); expect(internals.child).toBeUndefined(); });
it("shutdown during an in-flight host request leaves no dangling task once the handler settles", async () => { let releaseHandler: (() => void) | undefined; const handlerBlocked = new Promise<void>((resolve) => { releaseHandler = resolve; }); const { manager, internals } = configuredManager( (request, state) => { if (request.type !== "shutdown") return; state.handleEvent({ event: "done", id: request.id, status: "ok" }); state.child.exitCode = 0; state.child.emit("exit", 0, null); }, { "test.slow": async () => { await handlerBlocked; return { answer: 42 }; }, }, true, ); const snapshot = vi .spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot") .mockResolvedValue(null);
internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } }); const tracked = [...internals.inFlightHostRequests]; expect(tracked).toHaveLength(1);
await manager.shutdown(); expect(snapshot).not.toHaveBeenCalled(); expect(internals.child).toBeUndefined();
// The reply write fails against the torn-down child; the task must still settle. releaseHandler?.(); await Promise.all(tracked); expect(internals.inFlightHostRequests.size).toBe(0); });
it("shares one wire exchange across concurrent graceful shutdown calls", async () => { let exitKernel: (() => void) | undefined; const { manager, internals } = configuredManager((request, state) => { state.handleEvent({ event: "done", id: request.id, status: "ok" }); exitKernel = () => { state.child.exitCode = 0; state.child.emit("exit", 0, null); }; });
const owner = manager.shutdown(); await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); const concurrent = manager.shutdown(); exitKernel?.();
await expect(owner).resolves.toBe(true); await expect(concurrent).resolves.toBe(false); expect(internals.writeLine).toHaveBeenCalledOnce(); });
it("stands down when a hard teardown supersedes the snapshot flush", async () => { let releaseSnapshot: (() => void) | undefined; const snapshotBlocked = new Promise<null>((resolve) => { releaseSnapshot = () => resolve(null); }); const { manager, internals } = configuredManager(() => undefined, undefined, true); vi.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot").mockReturnValue( snapshotBlocked, ); const child = internals.child;
const shutdown = manager.shutdown({ snapshot: true }); await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); await manager.kill(); releaseSnapshot?.();
await expect(shutdown).resolves.toBe(false); expect(child.kill).toHaveBeenCalledOnce(); });
it("restart after a graceful shutdown starts the kernel again", async () => { const { manager, internals } = configuredManager((request, state) => { if (request.type !== "shutdown") return; state.handleEvent({ event: "done", id: request.id, status: "ok" }); // The runtime exits after acking shutdown; the exit handler must not // steal teardown ownership from the in-flight graceful shutdown(). state.child.exitCode = 0; state.child.emit("exit", 0, null); }); const start = vi.spyOn(manager, "start").mockResolvedValue(undefined);
await manager.restart();
expect(internals.child).toBeUndefined(); expect(start).toHaveBeenCalledTimes(1); });
it("restart does not resurrect a kernel a concurrent kill superseded", async () => { const { manager } = configuredManager((_request, state) => { // A concurrent kill supersedes the in-flight graceful shutdown. void manager.kill(); state.child.emit("exit", null, "SIGKILL"); }); const start = vi.spyOn(manager, "start").mockResolvedValue(undefined);
await manager.restart();
expect(start).not.toHaveBeenCalled(); });
it("kill() destroys stderr only while the child has not exited", async () => { // Alive child: it may ignore the signal and never emit 'exit', so the // post-exit drain would never bound the pipe — teardown must. const alive = configuredManager(() => {}); const aliveChild = alive.internals.child; await alive.manager.kill(); expect(aliveChild.stderr?.destroy).toHaveBeenCalledTimes(1);
// Exited child: the post-exit drain owns the stream; destroying it in // teardown would drop the kernel's buffered last words. const exited = configuredManager(() => {}); const exitedChild = exited.internals.child; exitedChild.exitCode = 42; await exited.manager.kill(); expect(exitedChild.stderr?.destroy).not.toHaveBeenCalled(); });
it("waits for the matching shutdown done and removes its waiter", async () => { const { manager, internals } = configuredManager(async (request, state) => { expect(request.type).toBe("shutdown"); expect(request.id).toBeTypeOf("string"); // An unrelated done must not release the shutdown waiter. state.handleEvent({ event: "done", id: "unrelated", status: "ok" }); expect(state.pendingDoneWaiters.size).toBe(1); state.handleEvent({ event: "done", id: request.id, status: "ok" }); queueMicrotask(() => { state.child.exitCode = 0; state.child.emit("exit", 0, null); }); });
await manager.shutdown();
expect(internals.pendingDoneWaiters.size).toBe(0); expect(internals.child).toBeUndefined(); });
it("keeps the kernel MCP close budget strictly inside the host shutdown deadline", () => { const source = readFileSync(resolve(__dirname, "../../../prime-agent-runtime/src/rlm/mcp.py"), "utf8"); const match = source.match(/^_SHUTDOWN_TIMEOUT = ([\d.]+)$/m); expect(match).not.toBeNull(); // +1s dispatch slack in mcp.py close(); the sum must undercut the host's 5s kill deadline. expect((Number(match![1]) + 1) * 1000).toBeLessThan(5000); });});