@oh-my-pi/pi-coding-agent 17.0.9 → 17.1.0

package/src/config/settings.ts

@@ -14,6 +14,7 @@
 import * as fs from "node:fs";
 import * as os from "node:os";
 import * as path from "node:path";
+import { configureCredentialRedaction } from "@oh-my-pi/pi-ai/providers/transform-messages";
 import { configureProviderMaxInFlightRequests } from "@oh-my-pi/pi-ai/stream";
 import {
 	getAgentDbPath,
@@ -33,7 +34,9 @@ import type { ModelRole } from "../config/model-roles";
 import { loadCapability } from "../discovery";
 import { isLightTheme, setAutoThemeMapping, setColorBlindMode, setSymbolPreset } from "../modes/theme/theme";
 import { AgentStorage } from "../session/agent-storage";
+import { AUTO_IMAGE_PROVIDER_ORDER, isImageProviderId } from "../tools/image-providers";
 import { type EditMode, normalizeEditMode } from "../utils/edit-mode";
+import { isSearchProviderId, SEARCH_PROVIDER_ORDER } from "../web/search/types";
 import { withFileLock } from "./file-lock";
 import {
 	type BashInterceptorRule,
@@ -361,6 +364,7 @@ export class Settings {
 		if (options.configFiles) configFiles.push(...options.configFiles);
 		this.#configFiles = configFiles.map(file => path.resolve(this.#cwd, expandTilde(file)));
 		this.#persist = !options.inMemory && options.readOnly !== true;
+		liveSettingsInstances.add(new WeakRef(this));
 
 		if (options.overrides) {
 			for (const [key, value] of Object.entries(options.overrides)) {
@@ -534,6 +538,23 @@ export class Settings {
 		}
 	}
 
+	/** Set once this instance is discarded; background saves become no-ops. */
+	#savesCancelled = false;
+
+	/**
+	 * Drop pending debounced saves and refuse any further background writes.
+	 * Used when an instance is being discarded (test teardown): an armed timer
+	 * or a chained in-flight save on a dropped instance would otherwise fire
+	 * later and race the successor's file locks.
+	 */
+	cancelPendingSaves(): void {
+		this.#savesCancelled = true;
+		clearTimeout(this.#saveTimer);
+		this.#saveTimer = undefined;
+		clearTimeout(this.#projectSaveTimer);
+		this.#projectSaveTimer = undefined;
+	}
+
 	/**
 	 * Flush any pending saves to disk.
 	 * Call before exit to ensure all changes are persisted.
@@ -1595,6 +1616,45 @@ export class Settings {
 		delete raw["mcp.discoveryMode"];
 		delete raw["mcp.discoveryDefaultServers"];
 
+		// providers.webSearch / providers.image (single preferred provider) →
+		// providers.webSearchOrder / providers.imageOrder (priority lists). A
+		// concrete legacy choice becomes the head of the new list with every
+		// remaining provider appended in its built-in order, so the old
+		// preference stays #1 and the fallback chain is written out explicitly.
+		// "auto" (or an unknown id) just drops the key — the default chain.
+		const providerPrefsObj = raw.providers as Record<string, unknown> | undefined;
+		const migrateProviderPreference = (
+			legacyKey: string,
+			orderKey: string,
+			expand: (value: string) => string[] | undefined,
+		): void => {
+			const flatLegacyKey = `providers.${legacyKey}`;
+			const legacy = providerPrefsObj?.[legacyKey] ?? raw[flatLegacyKey];
+			if (legacy === undefined) return;
+			const existingOrder = providerPrefsObj?.[orderKey] ?? raw[`providers.${orderKey}`];
+			const orderAlreadySet = Array.isArray(existingOrder) && existingOrder.length > 0;
+			if (!orderAlreadySet && typeof legacy === "string") {
+				const expanded = expand(legacy);
+				if (expanded) {
+					const root = providerPrefsObj ?? {};
+					root[orderKey] = expanded;
+					raw.providers = root;
+				}
+			}
+			if (providerPrefsObj) delete providerPrefsObj[legacyKey];
+			delete raw[flatLegacyKey];
+		};
+		migrateProviderPreference("webSearch", "webSearchOrder", value =>
+			value !== "auto" && isSearchProviderId(value)
+				? [value, ...SEARCH_PROVIDER_ORDER.filter(id => id !== value)]
+				: undefined,
+		);
+		migrateProviderPreference("image", "imageOrder", value =>
+			value !== "auto" && isImageProviderId(value)
+				? [value, ...AUTO_IMAGE_PROVIDER_ORDER.filter(id => id !== value)]
+				: undefined,
+		);
+
 		return raw;
 	}
 
@@ -1646,7 +1706,7 @@ export class Settings {
 	}
 
 	async #saveNow(): Promise<void> {
-		if (!this.#persist || !this.#configPath) return;
+		if (this.#savesCancelled || !this.#persist || !this.#configPath) return;
 		if (this.#modified.size === 0 && this.#modifiedGlobalModelRoles.size === 0) return;
 
 		const configPath = this.#configPath;
@@ -1750,7 +1810,7 @@ export class Settings {
 	}
 
 	async #saveProjectNow(): Promise<void> {
-		if (!this.#persist || this.#modifiedProjectModelRoles.size === 0) return;
+		if (this.#savesCancelled || !this.#persist || this.#modifiedProjectModelRoles.size === 0) return;
 
 		const projectConfigPath = path.join(this.#cwd, ".omp", "config.yml");
 		const modifiedModelRoles = [...this.#modifiedProjectModelRoles];
@@ -1919,6 +1979,9 @@ const SETTING_HOOKS: Partial<Record<SettingPath, SettingHook<any>>> = {
 	"providers.maxInFlightRequests": value => {
 		configureProviderMaxInFlightRequests(validateProviderMaxInFlightRequests(value));
 	},
+	"secrets.enabled": value => {
+		configureCredentialRedaction(value === true);
+	},
 	"hindsight.bankId": () => hindsightScopeSignal.fire(),
 	"hindsight.bankIdPrefix": () => hindsightScopeSignal.fire(),
 	"hindsight.scoping": () => hindsightScopeSignal.fire(),
@@ -1978,6 +2041,13 @@ export const onHindsightScopeChanged = (cb: () => void) => hindsightScopeSignal.
 // Global Singleton
 // ═══════════════════════════════════════════════════════════════════════════
 
+/**
+ * Weak registry of every constructed instance so `resetSettingsForTest` can
+ * disarm stray background saves on isolated instances too. WeakRefs never
+ * retain instances; the set is cleared on every test reset.
+ */
+const liveSettingsInstances = new Set<WeakRef<Settings>>();
+
 let globalInstance: Settings | null = null;
 let globalInstancePromise: Promise<Settings> | null = null;
 let boundSettingsInstance: Settings | null = null;
@@ -1997,10 +2067,19 @@ export function isSettingsInitialized(): boolean {
  * @internal
  */
 export function resetSettingsForTest(): void {
+	// Disarm every constructed instance's debounced saves — including isolated
+	// (non-singleton) instances: an armed timer or chained in-flight save on a
+	// dropped instance fires mid-way through the NEXT test and races its file
+	// locks/spies (cross-file pollution).
+	for (const ref of liveSettingsInstances) {
+		ref.deref()?.cancelPendingSaves();
+	}
+	liveSettingsInstances.clear();
 	globalInstance = null;
 	globalInstancePromise = null;
 	clearBoundSettingsMethods();
 	configureProviderMaxInFlightRequests(undefined);
+	configureCredentialRedaction(false);
 }
 
 /**

package/src/dap/client.ts

@@ -260,7 +260,7 @@ export class DapClient {
 		// socket connect fails, we must not leak the detached adapter process.
 		try {
 			await waitForCondition(() => isUnixSocketReady(socketPath), timeoutMs, proc);
-			const { readable, writeSink, socket } = await connectSocket({ unix: socketPath });
+			const { readable, writeSink, socket } = await connectSocket({ unix: socketPath }, timeoutMs);
 			const client = new DapClient(adapter, cwd, proc, { readable, writeSink, socket });
 			proc.exited.then(() => client.#handleProcessExit());
 			void client.#startMessageReader();
@@ -924,10 +924,20 @@ function socketToSink(socket: Bun.Socket<undefined>): DapWriteSink {
 	};
 }
 
-/** Connect to a unix domain socket and return DAP transport streams. */
-async function connectSocket(options: { unix: string }): Promise<SocketTransport> {
-	const { promise, resolve } = Promise.withResolvers<SocketTransport>();
+/**
+ * Connect to a unix domain socket and return DAP transport streams.
+ *
+ * Rejects (rather than hanging) when the connect fails — a stat-ready but dead
+ * socket returns ECONNREFUSED, a socket removed between the readiness stat and
+ * the connect returns ENOENT, a permission mismatch returns EACCES — and when
+ * neither `open` nor an error arrives within `timeoutMs` (e.g. a TOCTOU stall).
+ * `#spawnSocketUnix`'s catch then kills the detached adapter instead of leaking
+ * it. Exported so tests can drive the reject path deterministically.
+ */
+export async function connectSocket(options: { unix: string }, timeoutMs: number): Promise<SocketTransport> {
+	const { promise, resolve, reject } = Promise.withResolvers<SocketTransport>();
 	let streamController: ReadableStreamDefaultController<Uint8Array>;
+	let opened = false;
 
 	const readable = new ReadableStream<Uint8Array>({
 		start(controller) {
@@ -935,10 +945,21 @@ async function connectSocket(options: { unix: string }): Promise<SocketTransport
 		},
 	});
 
+	const timer = setTimeout(() => {
+		reject(new Error(`Timed out connecting to unix socket ${options.unix} after ${timeoutMs}ms`));
+	}, timeoutMs);
+	// A late socket callback after settle is a no-op; clearing the timer just
+	// stops it from keeping the event loop alive past the connect.
+	void promise.then(
+		() => clearTimeout(timer),
+		() => clearTimeout(timer),
+	);
+
 	Bun.connect({
 		unix: options.unix,
 		socket: {
 			open(socket) {
+				opened = true;
 				resolve({
 					readable,
 					writeSink: socketToSink(socket),
@@ -949,6 +970,9 @@ async function connectSocket(options: { unix: string }): Promise<SocketTransport
 				streamController.enqueue(new Uint8Array(data));
 			},
 			close() {
+				if (!opened) {
+					reject(new Error(`Unix socket ${options.unix} closed before opening`));
+				}
 				try {
 					streamController.close();
 				} catch {
@@ -956,6 +980,9 @@ async function connectSocket(options: { unix: string }): Promise<SocketTransport
 				}
 			},
 			error(_socket, err) {
+				if (!opened) {
+					reject(err);
+				}
 				try {
 					streamController.error(err);
 				} catch {
@@ -963,6 +990,12 @@ async function connectSocket(options: { unix: string }): Promise<SocketTransport
 				}
 			},
 		},
+	}).catch(err => {
+		// Bun.connect rejects the returned promise on synchronous connect
+		// failures (e.g. ENOENT) without always firing the `error` handler.
+		if (!opened) {
+			reject(err);
+		}
 	});
 
 	return promise;

package/src/edit/index.ts

@@ -3,7 +3,7 @@ import hashlineGrammar from "@oh-my-pi/hashline/grammar.lark" with { type: "text
 import hashlineDescription from "@oh-my-pi/hashline/prompt.md" with { type: "text" };
 import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import type { ToolExample } from "@oh-my-pi/pi-ai";
-import { prompt } from "@oh-my-pi/pi-utils";
+import { isEnoent, isEnotdir, prompt } from "@oh-my-pi/pi-utils";
 import { createLspWritethrough, flushLspWritethroughBatch, type WritethroughCallback, writethroughNoop } from "../lsp";
 import { DeferredDiagnostics } from "../lsp/deferred-diagnostics";
 import { getDiagnosticsLedger } from "../lsp/diagnostics-ledger";
@@ -12,7 +12,8 @@ import patchDescription from "../prompts/tools/patch.md" with { type: "text" };
 import replaceDescription from "../prompts/tools/replace.md" with { type: "text" };
 import type { ToolSession } from "../tools";
 import { truncateForPrompt } from "../tools/approval";
-import { isInternalUrlPath } from "../tools/path-utils";
+import { findUniqueWorkspaceSuffix, isInternalUrlPath } from "../tools/path-utils";
+import { resolvePlanPath } from "../tools/plan-mode-guard";
 import { type EditMode, normalizeEditMode, resolveEditMode } from "../utils/edit-mode";
 import { executeHashlineSingle, hashlineEditParamsSchema } from "./hashline";
 import { type ApplyPatchParams, applyPatchSchema, expandApplyPatchToEntries } from "./modes/apply-patch";
@@ -73,6 +74,24 @@ function resolveConfiguredEditMode(rawEditMode: string): EditMode | undefined {
 	return editMode;
 }
 
+async function resolveEditPath(
+	session: ToolSession,
+	authoredPath: string,
+	options: { mustExist: boolean; signal?: AbortSignal },
+): Promise<string> {
+	if (!options.mustExist || isInternalUrlPath(authoredPath)) return authoredPath;
+
+	try {
+		await Bun.file(resolvePlanPath(session, authoredPath)).stat();
+		return authoredPath;
+	} catch (error) {
+		if (!isEnoent(error) && !isEnotdir(error)) throw error;
+	}
+
+	const match = await findUniqueWorkspaceSuffix(authoredPath, session.cwd, options.signal);
+	return match?.displayPath ?? authoredPath;
+}
+
 function resolveAllowFuzzy(session: ToolSession, rawValue: string): boolean {
 	switch (rawValue) {
 		case "true":
@@ -517,7 +536,7 @@ export class EditTool implements AgentTool<TInput> {
 						note: "All entries in one call apply to the top-level `path`; use separate calls for different files.",
 					},
 				] satisfies readonly ToolExample<PatchParams>[],
-				execute: (
+				execute: async (
 					tool: EditTool,
 					params: EditParams,
 					signal: AbortSignal | undefined,
@@ -525,11 +544,15 @@ export class EditTool implements AgentTool<TInput> {
 					onUpdate?: (partialResult: AgentToolResult<EditToolDetails, TInput>) => void,
 				) => {
 					const { edits, path } = params as PatchParams;
+					const targetPath = await resolveEditPath(tool.session, path, {
+						mustExist: (edits[0]?.op ?? "update") !== "create",
+						signal,
+					});
 					const runs = (edits as PatchEditEntry[]).map(
 						entry => (br: LspBatchRequest | undefined) =>
 							executePatchSingle({
 								session: tool.session,
-								path,
+								path: targetPath,
 								params: entry,
 								signal,
 								batchRequest: br,
@@ -542,7 +565,7 @@ export class EditTool implements AgentTool<TInput> {
 								beginDeferredDiagnosticsForPath: p => tool.#deferredDiagnostics.begin(p),
 							}),
 					);
-					return executeSinglePathEntries(path, runs, batchRequest, onUpdate, tool.session.cwd, signal);
+					return executeSinglePathEntries(targetPath, runs, batchRequest, onUpdate, tool.session.cwd, signal);
 				},
 			},
 			apply_patch: {
@@ -564,14 +587,26 @@ export class EditTool implements AgentTool<TInput> {
 					onUpdate?: (partialResult: AgentToolResult<EditToolDetails, TInput>) => void,
 				) => {
 					const entries = expandApplyPatchToEntries(params as ApplyPatchParams);
+					// Resolve each authored path once per patch so paired hunks (e.g. delete
+					// then re-add of the same file) share the same workspace target.
+					const resolvedTargets = new Map<string, Promise<string>>();
+					const resolveOnce = (path: string, mustExist: boolean): Promise<string> => {
+						let pending = resolvedTargets.get(path);
+						if (!pending) {
+							pending = resolveEditPath(tool.session, path, { mustExist, signal });
+							resolvedTargets.set(path, pending);
+						}
+						return pending;
+					};
 					const perFile = entries.map(entry => {
 						const { path, ...patchParams } = entry;
 						return {
 							path,
-							run: (br: LspBatchRequest | undefined) =>
-								executePatchSingle({
+							run: async (br: LspBatchRequest | undefined) => {
+								const targetPath = await resolveOnce(path, patchParams.op !== "create");
+								return executePatchSingle({
 									session: tool.session,
-									path,
+									path: targetPath,
 									params: patchParams,
 									signal,
 									batchRequest: br,
@@ -579,7 +614,8 @@ export class EditTool implements AgentTool<TInput> {
 									fuzzyThreshold: tool.#fuzzyThreshold,
 									writethrough: tool.#writethrough,
 									beginDeferredDiagnosticsForPath: p => tool.#deferredDiagnostics.begin(p),
-								}),
+								});
+							},
 						};
 					});
 					return executeApplyPatchPerFile(perFile, batchRequest, tool.session.cwd, signal, onUpdate);
@@ -609,7 +645,7 @@ export class EditTool implements AgentTool<TInput> {
 			replace: {
 				description: () => prompt.render(replaceDescription),
 				parameters: replaceEditSchema,
-				execute: (
+				execute: async (
 					tool: EditTool,
 					params: EditParams,
 					signal: AbortSignal | undefined,
@@ -617,11 +653,12 @@ export class EditTool implements AgentTool<TInput> {
 					onUpdate?: (partialResult: AgentToolResult<EditToolDetails, TInput>) => void,
 				) => {
 					const { edits, path } = params as ReplaceParams;
+					const targetPath = await resolveEditPath(tool.session, path, { mustExist: true, signal });
 					const runs = (edits as ReplaceEditEntry[]).map(
 						entry => (br: LspBatchRequest | undefined) =>
 							executeReplaceSingle({
 								session: tool.session,
-								path,
+								path: targetPath,
 								params: entry,
 								signal,
 								batchRequest: br,
@@ -631,7 +668,7 @@ export class EditTool implements AgentTool<TInput> {
 								beginDeferredDiagnosticsForPath: p => tool.#deferredDiagnostics.begin(p),
 							}),
 					);
-					return executeSinglePathEntries(path, runs, batchRequest, onUpdate, tool.session.cwd, signal);
+					return executeSinglePathEntries(targetPath, runs, batchRequest, onUpdate, tool.session.cwd, signal);
 				},
 			},
 		}[this.mode];

package/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) {
-		throw new PythonExecutionCancelledError(false);
+): 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);
 	}
-	requireRemainingTimeoutMs(options.deadlineMs);
-	const next = await startKernel(cwd, options);
-	if (sessions.get(session.sessionKey) !== session) {
-		await next.shutdown().catch(() => undefined);
+	if (sessions.get(session.sessionKey) !== session || session.generation !== generation || session.kernel !== kernel) {
 		throw new PythonExecutionCancelledError(false);
 	}
-	session.kernel = next;
+
+	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);
+	}
+	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) {
-			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) {
+		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);
 		}
-		return await executeWithKernel(session.kernel, code, runOptions);
+		return await executeWithKernel(retryKernel, code, runOptions);
 	}
 }