refactor(coding-agent): extracted generic kernel session registry
- Python, Ruby and Julia each carried their own copy of the same session maps, acquire/reset/replace/dispose lifecycle and executeOnSession; one generic registry now owns it, parameterized by a per-language descriptor. - Julia additionally re-implemented seven executor-base helpers locally; those copies are gone and executor-base gained sparse managed-env and timeout resolver hooks so Julia's differing behavior survives unchanged. - All twelve exported entry points keep their names and signatures.
This commit is contained in:
can1357
@@ -116,9 +116,13 @@ export async function waitForPromiseWithCancellation<T>(
|
||||
promise: Promise<T>,
|
||||
options: { signal?: AbortSignal; deadlineMs?: number },
|
||||
cancelledErrorClass: CancelledErrorClass,
|
||||
timedOutResolver?: (error: unknown, signal?: AbortSignal) => boolean,
|
||||
): Promise<T> {
|
||||
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<T>(
|
||||
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<string, string>;
|
||||
}
|
||||
interface ManagedKernelEnvPolicy {
|
||||
sparse?: boolean;
|
||||
}
|
||||
|
||||
export function buildManagedKernelEnvPatch(options: ManagedKernelEnvOptions): Record<string, string | null> {
|
||||
export function buildManagedKernelEnvPatch(options: ManagedKernelEnvOptions): Record<string, string | null>;
|
||||
export function buildManagedKernelEnvPatch(
|
||||
options: ManagedKernelEnvOptions,
|
||||
policy: { sparse: true },
|
||||
): Record<string, string | undefined>;
|
||||
export function buildManagedKernelEnvPatch(
|
||||
options: ManagedKernelEnvOptions,
|
||||
policy?: ManagedKernelEnvPolicy,
|
||||
): KernelEnvPatch {
|
||||
const localRoots = options.localRoots;
|
||||
if (policy?.sparse) {
|
||||
const patch: Record<string, string | undefined> = {};
|
||||
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<string, string> | undefined {
|
||||
const patch = buildManagedKernelEnvPatch(options);
|
||||
export function buildManagedKernelEnv(
|
||||
options: ManagedKernelEnvOptions,
|
||||
policy?: ManagedKernelEnvPolicy,
|
||||
): Record<string, string> | undefined {
|
||||
const patch = policy?.sparse
|
||||
? buildManagedKernelEnvPatch(options, { sparse: true })
|
||||
: buildManagedKernelEnvPatch(options);
|
||||
const env: Record<string, string> = {};
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -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<string>;
|
||||
hasFallbackOwner: boolean;
|
||||
}
|
||||
|
||||
interface JuliaSession extends JuliaSessionOwners {
|
||||
sessionKey: string;
|
||||
sessionId: string;
|
||||
cwd: string;
|
||||
kernel: JuliaKernel;
|
||||
}
|
||||
|
||||
interface StartingJuliaSession extends JuliaSessionOwners {
|
||||
promise: Promise<JuliaSession>;
|
||||
}
|
||||
|
||||
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<string, JuliaSession>();
|
||||
const startingSessions = new Map<string, StartingJuliaSession>();
|
||||
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 = 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<JuliaExecutorOptions, "deadlineMs" | "timeoutMs">): 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<T>(
|
||||
async function waitForJuliaPromise<T>(
|
||||
promise: Promise<T>,
|
||||
options: Pick<JuliaExecutorOptions, "signal" | "deadlineMs">,
|
||||
): Promise<T> {
|
||||
if (options.signal?.aborted) {
|
||||
throw new JuliaExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal));
|
||||
}
|
||||
const cleanups: Array<() => void> = [];
|
||||
const { promise: cancelPromise, reject } = Promise.withResolvers<never>();
|
||||
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<string, string>;
|
||||
}): Record<string, string | undefined> {
|
||||
const patch: Record<string, string | undefined> = {};
|
||||
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<string, string>;
|
||||
}): Record<string, string> | undefined {
|
||||
const patch = buildKernelEnvPatch(options);
|
||||
const keys = Object.keys(patch);
|
||||
if (keys.length === 0) return undefined;
|
||||
const realEnv: Record<string, string> = {};
|
||||
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<JuliaKernel> {
|
||||
requireRemainingTimeoutMs(options.deadlineMs);
|
||||
const env: Record<string, string | undefined> = {};
|
||||
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<string, string | undefined> = 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<JuliaSession> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
}
|
||||
}
|
||||
|
||||
async function executeOnSession(code: string, cwd: string, options: JuliaExecutorOptions): Promise<JuliaResult> {
|
||||
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<JuliaKernel>
|
||||
>({
|
||||
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<void> {
|
||||
await sessionRegistry.disposeAll();
|
||||
}
|
||||
|
||||
export async function disposeJuliaKernelSessionsByOwner(ownerId: string): Promise<void> {
|
||||
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<JuliaResult> {
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -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<RegistryKernelShutdownResult>;
|
||||
}
|
||||
|
||||
export interface KernelSession<TKernel extends RegistryKernel> extends SessionOwners {
|
||||
sessionKey: string;
|
||||
sessionId: string;
|
||||
cwd: string;
|
||||
kernel: TKernel;
|
||||
}
|
||||
|
||||
interface StartingKernelSession<TSession> extends SessionOwners {
|
||||
promise: Promise<TSession>;
|
||||
}
|
||||
|
||||
export interface KernelSessionRegistryContext<
|
||||
TKernel extends RegistryKernel,
|
||||
TOptions extends KernelSessionRegistryOptions,
|
||||
TSession extends KernelSession<TKernel>,
|
||||
> {
|
||||
sessions: Map<string, TSession>;
|
||||
startKernel: (cwd: string, options: TOptions) => Promise<TKernel>;
|
||||
replaceSessionKernel: (session: TSession, cwd: string, options: TOptions) => Promise<TKernel>;
|
||||
}
|
||||
|
||||
interface KernelSessionRegistryDescriptor<
|
||||
TKernel extends RegistryKernel,
|
||||
TOptions extends KernelSessionRegistryOptions,
|
||||
TResult,
|
||||
TSession extends KernelSession<TKernel>,
|
||||
> {
|
||||
languageLabel: string;
|
||||
cancelledErrorClass: CancelledErrorClass;
|
||||
buildSessionKey: (sessionId: string, cwd: string, interpreter: string | undefined) => string;
|
||||
createSession: (session: KernelSession<TKernel>) => TSession;
|
||||
startKernel: (cwd: string, options: TOptions) => Promise<TKernel>;
|
||||
executeWithKernel: (kernel: TKernel, code: string, options: TOptions) => Promise<TResult>;
|
||||
waitForStartup?: (promise: Promise<TSession>, options: TOptions) => Promise<TSession>;
|
||||
replaceSessionKernel?: (
|
||||
session: TSession,
|
||||
cwd: string,
|
||||
options: TOptions,
|
||||
context: KernelSessionRegistryContext<TKernel, TOptions, TSession>,
|
||||
) => Promise<TKernel>;
|
||||
acquireLiveSessionKernel?: (
|
||||
session: TSession,
|
||||
cwd: string,
|
||||
options: TOptions,
|
||||
context: KernelSessionRegistryContext<TKernel, TOptions, TSession>,
|
||||
) => Promise<TKernel>;
|
||||
invalidateSession?: (session: TSession) => void;
|
||||
shutdownSession?: (session: TSession, resetting: boolean) => Promise<RegistryKernelShutdownResult>;
|
||||
clearResetsOnDisposeAll?: boolean;
|
||||
logBeforeReplacement?: boolean;
|
||||
isCancellation?: (error: unknown) => boolean;
|
||||
isTimedOutCancellation?: (error: unknown, signal?: AbortSignal) => boolean;
|
||||
validateKernel?: (session: TSession, kernel: TKernel) => boolean;
|
||||
}
|
||||
|
||||
interface KernelSessionRegistry<TOptions extends KernelSessionRegistryOptions, TResult> {
|
||||
disposeAll(): Promise<void>;
|
||||
disposeByOwner(ownerId: string): Promise<void>;
|
||||
executeOnSession(code: string, cwd: string, options: TOptions): Promise<TResult>;
|
||||
}
|
||||
|
||||
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<TKernel>,
|
||||
>(
|
||||
descriptor: KernelSessionRegistryDescriptor<TKernel, TOptions, TResult, TSession>,
|
||||
): KernelSessionRegistry<TOptions, TResult> {
|
||||
const sessions = new Map<string, TSession>();
|
||||
const startingSessions = new Map<string, StartingKernelSession<TSession>>();
|
||||
const resettingSessions = new Map<string, Promise<void>>();
|
||||
|
||||
const context: KernelSessionRegistryContext<TKernel, TOptions, TSession> = {
|
||||
sessions,
|
||||
startKernel: descriptor.startKernel,
|
||||
replaceSessionKernel,
|
||||
};
|
||||
|
||||
function waitForStartup(promise: Promise<TSession>, options: TOptions): Promise<TSession> {
|
||||
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<TSession> {
|
||||
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<TSession>;
|
||||
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<TKernel> {
|
||||
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<TKernel> {
|
||||
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<RegistryKernelShutdownResult> {
|
||||
return await (descriptor.shutdownSession?.(session, resetting) ?? session.kernel.shutdown());
|
||||
}
|
||||
|
||||
async function resetSession(sessionKey: string): Promise<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
const toShutdown: TSession[] = [];
|
||||
const startingToShutdown: StartingKernelSession<TSession>[] = [];
|
||||
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<TResult> {
|
||||
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 };
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user