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.
23 kB · 621 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622import { afterEach, describe, expect, it, vi } from "vitest";import { AGENT_MESSAGE_DISPLAY_MIME, type KernelSentAgentMessage, ReplKernelManager,} from "../src/core/kernel/index.js";
async function waitForCalls(mock: { mock: { calls: unknown[][] } }, count: number): Promise<void> { for (let i = 0; i < 20; i++) { if (mock.mock.calls.length >= count) { return; } await Promise.resolve(); } expect(mock.mock.calls.length).toBeGreaterThanOrEqual(count);}
type ReplInternals = { state: string; writeLine: (request: Record<string, unknown>) => Promise<void>; start: () => Promise<void>; handleEvent: (event: Record<string, unknown>) => void; activeExecution?: { requestId: string }; child?: { kill: (signal?: NodeJS.Signals | number) => boolean; pid?: number; stdin?: unknown };};
function runningManagerWith(writeLine: (request: Record<string, unknown>) => Promise<void>): { manager: ReplKernelManager; internals: ReplInternals; kernelKill: ReturnType<typeof vi.fn>;} { const manager = new ReplKernelManager({ cwd: process.cwd() }); const kernelKill = vi.fn((_signal?: NodeJS.Signals | number) => true); const internals = manager as unknown as ReplInternals; Object.assign(internals, { state: "running", writeLine, start: async () => {}, child: { kill: kernelKill, pid: undefined, stdin: undefined }, }); return { manager, internals, kernelKill };}
describe("ReplKernelManager abort handling", () => { afterEach(() => { vi.useRealTimers(); });
it("does not poison startup after a caller starts with an aborted signal", async () => { const manager = new ReplKernelManager({ cwd: process.cwd() }); let startCount = 0; Object.assign( manager as unknown as { doStart: () => Promise<void>; }, { doStart: async () => { startCount++; }, }, ); const controller = new AbortController(); controller.abort();
await expect(manager.start({ signal: controller.signal })).rejects.toThrow("Kernel startup aborted"); await expect(manager.start()).resolves.toBeUndefined(); expect(startCount).toBe(1); });
it("does not cancel shared startup when one waiting caller aborts", async () => { const manager = new ReplKernelManager({ cwd: process.cwd() }); let releaseStart: () => void = () => {}; let startCount = 0; Object.assign( manager as unknown as { doStart: () => Promise<void>; }, { doStart: async () => { startCount++; await new Promise<void>((resolve) => { releaseStart = resolve; }); }, }, ); const controller = new AbortController();
const firstStart = manager.start({ signal: controller.signal }); const secondStart = manager.start(); controller.abort();
await expect(firstStart).rejects.toThrow("Kernel startup aborted"); releaseStart(); await expect(secondStart).resolves.toBeUndefined(); expect(startCount).toBe(1); });
it("settles an aborted execution when the runtime never sends done", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals, kernelKill } = runningManagerWith(writeLine); const controller = new AbortController(); const lateSentAgentMessages: KernelSentAgentMessage[] = [];
const executePromise = manager.execute("while True: pass", { signal: controller.signal, onLateSentAgentMessage: (message) => lateSentAgentMessages.push(message), }); await waitForCalls(writeLine, 1); expect(writeLine).toHaveBeenCalledTimes(1); expect(writeLine.mock.calls[0]?.[0]).toMatchObject({ type: "execute" });
controller.abort(); await vi.advanceTimersByTimeAsync(1000);
await expect(executePromise).resolves.toMatchObject({ status: "aborted" }); // The abort listener interrupted the runtime out-of-band. expect(writeLine.mock.calls.some((call) => (call[0] as { type?: string }).type === "interrupt")).toBe(true); expect(kernelKill).not.toHaveBeenCalled();
const activeExecution = internals.activeExecution; expect(activeExecution).toBeDefined(); if (!activeExecution) { throw new Error("Expected active execution to remain until the runtime's done"); } // A late display event for the aborted cell still dispatches the sent message. internals.handleEvent({ event: "display", id: activeExecution.requestId, data: { [AGENT_MESSAGE_DISPLAY_MIME]: { id: "agentmsg-after-abort", message: "still sent", deliveryStatus: "delivered", target: { activeSessionId: "beta", sessionId: "session-beta" }, }, }, }); expect(lateSentAgentMessages).toEqual([ { id: "agentmsg-after-abort", message: "still sent", deliveryStatus: "delivered", target: { activeSessionId: "beta", sessionId: "session-beta" }, }, ]); // The next execute waits for the stale cell's done before sending. const secondExecutePromise = manager.execute("x = 1"); await Promise.resolve(); expect(writeLine.mock.calls.filter((call) => (call[0] as { type?: string }).type === "execute")).toHaveLength(1);
internals.handleEvent({ event: "done", id: activeExecution.requestId, status: "error" }); await vi.waitFor(() => { expect(writeLine.mock.calls.filter((call) => (call[0] as { type?: string }).type === "execute")).toHaveLength( 2, ); });
const secondExecution = internals.activeExecution; expect(secondExecution).toBeDefined(); if (!secondExecution) { throw new Error("Expected second execution to start after the stale cell settled"); } internals.handleEvent({ event: "done", id: secondExecution.requestId, status: "ok" }); await expect(secondExecutePromise).resolves.toMatchObject({ status: "ok" });
manager.disposeSync(); expect(kernelKill).toHaveBeenCalledWith("SIGTERM"); });
it("settles an aborted execution when the stdin write never resolves", async () => { vi.useFakeTimers(); const interruptWrites: Record<string, unknown>[] = []; const writeLine = vi.fn((request: Record<string, unknown>) => { if (request.type === "interrupt") { interruptWrites.push(request); return Promise.resolve(); } return new Promise<void>(() => {}); }); const { manager } = runningManagerWith(writeLine); const controller = new AbortController();
const executePromise = manager.execute("while True: pass", { signal: controller.signal }); await waitForCalls(writeLine, 1); controller.abort(); await vi.advanceTimersByTimeAsync(1000);
await expect(executePromise).resolves.toMatchObject({ status: "aborted" }); expect(interruptWrites.length).toBeGreaterThan(0); });
it("fails a later execution fast when the interrupted cell never settles", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager } = runningManagerWith(writeLine); const controller = new AbortController();
const executePromise = manager.execute("while True: pass", { signal: controller.signal }); await waitForCalls(writeLine, 1); controller.abort(); await vi.advanceTimersByTimeAsync(1000); await expect(executePromise).resolves.toMatchObject({ status: "aborted" });
const secondExecutePromise = manager.execute("x = 1"); const secondExecuteExpectation = expect(secondExecutePromise).rejects.toThrow( "The Python kernel is still running the previously interrupted cell", ); await Promise.resolve(); await vi.advanceTimersByTimeAsync(5000);
await secondExecuteExpectation; expect(writeLine.mock.calls.filter((call) => (call[0] as { type?: string }).type === "execute")).toHaveLength(1); expect(writeLine.mock.calls.some((call) => (call[0] as { type?: string }).type === "interrupt")).toBe(true); manager.disposeSync(); });
it("cancels a hung final snapshot execution before teardown", async () => { vi.useFakeTimers(); const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); let releaseQueue: () => void = () => {}; const previousExecution = new Promise<void>((resolve) => { releaseQueue = resolve; }); const executeInner = vi.fn( async ( _requestFields: Record<string, unknown>, _code: string, opts: { signal?: AbortSignal }, ): Promise<{ stdout: string; stderr: string; status: "aborted"; durationMs: number }> => await new Promise((resolve) => { opts.signal?.addEventListener( "abort", () => resolve({ stdout: "", stderr: "", status: "aborted", durationMs: 5000 }), { once: true }, ); }), ); const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; executionQueue: Promise<void>; executeInner: typeof executeInner; start: () => Promise<void>; cleanupResources: () => void; }, { state: "running", executionQueue: previousExecution, executeInner, start: async () => {}, cleanupResources }, );
const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); expect(executeInner).not.toHaveBeenCalled(); releaseQueue(); await waitForCalls(executeInner, 1); const signal = executeInner.mock.calls[0]?.[2].signal; expect(signal?.aborted).toBe(false); await vi.advanceTimersByTimeAsync(4999); expect(signal?.aborted).toBe(false); expect(cleanupResources).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1);
expect(signal?.aborted).toBe(true); await expect(disposal).resolves.toBe(true); expect(cleanupResources).toHaveBeenCalledOnce(); });
it("tears down when the final snapshot is blocked behind a hung execution", async () => { vi.useFakeTimers(); const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); const executeInner = vi.fn(); const interrupt = vi.fn(async () => {}); const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; executionQueue: Promise<void>; activeExecution: object; executeInner: typeof executeInner; interrupt: typeof interrupt; cleanupResources: () => void; }, { state: "running", executionQueue: new Promise<void>(() => {}), activeExecution: {}, executeInner, interrupt, cleanupResources, }, );
const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); expect(interrupt).toHaveBeenCalledOnce(); await vi.advanceTimersByTimeAsync(4999); expect(cleanupResources).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1);
await expect(disposal).resolves.toBe(true); expect(executeInner).not.toHaveBeenCalled(); expect(cleanupResources).toHaveBeenCalledOnce(); });
it("rejects executions enqueued during the final snapshot flush so dispose stays bounded", async () => { vi.useFakeTimers(); const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); let releaseQueue: () => void = () => {}; const previousExecution = new Promise<void>((resolve) => { releaseQueue = resolve; }); const executeInner = vi.fn( async ( _requestFields: Record<string, unknown>, _code: string, opts: { signal?: AbortSignal }, ): Promise<{ stdout: string; stderr: string; status: "aborted"; durationMs: number }> => await new Promise((resolve) => { opts.signal?.addEventListener( "abort", () => resolve({ stdout: "", stderr: "", status: "aborted", durationMs: 5000 }), { once: true }, ); }), ); const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; executionQueue: Promise<void>; executeInner: typeof executeInner; start: () => Promise<void>; cleanupResources: () => void; }, { state: "running", executionQueue: previousExecution, executeInner, start: async () => {}, cleanupResources }, );
const disposal = manager.shutdown({ snapshot: true, drainHostRequests: true }); // A cell arriving mid-flush must not splice ahead of the final snapshot. await expect(manager.execute("1 + 1")).rejects.toThrow("Kernel is shutting down"); releaseQueue(); await waitForCalls(executeInner, 1); expect(executeInner.mock.calls[0]?.[0].type).toBe("snapshot"); await vi.advanceTimersByTimeAsync(5000);
await expect(disposal).resolves.toBe(true); expect(executeInner).toHaveBeenCalledOnce(); expect(cleanupResources).toHaveBeenCalledOnce(); });
it("rejects an execution whose re-bootstrap wait crossed the start of the final flush", async () => { const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, bootstrapCode: "boot-code", }); let releaseBootstrap: () => void = () => {}; const bootstrapBlocked = new Promise<void>((resolve) => { releaseBootstrap = resolve; }); let releaseSnapshot: () => void = () => {}; const snapshotBlocked = new Promise<void>((resolve) => { releaseSnapshot = resolve; }); const calls: string[] = []; const ok = { stdout: "", stderr: "", status: "ok" as const, durationMs: 0 }; const executeInner = vi.fn(async (requestFields: Record<string, unknown> & { type: string }, code: string) => { if (requestFields.type === "snapshot") { calls.push("snapshot"); await snapshotBlocked; return { ...ok, doneFields: { saved: [], skipped: [], bytes: 0 } }; } if (code === "boot-code") { calls.push("bootstrap"); await bootstrapBlocked; return ok; } calls.push("cell"); return ok; }); const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; pendingRebootstrap: boolean; executeInner: typeof executeInner; start: () => Promise<void>; cleanupResources: () => void; }, { state: "running", pendingRebootstrap: true, executeInner, start: async () => {}, cleanupResources }, );
// Admitted before the flush: passes the first guard, then parks on the // in-flight lazy re-bootstrap. const cell = manager.execute("user-cell"); cell.catch(() => undefined); await waitForCalls(executeInner, 1); expect(calls).toEqual(["bootstrap"]);
// The teardown's final flush starts while the cell is still parked. const teardown = manager.shutdown({ snapshot: true }); await new Promise((resolve) => globalThis.setTimeout(resolve, 0)); releaseBootstrap();
// Un-bootstrapped no more, the cell must NOT splice between the flush's // captured queue and the final snapshot; the guard re-check rejects it. await expect(cell).rejects.toThrow("Kernel is shutting down"); releaseSnapshot(); await expect(teardown).resolves.toBe(true); expect(calls).toEqual(["bootstrap", "snapshot"]); });
it("settles an aborted execution promptly while the lazy re-bootstrap is still running", async () => { const manager = new ReplKernelManager({ cwd: process.cwd(), bootstrapCode: "boot-code" }); let releaseBootstrap: () => void = () => {}; const bootstrapBlocked = new Promise<void>((resolve) => { releaseBootstrap = resolve; }); const ok = { stdout: "", stderr: "", status: "ok" as const, durationMs: 0 }; const executeInner = vi.fn(async (_requestFields: Record<string, unknown> & { type: string }, code: string) => { if (code === "boot-code") { await bootstrapBlocked; } return ok; }); Object.assign( manager as unknown as { state: "running"; pendingRebootstrap: boolean; executeInner: typeof executeInner; start: () => Promise<void>; }, { state: "running", pendingRebootstrap: true, executeInner, start: async () => {} }, );
const controller = new AbortController(); const cell = manager.execute("user-cell", { signal: controller.signal }); await waitForCalls(executeInner, 1);
// Aborted mid-wait: the cell must not ride out the bootstrap bound. controller.abort(); await expect(cell).resolves.toMatchObject({ status: "aborted" }); releaseBootstrap(); });
it("rejects restart() during the final snapshot flush instead of deadlocking the teardown", async () => { const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); Object.assign( manager as unknown as { state: "running"; flushingSnapshotForDispose: boolean; start: () => Promise<void> }, { state: "running", flushingSnapshotForDispose: true, start: async () => {}, }, );
// Taking a queue slot here and joining the owning shutdown would leave the // flush's snapshot parked behind the slot forever (mutual wait). await expect(manager.restart()).rejects.toThrow("Kernel is shutting down"); });
it("joins concurrent teardowns into a single final snapshot flush", async () => { const manager = new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, }); const executeInner = vi.fn( async ( _requestFields: Record<string, unknown> & { type: string }, ): Promise<{ stdout: string; stderr: string; status: "ok"; durationMs: number }> => ({ stdout: "", stderr: "", status: "ok", durationMs: 0, }), ); const cleanupResources = vi.fn(); Object.assign( manager as unknown as { state: "running"; executionQueue: Promise<void>; executeInner: typeof executeInner; start: () => Promise<void>; cleanupResources: () => void; }, { state: "running", executionQueue: Promise.resolve(), executeInner, start: async () => {}, cleanupResources }, );
// A session dispose racing a signal-handler shutdown must share one final // flush instead of queueing a second snapshot behind the first. await Promise.all([ manager.shutdown({ snapshot: true, drainHostRequests: true }), manager.shutdown({ snapshot: true }), ]);
const snapshotRequests = executeInner.mock.calls.filter((call) => call[0].type === "snapshot"); expect(snapshotRequests).toHaveLength(1); expect(cleanupResources).toHaveBeenCalled(); });
it("routes null-id and stale-id stream events into backgroundOutput, not stdout", async () => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
const executePromise = manager.execute("print('own')"); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; expect(execution).toBeDefined(); if (!execution) { throw new Error("Expected an active execution"); }
internals.handleEvent({ event: "stdout", id: execution.requestId, text: "own\n" }); internals.handleEvent({ event: "stdout", id: null, text: "SECRET-null\n" }); internals.handleEvent({ event: "stdout", id: "stale-cell", text: "SECRET-stale\n" }); internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
const result = await executePromise; expect(result.stdout).toBe("own\n"); expect(result.stdout).not.toContain("SECRET"); expect(result.backgroundOutput).toBe("SECRET-null\nSECRET-stale\n"); manager.disposeSync(); });
it("carries between-cell background output into the next execution's result", async () => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
// No active execution: null-id output parks as pending background output. internals.handleEvent({ event: "stdout", id: null, text: "between-cells\n" });
const executePromise = manager.execute("x = 1"); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; expect(execution).toBeDefined(); if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
const result = await executePromise; expect(result.stdout).toBe(""); expect(result.backgroundOutput).toBe("between-cells\n"); manager.disposeSync(); });
it("dispose writes a protocol shutdown request before hard-killing the child", async () => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine); const killSignals: (NodeJS.Signals | number | undefined)[] = []; internals.child = { kill: (signal?: NodeJS.Signals | number) => { killSignals.push(signal); return true; }, pid: undefined, stdin: { destroyed: false, destroy: () => undefined }, };
await manager.shutdown({ snapshot: true, drainHostRequests: true }); const types = writeLine.mock.calls.map((call) => (call[0] as { type?: string }).type); expect(types).toContain("shutdown"); expect(killSignals).toContain("SIGTERM"); });
it("drops stale between-cell background output on kernel teardown", async () => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
internals.handleEvent({ event: "stdout", id: null, text: "pre-restart-leftover\n" }); manager.disposeSync();
// Simulate the restarted kernel: a new execution must start clean. Object.assign(internals, { state: "running", start: async () => {} }); const executePromise = manager.execute("x = 1"); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; expect(execution).toBeDefined(); if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
const result = await executePromise; expect(result.backgroundOutput).toBeUndefined(); manager.disposeSync(); });
it("marks between-cell background output as truncated once the pending cap is hit", async () => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
// Overflow the pending buffer while idle; nothing more arrives during the cell. internals.handleEvent({ event: "stdout", id: null, text: "x".repeat(70 * 1024) });
const executePromise = manager.execute("x = 1"); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; expect(execution).toBeDefined(); if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
const result = await executePromise; expect(result.backgroundOutput).toContain("background output truncated at"); expect(result.backgroundOutput?.length).toBeLessThan(70 * 1024); manager.disposeSync(); });});