import { 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();
});
});