diff --git a/packages/coding-agent/src/eval/executor-base.ts b/packages/coding-agent/src/eval/executor-base.ts index c89bac7bd..5f7ace0f9 100644 --- a/packages/coding-agent/src/eval/executor-base.ts +++ b/packages/coding-agent/src/eval/executor-base.ts @@ -116,9 +116,13 @@ export async function waitForPromiseWithCancellation( promise: Promise, options: { signal?: AbortSignal; deadlineMs?: number }, cancelledErrorClass: CancelledErrorClass, + timedOutResolver?: (error: unknown, signal?: AbortSignal) => boolean, ): Promise { if (options.signal?.aborted) { - throw new cancelledErrorClass(isTimedOutCancellation(options.signal.reason, cancelledErrorClass, options.signal)); + throw new cancelledErrorClass( + timedOutResolver?.(options.signal.reason, options.signal) ?? + isTimedOutCancellation(options.signal.reason, cancelledErrorClass, options.signal), + ); } const remainingMs = getRemainingTimeoutMs(options.deadlineMs); if (remainingMs !== undefined && remainingMs <= 0) { @@ -139,7 +143,8 @@ export async function waitForPromiseWithCancellation( finish(() => reject( new cancelledErrorClass( - isTimedOutCancellation(options.signal?.reason, cancelledErrorClass, options.signal), + timedOutResolver?.(options.signal?.reason, options.signal) ?? + isTimedOutCancellation(options.signal?.reason, cancelledErrorClass, options.signal), ), ), ); @@ -276,9 +281,32 @@ interface ManagedKernelEnvOptions { bridge?: { url: string; token: string }; localRoots?: Record; } +interface ManagedKernelEnvPolicy { + sparse?: boolean; +} -export function buildManagedKernelEnvPatch(options: ManagedKernelEnvOptions): Record { +export function buildManagedKernelEnvPatch(options: ManagedKernelEnvOptions): Record; +export function buildManagedKernelEnvPatch( + options: ManagedKernelEnvOptions, + policy: { sparse: true }, +): Record; +export function buildManagedKernelEnvPatch( + options: ManagedKernelEnvOptions, + policy?: ManagedKernelEnvPolicy, +): KernelEnvPatch { const localRoots = options.localRoots; + if (policy?.sparse) { + const patch: Record = {}; + if (options.sessionFile) patch.PI_SESSION_FILE = options.sessionFile; + if (options.artifactsDir) patch.PI_ARTIFACTS_DIR = options.artifactsDir; + if (options.bridge) { + patch.PI_TOOL_BRIDGE_URL = options.bridge.url; + patch.PI_TOOL_BRIDGE_TOKEN = options.bridge.token; + patch.PI_TOOL_BRIDGE_SESSION = options.bridgeSessionId ?? ""; + } + if (localRoots) patch.PI_EVAL_LOCAL_ROOTS = JSON.stringify(localRoots); + return patch; + } return { PI_SESSION_FILE: options.sessionFile ?? null, PI_ARTIFACTS_DIR: options.artifactsDir ?? null, @@ -289,13 +317,18 @@ export function buildManagedKernelEnvPatch(options: ManagedKernelEnvOptions): Re }; } -export function buildManagedKernelEnv(options: ManagedKernelEnvOptions): Record | undefined { - const patch = buildManagedKernelEnvPatch(options); +export function buildManagedKernelEnv( + options: ManagedKernelEnvOptions, + policy?: ManagedKernelEnvPolicy, +): Record | undefined { + const patch = policy?.sparse + ? buildManagedKernelEnvPatch(options, { sparse: true }) + : buildManagedKernelEnvPatch(options); const env: Record = {}; let hasKeys = false; for (const key of MANAGED_KERNEL_ENV_KEYS) { const value = patch[key]; - if (value !== null) { + if (typeof value === "string") { env[key] = value; hasKeys = true; } diff --git a/packages/coding-agent/src/eval/jl/executor.ts b/packages/coding-agent/src/eval/jl/executor.ts index 174b8803c..54dba0eca 100644 --- a/packages/coding-agent/src/eval/jl/executor.ts +++ b/packages/coding-agent/src/eval/jl/executor.ts @@ -2,11 +2,21 @@ 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, - resolveOwnerScopedSessionKey, + getExecutionDeadlineMs, + isCancellationError, + isTimedOutCancellation, + waitForPromiseWithCancellation, } from "../executor-base"; +import { + createKernelSessionRegistry, + type KernelSession, + normalizeKernelSessionCwd, + requireRemainingKernelTimeoutMs, +} from "../kernel-session-registry"; import { ensurePyToolBridge, type PyToolBridgeInfo } from "../py/tool-bridge"; import type { EvalDisplayOutput, EvalStatusEvent } from "../types"; import { @@ -58,22 +68,6 @@ export interface JuliaResult { stdinRequested: boolean; } -interface JuliaSessionOwners { - ownerIds: Set; - hasFallbackOwner: boolean; -} - -interface JuliaSession extends JuliaSessionOwners { - sessionKey: string; - sessionId: string; - cwd: string; - kernel: JuliaKernel; -} - -interface StartingJuliaSession extends JuliaSessionOwners { - promise: Promise; -} - class JuliaExecutionCancelledError extends Error { constructor(readonly timedOut: boolean) { super(timedOut ? "Julia execution timed out" : "Julia execution cancelled"); @@ -81,14 +75,6 @@ class JuliaExecutionCancelledError extends Error { } } -const sessions = new Map(); -const startingSessions = new Map(); -const resettingSessions = new Map>(); - -function normalizeSessionCwd(cwd: string): string { - return path.resolve(cwd); -} - function normalizeExplicitInterpreter(cwd: string, interpreter: string | undefined): string { if (interpreter === undefined) return ""; const resolved = resolveExplicitJuliaRuntime(interpreter, cwd, {}).juliaPath; @@ -99,85 +85,39 @@ function normalizeExplicitInterpreter(cwd: string, interpreter: string | undefin } } -function buildSessionKey(sessionId: string, cwd: string, interpreter: string | undefined): string { - const normalizedCwd = normalizeSessionCwd(cwd); - const normalizedInterpreter = normalizeExplicitInterpreter(normalizedCwd, interpreter); - return `${sessionId}::${normalizedCwd}::${normalizedInterpreter}`; +function isJuliaCancellationError(error: unknown): boolean { + return ( + isCancellationError(error, JuliaExecutionCancelledError) || + (!!error && + typeof error === "object" && + "name" in error && + (error.name === "AbortError" || error.name === "TimeoutError")) + ); } -function isCancellationError(error: unknown): boolean { - if (error instanceof JuliaExecutionCancelledError) return true; - if (error instanceof Error && (error.name === "AbortError" || error.name === "TimeoutError")) return true; - if ( - error && - typeof error === "object" && - "name" in error && - (error.name === "AbortError" || error.name === "TimeoutError") - ) - return true; - return false; +function isTimedOutJuliaCancellation(error: unknown, signal?: AbortSignal): boolean { + return ( + isTimedOutCancellation(error, JuliaExecutionCancelledError, signal) || + (!!error && typeof error === "object" && "name" in error && error.name === "TimeoutError") + ); } -function isTimedOutCancellation(error: unknown, signal?: AbortSignal): boolean { - if (error instanceof JuliaExecutionCancelledError) return error.timedOut; - if (error instanceof Error && error.name === "TimeoutError") return true; - if (error && typeof error === "object" && "name" in error && error.name === "TimeoutError") return true; - if (signal?.reason instanceof Error && signal.reason.name === "TimeoutError") return true; - return false; -} - -function getExecutionDeadlineMs(options?: Pick): number | undefined { - if (options?.deadlineMs !== undefined) return options.deadlineMs; - if (options?.timeoutMs !== undefined && options.timeoutMs > 0) return Date.now() + options.timeoutMs; - return undefined; -} - -function getRemainingTimeoutMs(deadlineMs?: number): number | undefined { - if (deadlineMs === undefined) return undefined; - return Math.max(0, deadlineMs - Date.now()); -} - -function requireRemainingTimeoutMs(deadlineMs?: number): number | undefined { - if (deadlineMs === undefined) return undefined; - const remaining = getRemainingTimeoutMs(deadlineMs); - if (remaining !== undefined && remaining <= 0) { - throw new JuliaExecutionCancelledError(true); - } - return remaining; -} - -async function waitForPromiseWithCancellation( +async function waitForJuliaPromise( promise: Promise, options: Pick, ): Promise { - if (options.signal?.aborted) { - throw new JuliaExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); - } - const cleanups: Array<() => void> = []; - const { promise: cancelPromise, reject } = Promise.withResolvers(); + const deadlineMs = + typeof options.deadlineMs === "number" && options.deadlineMs > Date.now() ? options.deadlineMs : undefined; + return await waitForPromiseWithCancellation( + promise, + { ...options, deadlineMs }, + JuliaExecutionCancelledError, + isTimedOutJuliaCancellation, + ); +} - if (options.signal) { - const onAbort = () => { - reject(new JuliaExecutionCancelledError(isTimedOutCancellation(options.signal?.reason, options.signal))); - }; - options.signal.addEventListener("abort", onAbort, { once: true }); - cleanups.push(() => options.signal?.removeEventListener("abort", onAbort)); - } - - const deadlineMs = options.deadlineMs; - if (typeof deadlineMs === "number" && deadlineMs > Date.now()) { - const timeout = setTimeout(() => { - reject(new JuliaExecutionCancelledError(true)); - }, deadlineMs - Date.now()); - timeout.unref?.(); - cleanups.push(() => clearTimeout(timeout)); - } - - try { - return await Promise.race([promise, cancelPromise]); - } finally { - for (const cleanup of cleanups) cleanup(); - } +function requireRemainingTimeoutMs(deadlineMs?: number): number | undefined { + return requireRemainingKernelTimeoutMs(deadlineMs, JuliaExecutionCancelledError); } function formatTimeoutAnnotation(timeoutMs?: number): string | undefined { @@ -198,55 +138,9 @@ function createCancelledJuliaResult(_timedOut: boolean, timeoutMs?: number): Jul return createCancelledKernelResult(output); } -function buildKernelEnvPatch(options: { - sessionFile?: string; - artifactsDir?: string; - bridge?: PyToolBridgeInfo; - bridgeSessionId?: string; - localRoots?: Record; -}): Record { - const patch: Record = {}; - if (options.sessionFile) patch.PI_SESSION_FILE = options.sessionFile; - if (options.artifactsDir) patch.PI_ARTIFACTS_DIR = options.artifactsDir; - if (options.bridge) { - patch.PI_TOOL_BRIDGE_URL = options.bridge.url; - patch.PI_TOOL_BRIDGE_TOKEN = options.bridge.token; - patch.PI_TOOL_BRIDGE_SESSION = options.bridgeSessionId ?? ""; - } - if (options.localRoots) { - patch.PI_EVAL_LOCAL_ROOTS = JSON.stringify(options.localRoots); - } - return patch; -} - -function buildKernelEnv(options: { - sessionFile?: string; - artifactsDir?: string; - bridge?: PyToolBridgeInfo; - bridgeSessionId?: string; - localRoots?: Record; -}): Record | undefined { - const patch = buildKernelEnvPatch(options); - const keys = Object.keys(patch); - if (keys.length === 0) return undefined; - const realEnv: Record = {}; - for (const key in patch) { - const val = patch[key]; - if (typeof val === "string") realEnv[key] = val; - } - return realEnv; -} - async function startKernel(cwd: string, options: JuliaExecutorOptions): Promise { requireRemainingTimeoutMs(options.deadlineMs); - const env: Record = {}; - const patch = buildKernelEnv(options); - if (patch) { - for (const key in patch) { - const value = patch[key]; - if (typeof value === "string") env[key] = value; - } - } + const env: Record = buildManagedKernelEnv(options, { sparse: true }) ?? {}; return await JuliaKernel.start({ cwd, interpreter: options.interpreter, @@ -256,163 +150,6 @@ async function startKernel(cwd: string, options: JuliaExecutorOptions): Promise< }); } -async function acquireSession( - sessionKey: string, - sessionId: string, - cwd: string, - options: JuliaExecutorOptions, -): Promise { - const existing = sessions.get(sessionKey); - if (existing) { - attachSessionOwner(existing, sessionId, options.kernelOwnerId); - return existing; - } - - const inFlight = startingSessions.get(sessionKey); - if (inFlight) { - attachSessionOwner(inFlight, sessionId, options.kernelOwnerId); - return await waitForPromiseWithCancellation(inFlight.promise, options); - } - - let startingSession!: StartingJuliaSession; - const startPromise = (async () => { - const kernel = await startKernel(cwd, options); - const session: JuliaSession = { - 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: startPromise, - }; - attachSessionOwner(startingSession, sessionId, options.kernelOwnerId); - startingSessions.set(sessionKey, startingSession); - try { - return await waitForPromiseWithCancellation(startPromise, options); - } finally { - if (startingSessions.get(sessionKey) === startingSession) startingSessions.delete(sessionKey); - } -} - -async function replaceSessionKernel(session: JuliaSession, cwd: string, options: JuliaExecutorOptions): Promise { - logger.warn("Julia subprocess died or is unresponsive; spawning fresh process", { - sessionKey: session.sessionKey, - }); - const oldKernel = session.kernel; - const remaining = getRemainingTimeoutMs(options.deadlineMs); - await oldKernel - .shutdown(remaining !== undefined ? { timeoutMs: Math.max(0, remaining) } : undefined) - .catch(() => undefined); - if (sessions.get(session.sessionKey) !== session) { - throw new JuliaExecutionCancelledError(false); - } - requireRemainingTimeoutMs(options.deadlineMs); - const nextKernel = await startKernel(cwd, options); - if (sessions.get(session.sessionKey) !== session) { - await nextKernel.shutdown().catch(() => undefined); - throw new JuliaExecutionCancelledError(false); - } - session.kernel = nextKernel; -} - -async function resetSession(sessionKey: string): Promise { - const session = sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.promise.catch(() => undefined)); - if (!session) return; - sessions.delete(sessionKey); - await session.kernel.shutdown({ timeoutMs: SHUTDOWN_GRACE_MS }).catch(() => undefined); -} - -export async function disposeAllJuliaKernelSessions(): Promise { - const pending = [...startingSessions.values()].map(starting => starting.promise); - startingSessions.clear(); - resettingSessions.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("Julia kernel shutdown not confirmed", { - sessionId: session.sessionId, - sessionKey: id, - cwd: session.cwd, - reason, - }); - if (!sessions.has(id)) sessions.set(id, session); - } -} - -export async function disposeJuliaKernelSessionsByOwner(ownerId: string): Promise { - const toShutdown: JuliaSession[] = []; - const startingToShutdown: StartingJuliaSession[] = []; - 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("Julia kernel shutdown not confirmed", { - sessionId: session.sessionId, - sessionKey: session.sessionKey, - cwd: session.cwd, - reason, - }); - if (!sessions.has(session.sessionKey)) sessions.set(session.sessionKey, session); - } -} - async function executeWithKernel( kernel: JuliaKernel, code: string, @@ -426,7 +163,7 @@ async function executeWithKernel( errorLogLabel: "Julia", isJulia: true, cancelledErrorClass: JuliaExecutionCancelledError, - buildKernelEnvPatch, + buildKernelEnvPatch: opts => buildManagedKernelEnvPatch(opts, { sparse: true }), formatKernelTimeoutAnnotation, formatTimeoutAnnotation, resolveDeadlineMs: opts => opts?.deadlineMs, @@ -434,10 +171,7 @@ async function executeWithKernel( } async function ensureKernelAvailable(cwd: string, options: JuliaExecutorOptions): Promise { - const availability = await waitForPromiseWithCancellation( - checkJuliaKernelAvailability(cwd, options.interpreter), - options, - ); + const availability = await waitForJuliaPromise(checkJuliaKernelAvailability(cwd, options.interpreter), options); if (!availability.ok) { throw new Error(availability.reason ?? "Julia kernel unavailable"); } @@ -454,65 +188,37 @@ async function ensureToolBridge(options: JuliaExecutorOptions): Promise { } } -async function executeOnSession(code: string, cwd: string, options: JuliaExecutorOptions): Promise { - 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 JuliaExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); - } - if (sessions.get(session.sessionKey) !== session) { - throw new JuliaExecutionCancelledError(false); - } - if (!session.kernel.isAlive()) { - await replaceSessionKernel(session, cwd, options); - if (sessions.get(session.sessionKey) !== session) { - throw new JuliaExecutionCancelledError(false); - } - } - const runOptions = { ...options, cwd }; - try { - return await executeWithKernel(session.kernel, code, runOptions); - } catch (err) { - if (isCancellationError(err) || options.signal?.aborted) throw err; - if (session.kernel.isAlive()) throw err; - if (sessions.get(session.sessionKey) !== session) { - throw new JuliaExecutionCancelledError(false); - } - await replaceSessionKernel(session, cwd, options); - if (sessions.get(session.sessionKey) !== session) { - throw new JuliaExecutionCancelledError(false); - } - return await executeWithKernel(session.kernel, code, runOptions); - } +const sessionRegistry = createKernelSessionRegistry< + JuliaKernel, + JuliaExecutorOptions, + JuliaResult, + KernelSession +>({ + languageLabel: "Julia", + cancelledErrorClass: JuliaExecutionCancelledError, + buildSessionKey: (sessionId, cwd, interpreter) => { + const normalizedCwd = normalizeKernelSessionCwd(cwd); + const normalizedInterpreter = normalizeExplicitInterpreter(normalizedCwd, interpreter); + return `${sessionId}::${normalizedCwd}::${normalizedInterpreter}`; + }, + createSession: session => session, + startKernel, + executeWithKernel, + waitForStartup: waitForJuliaPromise, + shutdownSession: (session, resetting) => + resetting ? session.kernel.shutdown({ timeoutMs: SHUTDOWN_GRACE_MS }) : session.kernel.shutdown(), + clearResetsOnDisposeAll: true, + logBeforeReplacement: true, + isCancellation: isJuliaCancellationError, + isTimedOutCancellation: isTimedOutJuliaCancellation, +}); + +export async function disposeAllJuliaKernelSessions(): Promise { + await sessionRegistry.disposeAll(); +} + +export async function disposeJuliaKernelSessionsByOwner(ownerId: string): Promise { + await sessionRegistry.disposeByOwner(ownerId); } export async function executeJuliaWithKernel( @@ -524,8 +230,13 @@ export async function executeJuliaWithKernel( } export async function executeJulia(code: string, options?: JuliaExecutorOptions): Promise { - const cwd = normalizeSessionCwd(options?.cwd ?? getProjectDir()); - const deadlineMs = getExecutionDeadlineMs(options); + const cwd = normalizeKernelSessionCwd(options?.cwd ?? getProjectDir()); + const deadlineMs = + options?.deadlineMs !== undefined + ? options.deadlineMs + : options?.timeoutMs !== undefined && options.timeoutMs > 0 + ? getExecutionDeadlineMs(options) + : undefined; const executionOptions: JuliaExecutorOptions = { ...(options ?? {}), cwd, @@ -536,15 +247,15 @@ export async function executeJulia(code: string, options?: JuliaExecutorOptions) requireRemainingTimeoutMs(deadlineMs); if (executionOptions.signal?.aborted) { throw new JuliaExecutionCancelledError( - isTimedOutCancellation(executionOptions.signal.reason, executionOptions.signal), + isTimedOutJuliaCancellation(executionOptions.signal.reason, executionOptions.signal), ); } 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) || executionOptions.signal?.aborted) { - return createCancelledJuliaResult(isTimedOutCancellation(err, executionOptions.signal)); + if (isJuliaCancellationError(err) || executionOptions.signal?.aborted) { + return createCancelledJuliaResult(isTimedOutJuliaCancellation(err, executionOptions.signal)); } throw err; } diff --git a/packages/coding-agent/src/eval/kernel-session-registry.ts b/packages/coding-agent/src/eval/kernel-session-registry.ts new file mode 100644 index 000000000..e5d8b0b8b --- /dev/null +++ b/packages/coding-agent/src/eval/kernel-session-registry.ts @@ -0,0 +1,398 @@ +import * as path from "node:path"; + +import { logger } from "@oh-my-pi/pi-utils"; +import { + attachSessionOwner, + type CancelledErrorClass, + getRemainingTimeoutMs, + isCancellationError, + isTimedOutCancellation, + resolveOwnerScopedSessionKey, + type SessionOwners, +} from "./executor-base"; + +interface KernelSessionRegistryOptions { + sessionId?: string; + kernelOwnerId?: string; + interpreter?: string; + reset?: boolean; + signal?: AbortSignal; + deadlineMs?: number; + bridge?: unknown; + bridgeSessionId?: string; +} + +interface RegistryKernelShutdownResult { + confirmed?: boolean; +} + +interface RegistryKernel { + isAlive(): boolean; + shutdown(options?: { timeoutMs: number }): Promise; +} + +export interface KernelSession extends SessionOwners { + sessionKey: string; + sessionId: string; + cwd: string; + kernel: TKernel; +} + +interface StartingKernelSession extends SessionOwners { + promise: Promise; +} + +export interface KernelSessionRegistryContext< + TKernel extends RegistryKernel, + TOptions extends KernelSessionRegistryOptions, + TSession extends KernelSession, +> { + sessions: Map; + startKernel: (cwd: string, options: TOptions) => Promise; + replaceSessionKernel: (session: TSession, cwd: string, options: TOptions) => Promise; +} + +interface KernelSessionRegistryDescriptor< + TKernel extends RegistryKernel, + TOptions extends KernelSessionRegistryOptions, + TResult, + TSession extends KernelSession, +> { + languageLabel: string; + cancelledErrorClass: CancelledErrorClass; + buildSessionKey: (sessionId: string, cwd: string, interpreter: string | undefined) => string; + createSession: (session: KernelSession) => TSession; + startKernel: (cwd: string, options: TOptions) => Promise; + executeWithKernel: (kernel: TKernel, code: string, options: TOptions) => Promise; + waitForStartup?: (promise: Promise, options: TOptions) => Promise; + replaceSessionKernel?: ( + session: TSession, + cwd: string, + options: TOptions, + context: KernelSessionRegistryContext, + ) => Promise; + acquireLiveSessionKernel?: ( + session: TSession, + cwd: string, + options: TOptions, + context: KernelSessionRegistryContext, + ) => Promise; + invalidateSession?: (session: TSession) => void; + shutdownSession?: (session: TSession, resetting: boolean) => Promise; + clearResetsOnDisposeAll?: boolean; + logBeforeReplacement?: boolean; + isCancellation?: (error: unknown) => boolean; + isTimedOutCancellation?: (error: unknown, signal?: AbortSignal) => boolean; + validateKernel?: (session: TSession, kernel: TKernel) => boolean; +} + +interface KernelSessionRegistry { + disposeAll(): Promise; + disposeByOwner(ownerId: string): Promise; + executeOnSession(code: string, cwd: string, options: TOptions): Promise; +} + +export function normalizeKernelSessionCwd(cwd: string): string { + return path.resolve(cwd); +} + +export function requireRemainingKernelTimeoutMs( + deadlineMs: number | undefined, + cancelledErrorClass: CancelledErrorClass, +): number | undefined { + const remainingMs = getRemainingTimeoutMs(deadlineMs); + if (remainingMs === undefined) return undefined; + if (remainingMs <= 0) { + throw new cancelledErrorClass(true); + } + return remainingMs; +} + +export function formatSessionTimeoutAnnotation(timeoutMs?: number): string { + if (timeoutMs === undefined) return "Command timed out"; + const secs = Math.max(1, Math.round(timeoutMs / 1000)); + return `Command timed out after ${secs} seconds`; +} + +export function formatSessionKernelTimeoutAnnotation(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.`; +} + +export function createKernelSessionRegistry< + TKernel extends RegistryKernel, + TOptions extends KernelSessionRegistryOptions, + TResult, + TSession extends KernelSession, +>( + descriptor: KernelSessionRegistryDescriptor, +): KernelSessionRegistry { + const sessions = new Map(); + const startingSessions = new Map>(); + const resettingSessions = new Map>(); + + const context: KernelSessionRegistryContext = { + sessions, + startKernel: descriptor.startKernel, + replaceSessionKernel, + }; + + function waitForStartup(promise: Promise, options: TOptions): Promise { + return descriptor.waitForStartup?.(promise, options) ?? promise; + } + + function isCurrent(session: TSession, kernel?: TKernel): boolean { + return ( + sessions.get(session.sessionKey) === session && + (kernel === undefined || descriptor.validateKernel?.(session, kernel) !== false) + ); + } + + async function acquireSession( + sessionKey: string, + sessionId: string, + cwd: string, + options: TOptions, + ): Promise { + 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 waitForStartup(starting.promise, options); + } + let startingSession!: StartingKernelSession; + const startup = (async () => { + const kernel = await descriptor.startKernel(cwd, options); + const session = descriptor.createSession({ + 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 waitForStartup(startup, options); + } finally { + if (startingSessions.get(sessionKey) === startingSession) startingSessions.delete(sessionKey); + } + } + + async function replaceSessionKernel(session: TSession, cwd: string, options: TOptions): Promise { + if (descriptor.replaceSessionKernel) { + return await descriptor.replaceSessionKernel(session, cwd, options, context); + } + if (descriptor.logBeforeReplacement) { + logger.warn(`${descriptor.languageLabel} subprocess died or is unresponsive; spawning fresh process`, { + sessionKey: session.sessionKey, + }); + } + 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 descriptor.cancelledErrorClass(false); + } + requireRemainingKernelTimeoutMs(options.deadlineMs, descriptor.cancelledErrorClass); + const next = await descriptor.startKernel(cwd, options); + if (sessions.get(session.sessionKey) !== session) { + await next.shutdown().catch(() => undefined); + throw new descriptor.cancelledErrorClass(false); + } + session.kernel = next; + return next; + } + + async function acquireLiveSessionKernel(session: TSession, cwd: string, options: TOptions): Promise { + if (descriptor.acquireLiveSessionKernel) { + return await descriptor.acquireLiveSessionKernel(session, cwd, options, context); + } + if (!isCurrent(session)) throw new descriptor.cancelledErrorClass(false); + if (!session.kernel.isAlive()) await replaceSessionKernel(session, cwd, options); + if (!isCurrent(session)) throw new descriptor.cancelledErrorClass(false); + return session.kernel; + } + + async function shutdownSession(session: TSession, resetting: boolean): Promise { + return await (descriptor.shutdownSession?.(session, resetting) ?? session.kernel.shutdown()); + } + + async function resetSession(sessionKey: string): Promise { + const existing = + sessions.get(sessionKey) ?? (await startingSessions.get(sessionKey)?.promise.catch(() => undefined)); + if (!existing) return; + descriptor.invalidateSession?.(existing); + sessions.delete(sessionKey); + await shutdownSession(existing, true).catch(() => undefined); + } + + async function disposeAll(): Promise { + const pending = [...startingSessions.values()].map(starting => starting.promise); + startingSessions.clear(); + if (descriptor.clearResetsOnDisposeAll) resettingSessions.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) { + descriptor.invalidateSession?.(session); + if (sessions.get(id) === session) sessions.delete(id); + } + const results = await Promise.allSettled(all.map(([, session]) => shutdownSession(session, false))); + 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(`${descriptor.languageLabel} kernel shutdown not confirmed`, { + sessionId: session.sessionId, + sessionKey: id, + cwd: session.cwd, + reason, + }); + if (!sessions.has(id)) sessions.set(id, session); + } + } + + async function disposeByOwner(ownerId: string): Promise { + const toShutdown: TSession[] = []; + const startingToShutdown: StartingKernelSession[] = []; + 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) { + descriptor.invalidateSession?.(session); + 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; + descriptor.invalidateSession?.(session); + if (sessions.get(session.sessionKey) === session) sessions.delete(session.sessionKey); + toShutdown.push(session); + } + const results = await Promise.allSettled(toShutdown.map(session => shutdownSession(session, false))); + 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(`${descriptor.languageLabel} kernel shutdown not confirmed`, { + sessionId: session.sessionId, + sessionKey: session.sessionKey, + cwd: session.cwd, + reason, + }); + if (!sessions.has(session.sessionKey)) sessions.set(session.sessionKey, session); + } + } + + async function executeOnSession(code: string, cwd: string, options: TOptions): Promise { + const sessionId = options.sessionId ?? `session:${cwd}`; + const sessionKey = resolveOwnerScopedSessionKey({ + baseKey: descriptor.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) { + const timedOut = + descriptor.isTimedOutCancellation?.(options.signal.reason, options.signal) ?? + isTimedOutCancellation(options.signal.reason, descriptor.cancelledErrorClass, options.signal); + throw new descriptor.cancelledErrorClass(timedOut); + } + const kernel = await acquireLiveSessionKernel(session, cwd, options); + if (!isCurrent(session, kernel)) throw new descriptor.cancelledErrorClass(false); + const runOptions = { ...options, cwd }; + try { + return await descriptor.executeWithKernel(kernel, code, runOptions); + } catch (err) { + if ( + descriptor.isCancellation?.(err) || + isCancellationError(err, descriptor.cancelledErrorClass) || + options.signal?.aborted + ) + throw err; + if (kernel.isAlive()) throw err; + let retryKernel: TKernel; + if (descriptor.acquireLiveSessionKernel) { + retryKernel = await acquireLiveSessionKernel(session, cwd, options); + } else { + if (!isCurrent(session, kernel)) throw new descriptor.cancelledErrorClass(false); + retryKernel = await replaceSessionKernel(session, cwd, options); + } + if (!isCurrent(session, retryKernel)) throw new descriptor.cancelledErrorClass(false); + return await descriptor.executeWithKernel(retryKernel, code, runOptions); + } + } + + return { disposeAll, disposeByOwner, executeOnSession }; +} diff --git a/packages/coding-agent/src/eval/py/executor.ts b/packages/coding-agent/src/eval/py/executor.ts index e0d895f92..ceef2356c 100644 --- a/packages/coding-agent/src/eval/py/executor.ts +++ b/packages/coding-agent/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; } -interface PythonSession { - sessionKey: string; - sessionId: string; - cwd: string; - kernel: PythonKernel; +interface PythonSession extends KernelSession { generation: number; replacement?: SessionKernelReplacement; - ownerIds: Set; - hasFallbackOwner: boolean; -} - -interface StartingPythonSession extends SessionOwners { - promise: Promise; -} - -const sessions = new Map(); -const startingSessions = new Map(); -const resettingSessions = new Map>(); - -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 { - 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, ): Promise { + 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, ): Promise { - 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 { - 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 { - 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 { - 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 { - 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({ + 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 { + await sessionRegistry.disposeAll(); +} + +export async function disposeKernelSessionsByOwner(ownerId: string): Promise { + 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 { - 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( diff --git a/packages/coding-agent/src/eval/rb/executor.ts b/packages/coding-agent/src/eval/rb/executor.ts index 740793962..17fc8f7e4 100644 --- a/packages/coding-agent/src/eval/rb/executor.ts +++ b/packages/coding-agent/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; - hasFallbackOwner: boolean; -} - -interface RubySession extends RubySessionOwners { - sessionKey: string; - sessionId: string; - cwd: string; - kernel: RubyKernel; -} - -interface StartingRubySession extends RubySessionOwners { - promise: Promise; -} - -const sessions = new Map(); -const startingSessions = new Map(); -const resettingSessions = new Map>(); - -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 { - 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 { - 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 { - 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 { - 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 { - 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 { } } -async function executeOnSession(code: string, cwd: string, options: RubyExecutorOptions): Promise { - 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 +>({ + 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 { + await sessionRegistry.disposeAll(); +} + +export async function disposeRubyKernelSessionsByOwner(ownerId: string): Promise { + 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 { - 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(