@gajae-code/coding-agent 0.5.2 → 0.5.4

package/src/eval/js/context-manager.ts

@@ -1,4 +1,5 @@
 import { isCompiledBinary, logger, Snowflake } from "@gajae-code/utils";
+import { registerResourceOwner } from "../../runtime/process-lifecycle";
 import type { ToolSession } from "../../tools";
 import { ToolAbortError, ToolError } from "../../tools/tool-errors";
 import { callSessionTool, type JsStatusEvent } from "./tool-bridge";
@@ -43,22 +44,66 @@ interface PendingRun {
 	settled: boolean;
 }
 
+interface ReadyDeferred {
+	promise: Promise<JsSession>;
+	resolve(session: JsSession): void;
+	reject(error: Error): void;
+}
+
+interface QueueDeferred {
+	promise: Promise<void>;
+	resolve(): void;
+	reject(error: Error): void;
+}
+
 interface JsSession {
 	sessionKey: string;
-	worker: WorkerHandle;
-	state: "alive" | "dead";
+	worker?: WorkerHandle;
+	state: "starting" | "alive" | "dead";
+	ownerId?: string;
 	pending: Map<string, PendingRun>;
 	queue: Promise<void>;
+	queuedWaiters: Set<(error: Error) => void>;
+	queueTail: QueueDeferred;
+	controllers: Set<AbortController>;
+	runSignal?: AbortSignal;
+	ready: ReadyDeferred;
+	unsubscribe?: () => void;
+	unregistered?: () => void;
 }
 
 const sessions = new Map<string, JsSession>();
+const sessionWaiters = new Map<string, Set<(error: Error) => void>>();
+let vmResourceCleanupRegistered = false;
+
+function ensureVmResourceCleanup(): void {
+	if (vmResourceCleanupRegistered) return;
+	vmResourceCleanupRegistered = true;
+	registerResourceOwner("js-vm-contexts", async () => {
+		try {
+			await disposeAllVmContexts();
+		} finally {
+			vmResourceCleanupRegistered = false;
+		}
+	});
+}
 const READY_TIMEOUT_MS_DEFAULT = 5_000;
 
+function getSessionWaiters(sessionKey: string): Set<(error: Error) => void> {
+	let waiters = sessionWaiters.get(sessionKey);
+	if (!waiters) {
+		waiters = new Set();
+		sessionWaiters.set(sessionKey, waiters);
+	}
+	return waiters;
+}
+
 export async function executeInVmContext(options: {
 	sessionKey: string;
 	sessionId: string;
 	cwd: string;
 	session: ToolSession;
+	ownerId?: string;
 	reset?: boolean;
 	code: string;
 	filename: string;
@@ -68,40 +113,84 @@ export async function executeInVmContext(options: {
 	if (options.reset) {
 		await resetVmContext(options.sessionKey);
 	}
-	const session = await acquireSession(
-		options.sessionKey,
-		{ cwd: options.cwd, sessionId: options.sessionId },
-		options.timeoutMs,
-	);
-	return await runQueued(session, () => runOnce(session, options));
+	const waiters = getSessionWaiters(options.sessionKey);
+	const { promise: contextResetPromise, reject: rejectContextReset } = Promise.withResolvers<never>();
+	contextResetPromise.catch(() => undefined);
+	waiters.add(rejectContextReset);
+	const runPromise = (async (): Promise<{ value: unknown }> => {
+		const session = await acquireSession(
+			options.sessionKey,
+			{ cwd: options.cwd, sessionId: options.sessionId },
+			options.ownerId,
+			options.timeoutMs,
+		);
+		return await runQueued(session, () => runOnce(session, options));
+	})();
+	try {
+		return await Promise.race([runPromise, contextResetPromise]);
+	} finally {
+		waiters.delete(rejectContextReset);
+	}
 }
 
 export async function resetVmContext(sessionKey: string): Promise<void> {
 	const session = sessions.get(sessionKey);
 	if (!session) return;
 	sessions.delete(sessionKey);
+	const waiters = sessionWaiters.get(sessionKey);
+	if (waiters) for (const reject of [...waiters]) reject(new ToolError("JS context reset"));
 	await killSession(session, new ToolError("JS context reset"));
 }
 
+export async function disposeVmContextsByOwner(ownerId: string): Promise<void> {
+	const owned = [...sessions.entries()].filter(
+		([sessionKey, session]) =>
+			session.ownerId === ownerId || sessionKey === ownerId || sessionKey === `js:${ownerId}`,
+	);
+	for (const [sessionKey, session] of owned) {
+		if (sessions.get(sessionKey) === session) sessions.delete(sessionKey);
+	}
+	await Promise.all(owned.map(([, session]) => killSession(session, new ToolError("JS context disposed"))));
+}
+
 export async function disposeAllVmContexts(): Promise<void> {
 	const all = [...sessions.values()];
 	sessions.clear();
 	await Promise.all(all.map(session => killSession(session, new ToolError("JS context disposed"))));
 }
 
+export function liveVmContextCount(): number {
+	return [...sessions.values()].filter(session => session.state !== "dead").length;
+}
+
 async function runQueued<T>(session: JsSession, work: () => Promise<T>): Promise<T> {
+	if (session.state !== "alive") throw new ToolError("JS worker is not alive");
 	const previous = session.queue;
-	const { promise, resolve } = Promise.withResolvers<void>();
-	session.queue = promise;
-	try {
-		await previous;
-	} catch {
-		// Previous run's failure must not poison this one.
-	}
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	const queueController = new AbortController();
+	const queueItem: QueueDeferred = { promise, resolve, reject };
+	const rejectWaiter = (error: Error): void => queueItem.reject(error);
+	session.queuedWaiters.add(rejectWaiter);
+	session.controllers.add(queueController);
+	session.queueTail = queueItem;
+	session.queue = (async () => {
+		await previous.catch(() => undefined);
+		await queueItem.promise;
+	})().catch(() => undefined);
 	try {
-		return await work();
+		await Promise.race([previous.catch(() => undefined), queueItem.promise, abortPromise(queueController.signal)]);
+		if (session.runSignal?.aborted) throw reasonToError(session.runSignal.reason, "JS worker is not alive");
+		if (session.state !== "alive") throw new ToolError("JS worker is not alive");
+		session.queuedWaiters.delete(rejectWaiter);
+		return await Promise.race([
+			work(),
+			queueItem.promise.then(() => new Promise<never>(() => {})),
+			abortPromise(queueController.signal),
+		]);
 	} finally {
-		resolve();
+		session.queuedWaiters.delete(rejectWaiter);
+		session.controllers.delete(queueController);
+		queueItem.resolve();
 	}
 }
 
@@ -118,6 +207,7 @@ async function runOnce(
 ): Promise<{ value: unknown }> {
 	const runId = `r-${Snowflake.next()}`;
 	const { promise, resolve, reject } = Promise.withResolvers<{ value: unknown }>();
+	const sessionSignal = session.runSignal;
 	const pending: PendingRun = {
 		runId,
 		runState: options.runState,
@@ -145,13 +235,19 @@ async function runOnce(
 	}
 
 	try {
-		session.worker.send({
-			type: "run",
-			runId,
-			code: options.code,
-			filename: options.filename,
-			snapshot: { cwd: options.cwd, sessionId: options.sessionId },
-		});
+		if (sessionSignal?.aborted) throw reasonToError(sessionSignal.reason, "JS worker is not alive");
+		if (
+			!safeSend(session, {
+				type: "run",
+				runId,
+				code: options.code,
+				filename: options.filename,
+				snapshot: { cwd: options.cwd, sessionId: options.sessionId },
+			})
+		) {
+			settleRunWithError(session, pending, new ToolError("JS worker send failed"));
+			return await promise;
+		}
 		return await promise;
 	} finally {
 		options.runState.signal?.removeEventListener("abort", onAbort);
@@ -159,46 +255,76 @@ async function runOnce(
 	}
 }
 
-async function acquireSession(sessionKey: string, snapshot: SessionSnapshot, timeoutMs?: number): Promise<JsSession> {
+async function acquireSession(
+	sessionKey: string,
+	snapshot: SessionSnapshot,
+	ownerId: string | undefined,
+	timeoutMs?: number,
+): Promise<JsSession> {
+	ensureVmResourceCleanup();
 	const existing = sessions.get(sessionKey);
-	if (existing && existing.state === "alive") return existing;
+	if (existing && existing.state !== "dead") return await existing.ready.promise;
 
-	const worker = await spawnJsWorker();
+	const { promise: ready, resolve: resolveSession, reject: rejectSession } = Promise.withResolvers<JsSession>();
+	ready.catch(() => undefined);
 	const session: JsSession = {
 		sessionKey,
-		worker,
-		state: "alive",
+		state: "starting",
+		ownerId,
 		pending: new Map(),
 		queue: Promise.resolve(),
+		queuedWaiters: new Set(),
+		queueTail: settledQueueDeferred(),
+		controllers: new Set(),
+		ready: { promise: ready, resolve: resolveSession, reject: rejectSession },
 	};
-	const { promise: readyPromise, resolve: resolveReady, reject: rejectReady } = Promise.withResolvers<void>();
-	let resolved = false;
-	const unsubscribe = worker.onMessage(msg => {
-		if (!resolved && msg.type === "ready") {
-			resolved = true;
-			resolveReady();
-			return;
-		}
-		if (!resolved && msg.type === "init-failed") {
-			resolved = true;
-			rejectReady(errorFromPayload(msg.error));
-			return;
-		}
-		handleSessionMessage(session, msg);
-	});
+	sessions.set(sessionKey, session);
+
+	let worker: WorkerHandle | undefined;
 	try {
-		// Cold-start can exceed 5s on slow hosts. Let the caller's per-cell timeout dominate so
-		// users can grant more headroom when they raise `timeout` on a cell.
+		worker = await spawnJsWorker();
+		const current = sessions.get(sessionKey);
+		if (current !== session) {
+			await worker.terminate().catch(() => undefined);
+			return (
+				(await current?.ready.promise) ?? Promise.reject(new ToolError("JS context replaced during initialization"))
+			);
+		}
+		session.worker = worker;
+		const { promise: readyPromise, resolve: resolveReady, reject: rejectReady } = Promise.withResolvers<void>();
+		let resolved = false;
+		session.unsubscribe = worker.onMessage(msg => {
+			if (!resolved && msg.type === "ready") {
+				resolved = true;
+				resolveReady();
+				return;
+			}
+			if (!resolved && msg.type === "init-failed") {
+				resolved = true;
+				rejectReady(errorFromPayload(msg.error));
+				return;
+			}
+			handleSessionMessage(session, msg);
+		});
 		const readyTimeoutMs = Math.max(READY_TIMEOUT_MS_DEFAULT, timeoutMs ?? 0);
 		await raceWithTimeout(readyPromise, readyTimeoutMs, "Timed out initializing JS eval worker");
+		if (sessions.get(sessionKey) !== session) {
+			await killSession(session, new ToolError("JS context replaced during initialization"));
+			return (
+				(await sessions.get(sessionKey)?.ready.promise) ??
+				Promise.reject(new ToolError("JS context replaced during initialization"))
+			);
+		}
+		worker.send({ type: "init", snapshot });
+		session.state = "alive";
+		session.ready.resolve(session);
+		return session;
 	} catch (error) {
-		unsubscribe();
-		await worker.terminate().catch(() => undefined);
+		if (sessions.get(sessionKey) === session) sessions.delete(sessionKey);
+		await killSession(session, error instanceof Error ? error : new Error(String(error)));
+		session.ready.reject(error instanceof Error ? error : new Error(String(error)));
 		throw error;
 	}
-	worker.send({ type: "init", snapshot });
-	sessions.set(sessionKey, session);
-	return session;
 }
 
 function handleSessionMessage(session: JsSession, msg: WorkerOutbound): void {
@@ -276,22 +402,54 @@ async function killSessionFor(session: JsSession, error: Error): Promise<void> {
 async function killSession(session: JsSession, error: Error): Promise<void> {
 	if (session.state === "dead") return;
 	session.state = "dead";
-	for (const pending of session.pending.values()) {
-		if (pending.settled) continue;
-		pending.settled = true;
-		for (const ctrl of pending.toolCalls.values()) ctrl.abort(error);
-		pending.reject(error);
-	}
+	const unsubscribe = session.unsubscribe;
+	session.unsubscribe = undefined;
+	unsubscribe?.();
+	session.ready.reject(error);
+	session.queueTail.reject(error);
+	for (const controller of [...session.controllers]) controller.abort(error);
+	session.controllers.clear();
+	session.runSignal = AbortSignal.abort(error);
+	for (const reject of [...session.queuedWaiters]) reject(error);
+	session.queuedWaiters.clear();
+	for (const pending of [...session.pending.values()]) settleRunWithError(session, pending, error);
 	session.pending.clear();
-	await session.worker.terminate().catch(() => undefined);
+	void session.worker?.terminate().catch(() => undefined);
+}
+
+function settleRunWithError(session: JsSession, pending: PendingRun, error: Error): void {
+	if (pending.settled) return;
+	pending.settled = true;
+	for (const ctrl of pending.toolCalls.values()) ctrl.abort(error);
+	pending.toolCalls.clear();
+	session.pending.delete(pending.runId);
+	pending.reject(error);
 }
 
-function safeSend(session: JsSession, msg: WorkerInbound): void {
-	if (session.state !== "alive") return;
+function settledQueueDeferred(): QueueDeferred {
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	resolve();
+	return { promise, resolve, reject };
+}
+
+function abortPromise(signal: AbortSignal): Promise<never> {
+	if (signal.aborted) return Promise.reject(reasonToError(signal.reason, "JS worker is not alive"));
+	const { promise, reject } = Promise.withResolvers<never>();
+	const onAbort = (): void => reject(reasonToError(signal.reason, "JS worker is not alive"));
+	signal.addEventListener("abort", onAbort, { once: true });
+	promise.finally(() => signal.removeEventListener("abort", onAbort)).catch(() => undefined);
+	return promise;
+}
+
+function safeSend(session: JsSession, msg: WorkerInbound): boolean {
+	if (session.state !== "alive") return false;
 	try {
-		session.worker.send(msg);
+		session.worker?.send(msg);
+		return true;
 	} catch (err) {
 		logger.debug("js worker send failed", { error: err instanceof Error ? err.message : String(err) });
+		void killSessionFor(session, err instanceof Error ? err : new Error(String(err)));
+		return false;
 	}
 }
 
@@ -352,10 +510,15 @@ async function spawnJsWorker(): Promise<WorkerHandle> {
 			: new Worker(new URL("./worker-entry.ts", import.meta.url).href, { type: "module" });
 		return wrapBunWorker(worker);
 	} catch (err) {
-		logger.warn("Bun Worker spawn failed; using inline JS eval worker (no sync-loop guard)", {
-			error: err instanceof Error ? err.message : String(err),
-		});
-		return spawnInlineWorker();
+		if (process.env.GAJAE_CODE_JS_EVAL_INLINE_WORKER === "1") {
+			logger.warn("Bun Worker spawn failed; using test-only inline JS eval worker", {
+				error: err instanceof Error ? err.message : String(err),
+			});
+			return spawnInlineWorker();
+		}
+		throw new ToolError(
+			`JS eval worker is unavailable and inline fallback is disabled because it cannot interrupt synchronous user code: ${err instanceof Error ? err.message : String(err)}`,
+		);
 	}
 }
 

package/src/eval/js/executor.ts

@@ -10,6 +10,7 @@ export interface JsExecutorOptions {
 	onChunk?: (chunk: string) => Promise<void> | void;
 	signal?: AbortSignal;
 	sessionId: string;
+	ownerId?: string;
 	reset?: boolean;
 	sessionFile?: string;
 	artifactPath?: string;
@@ -68,6 +69,7 @@ export async function executeJs(code: string, options: JsExecutorOptions): Promi
 		await executeInVmContext({
 			sessionKey: options.sessionId,
 			sessionId: options.sessionId,
+			ownerId: options.ownerId,
 			cwd: options.cwd ?? options.session.cwd,
 			session: options.session,
 			reset: options.reset,

package/src/eval/js/index.ts

@@ -23,6 +23,7 @@ export default {
 			deadlineMs: opts.deadlineMs,
 			signal: opts.signal,
 			sessionId: namespaceSessionId(opts.sessionId),
+			ownerId: opts.kernelOwnerId,
 			sessionFile: opts.sessionFile,
 			reset: opts.reset,
 			artifactPath: opts.artifactPath,

package/src/eval/js/worker-core.ts

@@ -12,6 +12,11 @@ interface ActiveRun {
 	pendingTools: Map<string, PendingTool>;
 }
 
+function rejectPendingTools(active: ActiveRun, error: Error): void {
+	for (const pending of active.pendingTools.values()) pending.reject(error);
+	active.pendingTools.clear();
+}
+
 function errorPayload(error: unknown): RunErrorPayload {
 	if (error instanceof Error) {
 		return {
@@ -99,7 +104,8 @@ export class WorkerCore {
 	async #runOne(runId: string, code: string, filename: string, snapshot: SessionSnapshot): Promise<void> {
 		const runtime = this.#ensureRuntime(snapshot);
 		runtime.setCwd(snapshot.cwd);
-		this.#active = { runId, pendingTools: new Map() };
+		const active: ActiveRun = { runId, pendingTools: new Map() };
+		this.#active = active;
 		try {
 			const value = await runtime.run(code, filename);
 			runtime.displayValue(value);
@@ -107,7 +113,8 @@ export class WorkerCore {
 		} catch (error) {
 			this.#transport.send({ type: "result", runId, ok: false, error: errorPayload(error) });
 		} finally {
-			this.#active = null;
+			rejectPendingTools(active, new ToolError("JS run ended"));
+			if (this.#active === active) this.#active = null;
 		}
 	}
 
@@ -132,10 +139,7 @@ export class WorkerCore {
 	#close(): void {
 		const active = this.#active;
 		if (active) {
-			for (const pending of active.pendingTools.values()) {
-				pending.reject(new ToolError("JS worker closed"));
-			}
-			active.pendingTools.clear();
+			rejectPendingTools(active, new ToolError("JS worker closed"));
 		}
 		this.#active = null;
 		this.#runtime = null;

package/src/eval/py/executor.ts

@@ -108,7 +108,18 @@ interface PythonSession {
 	queue: Promise<void>;
 }
 
-const sessions = new Map<string, PythonSession>();
+interface InitializingPythonSession {
+	sessionId: string;
+	promise: Promise<PythonSession>;
+}
+
+const sessions = new Map<string, PythonSession | InitializingPythonSession>();
+
+function isInitializingSession(
+	session: PythonSession | InitializingPythonSession,
+): session is InitializingPythonSession {
+	return "promise" in session;
+}
 
 // ---------------------------------------------------------------------------
 // Cancellation plumbing
@@ -288,20 +299,48 @@ function attachOwner(session: PythonSession, sessionId: string, ownerId: string
 async function acquireSession(sessionId: string, cwd: string, options: PythonExecutorOptions): Promise<PythonSession> {
 	const existing = sessions.get(sessionId);
 	if (existing) {
-		attachOwner(existing, sessionId, options.kernelOwnerId);
-		return existing;
+		const session = isInitializingSession(existing)
+			? await waitForPromiseWithCancellation(existing.promise, options)
+			: existing;
+		attachOwner(session, sessionId, options.kernelOwnerId);
+		return session;
 	}
-	const kernel = await startKernel(cwd, options);
-	const session: PythonSession = {
+
+	const initializing: InitializingPythonSession = {
 		sessionId,
-		kernel,
-		ownerIds: new Set(),
-		hasFallbackOwner: false,
-		queue: Promise.resolve(),
+		promise: Promise.resolve().then(async () => {
+			const kernel = await startKernel(cwd, options);
+			const current = sessions.get(sessionId);
+			if (current !== initializing) {
+				await kernel.shutdown().catch(() => undefined);
+				const winner = current
+					? isInitializingSession(current)
+						? await waitForPromiseWithCancellation(current.promise, options)
+						: current
+					: undefined;
+				if (winner) return winner;
+				throw new PythonExecutionCancelledError(false);
+			}
+			const session: PythonSession = {
+				sessionId,
+				kernel,
+				ownerIds: new Set(),
+				hasFallbackOwner: false,
+				queue: Promise.resolve(),
+			};
+			sessions.set(sessionId, session);
+			return session;
+		}),
 	};
-	attachOwner(session, sessionId, options.kernelOwnerId);
-	sessions.set(sessionId, session);
-	return session;
+	sessions.set(sessionId, initializing);
+	try {
+		const session = await waitForPromiseWithCancellation(initializing.promise, options);
+		attachOwner(session, sessionId, options.kernelOwnerId);
+		return session;
+	} catch (err) {
+		if (sessions.get(sessionId) === initializing) sessions.delete(sessionId);
+		throw err;
+	}
 }
 
 async function replaceSessionKernel(
@@ -330,7 +369,8 @@ async function resetSession(sessionId: string): Promise<void> {
 	const existing = sessions.get(sessionId);
 	if (!existing) return;
 	sessions.delete(sessionId);
-	await existing.kernel.shutdown().catch(() => undefined);
+	const session = isInitializingSession(existing) ? await existing.promise.catch(() => undefined) : existing;
+	await session?.kernel.shutdown().catch(() => undefined);
 }
 
 async function runQueued<T>(
@@ -363,11 +403,20 @@ export async function disposeAllKernelSessions(): Promise<void> {
 	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 resolved = await Promise.all(
+		all.map(
+			async ([id, entry]) =>
+				[id, isInitializingSession(entry) ? await entry.promise.catch(() => undefined) : entry] as const,
+		),
+	);
+	const shutdownTargets = resolved.filter(
+		(entry): entry is readonly [string, PythonSession] => entry[1] !== undefined,
+	);
+	const results = await Promise.allSettled(shutdownTargets.map(([, session]) => session.kernel.shutdown()));
+	for (let i = 0; i < shutdownTargets.length; i += 1) {
+		const [id, session] = shutdownTargets[i];
 		const result = results[i];
-		if (result.status === "fulfilled" && result.value?.confirmed !== false) continue;
+		if (result.status === "fulfilled" && result.value.confirmed) continue;
 		const reason = result.status === "rejected" ? result.reason : "not confirmed";
 		logger.warn("Python kernel shutdown not confirmed", { sessionId: id, reason });
 		if (!sessions.has(id)) sessions.set(id, session);
@@ -377,7 +426,7 @@ export async function disposeAllKernelSessions(): Promise<void> {
 export async function disposeKernelSessionsByOwner(ownerId: string): Promise<void> {
 	const toShutdown: PythonSession[] = [];
 	for (const session of [...sessions.values()]) {
-		if (!session.ownerIds.has(ownerId)) continue;
+		if (isInitializingSession(session) || !session.ownerIds.has(ownerId)) continue;
 		if (session.ownerIds.size === 1) {
 			toShutdown.push(session);
 			continue;
@@ -391,7 +440,7 @@ export async function disposeKernelSessionsByOwner(ownerId: string): Promise<voi
 	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) {
+		if (result.status === "fulfilled" && result.value.confirmed) {
 			session.ownerIds.delete(ownerId);
 			continue;
 		}