fix(coding-agent): scoped python kernel cleanup to sessions
scoped retained-kernel ownership to agent sessions and cleaned it up on session disposal cleaned up warmed python owners on session startup failure and rejected new direct and tool-based python starts during disposal, including async hook, preflight, and warmup races
This commit is contained in:
Carl
@@ -0,0 +1,803 @@
|
||||
import { afterEach, describe, expect, it, vi } from "bun:test";
|
||||
import * as fs from "node:fs";
|
||||
import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-ai";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import * as pythonExecutor from "@oh-my-pi/pi-coding-agent/ipy/executor";
|
||||
import type { PreludeHelper, PythonKernel as PythonKernelInstance } from "@oh-my-pi/pi-coding-agent/ipy/kernel";
|
||||
import * as pythonKernel from "@oh-my-pi/pi-coding-agent/ipy/kernel";
|
||||
import * as memories from "@oh-my-pi/pi-coding-agent/memories";
|
||||
import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { Snowflake } from "@oh-my-pi/pi-utils";
|
||||
|
||||
const OK_EXECUTION = { status: "ok", cancelled: false, timedOut: false, stdinRequested: false } as const;
|
||||
|
||||
class FakeKernel {
|
||||
executeCalls: string[] = [];
|
||||
shutdownCalls = 0;
|
||||
alive = true;
|
||||
blockedCode: string | undefined;
|
||||
blockedExecution: Promise<typeof OK_EXECUTION> | undefined;
|
||||
blockedExecutionStarted: (() => void) | undefined;
|
||||
blockedExecutionReject: ((error: Error) => void) | undefined;
|
||||
abortBlockedExecution = true;
|
||||
|
||||
isAlive(): boolean {
|
||||
return this.alive;
|
||||
}
|
||||
|
||||
async execute(code: string, options?: { signal?: AbortSignal }): Promise<typeof OK_EXECUTION> {
|
||||
this.executeCalls.push(code);
|
||||
if (code === this.blockedCode && this.blockedExecution) {
|
||||
this.blockedExecutionStarted?.();
|
||||
if (!this.abortBlockedExecution || !options?.signal) {
|
||||
return await this.blockedExecution;
|
||||
}
|
||||
return await Promise.race([
|
||||
this.blockedExecution,
|
||||
new Promise<typeof OK_EXECUTION>((_, reject) => {
|
||||
const onAbort = () => reject(new DOMException("Aborted", "AbortError"));
|
||||
if (options.signal?.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
options.signal?.addEventListener("abort", onAbort, { once: true });
|
||||
}),
|
||||
]);
|
||||
}
|
||||
return OK_EXECUTION;
|
||||
}
|
||||
|
||||
async ping(): Promise<boolean> {
|
||||
return this.alive;
|
||||
}
|
||||
|
||||
shutdown = vi.fn(async () => {
|
||||
this.shutdownCalls += 1;
|
||||
this.alive = false;
|
||||
this.blockedExecutionReject?.(new Error("Kernel shut down during execution"));
|
||||
return { confirmed: true };
|
||||
});
|
||||
}
|
||||
|
||||
const getModel = () => {
|
||||
const model = getBundledModel("anthropic", "claude-sonnet-4-5");
|
||||
if (!model) throw new Error("Expected bundled model");
|
||||
return model;
|
||||
};
|
||||
|
||||
const createTempProject = () => {
|
||||
const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-agent-session-python-cleanup-${Snowflake.next()}-`));
|
||||
const cwd = path.join(tempDir, "project");
|
||||
fs.mkdirSync(cwd, { recursive: true });
|
||||
return { tempDir, cwd };
|
||||
};
|
||||
|
||||
const createSession = async (
|
||||
tempDir: string,
|
||||
cwd: string,
|
||||
options: { extensions?: ExtensionFactory[]; sessionManager?: SessionManager } = {},
|
||||
) =>
|
||||
(
|
||||
await createAgentSession({
|
||||
cwd,
|
||||
agentDir: tempDir,
|
||||
sessionManager: options.sessionManager ?? SessionManager.inMemory(cwd),
|
||||
settings: Settings.isolated({ "python.kernelMode": "session" }),
|
||||
model: getModel(),
|
||||
disableExtensionDiscovery: true,
|
||||
extensions: options.extensions,
|
||||
skills: [],
|
||||
contextFiles: [],
|
||||
promptTemplates: [],
|
||||
slashCommands: [],
|
||||
enableMCP: false,
|
||||
enableLsp: false,
|
||||
forcePythonWarmup: true,
|
||||
toolNames: ["python"],
|
||||
})
|
||||
).session;
|
||||
|
||||
const stubPythonWarmup = () =>
|
||||
vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] });
|
||||
|
||||
const createWarmupKernel = (docs: PreludeHelper[] = []) => ({
|
||||
introspectPrelude: vi.fn().mockResolvedValue(docs),
|
||||
execute: vi.fn(async () => OK_EXECUTION),
|
||||
ping: vi.fn(async () => true),
|
||||
isAlive: () => true,
|
||||
shutdown: vi.fn(async () => ({ confirmed: true })),
|
||||
});
|
||||
|
||||
describe("AgentSession python cleanup", () => {
|
||||
const tempDirs: string[] = [];
|
||||
|
||||
afterEach(async () => {
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
vi.restoreAllMocks();
|
||||
await pythonExecutor.disposeAllKernelSessions();
|
||||
for (const tempDir of tempDirs.splice(0)) {
|
||||
fs.rmSync(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("cleans up warmed Python owners when createAgentSession fails before session construction", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
|
||||
const unrelatedKernel = createWarmupKernel();
|
||||
const warmedKernel = createWarmupKernel([
|
||||
{ name: "helper", signature: "helper()", docstring: "doc", category: "general" },
|
||||
]);
|
||||
const unrelatedCwd = path.join(tempDir, "unrelated-before");
|
||||
const throwingExtension: ExtensionFactory = () => {
|
||||
throw new Error("Extension init failed");
|
||||
};
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
const startSpy = vi
|
||||
.spyOn(pythonKernel.PythonKernel, "start")
|
||||
.mockResolvedValueOnce(unrelatedKernel as unknown as PythonKernelInstance)
|
||||
.mockResolvedValueOnce(warmedKernel as unknown as PythonKernelInstance);
|
||||
|
||||
await pythonExecutor.executePython("print('unrelated before')", {
|
||||
cwd: unrelatedCwd,
|
||||
sessionId: "unrelated-before-session",
|
||||
kernelMode: "session",
|
||||
kernelOwnerId: "other-owner",
|
||||
});
|
||||
|
||||
await expect(
|
||||
createAgentSession({
|
||||
cwd,
|
||||
agentDir: tempDir,
|
||||
sessionManager: SessionManager.inMemory(cwd),
|
||||
settings: Settings.isolated({ "python.kernelMode": "session" }),
|
||||
model: getModel(),
|
||||
disableExtensionDiscovery: true,
|
||||
extensions: [throwingExtension],
|
||||
skills: [],
|
||||
contextFiles: [],
|
||||
promptTemplates: [],
|
||||
slashCommands: [],
|
||||
enableMCP: false,
|
||||
enableLsp: false,
|
||||
forcePythonWarmup: true,
|
||||
toolNames: ["python"],
|
||||
}),
|
||||
).rejects.toThrow("Extension init failed");
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(2);
|
||||
expect(warmedKernel.introspectPrelude).toHaveBeenCalledTimes(1);
|
||||
expect(warmedKernel.shutdown).toHaveBeenCalledTimes(1);
|
||||
expect(unrelatedKernel.shutdown).not.toHaveBeenCalled();
|
||||
|
||||
await pythonExecutor.executePython("print('still alive before')", {
|
||||
cwd: unrelatedCwd,
|
||||
sessionId: "unrelated-before-session",
|
||||
kernelMode: "session",
|
||||
kernelOwnerId: "other-owner",
|
||||
});
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(2);
|
||||
expect(unrelatedKernel.execute).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("cleans up warmed Python owners when createAgentSession fails after session construction", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
|
||||
const unrelatedKernel = createWarmupKernel();
|
||||
const warmedKernel = createWarmupKernel([
|
||||
{ name: "helper", signature: "helper()", docstring: "doc", category: "general" },
|
||||
]);
|
||||
const unrelatedCwd = path.join(tempDir, "unrelated-after");
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
const startSpy = vi
|
||||
.spyOn(pythonKernel.PythonKernel, "start")
|
||||
.mockResolvedValueOnce(unrelatedKernel as unknown as PythonKernelInstance)
|
||||
.mockResolvedValueOnce(warmedKernel as unknown as PythonKernelInstance);
|
||||
vi.spyOn(memories, "startMemoryStartupTask").mockImplementation(() => {
|
||||
throw new Error("Memory startup failed");
|
||||
});
|
||||
|
||||
await pythonExecutor.executePython("print('unrelated after')", {
|
||||
cwd: unrelatedCwd,
|
||||
sessionId: "unrelated-after-session",
|
||||
kernelMode: "session",
|
||||
kernelOwnerId: "other-owner",
|
||||
});
|
||||
|
||||
await expect(
|
||||
createAgentSession({
|
||||
cwd,
|
||||
agentDir: tempDir,
|
||||
sessionManager: SessionManager.inMemory(cwd),
|
||||
settings: Settings.isolated({ "python.kernelMode": "session" }),
|
||||
model: getModel(),
|
||||
disableExtensionDiscovery: true,
|
||||
skills: [],
|
||||
contextFiles: [],
|
||||
promptTemplates: [],
|
||||
slashCommands: [],
|
||||
enableMCP: false,
|
||||
enableLsp: false,
|
||||
forcePythonWarmup: true,
|
||||
toolNames: ["python"],
|
||||
}),
|
||||
).rejects.toThrow("Memory startup failed");
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(2);
|
||||
expect(warmedKernel.introspectPrelude).toHaveBeenCalledTimes(1);
|
||||
expect(warmedKernel.shutdown).toHaveBeenCalledTimes(1);
|
||||
expect(unrelatedKernel.shutdown).not.toHaveBeenCalled();
|
||||
|
||||
await pythonExecutor.executePython("print('still alive after')", {
|
||||
cwd: unrelatedCwd,
|
||||
sessionId: "unrelated-after-session",
|
||||
kernelMode: "session",
|
||||
kernelOwnerId: "other-owner",
|
||||
});
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(2);
|
||||
expect(unrelatedKernel.execute).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("waits for active SDK session Python work before releasing a shared retained kernel", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
stubPythonWarmup();
|
||||
|
||||
const kernel = new FakeKernel();
|
||||
const blockedExecution = Promise.withResolvers<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
let blockedExecutionSettled = false;
|
||||
blockedExecution.promise.then(
|
||||
() => {
|
||||
blockedExecutionSettled = true;
|
||||
},
|
||||
() => {
|
||||
blockedExecutionSettled = true;
|
||||
},
|
||||
);
|
||||
kernel.blockedCode = "print('first')";
|
||||
kernel.blockedExecution = blockedExecution.promise;
|
||||
kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve();
|
||||
kernel.blockedExecutionReject = error => blockedExecution.reject(error);
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
const startSpy = vi
|
||||
.spyOn(pythonKernel.PythonKernel, "start")
|
||||
.mockResolvedValue(kernel as unknown as PythonKernelInstance);
|
||||
const firstSession = await createSession(tempDir, cwd);
|
||||
const secondSession = await createSession(tempDir, cwd);
|
||||
expect(startSpy).toHaveBeenCalledTimes(0);
|
||||
let firstDisposed = false;
|
||||
|
||||
try {
|
||||
const firstExecution = firstSession.executePython("print('first')");
|
||||
let firstExecutionSettled = false;
|
||||
const observedFirstExecution = firstExecution.finally(() => {
|
||||
firstExecutionSettled = true;
|
||||
});
|
||||
await blockedExecutionStarted.promise;
|
||||
|
||||
const disposeFirst = firstSession.dispose().then(() => {
|
||||
expect(blockedExecutionSettled).toBe(true);
|
||||
expect(firstExecutionSettled).toBe(true);
|
||||
firstDisposed = true;
|
||||
});
|
||||
await Bun.sleep(0);
|
||||
expect(firstDisposed).toBe(false);
|
||||
expect(blockedExecutionSettled).toBe(false);
|
||||
expect(firstExecutionSettled).toBe(false);
|
||||
|
||||
const secondExecution = secondSession.executePython("print('second')");
|
||||
await Bun.sleep(0);
|
||||
|
||||
expect(firstDisposed).toBe(false);
|
||||
expect(blockedExecutionSettled).toBe(false);
|
||||
expect(firstExecutionSettled).toBe(false);
|
||||
expect(kernel.shutdownCalls).toBe(0);
|
||||
|
||||
blockedExecution.resolve(OK_EXECUTION);
|
||||
await Promise.all([observedFirstExecution, secondExecution, disposeFirst]);
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(1);
|
||||
expect(kernel.shutdownCalls).toBe(0);
|
||||
expect(kernel.executeCalls).toEqual(["print('first')", "print('second')"]);
|
||||
|
||||
await secondSession.executePython("print('third')");
|
||||
|
||||
expect(startSpy).toHaveBeenCalledTimes(1);
|
||||
expect(kernel.executeCalls).toEqual(["print('first')", "print('second')", "print('third')"]);
|
||||
} finally {
|
||||
if (!firstDisposed) {
|
||||
await firstSession.dispose();
|
||||
}
|
||||
await secondSession.dispose();
|
||||
}
|
||||
|
||||
expect(kernel.shutdownCalls).toBe(1);
|
||||
});
|
||||
|
||||
it("aborts tracked Python tool warmup during session dispose before executePython starts", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
|
||||
const blockedWarmupStarted = Promise.withResolvers<void>();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "tool ok",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 7,
|
||||
outputLines: 1,
|
||||
outputBytes: 7,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
let warmupCallCount = 0;
|
||||
const warmupSpy = vi
|
||||
.spyOn(pythonExecutor, "warmPythonEnvironment")
|
||||
.mockImplementation(async (_cwd, _sessionId, _useSharedGateway, _sessionFile, _kernelOwnerId, signal) => {
|
||||
warmupCallCount += 1;
|
||||
if (warmupCallCount === 1) {
|
||||
return { ok: true, docs: [] };
|
||||
}
|
||||
blockedWarmupStarted.resolve();
|
||||
return await new Promise<{ ok: boolean; reason?: string; docs: [] }>(resolve => {
|
||||
const onAbort = () => resolve({ ok: false, reason: "Warmup aborted", docs: [] });
|
||||
if (signal?.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
signal?.addEventListener("abort", onAbort, { once: true });
|
||||
});
|
||||
});
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
const pythonTool = session.getToolByName("python");
|
||||
expect(pythonTool).toBeDefined();
|
||||
let toolExecutionSettled = false;
|
||||
const toolExecution = pythonTool!
|
||||
.execute("call-id", { cells: [{ code: "print('tool')" }] }, undefined, undefined, undefined)
|
||||
.finally(() => {
|
||||
toolExecutionSettled = true;
|
||||
});
|
||||
await blockedWarmupStarted.promise;
|
||||
|
||||
let disposed = false;
|
||||
const disposeSession = session.dispose().then(() => {
|
||||
disposed = true;
|
||||
});
|
||||
await Bun.sleep(0);
|
||||
|
||||
expect(disposed).toBe(false);
|
||||
expect(toolExecutionSettled).toBe(false);
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(2);
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
|
||||
await expect(toolExecution).rejects.toThrow("Operation aborted");
|
||||
await disposeSession;
|
||||
|
||||
expect(disposed).toBe(true);
|
||||
expect(toolExecutionSettled).toBe(true);
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("rejects Python tool starts when warmup finishes after dispose begins", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
|
||||
const blockedWarmupStarted = Promise.withResolvers<void>();
|
||||
const releaseWarmup = Promise.withResolvers<void>();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
let warmupCallCount = 0;
|
||||
const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockImplementation(async () => {
|
||||
warmupCallCount += 1;
|
||||
if (warmupCallCount === 1) {
|
||||
return { ok: true, docs: [] };
|
||||
}
|
||||
blockedWarmupStarted.resolve();
|
||||
await releaseWarmup.promise;
|
||||
return { ok: true, docs: [] };
|
||||
});
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
const pythonTool = session.getToolByName("python");
|
||||
expect(pythonTool).toBeDefined();
|
||||
const toolExecution = pythonTool!.execute(
|
||||
"call-id",
|
||||
{ cells: [{ code: "print('tool after warmup')" }] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
);
|
||||
await blockedWarmupStarted.promise;
|
||||
const disposeSession = session.dispose();
|
||||
releaseWarmup.resolve();
|
||||
await expect(toolExecution).rejects.toThrow(
|
||||
"Python execution is unavailable while session disposal is in progress",
|
||||
);
|
||||
await disposeSession;
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(2);
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
it("aborts tracked Python tool execution during session dispose after warmup completes", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
|
||||
const helper: PreludeHelper = {
|
||||
name: "helper",
|
||||
signature: "helper()",
|
||||
docstring: "doc",
|
||||
category: "general",
|
||||
};
|
||||
let preludeDocs: PreludeHelper[] = [];
|
||||
vi.spyOn(pythonExecutor, "getPreludeDocs").mockImplementation(() => preludeDocs);
|
||||
const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockImplementation(async () => {
|
||||
preludeDocs = [helper];
|
||||
return { ok: true, docs: preludeDocs };
|
||||
});
|
||||
const blockedExecuteStarted = Promise.withResolvers<void>();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockImplementation(async (_code, options) => {
|
||||
const signal = options?.signal;
|
||||
if (!signal) {
|
||||
throw new Error("Expected abort signal");
|
||||
}
|
||||
blockedExecuteStarted.resolve();
|
||||
return await new Promise(resolve => {
|
||||
const onAbort = () =>
|
||||
resolve({
|
||||
output: "Command aborted",
|
||||
exitCode: undefined,
|
||||
cancelled: true,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 15,
|
||||
outputLines: 1,
|
||||
outputBytes: 15,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
if (signal.aborted) {
|
||||
onAbort();
|
||||
return;
|
||||
}
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
});
|
||||
});
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
const pythonTool = session.getToolByName("python");
|
||||
expect(pythonTool).toBeDefined();
|
||||
let toolExecutionSettled = false;
|
||||
const toolExecution = pythonTool!
|
||||
.execute("call-id", { cells: [{ code: "print('tool')" }] }, undefined, undefined, undefined)
|
||||
.finally(() => {
|
||||
toolExecutionSettled = true;
|
||||
});
|
||||
await blockedExecuteStarted.promise;
|
||||
|
||||
let disposed = false;
|
||||
const disposeSession = session.dispose().then(() => {
|
||||
disposed = true;
|
||||
});
|
||||
await Bun.sleep(0);
|
||||
|
||||
expect(disposed).toBe(false);
|
||||
expect(toolExecutionSettled).toBe(false);
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(1);
|
||||
expect(executeSpy).toHaveBeenCalledTimes(1);
|
||||
|
||||
const [toolResult] = await Promise.all([toolExecution, disposeSession]);
|
||||
|
||||
expect(disposed).toBe(true);
|
||||
expect(toolExecutionSettled).toBe(true);
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(1);
|
||||
expect(executeSpy).toHaveBeenCalledTimes(1);
|
||||
expect(toolResult.details?.isError).toBe(true);
|
||||
expect(toolResult.content).toContainEqual(
|
||||
expect.objectContaining({ type: "text", text: expect.stringContaining("Command aborted") }),
|
||||
);
|
||||
});
|
||||
|
||||
it("detaches retained kernel ownership even when dispose times out waiting for Python work", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
stubPythonWarmup();
|
||||
|
||||
const kernel = new FakeKernel();
|
||||
const blockedExecution = Promise.withResolvers<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
kernel.blockedCode = "print('blocked')";
|
||||
kernel.blockedExecution = blockedExecution.promise;
|
||||
kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve();
|
||||
kernel.abortBlockedExecution = false;
|
||||
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
const startSpy = vi
|
||||
.spyOn(pythonKernel.PythonKernel, "start")
|
||||
.mockResolvedValue(kernel as unknown as PythonKernelInstance);
|
||||
|
||||
const firstSession = await createSession(tempDir, cwd);
|
||||
const secondSession = await createSession(tempDir, cwd);
|
||||
|
||||
await secondSession.executePython("print('owner-b warmup')");
|
||||
const firstExecution = firstSession.executePython("print('blocked')");
|
||||
await blockedExecutionStarted.promise;
|
||||
let firstExecutionSettled = false;
|
||||
void firstExecution.finally(() => {
|
||||
firstExecutionSettled = true;
|
||||
});
|
||||
|
||||
let firstDisposed = false;
|
||||
const disposeFirst = firstSession.dispose().then(() => {
|
||||
firstDisposed = true;
|
||||
});
|
||||
await disposeFirst;
|
||||
|
||||
expect(firstDisposed).toBe(true);
|
||||
expect(firstExecutionSettled).toBe(false);
|
||||
expect(kernel.shutdownCalls).toBe(0);
|
||||
expect(startSpy).toHaveBeenCalledTimes(1);
|
||||
|
||||
blockedExecution.resolve(OK_EXECUTION);
|
||||
await expect(firstExecution).resolves.toMatchObject({
|
||||
cancelled: false,
|
||||
exitCode: 0,
|
||||
stdinRequested: false,
|
||||
});
|
||||
await secondSession.executePython("print('owner-b after detach')");
|
||||
expect(startSpy).toHaveBeenCalledTimes(1);
|
||||
expect(kernel.executeCalls).toEqual([
|
||||
"print('owner-b warmup')",
|
||||
"print('blocked')",
|
||||
"print('owner-b after detach')",
|
||||
]);
|
||||
await secondSession.dispose();
|
||||
|
||||
expect(kernel.shutdownCalls).toBe(1);
|
||||
}, 10000);
|
||||
|
||||
it("rejects direct session Python starts once dispose begins", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
stubPythonWarmup();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
const disposeSession = session.dispose();
|
||||
await expect(session.executePython("print('late')")).rejects.toThrow(
|
||||
"Python execution is unavailable while session disposal is in progress",
|
||||
);
|
||||
await disposeSession;
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("rejects direct session Python starts after an async user_python hook yields during dispose", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
const hookStarted = Promise.withResolvers<void>();
|
||||
const releaseHook = Promise.withResolvers<void>();
|
||||
const hookExtension: ExtensionFactory = api => {
|
||||
api.on("user_python", async () => {
|
||||
hookStarted.resolve();
|
||||
await releaseHook.promise;
|
||||
return undefined;
|
||||
});
|
||||
};
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
|
||||
const session = await createSession(tempDir, cwd, { extensions: [hookExtension] });
|
||||
const execution = session.executePython("print('late after hook')");
|
||||
await hookStarted.promise;
|
||||
await session.dispose();
|
||||
releaseHook.resolve();
|
||||
await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress");
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("rejects async user_python hook results after dispose begins", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
const hookStarted = Promise.withResolvers<void>();
|
||||
const releaseHook = Promise.withResolvers<void>();
|
||||
const hookExtension: ExtensionFactory = api => {
|
||||
api.on("user_python", async () => {
|
||||
hookStarted.resolve();
|
||||
await releaseHook.promise;
|
||||
return {
|
||||
result: {
|
||||
output: "hooked late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 11,
|
||||
outputLines: 1,
|
||||
outputBytes: 11,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
},
|
||||
};
|
||||
});
|
||||
};
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
|
||||
const session = await createSession(tempDir, cwd, { extensions: [hookExtension] });
|
||||
const execution = session.executePython("print('late hook result')");
|
||||
await hookStarted.promise;
|
||||
await session.dispose();
|
||||
releaseHook.resolve();
|
||||
await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress");
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
expect(session.messages.some(message => message.role === "pythonExecution")).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects Python tool starts once dispose begins", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
const warmupSpy = stubPythonWarmup();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
const pythonTool = session.getToolByName("python");
|
||||
expect(pythonTool).toBeDefined();
|
||||
const disposeSession = session.dispose();
|
||||
await expect(
|
||||
pythonTool!.execute("call-id", { cells: [{ code: "print('late')" }] }, undefined, undefined, undefined),
|
||||
).rejects.toThrow("Python execution is unavailable while session disposal is in progress");
|
||||
await disposeSession;
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(1);
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("rejects Python tool starts that reach async preflight after dispose begins", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
const warmupSpy = stubPythonWarmup();
|
||||
const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({
|
||||
output: "late",
|
||||
exitCode: 0,
|
||||
cancelled: false,
|
||||
truncated: false,
|
||||
totalLines: 1,
|
||||
totalBytes: 4,
|
||||
outputLines: 1,
|
||||
outputBytes: 4,
|
||||
displayOutputs: [],
|
||||
stdinRequested: false,
|
||||
});
|
||||
const artifactStarted = Promise.withResolvers<void>();
|
||||
const releaseArtifact = Promise.withResolvers<void>();
|
||||
const sessionManager = SessionManager.inMemory(cwd);
|
||||
vi.spyOn(sessionManager, "allocateArtifactPath").mockImplementation(async () => {
|
||||
artifactStarted.resolve();
|
||||
await releaseArtifact.promise;
|
||||
return {};
|
||||
});
|
||||
|
||||
const session = await createSession(tempDir, cwd, { sessionManager });
|
||||
const pythonTool = session.getToolByName("python");
|
||||
expect(pythonTool).toBeDefined();
|
||||
const execution = pythonTool!.execute(
|
||||
"call-id",
|
||||
{ cells: [{ code: "print('late after artifact')" }] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
);
|
||||
await artifactStarted.promise;
|
||||
const disposeSession = session.dispose();
|
||||
releaseArtifact.resolve();
|
||||
await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress");
|
||||
await disposeSession;
|
||||
expect(warmupSpy).toHaveBeenCalledTimes(1);
|
||||
expect(executeSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("aborts every active Python execution owned by the session during dispose", async () => {
|
||||
const { tempDir, cwd } = createTempProject();
|
||||
tempDirs.push(tempDir);
|
||||
pythonExecutor.resetPreludeDocsCache();
|
||||
stubPythonWarmup();
|
||||
|
||||
const kernel = new FakeKernel();
|
||||
const blockedExecution = Promise.withResolvers<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
kernel.blockedCode = "print('first')";
|
||||
kernel.blockedExecution = blockedExecution.promise;
|
||||
kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve();
|
||||
|
||||
vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true });
|
||||
vi.spyOn(pythonKernel.PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance);
|
||||
|
||||
const session = await createSession(tempDir, cwd);
|
||||
|
||||
const firstExecution = session.executePython("print('first')");
|
||||
await blockedExecutionStarted.promise;
|
||||
const secondExecution = session.executePython("print('second')");
|
||||
await Bun.sleep(0);
|
||||
|
||||
await session.dispose();
|
||||
const [firstResult, secondResult] = await Promise.all([firstExecution, secondExecution]);
|
||||
|
||||
expect(firstResult.cancelled).toBe(true);
|
||||
expect(secondResult.cancelled).toBe(true);
|
||||
expect(kernel.executeCalls).toEqual(["print('first')"]);
|
||||
expect(kernel.shutdownCalls).toBe(1);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user