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

packages/coding-agent/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- Fixed concurrent dead-kernel replacement in persistent Python sessions starting multiple generations and orphaning the losing process; callers now share one generation-scoped replacement, while resets and disposal invalidate and drain stale replacements ([#6367](https://github.com/can1357/oh-my-pi/issues/6367)).
 - Fixed credential-shaped tokens (GitHub/GitLab/OpenAI/Anthropic key patterns) being redacted from outbound provider requests even with `secrets.enabled` off; the pattern redaction now follows the `secrets.enabled` ("Hide Secrets") setting like the secret obfuscator.
 - Fixed Ctrl-clicking a wrapped OAuth authorization URL opening only the clicked row's truncated fragment by preserving the complete hyperlink target on every rendered row.
 - Fixed used-only absolute usage amounts across output surfaces: CLI now renders `$123.45 used`; the TUI shows a neutral, width-bounded amount instead of a pending/dotted/account-count placeholder; and ACP preserves `123.45 usd used` while suppressing duplicate window suffixes such as `— extra`. ([#5575](https://github.com/can1357/oh-my-pi/issues/5575))

packages/coding-agent/src/eval/py/executor.ts

@@ -21,6 +21,7 @@ import {
 	type KernelDisplayOutput,
 	type KernelExecuteOptions,
 	type KernelExecuteResult,
+	type KernelShutdownResult,
 	PythonKernel,
 } from "./kernel";
 import { resolveExplicitPythonRuntime } from "./runtime";
@@ -136,11 +137,19 @@ export interface PythonResult {
 // the same tuple; the kernel stays alive until the last owner detaches.
 // ---------------------------------------------------------------------------
 
+interface SessionKernelReplacement {
+	generation: number;
+	deadlineMs?: number;
+	promise: Promise<PythonKernel>;
+}
+
 interface PythonSession {
 	sessionKey: string;
 	sessionId: string;
 	cwd: string;
 	kernel: PythonKernel;
+	generation: number;
+	replacement?: SessionKernelReplacement;
 	ownerIds: Set<string>;
 	hasFallbackOwner: boolean;
 }
@@ -254,6 +263,7 @@ async function acquireSession(
 			sessionId,
 			cwd,
 			kernel,
+			generation: 0,
 			ownerIds: new Set(),
 			hasFallbackOwner: false,
 		};
@@ -272,31 +282,99 @@ async function acquireSession(
 
 async function replaceSessionKernel(
 	session: PythonSession,
+	kernel: PythonKernel,
+	generation: number,
 	cwd: string,
 	options: PythonExecutorOptions,
-): Promise<void> {
-	const old = session.kernel;
-	const remaining = getRemainingTimeoutMs(options.deadlineMs);
-	await old
-		.shutdown(remaining !== undefined ? { timeoutMs: Math.max(0, remaining) } : undefined)
-		.catch(() => undefined);
-	if (sessions.get(session.sessionKey) !== session) {
+): Promise<PythonKernel> {
+	const inFlight = session.replacement;
+	if (inFlight?.generation === generation) {
+		if (
+			inFlight.deadlineMs !== undefined &&
+			(options.deadlineMs === undefined || options.deadlineMs > inFlight.deadlineMs)
+		) {
+			inFlight.deadlineMs = options.deadlineMs;
+		}
+		return await waitForPromiseWithCancellation(inFlight.promise, options, PythonExecutionCancelledError);
+	}
+	if (sessions.get(session.sessionKey) !== session || session.generation !== generation || session.kernel !== kernel) {
 		throw new PythonExecutionCancelledError(false);
 	}
-	requireRemainingTimeoutMs(options.deadlineMs);
-	const next = await startKernel(cwd, options);
-	if (sessions.get(session.sessionKey) !== session) {
-		await next.shutdown().catch(() => undefined);
-		throw new PythonExecutionCancelledError(false);
+
+	const deferred = Promise.withResolvers<PythonKernel>();
+	const replacement: SessionKernelReplacement = {
+		generation,
+		deadlineMs: options.deadlineMs,
+		promise: deferred.promise,
+	};
+	session.replacement = replacement;
+	void (async () => {
+		try {
+			const remaining = getRemainingTimeoutMs(options.deadlineMs);
+			await kernel
+				.shutdown(remaining !== undefined ? { timeoutMs: Math.max(0, remaining) } : undefined)
+				.catch(() => undefined);
+			if (replacement.deadlineMs !== undefined && replacement.deadlineMs <= Date.now()) {
+				throw new PythonExecutionCancelledError(true);
+			}
+			if (
+				sessions.get(session.sessionKey) !== session ||
+				session.generation !== generation ||
+				session.kernel !== kernel
+			) {
+				throw new PythonExecutionCancelledError(false);
+			}
+			const next = await startKernel(cwd, {
+				...options,
+				signal: undefined,
+				deadlineMs: undefined,
+			});
+			if (
+				sessions.get(session.sessionKey) !== session ||
+				session.generation !== generation ||
+				session.kernel !== kernel
+			) {
+				await next.shutdown().catch(() => undefined);
+				throw new PythonExecutionCancelledError(false);
+			}
+			session.kernel = next;
+			session.generation += 1;
+			deferred.resolve(next);
+		} catch (err) {
+			deferred.reject(err);
+		} finally {
+			if (session.replacement === replacement) session.replacement = undefined;
+		}
+	})();
+	return await waitForPromiseWithCancellation(deferred.promise, options, PythonExecutionCancelledError);
+}
+
+async function shutdownInvalidatedSession(session: PythonSession): Promise<KernelShutdownResult> {
+	const replacement = session.replacement;
+	if (replacement) await replacement.promise.catch(() => undefined);
+	return await session.kernel.shutdown();
+}
+
+async function acquireLiveSessionKernel(
+	session: PythonSession,
+	cwd: string,
+	options: PythonExecutorOptions,
+): Promise<PythonKernel> {
+	while (sessions.get(session.sessionKey) === session) {
+		const kernel = session.kernel;
+		const generation = session.generation;
+		if (kernel.isAlive()) return kernel;
+		await replaceSessionKernel(session, kernel, generation, cwd, options);
 	}
-	session.kernel = next;
+	throw new PythonExecutionCancelledError(false);
 }
 
 async function resetSession(sessionKey: string): Promise<void> {
 	const existing = sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.catch(() => undefined));
 	if (!existing) return;
+	existing.generation += 1;
 	sessions.delete(sessionKey);
-	await existing.kernel.shutdown().catch(() => undefined);
+	await shutdownInvalidatedSession(existing).catch(() => undefined);
 }
 
 // ---------------------------------------------------------------------------
@@ -315,9 +393,10 @@ export async function disposeAllKernelSessions(): Promise<void> {
 		}
 	}
 	for (const [id, session] of all) {
+		session.generation += 1;
 		if (sessions.get(id) === session) sessions.delete(id);
 	}
-	const results = await Promise.allSettled(all.map(([, session]) => session.kernel.shutdown()));
+	const results = await Promise.allSettled(all.map(([, session]) => shutdownInvalidatedSession(session)));
 	for (let i = 0; i < all.length; i += 1) {
 		const [id, session] = all[i];
 		const result = results[i];
@@ -344,9 +423,10 @@ export async function disposeKernelSessionsByOwner(ownerId: string): Promise<voi
 		session.ownerIds.delete(ownerId);
 	}
 	for (const session of toShutdown) {
+		session.generation += 1;
 		if (sessions.get(session.sessionKey) === session) sessions.delete(session.sessionKey);
 	}
-	const results = await Promise.allSettled(toShutdown.map(session => session.kernel.shutdown()));
+	const results = await Promise.allSettled(toShutdown.map(session => shutdownInvalidatedSession(session)));
 	for (let i = 0; i < toShutdown.length; i += 1) {
 		const session = toShutdown[i];
 		const result = results[i];
@@ -458,31 +538,21 @@ async function executeOnSession(code: string, cwd: string, options: PythonExecut
 			isTimedOutCancellation(options.signal.reason, PythonExecutionCancelledError, options.signal),
 		);
 	}
-	if (sessions.get(session.sessionKey) !== session) {
+	const kernel = await acquireLiveSessionKernel(session, cwd, options);
+	if (sessions.get(session.sessionKey) !== session || session.kernel !== kernel) {
 		throw new PythonExecutionCancelledError(false);
 	}
-	if (!session.kernel.isAlive()) {
-		await replaceSessionKernel(session, cwd, options);
-		if (sessions.get(session.sessionKey) !== session) {
-			throw new PythonExecutionCancelledError(false);
-		}
-	}
 	const runOptions = { ...options, cwd };
 	try {
-		return await executeWithKernel(session.kernel, code, runOptions);
+		return await executeWithKernel(kernel, code, runOptions);
 	} catch (err) {
 		if (isCancellationError(err, PythonExecutionCancelledError) || options.signal?.aborted) throw err;
-		if (session.kernel.isAlive()) throw err;
-		if (sessions.get(session.sessionKey) !== session) {
+		if (kernel.isAlive()) throw err;
+		const retryKernel = await acquireLiveSessionKernel(session, cwd, options);
+		if (sessions.get(session.sessionKey) !== session || session.kernel !== retryKernel) {
 			throw new PythonExecutionCancelledError(false);
 		}
-		// Shared kernels are keyed by cwd, so a dead kernel can be recreated in place
-		// without risking cross-directory state bleed.
-		await replaceSessionKernel(session, cwd, options);
-		if (sessions.get(session.sessionKey) !== session) {
-			throw new PythonExecutionCancelledError(false);
-		}
-		return await executeWithKernel(session.kernel, code, runOptions);
+		return await executeWithKernel(retryKernel, code, runOptions);
 	}
 }
 

packages/coding-agent/test/core/python-executor.lifecycle.test.ts

@@ -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);