fix(eval): coalesced python kernel replacement

Tracked retained Python kernel generations and shared one replacement promise per dead generation. Reset and disposal now invalidate and drain replacement work before allowing a new session to take ownership. Added deterministic fake-kernel coverage for concurrent callers, cancellation, reset, owner/global disposal, and independent cwd keys. Fixes #6367
2026-07-23 19:00:08 +00:00
parent 130578aced
commit 03489d1ebe
3 changed files with 326 additions and 34 deletions
@@ -1,5 +1,9 @@
 import { afterEach, describe, expect, it } from "bun:test";
-import { disposeAllKernelSessions, executePython } from "@oh-my-pi/pi-coding-agent/eval/py/executor";
+import {
+	disposeAllKernelSessions,
+	disposeKernelSessionsByOwner,
+	executePython,
+} from "@oh-my-pi/pi-coding-agent/eval/py/executor";
 import {
 	type KernelExecuteOptions,
 	type KernelExecuteResult,
@@ -29,6 +33,10 @@ class FakeKernel {
 		return this.#alive;
 	}

+	markDead(): void {
+		this.#alive = false;
+	}
+
 	async execute(code: string, options?: KernelExecuteOptions): Promise<KernelExecuteResult> {
 		this.executeCalls.push(code);
 		this.#onExecute?.(options);
@@ -53,6 +61,10 @@ const okResult: KernelExecuteResult = {
 	stdinRequested: false,
 };

+async function flushMicrotasks(): Promise<void> {
+	for (let i = 0; i < 5; i += 1) await Promise.resolve();
+}
+
 describe("executePython session lifecycle", () => {
 	const originalStart = PythonKernel.start;

@@ -95,6 +107,215 @@ describe("executePython session lifecycle", () => {
 		expect(liveKernel.executeCalls).toEqual(["print('restart')"]);
 	});

+	it("coalesces concurrent replacement of one dead session generation", async () => {
+		const deadKernel = new FakeKernel(okResult);
+		const replacementOne = new FakeKernel(okResult);
+		const replacementTwo = new FakeKernel(okResult);
+		const replacementStarted = Promise.withResolvers<void>();
+		const releaseReplacement = Promise.withResolvers<void>();
+		const replacements = [replacementOne, replacementTwo];
+		let startCount = 0;
+
+		PythonKernel.start = async () => {
+			startCount += 1;
+			if (startCount === 1) return deadKernel as unknown as PythonKernel;
+			replacementStarted.resolve();
+			await releaseReplacement.promise;
+			return replacements.shift() as unknown as PythonKernel;
+		};
+
+		await executePython("print('setup')", { sessionId: "session-concurrent-restart" });
+		deadKernel.markDead();
+
+		const first = executePython("print('first')", { sessionId: "session-concurrent-restart" });
+		await replacementStarted.promise;
+		const second = executePython("print('second')", { sessionId: "session-concurrent-restart" });
+		await Promise.resolve();
+		await Promise.resolve();
+
+		expect(startCount).toBe(2);
+
+		releaseReplacement.resolve();
+		await Promise.all([first, second]);
+
+		expect(replacementOne.executeCalls).toEqual(["print('first')", "print('second')"]);
+		expect(replacementTwo.executeCalls).toEqual([]);
+
+		await disposeAllKernelSessions();
+		expect(replacementOne.shutdownCalls).toBe(1);
+		expect(replacementTwo.shutdownCalls).toBe(0);
+	});
+
+	it("keeps a shared replacement alive when one caller cancels", async () => {
+		const deadKernel = new FakeKernel(okResult);
+		const replacement = new FakeKernel(okResult);
+		const replacementStarted = Promise.withResolvers<void>();
+		const releaseReplacement = Promise.withResolvers<void>();
+		let startCount = 0;
+
+		PythonKernel.start = async () => {
+			startCount += 1;
+			if (startCount === 1) return deadKernel as unknown as PythonKernel;
+			replacementStarted.resolve();
+			await releaseReplacement.promise;
+			return replacement as unknown as PythonKernel;
+		};
+
+		await executePython("print('setup')", { sessionId: "session-cancelled-restart" });
+		deadKernel.markDead();
+
+		const abortController = new AbortController();
+		const cancelled = executePython("print('cancelled')", {
+			sessionId: "session-cancelled-restart",
+			signal: abortController.signal,
+		});
+		await replacementStarted.promise;
+		const retained = executePython("print('retained')", { sessionId: "session-cancelled-restart" });
+		await flushMicrotasks();
+		abortController.abort(Object.assign(new Error("replacement wait cancelled"), { name: "AbortError" }));
+
+		expect((await cancelled).cancelled).toBe(true);
+		expect(startCount).toBe(2);
+
+		releaseReplacement.resolve();
+		expect((await retained).cancelled).toBe(false);
+		expect(replacement.executeCalls).toEqual(["print('retained')"]);
+	});
+
+	it("invalidates an in-flight replacement before resetting to a fresh generation", async () => {
+		const deadKernel = new FakeKernel(okResult);
+		const staleReplacement = new FakeKernel(okResult);
+		const freshKernel = new FakeKernel(okResult);
+		const replacementStarted = Promise.withResolvers<void>();
+		const releaseReplacement = Promise.withResolvers<void>();
+		let startCount = 0;
+
+		PythonKernel.start = async () => {
+			startCount += 1;
+			if (startCount === 1) return deadKernel as unknown as PythonKernel;
+			if (startCount === 2) {
+				replacementStarted.resolve();
+				await releaseReplacement.promise;
+				return staleReplacement as unknown as PythonKernel;
+			}
+			return freshKernel as unknown as PythonKernel;
+		};
+
+		await executePython("print('setup')", { sessionId: "session-reset-replacement" });
+		deadKernel.markDead();
+
+		const obsolete = executePython("print('obsolete')", { sessionId: "session-reset-replacement" });
+		await replacementStarted.promise;
+		const reset = executePython("print('reset')", {
+			sessionId: "session-reset-replacement",
+			reset: true,
+		});
+		await flushMicrotasks();
+		releaseReplacement.resolve();
+
+		expect((await obsolete).cancelled).toBe(true);
+		expect((await reset).cancelled).toBe(false);
+		expect(staleReplacement.executeCalls).toEqual([]);
+		expect(staleReplacement.shutdownCalls).toBe(1);
+		expect(freshKernel.executeCalls).toEqual(["print('reset')"]);
+
+		await executePython("print('later')", { sessionId: "session-reset-replacement" });
+		expect(startCount).toBe(3);
+		expect(freshKernel.executeCalls).toEqual(["print('reset')", "print('later')"]);
+	});
+
+	it("drains replacements invalidated by owner and global disposal", async () => {
+		const ownerKernel = new FakeKernel(okResult);
+		const globalKernel = new FakeKernel(okResult);
+		const ownerReplacement = new FakeKernel(okResult);
+		const globalReplacement = new FakeKernel(okResult);
+		const replacementsStarted = Promise.withResolvers<void>();
+		const releaseReplacements = Promise.withResolvers<void>();
+		const initialKernels = [ownerKernel, globalKernel];
+		const replacementKernels = [ownerReplacement, globalReplacement];
+		let replacementStartCount = 0;
+
+		PythonKernel.start = async () => {
+			const initial = initialKernels.shift();
+			if (initial) return initial as unknown as PythonKernel;
+			replacementStartCount += 1;
+			if (replacementStartCount === 2) replacementsStarted.resolve();
+			await releaseReplacements.promise;
+			return replacementKernels.shift() as unknown as PythonKernel;
+		};
+
+		await executePython("print('owner setup')", {
+			sessionId: "session-owner-disposal-replacement",
+			kernelOwnerId: "replacement-owner",
+		});
+		await executePython("print('global setup')", { sessionId: "session-global-disposal-replacement" });
+		ownerKernel.markDead();
+		globalKernel.markDead();
+
+		const ownerExecution = executePython("print('owner obsolete')", {
+			sessionId: "session-owner-disposal-replacement",
+			kernelOwnerId: "replacement-owner",
+		});
+		const globalExecution = executePython("print('global obsolete')", {
+			sessionId: "session-global-disposal-replacement",
+		});
+		await replacementsStarted.promise;
+
+		const ownerDisposal = disposeKernelSessionsByOwner("replacement-owner");
+		await flushMicrotasks();
+		const globalDisposal = disposeAllKernelSessions();
+		await flushMicrotasks();
+		releaseReplacements.resolve();
+
+		expect((await ownerExecution).cancelled).toBe(true);
+		expect((await globalExecution).cancelled).toBe(true);
+		await Promise.all([ownerDisposal, globalDisposal]);
+		expect(ownerReplacement.executeCalls).toEqual([]);
+		expect(globalReplacement.executeCalls).toEqual([]);
+		expect(ownerReplacement.shutdownCalls).toBe(1);
+		expect(globalReplacement.shutdownCalls).toBe(1);
+	});
+
+	it("keeps replacement coordination independent across normalized cwd keys", async () => {
+		const deadOne = new FakeKernel(okResult);
+		const deadTwo = new FakeKernel(okResult);
+		const replacementOne = new FakeKernel(okResult);
+		const replacementTwo = new FakeKernel(okResult);
+		const kernels = [deadOne, deadTwo, replacementOne, replacementTwo];
+		let startCount = 0;
+
+		PythonKernel.start = async () => {
+			startCount += 1;
+			return kernels.shift() as unknown as PythonKernel;
+		};
+
+		await executePython("print('setup one')", {
+			cwd: "/tmp/replacement-key-one",
+			sessionId: "session-independent-replacement",
+		});
+		await executePython("print('setup two')", {
+			cwd: "/tmp/replacement-key-two",
+			sessionId: "session-independent-replacement",
+		});
+		deadOne.markDead();
+		deadTwo.markDead();
+
+		await Promise.all([
+			executePython("print('one')", {
+				cwd: "/tmp/replacement-key-one",
+				sessionId: "session-independent-replacement",
+			}),
+			executePython("print('two')", {
+				cwd: "/tmp/replacement-key-two",
+				sessionId: "session-independent-replacement",
+			}),
+		]);
+
+		expect(startCount).toBe(4);
+		expect(replacementOne.executeCalls).toEqual(["print('one')"]);
+		expect(replacementTwo.executeCalls).toEqual(["print('two')"]);
+	});
+
 	it("resets the session kernel when requested", async () => {
 		const firstKernel = new FakeKernel(okResult);
 		const secondKernel = new FakeKernel(okResult);