@oh-my-pi/pi-coding-agent 17.2.10 → 17.2.12

package/src/eval/py/executor.ts

@@ -1,10 +1,8 @@
 import * as fs from "node:fs";
-import * as path from "node:path";
 
 import { getProjectDir, logger } from "@oh-my-pi/pi-utils";
 import type { ToolSession } from "../../tools";
 import {
-	attachSessionOwner,
 	buildManagedKernelEnv,
 	buildManagedKernelEnvPatch,
 	createCancelledKernelResult,
@@ -13,11 +11,18 @@ import {
 	getRemainingTimeoutMs,
 	isCancellationError,
 	isTimedOutCancellation,
-	resolveOwnerScopedSessionKey,
-	type SessionOwners,
 	waitForPromiseWithCancellation,
 } from "../executor-base";
 import type { JsStatusEvent } from "../js/shared/types";
+import {
+	createKernelSessionRegistry,
+	formatSessionKernelTimeoutAnnotation,
+	formatSessionTimeoutAnnotation,
+	type KernelSession,
+	type KernelSessionRegistryContext,
+	normalizeKernelSessionCwd,
+	requireRemainingKernelTimeoutMs,
+} from "../kernel-session-registry";
 import {
 	checkPythonKernelAvailability,
 	type KernelDisplayOutput,
@@ -145,27 +150,9 @@ interface SessionKernelReplacement {
 	promise: Promise<PythonKernel>;
 }
 
-interface PythonSession {
-	sessionKey: string;
-	sessionId: string;
-	cwd: string;
-	kernel: PythonKernel;
+interface PythonSession extends KernelSession<PythonKernel> {
 	generation: number;
 	replacement?: SessionKernelReplacement;
-	ownerIds: Set<string>;
-	hasFallbackOwner: boolean;
-}
-
-interface StartingPythonSession extends SessionOwners {
-	promise: Promise<PythonSession>;
-}
-
-const sessions = new Map<string, PythonSession>();
-const startingSessions = new Map<string, StartingPythonSession>();
-const resettingSessions = new Map<string, Promise<void>>();
-
-function normalizeSessionCwd(cwd: string): string {
-	return path.resolve(cwd);
 }
 
 function normalizeExplicitInterpreter(cwd: string, interpreter: string | undefined): string {
@@ -178,11 +165,6 @@ function normalizeExplicitInterpreter(cwd: string, interpreter: string | undefin
 	}
 }
 
-function buildSessionKey(sessionId: string, cwd: string, interpreter: string | undefined): string {
-	const normalizedCwd = normalizeSessionCwd(cwd);
-	return `${sessionId}\0${normalizedCwd}\0${normalizeExplicitInterpreter(normalizedCwd, interpreter)}`;
-}
-
 // ---------------------------------------------------------------------------
 // Cancellation plumbing
 // ---------------------------------------------------------------------------
@@ -198,32 +180,16 @@ class PythonExecutionCancelledError extends Error {
 }
 
 function requireRemainingTimeoutMs(deadlineMs?: number): number | undefined {
-	const remainingMs = getRemainingTimeoutMs(deadlineMs);
-	if (remainingMs === undefined) return undefined;
-	if (remainingMs <= 0) {
-		throw new PythonExecutionCancelledError(true);
-	}
-	return remainingMs;
+	return requireRemainingKernelTimeoutMs(deadlineMs, PythonExecutionCancelledError);
 }
 
 // ---------------------------------------------------------------------------
 // Result formatting
 // ---------------------------------------------------------------------------
 
-function formatTimeoutAnnotation(timeoutMs?: number): string | undefined {
-	if (timeoutMs === undefined) return "Command timed out";
-	const secs = Math.max(1, Math.round(timeoutMs / 1000));
-	return `Command timed out after ${secs} seconds`;
-}
+const formatTimeoutAnnotation = formatSessionTimeoutAnnotation;
 
-function formatKernelTimeoutAnnotation(timeoutMs: number | undefined, kernelKilled: boolean): string {
-	const secs = timeoutMs === undefined ? undefined : Math.max(1, Math.round(timeoutMs / 1000));
-	if (kernelKilled) {
-		return "eval cell timed out and the kernel was unresponsive to interrupt; the kernel has been killed and will be recreated on the next call.";
-	}
-	const duration = secs === undefined ? "the configured timeout" : `${secs}s`;
-	return `eval cell timed out after ${duration}; kernel interrupted but remains running. Reset the kernel via { reset: true } if state appears corrupted.`;
-}
+const formatKernelTimeoutAnnotation = formatSessionKernelTimeoutAnnotation;
 
 function createCancelledPythonResult(timedOut: boolean, timeoutMs?: number): PythonResult {
 	const output = timedOut ? (formatTimeoutAnnotation(timeoutMs) ?? "Command timed out") : "";
@@ -245,63 +211,14 @@ async function startKernel(cwd: string, options: PythonExecutorOptions): Promise
 	});
 }
 
-async function acquireSession(
-	sessionKey: string,
-	sessionId: string,
-	cwd: string,
-	options: PythonExecutorOptions,
-): Promise<PythonSession> {
-	const existing = sessions.get(sessionKey);
-	if (existing) {
-		attachSessionOwner(existing, sessionId, options.kernelOwnerId);
-		return existing;
-	}
-	const starting = startingSessions.get(sessionKey);
-	if (starting) {
-		attachSessionOwner(starting, sessionId, options.kernelOwnerId);
-		return await starting.promise;
-	}
-	let startingSession!: StartingPythonSession;
-	const startup = (async () => {
-		const kernel = await startKernel(cwd, options);
-		const session: PythonSession = {
-			sessionKey,
-			sessionId,
-			cwd,
-			kernel,
-			generation: 0,
-			ownerIds: new Set(startingSession.ownerIds),
-			hasFallbackOwner: startingSession.hasFallbackOwner,
-		};
-		// Publish only while this startup still owns the key: owner disposal or
-		// a concurrent dispose-all may have already reaped the starting record,
-		// and publishing here would resurrect a kernel that was just torn down.
-		if (startingSessions.get(sessionKey) === startingSession) {
-			sessions.set(sessionKey, session);
-		}
-		return session;
-	})();
-	startingSession = {
-		ownerIds: new Set(),
-		hasFallbackOwner: false,
-		promise: startup,
-	};
-	attachSessionOwner(startingSession, sessionId, options.kernelOwnerId);
-	startingSessions.set(sessionKey, startingSession);
-	try {
-		return await startup;
-	} finally {
-		if (startingSessions.get(sessionKey) === startingSession) startingSessions.delete(sessionKey);
-	}
-}
-
 async function replaceSessionKernel(
 	session: PythonSession,
-	kernel: PythonKernel,
-	generation: number,
 	cwd: string,
 	options: PythonExecutorOptions,
+	context: KernelSessionRegistryContext<PythonKernel, PythonExecutorOptions, PythonSession>,
 ): Promise<PythonKernel> {
+	const kernel = session.kernel;
+	const generation = session.generation;
 	const inFlight = session.replacement;
 	if (inFlight?.generation === generation) {
 		if (
@@ -312,7 +229,11 @@ async function replaceSessionKernel(
 		}
 		return await waitForPromiseWithCancellation(inFlight.promise, options, PythonExecutionCancelledError);
 	}
-	if (sessions.get(session.sessionKey) !== session || session.generation !== generation || session.kernel !== kernel) {
+	if (
+		context.sessions.get(session.sessionKey) !== session ||
+		session.generation !== generation ||
+		session.kernel !== kernel
+	) {
 		throw new PythonExecutionCancelledError(false);
 	}
 
@@ -333,7 +254,7 @@ async function replaceSessionKernel(
 				throw new PythonExecutionCancelledError(true);
 			}
 			if (
-				sessions.get(session.sessionKey) !== session ||
+				context.sessions.get(session.sessionKey) !== session ||
 				session.generation !== generation ||
 				session.kernel !== kernel
 			) {
@@ -345,7 +266,7 @@ async function replaceSessionKernel(
 				deadlineMs: undefined,
 			});
 			if (
-				sessions.get(session.sessionKey) !== session ||
+				context.sessions.get(session.sessionKey) !== session ||
 				session.generation !== generation ||
 				session.kernel !== kernel
 			) {
@@ -374,111 +295,16 @@ async function acquireLiveSessionKernel(
 	session: PythonSession,
 	cwd: string,
 	options: PythonExecutorOptions,
+	context: KernelSessionRegistryContext<PythonKernel, PythonExecutorOptions, PythonSession>,
 ): Promise<PythonKernel> {
-	while (sessions.get(session.sessionKey) === session) {
+	while (context.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);
+		await context.replaceSessionKernel(session, cwd, options);
 	}
 	throw new PythonExecutionCancelledError(false);
 }
 
-async function resetSession(sessionKey: string): Promise<void> {
-	const existing =
-		sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.promise.catch(() => undefined));
-	if (!existing) return;
-	existing.generation += 1;
-	sessions.delete(sessionKey);
-	await shutdownInvalidatedSession(existing).catch(() => undefined);
-}
-
-// ---------------------------------------------------------------------------
-// Public dispose entry points
-// ---------------------------------------------------------------------------
-
-export async function disposeAllKernelSessions(): Promise<void> {
-	const pending = [...startingSessions.values()].map(starting => starting.promise);
-	startingSessions.clear();
-	const started = await Promise.allSettled(pending);
-	const all = [...sessions.entries()];
-	for (const result of started) {
-		if (result.status !== "fulfilled") continue;
-		if (!all.some(([, session]) => session === result.value)) {
-			all.push([result.value.sessionKey, result.value]);
-		}
-	}
-	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]) => shutdownInvalidatedSession(session)));
-	for (let i = 0; i < all.length; i += 1) {
-		const [id, session] = all[i];
-		const result = results[i];
-		if (result.status === "fulfilled" && result.value?.confirmed !== false) continue;
-		const reason = result.status === "rejected" ? result.reason : "not confirmed";
-		logger.warn("Python kernel shutdown not confirmed", {
-			sessionId: session.sessionId,
-			sessionKey: id,
-			cwd: session.cwd,
-			reason,
-		});
-		if (!sessions.has(id)) sessions.set(id, session);
-	}
-}
-
-export async function disposeKernelSessionsByOwner(ownerId: string): Promise<void> {
-	const toShutdown: PythonSession[] = [];
-	for (const session of [...sessions.values()]) {
-		if (!session.ownerIds.has(ownerId)) continue;
-		if (session.ownerIds.size === 1) {
-			toShutdown.push(session);
-			continue;
-		}
-		session.ownerIds.delete(ownerId);
-	}
-	const startingToShutdown: StartingPythonSession[] = [];
-	for (const [sessionKey, starting] of [...startingSessions.entries()]) {
-		if (sessions.has(sessionKey) || !starting.ownerIds.has(ownerId)) continue;
-		if (starting.ownerIds.size === 1) {
-			startingSessions.delete(sessionKey);
-			startingToShutdown.push(starting);
-			continue;
-		}
-		starting.ownerIds.delete(ownerId);
-	}
-	for (const session of toShutdown) {
-		session.generation += 1;
-		if (sessions.get(session.sessionKey) === session) sessions.delete(session.sessionKey);
-	}
-	const started = await Promise.allSettled(startingToShutdown.map(starting => starting.promise));
-	for (const result of started) {
-		if (result.status !== "fulfilled") continue;
-		const session = result.value;
-		session.generation += 1;
-		if (sessions.get(session.sessionKey) === session) sessions.delete(session.sessionKey);
-		toShutdown.push(session);
-	}
-	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];
-		if (result.status === "fulfilled" && result.value?.confirmed !== false) {
-			session.ownerIds.delete(ownerId);
-			continue;
-		}
-		const reason = result.status === "rejected" ? result.reason : "not confirmed";
-		logger.warn("Python kernel shutdown not confirmed", {
-			sessionId: session.sessionId,
-			sessionKey: session.sessionKey,
-			cwd: session.cwd,
-			reason,
-		});
-		if (!sessions.has(session.sessionKey)) sessions.set(session.sessionKey, session);
-	}
-}
-
 // ---------------------------------------------------------------------------
 // Execution
 // ---------------------------------------------------------------------------
@@ -535,65 +361,31 @@ async function executePerCall(code: string, cwd: string, options: PythonExecutor
 	}
 }
 
-async function executeOnSession(code: string, cwd: string, options: PythonExecutorOptions): Promise<PythonResult> {
-	const sessionId = options.sessionId ?? `session:${cwd}`;
-	const sessionKey = resolveOwnerScopedSessionKey({
-		baseKey: buildSessionKey(sessionId, cwd, options.interpreter),
-		ownerId: options.kernelOwnerId,
-		reset: options.reset === true,
-		hasSession: key => sessions.has(key) || startingSessions.has(key),
-		getOwners: key => sessions.get(key) ?? startingSessions.get(key),
-	});
-	if (options.bridge && !options.bridgeSessionId) {
-		options.bridgeSessionId = sessionId;
-	}
-	if (options.reset) {
-		// Coalesce concurrent resets: if another reset is in flight for this
-		// session, await it instead of throwing — the caller's intent ("start
-		// from a clean kernel") is satisfied once that reset settles.
-		const inFlight = resettingSessions.get(sessionKey);
-		if (inFlight) await inFlight.catch(() => undefined);
-		else {
-			const resetPromise = resetSession(sessionKey);
-			resettingSessions.set(
-				sessionKey,
-				resetPromise.then(() => undefined),
-			);
-			try {
-				await resetPromise;
-			} finally {
-				resettingSessions.delete(sessionKey);
-			}
-		}
-	} else {
-		// A reset already in progress is an internal coordination state, not a
-		// user-visible failure. Wait for it to clear, then proceed with the
-		// requested execution on the freshly-restarted kernel.
-		const inFlight = resettingSessions.get(sessionKey);
-		if (inFlight) await inFlight.catch(() => undefined);
-	}
-	const session = await acquireSession(sessionKey, sessionId, cwd, options);
-	if (options.signal?.aborted) {
-		throw new PythonExecutionCancelledError(
-			isTimedOutCancellation(options.signal.reason, PythonExecutionCancelledError, options.signal),
-		);
-	}
-	const kernel = await acquireLiveSessionKernel(session, cwd, options);
-	if (sessions.get(session.sessionKey) !== session || session.kernel !== kernel) {
-		throw new PythonExecutionCancelledError(false);
-	}
-	const runOptions = { ...options, cwd };
-	try {
-		return await executeWithKernel(kernel, code, runOptions);
-	} catch (err) {
-		if (isCancellationError(err, PythonExecutionCancelledError) || options.signal?.aborted) throw err;
-		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(retryKernel, code, runOptions);
-	}
+const sessionRegistry = createKernelSessionRegistry<PythonKernel, PythonExecutorOptions, PythonResult, PythonSession>({
+	languageLabel: "Python",
+	cancelledErrorClass: PythonExecutionCancelledError,
+	buildSessionKey: (sessionId, cwd, interpreter) => {
+		const normalizedCwd = normalizeKernelSessionCwd(cwd);
+		return `${sessionId}\0${normalizedCwd}\0${normalizeExplicitInterpreter(normalizedCwd, interpreter)}`;
+	},
+	createSession: session => ({ ...session, generation: 0 }),
+	startKernel,
+	executeWithKernel,
+	replaceSessionKernel,
+	acquireLiveSessionKernel,
+	invalidateSession: session => {
+		session.generation += 1;
+	},
+	shutdownSession: session => shutdownInvalidatedSession(session),
+	validateKernel: (session, kernel) => session.kernel === kernel,
+});
+
+export async function disposeAllKernelSessions(): Promise<void> {
+	await sessionRegistry.disposeAll();
+}
+
+export async function disposeKernelSessionsByOwner(ownerId: string): Promise<void> {
+	await sessionRegistry.disposeByOwner(ownerId);
 }
 
 export async function executePythonWithKernel(
@@ -605,7 +397,7 @@ export async function executePythonWithKernel(
 }
 
 export async function executePython(code: string, options?: PythonExecutorOptions): Promise<PythonResult> {
-	const cwd = normalizeSessionCwd(options?.cwd ?? getProjectDir());
+	const cwd = normalizeKernelSessionCwd(options?.cwd ?? getProjectDir());
 	const deadlineMs = getExecutionDeadlineMs(options);
 	const executionOptions: PythonExecutorOptions = {
 		...(options ?? {}),
@@ -631,7 +423,7 @@ export async function executePython(code: string, options?: PythonExecutorOption
 		if (kernelMode === "per-call") {
 			return await executePerCall(code, cwd, executionOptions);
 		}
-		return await executeOnSession(code, cwd, executionOptions);
+		return await sessionRegistry.executeOnSession(code, cwd, executionOptions);
 	} catch (err) {
 		if (isCancellationError(err, PythonExecutionCancelledError) || executionOptions.signal?.aborted) {
 			return createCancelledPythonResult(