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.
5.8 kB · 167 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { afterEach, describe, expect, it, vi } from "vitest";import { ReplKernelManager } from "../src/core/kernel/index.js";import { kernelCwdPathIn, manifestPathIn, readKernelCwd, snapshotPathIn, writeKernelCwd,} from "../src/core/kernel/state-snapshot.js";
describe("kernel state snapshot paths", () => { it("places snapshot + manifest inside the session artifact directory", () => { const artifactDir = "/home/u/.prime/agent/session-artifacts/abc-123"; expect(snapshotPathIn(artifactDir)).toBe(join(artifactDir, "kernel-state.dill")); expect(manifestPathIn(artifactDir)).toBe(join(artifactDir, "kernel-state.json")); expect(kernelCwdPathIn(artifactDir)).toBe(join(artifactDir, "kernel-cwd")); });});
describe("recorded kernel working directory", () => { const dirs: string[] = [];
function artifactDir(): string { const dir = mkdtempSync(join(tmpdir(), `pi-kernel-cwd-${process.pid}-`)); dirs.push(dir); return dir; }
afterEach(() => { while (dirs.length > 0) rmSync(dirs.pop() as string, { recursive: true, force: true }); });
it("round-trips a directory that exists", () => { const dir = artifactDir(); const work = join(dir, "work"); mkdirSync(work);
writeKernelCwd(dir, work);
expect(readKernelCwd(dir)).toEqual({ cwd: work }); });
it("reports a recorded directory that is gone instead of using it", () => { const dir = artifactDir(); const gone = join(dir, "gone"); mkdirSync(gone); writeKernelCwd(dir, gone); rmSync(gone, { recursive: true, force: true });
expect(readKernelCwd(dir)).toEqual({ missing: gone }); });
it("ignores a relative path, an empty record, and a missing record", () => { const dir = artifactDir(); writeKernelCwd(dir, "relative/path"); expect(readKernelCwd(dir)).toEqual({});
writeFileSync(kernelCwdPathIn(dir), "\n"); expect(readKernelCwd(dir)).toEqual({});
rmSync(kernelCwdPathIn(dir)); expect(readKernelCwd(dir)).toEqual({}); });});
type ReplInternals = { state: string; writeLine: (request: Record<string, unknown>) => Promise<void>; start: () => Promise<void>; handleEvent: (event: Record<string, unknown>) => void; activeExecution?: { requestId: string }; scheduleSnapshot: () => void;};
/** A manager whose kernel is a stub that answers state requests by hand. */function stubManager(): { manager: ReplKernelManager; internals: ReplInternals; requests: Record<string, unknown>[]; requestTypes: () => (string | undefined)[];} { const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); const requests: Record<string, unknown>[] = []; const internals = manager as unknown as ReplInternals; Object.assign(internals, { state: "running", start: async () => {}, writeLine: async (request: Record<string, unknown>) => { requests.push(request); if (request.type === "restore" && typeof request.id === "string") { internals.handleEvent({ event: "done", id: request.id, status: "ok", restored: ["alpha"], failed: [] }); } }, }); return { manager, internals, requests, requestTypes: () => requests.map((request) => String(request.type)) };}
async function settleExecution(internals: ReplInternals): Promise<void> { for (let i = 0; i < 20 && !internals.activeExecution; i += 1) { await Promise.resolve(); }}
describe("kernel state snapshots across a cell deadline", () => { afterEach(() => { vi.useRealTimers(); });
it("never snapshots the timed-out cell's namespace", async () => { vi.useFakeTimers(); const { manager, internals, requestTypes } = stubManager(); const scheduleSnapshot = vi.fn(); Object.assign(internals, { scheduleSnapshot });
const executePromise = manager.execute("while True: pass", { timeoutMs: 1000 }); await settleExecution(internals); const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "stdout", id: execution.requestId, text: "before\n" });
// Deadline at 1000ms, grace at 2000ms: the cell is reported as aborted. await vi.advanceTimersByTimeAsync(2000); const result = await executePromise; expect(result.status).toBe("aborted"); expect(result.timedOut).toBe(true); expect(result.stdout).toBe("before\n");
// execute() schedules the post-cell snapshot only on status "ok": the on-disk // snapshot stays the last successful cell's, and the next start restores it. expect(scheduleSnapshot).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(60_000); expect(requestTypes().filter((type) => type === "snapshot")).toEqual([]);
// The next kernel over the same snapshot files restores the last successful // cell's namespace; the timed-out cell never contributed one. internals.handleEvent({ event: "done", id: execution.requestId, status: "error" }); await vi.advanceTimersByTimeAsync(0); expect((await manager.restoreState())?.restored).toEqual(["alpha"]); manager.disposeSync(); });
it("schedules the snapshot for a cell that succeeded", async () => { const { manager, internals } = stubManager(); const scheduleSnapshot = vi.fn(); Object.assign(internals, { scheduleSnapshot });
const executePromise = manager.execute("x = 41"); await settleExecution(internals); const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
await expect(executePromise).resolves.toMatchObject({ status: "ok" }); expect(scheduleSnapshot).toHaveBeenCalledTimes(1); manager.disposeSync(); });});