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.
18 kB · 482 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483import { chmodSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { cleanupSessionResources } from "@earendil-works/pi-ai";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import type { ExtensionContext } from "../src/core/extensions/types.js";import type { KernelBootstrapProgressHandler } from "../src/core/kernel/bootstrap.js";import { type ExecuteResult, KernelBusyAfterInterruptError, type KernelClient, ReplKernelManager,} from "../src/core/kernel/index.js";import { createIpythonToolDefinition, IpythonKernelProvisioner } from "../src/core/tools/ipython.js";
let tempDir = "";
function writeFakePython(opts: { sleepSeconds?: number } = {}): { python: string; countRuns: () => number } { const python = join(tempDir, "python"); const countFile = join(tempDir, "runs"); writeFileSync( python, [ "#!/bin/sh", `echo run >> "${countFile}"`, ...(opts.sleepSeconds ? [`sleep ${opts.sleepSeconds}`] : []), "exit 42", "", ].join("\n"), ); chmodSync(python, 0o755); const countRuns = () => { try { return readFileSync(countFile, "utf8").split("\n").filter(Boolean).length; } catch { return 0; } }; return { python, countRuns };}
function okExecuteResult(): ExecuteResult { return { stdout: "ok", stderr: "", status: "ok", durationMs: 1 };}
function createBusyKernelContext( select: (title: string, options: string[]) => Promise<string | undefined>, options: { throwWorkingMessage?: boolean } = {},): { ctx: ExtensionContext; setWorkingMessage: ReturnType<typeof vi.fn>;} { const setWorkingMessage = vi.fn(() => { if (options.throwWorkingMessage) { throw new Error("stale UI context"); } }); const ctx = { hasUI: true, ui: { select, setWorkingMessage, }, } as unknown as ExtensionContext; return { ctx, setWorkingMessage };}
function writeFakeReplRuntime( markerPath: string, options: { gatedExecute?: { startedPath: string; gatePath: string } } = {},): string { const python = join(tempDir, "python-repl"); const refuse = `emit({ event: "error", id: request.id, ename: "RuntimeError", evalue: "bootstrap refused", traceback: [] }); emit({ event: "done", id: request.id, status: "error" });`; const executeBranch = options.gatedExecute ? `fs.writeFileSync(${JSON.stringify(options.gatedExecute.startedPath)}, "1"); const gate = setInterval(() => { if (!fs.existsSync(${JSON.stringify(options.gatedExecute.gatePath)})) return; clearInterval(gate); ${refuse} }, 10);` : refuse; writeFileSync( python, `#!/usr/bin/env nodeconst fs = require("node:fs");const readline = require("node:readline");const emit = (event) => process.stdout.write(JSON.stringify(event) + "\\n");emit({ event: "ready", protocol: 3, python: process.version });const input = readline.createInterface({ input: process.stdin });input.on("line", (line) => { const request = JSON.parse(line); if (request.type === "restore") { emit({ event: "done", id: request.id, status: "ok", restored: [], failed: [] }); return; } if (request.type === "snapshot") { setTimeout(() => { fs.writeFileSync(${JSON.stringify(markerPath)}, "1"); emit({ event: "done", id: request.id, status: "ok", saved: [], skipped: [], bytes: 0 }); }, 150); return; } if (request.type === "execute") { ${executeBranch} return; } if (request.type === "shutdown") { emit({ event: "done", id: request.id, status: "ok" }); process.exit(0); }});`, ); chmodSync(python, 0o755); return python;}
describe("IpythonKernelProvisioner", () => { beforeEach(() => { tempDir = mkdtempSync(join(tmpdir(), "prime-agent-provisioner-")); });
afterEach(() => { if (tempDir) { rmSync(tempDir, { recursive: true, force: true }); tempDir = ""; } });
it("does not surface a failed startup before the kernel's final snapshot flush finished", async () => { const marker = join(tempDir, "snapshot-flushed"); const snapshotDir = join(tempDir, "snapshots"); mkdirSync(snapshotDir, { recursive: true }); const python = writeFakeReplRuntime(marker); const provisioner = new IpythonKernelProvisioner(tempDir, { python, snapshotDir }); try { await expect(provisioner.ensure()).rejects.toThrow(/Failed to initialize rlm runtime/); // The failed kernel's teardown (final snapshot flush included) completed // before the failure surfaced, so a replacement provisioner gated on this // one cannot race the still-flushing kernel over the same snapshot files. expect(existsSync(marker)).toBe(true); } finally { await provisioner.dispose(); } });
it("memoizes concurrent ensure() calls into one startup", async () => { const { python, countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python });
const [a, b] = await Promise.allSettled([provisioner.ensure(), provisioner.ensure()]); expect(a.status).toBe("rejected"); expect(b.status).toBe("rejected"); expect(countRuns()).toBe(1); });
it("retries after a failed startup instead of caching the rejection", async () => { const { python, countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python });
await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/); await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/); expect(countRuns()).toBe(2); });
it("prewarm() swallows the failure and the next ensure() starts fresh", async () => { const { python, countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python });
provisioner.prewarm(); expect(provisioner.manager).toBeUndefined();
// Once the prewarm startup settles, ensure() must launch a second attempt. await vi.waitFor(async () => { await expect(provisioner.ensure()).rejects.toThrow(); expect(countRuns()).toBeGreaterThanOrEqual(2); }); });
it("replays the current startup stage to listeners attaching mid-flight", async () => { const { python } = writeFakePython({ sleepSeconds: 1 }); const provisioner = new IpythonKernelProvisioner(tempDir, { python });
provisioner.prewarm(); const messages: string[] = []; const joined = provisioner.ensure((message) => messages.push(message)); expect(messages).toContain("Starting Python kernel..."); await expect(joined).rejects.toThrow(); });
it("dispose() settles a startup that is still in flight", async () => { const { python } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python });
provisioner.prewarm(); await provisioner.dispose(); expect(provisioner.manager).toBeUndefined(); });
it("skips the snapshot when dispose({ snapshot: false }) aborts a startup in flight", async () => { const marker = join(tempDir, "snapshot-flushed"); const executeStarted = join(tempDir, "execute-started"); const executeGate = join(tempDir, "execute-gate"); const snapshotDir = join(tempDir, "snapshots"); mkdirSync(snapshotDir, { recursive: true }); const python = writeFakeReplRuntime(marker, { gatedExecute: { startedPath: executeStarted, gatePath: executeGate }, }); const provisioner = new IpythonKernelProvisioner(tempDir, { python, snapshotDir });
const started = provisioner.ensure().catch(() => undefined); await vi.waitFor(() => expect(existsSync(executeStarted)).toBe(true)); const disposed = provisioner.dispose({ snapshot: false }); writeFileSync(executeGate, "1"); await Promise.all([disposed, started]); expect(existsSync(marker)).toBe(false); });
it("dispose({ snapshot: false }) skips the kernel's final snapshot flush", async () => { const { python } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python }); const shutdown = vi.fn(async () => {}); Reflect.set(provisioner, "managerPromise", Promise.resolve({ shutdown }));
await provisioner.dispose({ snapshot: false }); expect(shutdown).toHaveBeenCalledWith({ snapshot: false, drainHostRequests: true }); });
it("dispose() before the boot slot prevents the kernel from spawning", async () => { const { python, countRuns } = writeFakePython(); let release: () => void = () => {}; const gate = new Promise<void>((r) => { release = r; }); const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate });
const started = provisioner.ensure().catch(() => {}); const disposed = provisioner.dispose(); // aborts while the boot waits on readyGate release(); await Promise.all([started, disposed]); expect(countRuns()).toBe(0); // disposed boot must never spawn a kernel });
it("aborting the startup owner before the boot slot prevents the kernel from spawning", async () => { const { python, countRuns } = writeFakePython(); let release: () => void = () => {}; const gate = new Promise<void>((r) => { release = r; }); const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate }); const controller = new AbortController();
const started = provisioner.ensure(undefined, controller.signal); controller.abort(); await expect(started).rejects.toThrow("Python execution aborted"); release(); await new Promise((r) => setTimeout(r, 50));
expect(countRuns()).toBe(0); expect(provisioner.manager).toBeUndefined(); });
it("waits for readyGate before starting the kernel", async () => { const { python, countRuns } = writeFakePython(); let release: () => void = () => {}; const gate = new Promise<void>((r) => { release = r; }); const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate });
const started = provisioner.ensure().catch(() => {}); await new Promise((r) => setTimeout(r, 50)); expect(countRuns()).toBe(0); // gated: must not spawn the kernel yet
release(); await started; expect(countRuns()).toBe(1); });
it("namespace maintenance returns null when no kernel is running", async () => { const provisioner = new IpythonKernelProvisioner(tempDir, {}); expect(await provisioner.listNamespaceNames()).toBeNull(); expect(await provisioner.pruneOversizedVariables()).toBeNull(); });
it("does not dispose a running kernel when an ensure caller is aborted", async () => { const provisioner = new IpythonKernelProvisioner(tempDir, {}); const dispose = vi.fn(async () => {}); const manager = { dispose, isRunning: true } as unknown as KernelClient; Object.assign( provisioner as unknown as { managerPromise: Promise<KernelClient>; startedManager: KernelClient; }, { managerPromise: Promise.resolve(manager), startedManager: manager, }, ); const controller = new AbortController(); controller.abort();
await expect(provisioner.ensure(undefined, controller.signal)).rejects.toThrow("Python execution aborted"); expect(dispose).not.toHaveBeenCalled(); expect(provisioner.manager).toBe(manager); });
function primeKernelMemo(provisioner: IpythonKernelProvisioner, manager: KernelClient) { Object.assign(provisioner as unknown as { managerPromise: Promise<KernelClient>; startedManager: KernelClient }, { managerPromise: Promise.resolve(manager), startedManager: manager, }); }
it("drops a dead kernel memo so ensure() restarts instead of reusing it", async () => { const { python, countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, { python }); primeKernelMemo(provisioner, { isRunning: false, isDefunct: true } as unknown as KernelClient);
await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/); expect(countRuns()).toBe(1); });
it("keeps the memo for a kernel that is repairing itself, not defunct", async () => { const { countRuns } = writeFakePython(); const provisioner = new IpythonKernelProvisioner(tempDir, {}); const repairing = { isRunning: false, isDefunct: false } as unknown as KernelClient; primeKernelMemo(provisioner, repairing);
// A second provisioner kernel during protocol repair would split the snapshot dir. await expect(provisioner.ensure()).resolves.toBe(repairing); expect(countRuns()).toBe(0); });
it("removes startup progress listeners when an ensure caller is aborted", async () => { const provisioner = new IpythonKernelProvisioner(tempDir, {}); Object.assign( provisioner as unknown as { managerPromise: Promise<KernelClient>; }, { managerPromise: new Promise<KernelClient>(() => {}), }, ); const controller = new AbortController(); const onProgress = vi.fn();
const ensurePromise = provisioner.ensure(onProgress, controller.signal).catch(() => undefined); controller.abort(); await ensurePromise;
const internals = provisioner as unknown as { startupListeners: Set<KernelBootstrapProgressHandler>; }; expect(internals.startupListeners.has(onProgress)).toBe(false); });
it("surfaces backgroundOutput in details without changing model content", async () => { const execute = vi .fn<KernelClient["execute"]>() .mockResolvedValueOnce({ ...okExecuteResult(), backgroundOutput: "bg-line" }); const manager = { execute } as unknown as KernelClient; const ensure = vi.fn(async () => manager); const kill = vi.fn(async () => {}); const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner });
const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, {} as ExtensionContext);
expect(result.details.backgroundOutput).toBe("bg-line"); expect(result.content).toEqual([{ type: "text", text: "ok\n[background output (unattributed)]\nbg-line" }]); });
it("lets the user wait when an interrupted kernel is still busy", async () => { const execute = vi .fn<KernelClient["execute"]>() .mockRejectedValueOnce(new KernelBusyAfterInterruptError()) .mockResolvedValueOnce(okExecuteResult()); const manager = { execute } as unknown as KernelClient; const ensure = vi.fn(async () => manager); const kill = vi.fn(async () => {}); const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner; const select = vi.fn(async () => "Wait and preserve state"); const { ctx, setWorkingMessage } = createBusyKernelContext(select, { throwWorkingMessage: true }); const tool = createIpythonToolDefinition(tempDir, { provisioner });
const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, ctx);
expect(result.details.status).toBe("ok"); expect(result.details.kernelRestarted).toBe(false); expect(ensure).toHaveBeenCalledTimes(2); expect(kill).not.toHaveBeenCalled(); expect(select).toHaveBeenCalledWith( expect.stringContaining("previous cell has not stopped"), ["Wait and preserve state", "Kill kernel and restart"], { signal: undefined, }, ); expect(setWorkingMessage).toHaveBeenCalledWith("Waiting for Python kernel..."); expect(setWorkingMessage).toHaveBeenLastCalledWith(undefined); });
it("lets the user kill and restart a busy interrupted kernel", async () => { const busyManager = { execute: vi.fn<KernelClient["execute"]>().mockRejectedValueOnce(new KernelBusyAfterInterruptError()), } as unknown as KernelClient; const freshManager = { execute: vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult()), } as unknown as KernelClient; const ensure = vi.fn(async () => { return ensure.mock.calls.length === 1 ? busyManager : freshManager; }); const kill = vi.fn(async () => {}); const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner; const select = vi.fn(async () => "Kill kernel and restart"); const { ctx, setWorkingMessage } = createBusyKernelContext(select, { throwWorkingMessage: true }); const tool = createIpythonToolDefinition(tempDir, { provisioner });
const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, ctx); const text = result.content[0]?.type === "text" ? result.content[0].text : "";
expect(result.details.status).toBe("ok"); expect(result.details.kernelRestarted).toBe(true); expect(text).toContain("<ipython_kernel_reset>"); expect(text).toContain("Variables, imports, async tasks, and open resources"); expect(text).toContain("ok"); expect(ensure).toHaveBeenCalledTimes(2); expect(kill).toHaveBeenCalledTimes(1); expect(freshManager.execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ signal: undefined })); expect(setWorkingMessage).toHaveBeenCalledWith("Restarting Python kernel..."); expect(setWorkingMessage).toHaveBeenLastCalledWith(undefined); });
it("does not delete the on-disk snapshot (the kernel survives compaction)", async () => { const snapshotDir = join(tempDir, "artifacts"); const provisioner = new IpythonKernelProvisioner(tempDir, { snapshotDir }); const dill = join(snapshotDir, "kernel-state.dill"); const manifest = join(snapshotDir, "kernel-state.json"); mkdirSync(snapshotDir, { recursive: true }); writeFileSync(dill, "payload"); writeFileSync(manifest, "{}");
// listing the namespace must never touch the on-disk snapshot await provisioner.listNamespaceNames();
expect(existsSync(dill)).toBe(true); expect(existsSync(manifest)).toBe(true); });});
describe("ReplKernelManager session cleanup during startup", () => { beforeEach(() => { tempDir = mkdtempSync(join(tmpdir(), "prime-agent-kernel-cleanup-")); });
afterEach(() => { if (tempDir) { rmSync(tempDir, { recursive: true, force: true }); tempDir = ""; } });
it("disposes a kernel that is still booting when its session is cleaned up", async () => { const python = join(tempDir, "python"); // Never emits the ready line - stays in the booting phase until killed. writeFileSync(python, ["#!/bin/sh", "sleep 30", ""].join("\n")); chmodSync(python, 0o755); const sessionId = `provisioner-test-${Date.now()}`; const manager = new ReplKernelManager({ python, cwd: tempDir, sessionId });
try { const startup = manager.start(); cleanupSessionResources(sessionId); await expect(startup).rejects.toThrow(/Kernel exited before ready|disposed during startup/); expect(manager.isRunning).toBe(false); } finally { await manager.shutdown({ snapshot: true, drainHostRequests: true }); } });});