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.
28 kB · 700 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701import { afterEach, describe, expect, it, vi } from "vitest";import { AGENT_MESSAGE_DISPLAY_MIME, type KernelSentAgentMessage, ReplKernelManager,} from "../src/core/kernel/index.js";
type ExecuteResult = Awaited<ReturnType<ReplKernelManager["execute"]>>;
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 };};
/** The manager's queue/teardown internals are private; these tests drive them directly. */function stub(manager: ReplKernelManager, overrides: Record<string, unknown>): void { Object.assign(manager as unknown as Record<string, unknown>, overrides);}
const noStart = { start: async (): Promise<void> => {} };
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); stub(manager, { state: "running", writeLine, ...noStart, child: { kill: kernelKill, pid: undefined, stdin: undefined }, }); return { manager, internals: manager as unknown as ReplInternals, kernelKill };}
function requireExecution(internals: ReplInternals): { requestId: string } { const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); } return execution;}
const OK_RESULT = { stdout: "", stderr: "", status: "ok" as const, durationMs: 0 };
/** Snapshot execution that never finishes on its own: it only settles when aborted. */function abortOnlyExecuteInner() { return vi.fn( async ( _requestFields: Record<string, unknown> & { type: string }, _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 }, ); }), );}
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; stub(manager, { 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; stub(manager, { 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.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();
// The stale cell stays active until the runtime's done arrives. const activeExecution = requireExecution(internals); // A late display event for the aborted cell still dispatches the sent message. const sentMessage = { id: "agentmsg-after-abort", message: "still sent", deliveryStatus: "delivered", target: { activeSessionId: "beta", sessionId: "session-beta" }, }; internals.handleEvent({ event: "display", id: activeExecution.requestId, data: { [AGENT_MESSAGE_DISPLAY_MIME]: sentMessage }, }); expect(lateSentAgentMessages).toEqual([sentMessage]);
// The next execute waits for the stale cell's done before sending. const executeCount = (): number => writeLine.mock.calls.filter((call) => (call[0] as { type?: string }).type === "execute").length; const secondExecutePromise = manager.execute("x = 1"); await Promise.resolve(); expect(executeCount()).toBe(1);
internals.handleEvent({ event: "done", id: activeExecution.requestId, status: "error" }); await vi.waitFor(() => { expect(executeCount()).toBe(2); });
internals.handleEvent({ event: "done", id: requireExecution(internals).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); manager.disposeSync(); });
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 executeInner = vi.fn(async (_requestFields: Record<string, unknown> & { type: string }, code: string) => { if (code === "boot-code") { await bootstrapBlocked; } return OK_RESULT; }); stub(manager, { state: "running", pendingRebootstrap: true, executeInner, ...noStart });
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(); });
describe("final snapshot flush", () => { const snapshotManager = (): ReplKernelManager => new ReplKernelManager({ cwd: process.cwd(), snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" }, });
it("cancels a hung final snapshot execution before teardown", async () => { vi.useFakeTimers(); const manager = snapshotManager(); let releaseQueue: () => void = () => {}; const previousExecution = new Promise<void>((resolve) => { releaseQueue = resolve; }); const executeInner = abortOnlyExecuteInner(); const cleanupResources = vi.fn(); stub(manager, { state: "running", executionQueue: previousExecution, executeInner, cleanupResources, ...noStart, });
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 = snapshotManager(); const executeInner = vi.fn(); const interrupt = vi.fn(async () => {}); const cleanupResources = vi.fn(); stub(manager, { 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 = snapshotManager(); let releaseQueue: () => void = () => {}; const previousExecution = new Promise<void>((resolve) => { releaseQueue = resolve; }); const executeInner = abortOnlyExecuteInner(); const cleanupResources = vi.fn(); stub(manager, { state: "running", executionQueue: previousExecution, executeInner, cleanupResources, ...noStart, });
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 executeInner = vi.fn(async (requestFields: Record<string, unknown> & { type: string }, code: string) => { if (requestFields.type === "snapshot") { calls.push("snapshot"); await snapshotBlocked; return { ...OK_RESULT, doneFields: { saved: [], skipped: [], bytes: 0 } }; } calls.push(code === "boot-code" ? "bootstrap" : "cell"); if (code === "boot-code") await bootstrapBlocked; return OK_RESULT; }); const cleanupResources = vi.fn(); stub(manager, { state: "running", pendingRebootstrap: true, executeInner, cleanupResources, ...noStart });
// 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("rejects restart() during the final snapshot flush instead of deadlocking the teardown", async () => { const manager = snapshotManager(); // Taking a queue slot here and joining the owning shutdown would leave the // flush's snapshot parked behind the slot forever (mutual wait). stub(manager, { state: "running", flushingSnapshotForDispose: true, ...noStart });
await expect(manager.restart()).rejects.toThrow("Kernel is shutting down"); });
it("joins concurrent teardowns into a single final snapshot flush", async () => { const manager = snapshotManager(); const executeInner = vi.fn(async (_requestFields: Record<string, unknown> & { type: string }) => OK_RESULT); const cleanupResources = vi.fn(); stub(manager, { state: "running", executionQueue: Promise.resolve(), executeInner, cleanupResources, ...noStart, });
// 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 }), ]);
expect(executeInner.mock.calls.filter((call) => call[0].type === "snapshot")).toHaveLength(1); expect(cleanupResources).toHaveBeenCalled(); });
it("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 });
expect(writeLine.mock.calls.map((call) => (call[0] as { type?: string }).type)).toContain("shutdown"); expect(killSignals).toContain("SIGTERM"); }); });
it("interrupts a cell that outlives its deadline and reports it as timedOut", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine); const requestsOfType = (type: string) => writeLine.mock.calls.map((call) => call[0] as { type?: string; id?: string }).filter((r) => r.type === type);
const executePromise = manager.execute("while True: pass", { timeoutMs: 3000 }); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); }
await vi.advanceTimersByTimeAsync(2999); expect(requestsOfType("interrupt")).toEqual([]); await vi.advanceTimersByTimeAsync(1);
// One out-of-band interrupt, addressed to this cell. const interrupts = requestsOfType("interrupt"); expect(interrupts).toHaveLength(1); expect(interrupts[0]).toMatchObject({ id: execution.requestId });
// The deadline settles the caller after the grace; the runtime's done has not arrived. await vi.advanceTimersByTimeAsync(1000); await expect(executePromise).resolves.toMatchObject({ status: "aborted", timedOut: true }); expect(internals.activeExecution).toBeDefined(); manager.disposeSync(); });
it("does not charge a queued cell for the cell in front of it", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine); const requestsOfType = (type: string) => writeLine.mock.calls.map((call) => call[0] as { type?: string; id?: string }).filter((r) => r.type === type);
const first = manager.execute("while True: pass", { timeoutMs: 60_000 }); await waitForCalls(writeLine, 1); const firstExecution = internals.activeExecution; if (!firstExecution) { throw new Error("Expected an active execution"); }
// Queued behind the first cell: it neither starts nor starts its clock. const second = manager.execute("x = 1", { timeoutMs: 1000 }); await vi.advanceTimersByTimeAsync(5000); expect(requestsOfType("execute")).toHaveLength(1); expect(requestsOfType("interrupt")).toEqual([]);
internals.handleEvent({ event: "done", id: firstExecution.requestId, status: "ok" }); await expect(first).resolves.toMatchObject({ status: "ok" });
// The slot is free: the queued cell starts and gets its own 1000ms clock. await vi.advanceTimersByTimeAsync(0); const secondExecution = internals.activeExecution; if (!secondExecution) { throw new Error("Expected the queued execution to start"); } expect(secondExecution.requestId).not.toBe(firstExecution.requestId); expect(requestsOfType("execute")).toHaveLength(2); await vi.advanceTimersByTimeAsync(999); expect(requestsOfType("interrupt")).toEqual([]); await vi.advanceTimersByTimeAsync(1); expect(requestsOfType("interrupt")).toMatchObject([{ id: secondExecution.requestId }]);
internals.handleEvent({ event: "done", id: secondExecution.requestId, status: "error" }); await expect(second).resolves.toMatchObject({ status: "aborted", timedOut: true }); manager.disposeSync(); });
it("leaves a cell that finishes before its deadline untimed and un-interrupted", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine); const requestsOfType = (type: string) => writeLine.mock.calls.map((call) => call[0] as { type?: string; id?: string }).filter((r) => r.type === type);
const first = manager.execute("x = 1", { timeoutMs: 3000 }); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); } internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" }); await expect(first).resolves.toMatchObject({ status: "ok" }); expect((await first).timedOut).toBeUndefined();
// The cleared timer must never interrupt the next cell. const second = manager.execute("y = 2"); await waitForCalls(writeLine, 2); const secondExecution = internals.activeExecution; if (!secondExecution) { throw new Error("Expected a second execution"); } await vi.advanceTimersByTimeAsync(60_000); expect(requestsOfType("interrupt")).toEqual([]); internals.handleEvent({ event: "done", id: secondExecution.requestId, status: "ok" }); await expect(second).resolves.toMatchObject({ status: "ok" }); manager.disposeSync(); });
it("arms no deadline when the caller supplies none", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine); const requestsOfType = (type: string) => writeLine.mock.calls.map((call) => call[0] as { type?: string; id?: string }).filter((r) => r.type === type);
// The runtime bootstrap is deliberately unbounded: no timeoutMs, no timer. const bootstrap = manager.execute("import time\ntime.sleep(600)"); await waitForCalls(writeLine, 1); const bootstrapExecution = internals.activeExecution; if (!bootstrapExecution) { throw new Error("Expected an active execution"); } await vi.advanceTimersByTimeAsync(10 * 60_000); expect(requestsOfType("interrupt")).toEqual([]); internals.handleEvent({ event: "done", id: bootstrapExecution.requestId, status: "ok" }); const bootstrapResult = await bootstrap; expect(bootstrapResult.status).toBe("ok"); expect(bootstrapResult.timedOut).toBeUndefined();
// State ops keep their own explicit bound instead of inheriting the cell deadline. const stateOp = manager.execute("pass", { internal: true }); await waitForCalls(writeLine, 2); const stateOpExecution = internals.activeExecution; if (!stateOpExecution) { throw new Error("Expected an internal execution"); } await vi.advanceTimersByTimeAsync(10 * 60_000); expect(requestsOfType("interrupt")).toEqual([]); internals.handleEvent({ event: "done", id: stateOpExecution.requestId, status: "ok" }); await expect(stateOp).resolves.toMatchObject({ status: "ok" }); manager.disposeSync(); });
it("keeps a caller abort distinguishable from a deadline", 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, timeoutMs: 60_000, }); await waitForCalls(writeLine, 1); controller.abort(); await vi.advanceTimersByTimeAsync(1000);
const aborted = await executePromise; expect(aborted.status).toBe("aborted"); expect(aborted.timedOut).toBeUndefined(); manager.disposeSync(); });
it("reports a cell that never settles and leaves the kill reachable", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals, kernelKill } = runningManagerWith(writeLine); const requestsOfType = (type: string) => writeLine.mock.calls.map((call) => call[0] as { type?: string; id?: string }).filter((r) => r.type === type);
const executePromise = manager.execute("while True: pass", { timeoutMs: 1000 }); await waitForCalls(writeLine, 1); await vi.advanceTimersByTimeAsync(2000); await expect(executePromise).resolves.toMatchObject({ status: "aborted", timedOut: true });
// The runtime never sent done: the settle wait reports that instead of throwing. const settled = manager.waitForInterruptedCellToSettle(1000); await vi.advanceTimersByTimeAsync(1000); await expect(settled).resolves.toBe(false); expect(requestsOfType("interrupt").length).toBeGreaterThanOrEqual(2); expect(internals.activeExecution).toBeDefined();
await manager.kill(); expect(kernelKill).toHaveBeenCalledWith("SIGKILL");
// A killed kernel is terminal, and there is nothing left to settle. await expect(manager.waitForInterruptedCellToSettle(1000)).resolves.toBe(true); });
it("resolves the settle wait as soon as the interrupted cell's done arrives", async () => { vi.useFakeTimers(); const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
const executePromise = manager.execute("while True: pass", { timeoutMs: 1000 }); await waitForCalls(writeLine, 1); const execution = internals.activeExecution; if (!execution) { throw new Error("Expected an active execution"); } await vi.advanceTimersByTimeAsync(2000); await expect(executePromise).resolves.toMatchObject({ status: "aborted", timedOut: true });
const settled = manager.waitForInterruptedCellToSettle(5000); internals.handleEvent({ event: "done", id: execution.requestId, status: "error" }); await expect(settled).resolves.toBe(true); manager.disposeSync(); });
// Output that is not owned by the running cell must never leak into its stdout. it.each<{ label: string; before?: (internals: ReplInternals) => void; during?: (internals: ReplInternals, requestId: string) => void; check: (result: ExecuteResult) => void; }>([ { label: "routes null-id and stale-id stream events into backgroundOutput, not stdout", during: (internals, requestId) => { internals.handleEvent({ event: "stdout", id: 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" }); }, check: (result) => { expect(result.stdout).toBe("own\n"); expect(result.backgroundOutput).toBe("SECRET-null\nSECRET-stale\n"); }, }, { label: "carries between-cell background output into the next execution's result", before: (internals) => internals.handleEvent({ event: "stdout", id: null, text: "between-cells\n" }), check: (result) => { expect(result.stdout).toBe(""); expect(result.backgroundOutput).toBe("between-cells\n"); }, }, { label: "marks between-cell background output as truncated once the pending cap is hit", before: (internals) => internals.handleEvent({ event: "stdout", id: null, text: "x".repeat(70 * 1024) }), check: (result) => { expect(result.backgroundOutput).toContain("background output truncated at"); expect(result.backgroundOutput?.length).toBeLessThan(70 * 1024); }, }, ])("$label", async ({ before, during, check }) => { const writeLine = vi.fn(async (_request: Record<string, unknown>) => {}); const { manager, internals } = runningManagerWith(writeLine);
before?.(internals); const executePromise = manager.execute("print('own')"); await waitForCalls(writeLine, 1); const execution = requireExecution(internals); during?.(internals, execution.requestId); internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });
check(await executePromise); manager.disposeSync(); });});