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

package/src/eval/rb/executor.ts

@@ -1,22 +1,26 @@
 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,
 	executeWithKernelBase,
 	getExecutionDeadlineMs,
-	getRemainingTimeoutMs,
 	isCancellationError,
 	isTimedOutCancellation,
-	resolveOwnerScopedSessionKey,
 	waitForPromiseWithCancellation,
 } from "../executor-base";
 import type { JsStatusEvent } from "../js/shared/types";
+import {
+	createKernelSessionRegistry,
+	formatSessionKernelTimeoutAnnotation,
+	formatSessionTimeoutAnnotation,
+	type KernelSession,
+	normalizeKernelSessionCwd,
+	requireRemainingKernelTimeoutMs,
+} from "../kernel-session-registry";
 import { ensurePyToolBridge } from "../py/tool-bridge";
 import {
 	checkRubyKernelAvailability,
@@ -97,39 +101,6 @@ export interface RubyResult {
 	stdinRequested: boolean;
 }
 
-// ---------------------------------------------------------------------------
-// Session bookkeeping
-//
-// One RubyKernel subprocess per (session id, cwd, interpreter) tuple. The
-// runner mutates process-global cwd/$LOAD_PATH/ENV during execution, so
-// cross-directory work must never share a live kernel. Multiple agent owners can
-// register against the same tuple; the kernel stays alive until the last owner detaches.
-// ---------------------------------------------------------------------------
-
-interface RubySessionOwners {
-	ownerIds: Set<string>;
-	hasFallbackOwner: boolean;
-}
-
-interface RubySession extends RubySessionOwners {
-	sessionKey: string;
-	sessionId: string;
-	cwd: string;
-	kernel: RubyKernel;
-}
-
-interface StartingRubySession extends RubySessionOwners {
-	promise: Promise<RubySession>;
-}
-
-const sessions = new Map<string, RubySession>();
-const startingSessions = new Map<string, StartingRubySession>();
-const resettingSessions = new Map<string, Promise<void>>();
-
-function normalizeSessionCwd(cwd: string): string {
-	return path.resolve(cwd);
-}
-
 function normalizeExplicitInterpreter(cwd: string, interpreter: string | undefined): string {
 	if (interpreter === undefined) return "";
 	const resolved = resolveExplicitRubyRuntime(interpreter, cwd, {}).rubyPath;
@@ -140,11 +111,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
 // ---------------------------------------------------------------------------
@@ -160,32 +126,16 @@ class RubyExecutionCancelledError extends Error {
 }
 
 function requireRemainingTimeoutMs(deadlineMs?: number): number | undefined {
-	const remainingMs = getRemainingTimeoutMs(deadlineMs);
-	if (remainingMs === undefined) return undefined;
-	if (remainingMs <= 0) {
-		throw new RubyExecutionCancelledError(true);
-	}
-	return remainingMs;
+	return requireRemainingKernelTimeoutMs(deadlineMs, RubyExecutionCancelledError);
 }
 
 // ---------------------------------------------------------------------------
 // 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 createCancelledRubyResult(timedOut: boolean, timeoutMs?: number): RubyResult {
 	const output = timedOut ? (formatTimeoutAnnotation(timeoutMs) ?? "Command timed out") : "";
@@ -207,161 +157,6 @@ async function startKernel(cwd: string, options: RubyExecutorOptions): Promise<R
 	});
 }
 
-async function acquireSession(
-	sessionKey: string,
-	sessionId: string,
-	cwd: string,
-	options: RubyExecutorOptions,
-): Promise<RubySession> {
-	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!: StartingRubySession;
-	const startup = (async () => {
-		const kernel = await startKernel(cwd, options);
-		const session: RubySession = {
-			sessionKey,
-			sessionId,
-			cwd,
-			kernel,
-			ownerIds: new Set(startingSession.ownerIds),
-			hasFallbackOwner: startingSession.hasFallbackOwner,
-		};
-		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: RubySession, cwd: string, options: RubyExecutorOptions): 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 RubyExecutionCancelledError(false);
-	}
-	requireRemainingTimeoutMs(options.deadlineMs);
-	const next = await startKernel(cwd, options);
-	if (sessions.get(session.sessionKey) !== session) {
-		await next.shutdown().catch(() => undefined);
-		throw new RubyExecutionCancelledError(false);
-	}
-	session.kernel = next;
-}
-
-async function resetSession(sessionKey: string): Promise<void> {
-	const existing =
-		sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.promise.catch(() => undefined));
-	if (!existing) return;
-	sessions.delete(sessionKey);
-	await existing.kernel.shutdown().catch(() => undefined);
-}
-
-// ---------------------------------------------------------------------------
-// Public dispose entry points
-// ---------------------------------------------------------------------------
-
-export async function disposeAllRubyKernelSessions(): 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) {
-		if (sessions.get(id) === session) sessions.delete(id);
-	}
-	const results = await Promise.allSettled(all.map(([, session]) => session.kernel.shutdown()));
-	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("Ruby kernel shutdown not confirmed", {
-			sessionId: session.sessionId,
-			sessionKey: id,
-			cwd: session.cwd,
-			reason,
-		});
-		if (!sessions.has(id)) sessions.set(id, session);
-	}
-}
-
-export async function disposeRubyKernelSessionsByOwner(ownerId: string): Promise<void> {
-	const toShutdown: RubySession[] = [];
-	const startingToShutdown: StartingRubySession[] = [];
-	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);
-	}
-	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) {
-		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;
-		if (sessions.get(session.sessionKey) === session) sessions.delete(session.sessionKey);
-		toShutdown.push(session);
-	}
-	const results = await Promise.allSettled(toShutdown.map(session => session.kernel.shutdown()));
-	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("Ruby 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
 // ---------------------------------------------------------------------------
@@ -406,67 +201,29 @@ async function ensureToolBridge(options: RubyExecutorOptions): Promise<void> {
 	}
 }
 
-async function executeOnSession(code: string, cwd: string, options: RubyExecutorOptions): Promise<RubyResult> {
-	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) {
-		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 {
-		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 RubyExecutionCancelledError(
-			isTimedOutCancellation(options.signal.reason, RubyExecutionCancelledError, options.signal),
-		);
-	}
-	if (sessions.get(session.sessionKey) !== session) {
-		throw new RubyExecutionCancelledError(false);
-	}
-	if (!session.kernel.isAlive()) {
-		await replaceSessionKernel(session, cwd, options);
-		if (sessions.get(session.sessionKey) !== session) {
-			throw new RubyExecutionCancelledError(false);
-		}
-	}
-	const runOptions = { ...options, cwd };
-	try {
-		return await executeWithKernel(session.kernel, code, runOptions);
-	} catch (err) {
-		if (isCancellationError(err, RubyExecutionCancelledError) || options.signal?.aborted) throw err;
-		if (session.kernel.isAlive()) throw err;
-		if (sessions.get(session.sessionKey) !== session) {
-			throw new RubyExecutionCancelledError(false);
-		}
-		await replaceSessionKernel(session, cwd, options);
-		if (sessions.get(session.sessionKey) !== session) {
-			throw new RubyExecutionCancelledError(false);
-		}
-		return await executeWithKernel(session.kernel, code, runOptions);
-	}
+const sessionRegistry = createKernelSessionRegistry<
+	RubyKernel,
+	RubyExecutorOptions,
+	RubyResult,
+	KernelSession<RubyKernel>
+>({
+	languageLabel: "Ruby",
+	cancelledErrorClass: RubyExecutionCancelledError,
+	buildSessionKey: (sessionId, cwd, interpreter) => {
+		const normalizedCwd = normalizeKernelSessionCwd(cwd);
+		return `${sessionId}\0${normalizedCwd}\0${normalizeExplicitInterpreter(normalizedCwd, interpreter)}`;
+	},
+	createSession: session => session,
+	startKernel,
+	executeWithKernel,
+});
+
+export async function disposeAllRubyKernelSessions(): Promise<void> {
+	await sessionRegistry.disposeAll();
+}
+
+export async function disposeRubyKernelSessionsByOwner(ownerId: string): Promise<void> {
+	await sessionRegistry.disposeByOwner(ownerId);
 }
 
 export async function executeRubyWithKernel(
@@ -478,7 +235,7 @@ export async function executeRubyWithKernel(
 }
 
 export async function executeRuby(code: string, options?: RubyExecutorOptions): Promise<RubyResult> {
-	const cwd = normalizeSessionCwd(options?.cwd ?? getProjectDir());
+	const cwd = normalizeKernelSessionCwd(options?.cwd ?? getProjectDir());
 	const deadlineMs = getExecutionDeadlineMs(options);
 	const executionOptions: RubyExecutorOptions = {
 		...(options ?? {}),
@@ -499,7 +256,7 @@ export async function executeRuby(code: string, options?: RubyExecutorOptions):
 		}
 		await ensureKernelAvailable(cwd, executionOptions);
 		await ensureToolBridge(executionOptions);
-		return await executeOnSession(code, cwd, executionOptions);
+		return await sessionRegistry.executeOnSession(code, cwd, executionOptions);
 	} catch (err) {
 		if (isCancellationError(err, RubyExecutionCancelledError) || executionOptions.signal?.aborted) {
 			return createCancelledRubyResult(

package/src/export/share.ts

@@ -24,7 +24,8 @@ import type { AssistantMessage, ImageContent, TextContent } from "@oh-my-pi/pi-a
 import { $which, logger } from "@oh-my-pi/pi-utils";
 import { DEFAULT_SHARE_URL } from "@oh-my-pi/pi-wire";
 import { $ } from "bun";
-import { obfuscateToolArguments, type SecretObfuscator } from "../secrets/obfuscator";
+import { obfuscateToolArguments } from "../secrets/message-transform";
+import type { SecretObfuscator } from "../secrets/obfuscator";
 import { type SessionEntry, type SessionHeader, TITLE_CHANGE_ENTRY_TYPE } from "../session/session-entries";
 import type { SessionManager } from "../session/session-manager";
 import type { OutputMeta } from "../tools/output-meta";