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.
35 kB · 873 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874import { 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 { Value } from "typebox/value";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 };}
/** The declared AgentToolResult shape omits `isError`; the tool still sets it. */function isErrorResult(result: { content: unknown }): boolean { return (result as { isError?: boolean }).isError === true;}
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;}
/** One Python skill's runtime info, as the session passes it to the provisioner. */function pythonSkillInfo(name: string): { name: string; importName: string; packagePath: string; pyprojectPath: string;} { const importName = name.replaceAll("-", "_"); return { name, importName, packagePath: join(tempDir, name), pyprojectPath: join(tempDir, name, "pyproject.toml"), };}
/** * A runtime stand-in that accepts the bootstrap and answers `signatures`, recording every * request it receives. Without the capability it answers the way an older runtime does. */function writeFakeDigestRuntime(options: { capabilities?: string[]; digest?: string } = {}): { python: string; requested: () => string[][];} { const capable = options.capabilities?.includes("signatures") ?? false; const python = join(tempDir, "python-digest"); const requestLog = join(tempDir, "digest-requests"); const ready = { event: "ready", protocol: 3, python: "3.13.11", capabilities: options.capabilities ?? [] }; const digest = options.digest ?? "# skill surface (generated)\n"; const answer = capable ? `emit({ event: "done", id: request.id, status: "ok", calls: 1, failed: [], digest: ${JSON.stringify(digest)} });` : `emit({ event: "error", id: null, ename: "ProtocolError", evalue: "unknown request type: signatures", traceback: [] }); emit({ event: "done", id: request.id, status: "error", reason: "unknown request type: signatures" });`; 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(${JSON.stringify(ready)});const input = readline.createInterface({ input: process.stdin });input.on("line", (line) => { const request = JSON.parse(line); if (request.type === "signatures") { fs.appendFileSync(${JSON.stringify(requestLog)}, JSON.stringify(request.skills) + "\\n"); ${answer} return; } if (request.type === "execute") { emit({ event: "done", id: request.id, status: "ok" }); return; } if (request.type === "snapshot") { emit({ event: "done", id: request.id, status: "ok", saved: [], skipped: [], bytes: 0 }); return; } if (request.type === "restore") { emit({ event: "done", id: request.id, status: "ok", restored: [], failed: [] }); return; } if (request.type === "shutdown") { emit({ event: "done", id: request.id, status: "ok" }); process.exit(0); }});`, ); chmodSync(python, 0o755); const requested = () => { try { return readFileSync(requestLog, "utf8") .split("\n") .filter(Boolean) .map((line) => JSON.parse(line) as string[]); } catch { return []; // No request has been recorded yet. } }; return { python, requested };}
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("reports a timed-out cell's partial output and its deadline notice", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({ ...okExecuteResult(), stdout: "half-done\n", status: "aborted", timedOut: true, durationMs: 60_100, }); const settle = vi.fn(async () => true); const manager = { execute, waitForInterruptedCellToSettle: settle } 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, cellTimeoutMs: 60_000 });
const result = await tool.execute( "tool-call", { code: "print('half-done')" }, undefined, undefined, {} as ExtensionContext, ); const text = result.content[0]?.type === "text" ? result.content[0].text : "";
expect(text).toContain("half-done"); expect(text).toContain("<ipython_cell_timeout>"); expect(text).toContain( "Cell deadline 60s hit at 60.1s; the cell was interrupted. Nothing from before it is lost: earlier cells' variables and imports survive, and everything this cell printed or already wrote is kept.", ); expect(text).toContain("cell: for _ in range(60): await asyncio.sleep(5)"); expect(text).toContain( 'If you were waiting on work, do not poll inside a cell: start it detached and end your turn -- h = bash("nix build .#prime-agent") returns at once, and its completion notice wakes you with the result.', ); expect(text).toContain( "If it was one genuinely long call (cargo/nix build, test suite, training), re-run it with a larger deadline: timeout=600 (seconds, max 3600).", ); expect(text).toContain("</ipython_cell_timeout>"); expect(text).not.toContain("The kernel was killed and restarted"); expect(result.details.timedOut).toBe(true); expect(result.details.timeoutMs).toEqual({ requested: 60_000, effective: 60_000 }); expect(result.details.status).toBe("aborted"); expect(result.details.kernelRestarted).toBe(false); expect(isErrorResult(result)).toBe(true); expect(execute).toHaveBeenCalledWith("print('half-done')", expect.objectContaining({ timeoutMs: 60_000 })); expect(settle).toHaveBeenCalledTimes(1); expect(kill).not.toHaveBeenCalled(); });
it("kills the kernel when the interrupted cell never settles, and says so twice", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({ ...okExecuteResult(), status: "aborted", timedOut: true, durationMs: 60_100, }); const settle = vi.fn(async () => false); const manager = { execute, waitForInterruptedCellToSettle: settle } 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, cellTimeoutMs: 60_000 });
const result = await tool.execute("tool-call", { code: "while True: pass" }, undefined, undefined, { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() }, } as unknown as ExtensionContext); const text = result.content[0]?.type === "text" ? result.content[0].text : "";
expect(kill).toHaveBeenCalledTimes(1); expect(ensure).toHaveBeenCalledTimes(1); expect(execute).toHaveBeenCalledTimes(1); expect(result.details.kernelRestarted).toBe(true); expect(result.details.timedOut).toBe(true); expect(text).toContain("<ipython_kernel_reset>"); expect(text).toContain("<ipython_cell_timeout>"); expect(text).toContain("The kernel was killed and restarted because the cell did not stop."); expect(isErrorResult(result)).toBe(true); });
it("never asks the user before escalating a deadline", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({ ...okExecuteResult(), status: "aborted", timedOut: true, durationMs: 60_100, }); const manager = { execute, waitForInterruptedCellToSettle: vi.fn(async () => false), } as unknown as KernelClient; const kill = vi.fn(async () => {}); const provisioner = { ensure: vi.fn(async () => manager), kill } as unknown as IpythonKernelProvisioner; const select = vi.fn(async () => "Wait and preserve state"); const { ctx, setWorkingMessage } = createBusyKernelContext(select); const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });
await tool.execute("tool-call", { code: "while True: pass" }, undefined, undefined, ctx);
// A watchdog that asks a human is not a watchdog: a child session has no UI. expect(select).not.toHaveBeenCalled(); expect(kill).toHaveBeenCalledTimes(1); expect(setWorkingMessage).toHaveBeenCalledWith("Restarting Python kernel..."); });
it("clamps an over-ceiling timeout, names both numbers, and suggests nothing larger", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({ ...okExecuteResult(), status: "aborted", timedOut: true, durationMs: 3_600_500, }); const manager = { execute, waitForInterruptedCellToSettle: vi.fn(async () => true), } as unknown as KernelClient; const provisioner = { ensure: vi.fn(async () => manager), kill: vi.fn(async () => {}), } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });
const result = await tool.execute( "tool-call", { code: "while True: pass", timeout: 7200 }, undefined, undefined, { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() }, } as unknown as ExtensionContext, ); const text = result.content[0]?.type === "text" ? result.content[0].text : "";
expect(execute).toHaveBeenCalledWith("while True: pass", expect.objectContaining({ timeoutMs: 3_600_000 })); expect(result.details.timeoutMs).toEqual({ requested: 7_200_000, effective: 3_600_000 }); expect(text).toContain("timeout=7200s exceeds the 3600s ceiling; this cell used 3600s."); expect(text).toContain("Cell deadline 3600s hit at 3600.5s"); expect(text).not.toContain("re-run it with a larger deadline"); expect(text).toContain("start it detached and end your turn"); });
it("reports the clamp even when the clamped cell finishes in time", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({ ...okExecuteResult(), status: "ok", }); const manager = { execute, waitForInterruptedCellToSettle: vi.fn(async () => true), } as unknown as KernelClient; const provisioner = { ensure: vi.fn(async () => manager), kill: vi.fn(async () => {}), } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });
const result = await tool.execute("tool-call", { code: "1 + 1", timeout: 7200 }, undefined, undefined, { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() }, } as unknown as ExtensionContext); const text = result.content[0]?.type === "text" ? result.content[0].text : "";
expect(text).toContain("timeout=7200s exceeds the 3600s ceiling; this cell used 3600s."); expect(text).not.toContain("<ipython_cell_timeout>"); expect(result.details.timedOut).toBeUndefined(); });
it("drops the larger-deadline suggestion when the same code times out again", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValue({ ...okExecuteResult(), status: "aborted", timedOut: true, durationMs: 60_100, }); const manager = { execute, waitForInterruptedCellToSettle: vi.fn(async () => true), } as unknown as KernelClient; const provisioner = { ensure: vi.fn(async () => manager), kill: vi.fn(async () => {}), } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 }); const ctx = { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() } } as unknown as ExtensionContext;
const first = await tool.execute("tool-call-1", { code: "while True: pass" }, undefined, undefined, ctx); const second = await tool.execute("tool-call-2", { code: "while True: pass" }, undefined, undefined, ctx); const other = await tool.execute( "tool-call-3", { code: "while True: sleep_forever()" }, undefined, undefined, ctx, ); const textOf = (result: { content: unknown }) => (result.content as { type: string; text?: string }[])[0]?.text ?? "";
expect(textOf(first)).toContain("timeout=600 (seconds, max 3600)"); // The same code again: the suggestion already failed once, so teach detach only. expect(textOf(second)).not.toContain("timeout=600"); expect(textOf(second)).not.toContain("re-run it with a larger deadline"); expect(textOf(second)).toContain("start it detached and end your turn"); // Different code: a fresh first timeout, so the suggestion is back. expect(textOf(other)).toContain("timeout=600 (seconds, max 3600)"); });
it("rejects a non-positive or over-ceiling timeout before any kernel call", async () => { const execute = vi.fn<KernelClient["execute"]>(); const ensure = vi.fn(async () => ({ execute }) as unknown as KernelClient); const provisioner = { ensure, kill: vi.fn(async () => {}) } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner });
// The same validator the agent loop runs before it invokes the tool. expect(Value.Check(tool.parameters, { code: "1" })).toBe(true); expect(Value.Check(tool.parameters, { code: "1", timeout: 3600 })).toBe(true); expect(Value.Check(tool.parameters, { code: "1", timeout: 0 })).toBe(false); expect(Value.Check(tool.parameters, { code: "1", timeout: -5 })).toBe(false); expect(Value.Check(tool.parameters, { code: "1", timeout: 3601 })).toBe(false); expect(ensure).not.toHaveBeenCalled(); expect(execute).not.toHaveBeenCalled(); });
it("passes the session cell deadline to the manager and publishes it in the description", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult()); const manager = { execute } as unknown as KernelClient; const provisioner = { ensure: vi.fn(async () => manager), kill: vi.fn(async () => {}), } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 120_000 });
expect(tool.description).toContain("Cells stop at 120s by default (max 3600s)"); const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() }, } as unknown as ExtensionContext);
expect(execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ timeoutMs: 120_000 })); expect(result.details.timedOut).toBeUndefined(); expect(result.content).toEqual([{ type: "text", text: "ok" }]); });
it("leaves the cell unbounded when the session turns the deadline off", async () => { const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult()); const manager = { execute } as unknown as KernelClient; const provisioner = { ensure: vi.fn(async () => manager), kill: vi.fn(async () => {}), } as unknown as IpythonKernelProvisioner; const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: "off" });
expect(tool.description).toContain("This session sets no cell deadline"); const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() }, } as unknown as ExtensionContext);
expect(execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ timeoutMs: undefined })); expect(result.details.timedOut).toBeUndefined(); });
it("caches the generated skill digest after the runtime bootstrap", async () => { const { python, requested } = writeFakeDigestRuntime({ capabilities: ["signatures"], digest: "# skill surface (generated)\ngit.status(cwd?) -> str\n", }); const onSkillDigest = vi.fn(); const provisioner = new IpythonKernelProvisioner(tempDir, { python, pythonSkills: [pythonSkillInfo("git")], onSkillDigest, });
try { await provisioner.ensure(); expect(provisioner.skillDigest).toContain("# skill surface"); expect(onSkillDigest).toHaveBeenCalledWith(provisioner.skillDigest); // The digest indexes the modules the bootstrap publishes, so the request names them. expect(requested()).toEqual([["git"]]); } finally { await provisioner.dispose(); } });
it("publishes no digest and sends nothing to a runtime without the capability", async () => { const { python, requested } = writeFakeDigestRuntime(); const onSkillDigest = vi.fn(); const provisioner = new IpythonKernelProvisioner(tempDir, { python, pythonSkills: [pythonSkillInfo("git")], onSkillDigest, });
try { // An older runtime still starts; the missing capability is not a startup failure. await provisioner.ensure(); expect(provisioner.skillDigest).toBeUndefined(); expect(onSkillDigest).not.toHaveBeenCalled(); expect(requested()).toEqual([]); } finally { await provisioner.dispose(); } });
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 }); } });});