@oh-my-pi/pi-coding-agent 12.17.0 → 12.18.0

package/src/exec/bash-executor.ts

@@ -34,6 +34,8 @@ export interface BashResult {
 	artifactId?: string;
 }
 
+const HARD_TIMEOUT_GRACE_MS = 5_000;
+
 const shellSessions = new Map<string, Shell>();
 
 export async function executeBash(command: string, options?: BashExecutorOptions): Promise<BashResult> {
@@ -65,74 +67,106 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 		};
 	}
 
+	const sessionKey = buildSessionKey(shell, prefix, snapshotPath, shellEnv, options?.sessionKey);
+	let shellSession = shellSessions.get(sessionKey);
+	if (!shellSession) {
+		shellSession = new Shell({ sessionEnv: shellEnv, snapshotPath: snapshotPath ?? undefined });
+		shellSessions.set(sessionKey, shellSession);
+	}
+	const signal = options?.signal;
+	const abortHandler = () => {
+		shellSession.abort(signal?.reason instanceof Error ? signal.reason.message : undefined);
+	};
+	if (signal) {
+		signal.addEventListener("abort", abortHandler, { once: true });
+	}
+
+	let hardTimeoutTimer: NodeJS.Timeout | undefined;
+	const hardTimeoutDeferred = Promise.withResolvers<"hard-timeout">();
+	const baseTimeoutMs = Math.max(1_000, options?.timeout ?? 300_000);
+	const hardTimeoutMs = baseTimeoutMs + HARD_TIMEOUT_GRACE_MS;
+	hardTimeoutTimer = setTimeout(() => {
+		shellSession.abort(`Hard timeout after ${Math.round(hardTimeoutMs / 1000)}s`);
+		hardTimeoutDeferred.resolve("hard-timeout");
+	}, hardTimeoutMs);
+
+	let resetSession = false;
+
 	try {
-		const sessionKey = buildSessionKey(shell, prefix, snapshotPath, shellEnv, options?.sessionKey);
-		let shellSession = shellSessions.get(sessionKey);
-		if (!shellSession) {
-			shellSession = new Shell({ sessionEnv: shellEnv, snapshotPath: snapshotPath ?? undefined });
-			shellSessions.set(sessionKey, shellSession);
+		const runPromise = shellSession.run(
+			{
+				command: finalCommand,
+				cwd: options?.cwd,
+				env: options?.env,
+				timeoutMs: options?.timeout,
+				signal,
+			},
+			(err, chunk) => {
+				if (!err) {
+					enqueueChunk(chunk);
+				}
+			},
+		);
+
+		const winner = await Promise.race([
+			runPromise.then(result => ({ kind: "result" as const, result })),
+			hardTimeoutDeferred.promise.then(() => ({ kind: "hard-timeout" as const })),
+		]);
+
+		await pendingChunks;
+
+		if (winner.kind === "hard-timeout") {
+			resetSession = true;
+			return {
+				exitCode: undefined,
+				cancelled: true,
+				...(await sink.dump(`Command exceeded hard timeout after ${Math.round(hardTimeoutMs / 1000)} seconds`)),
+			};
 		}
 
-		const signal = options?.signal;
-		const abortHandler = () => {
-			shellSession.abort(signal?.reason instanceof Error ? signal.reason.message : undefined);
-		};
-		if (signal) {
-			signal.addEventListener("abort", abortHandler, { once: true });
+		// Handle timeout
+		if (winner.result.timedOut) {
+			const annotation = options?.timeout
+				? `Command timed out after ${Math.round(options.timeout / 1000)} seconds`
+				: "Command timed out";
+			resetSession = true;
+			return {
+				exitCode: undefined,
+				cancelled: true,
+				...(await sink.dump(annotation)),
+			};
 		}
 
-		try {
-			const result = await shellSession.run(
-				{
-					command: finalCommand,
-					cwd: options?.cwd,
-					env: options?.env,
-					timeoutMs: options?.timeout,
-					signal,
-				},
-				(err, chunk) => {
-					if (!err) {
-						enqueueChunk(chunk);
-					}
-				},
-			);
-
-			await pendingChunks;
-
-			// Handle timeout
-			if (result.timedOut) {
-				const annotation = options?.timeout
-					? `Command timed out after ${Math.round(options.timeout / 1000)} seconds`
-					: "Command timed out";
-				return {
-					exitCode: undefined,
-					cancelled: true,
-					...(await sink.dump(annotation)),
-				};
-			}
-
-			// Handle cancellation
-			if (result.cancelled) {
-				return {
-					exitCode: undefined,
-					cancelled: true,
-					...(await sink.dump("Command cancelled")),
-				};
-			}
-
-			// Normal completion
+		// Handle cancellation
+		if (winner.result.cancelled) {
+			resetSession = true;
 			return {
-				exitCode: result.exitCode,
-				cancelled: false,
-				...(await sink.dump()),
+				exitCode: undefined,
+				cancelled: true,
+				...(await sink.dump("Command cancelled")),
 			};
-		} finally {
-			if (signal) {
-				signal.removeEventListener("abort", abortHandler);
-			}
 		}
+
+		// Normal completion
+		return {
+			exitCode: winner.result.exitCode,
+			cancelled: false,
+			...(await sink.dump()),
+		};
+	} catch (err) {
+		resetSession = true;
+		throw err;
 	} finally {
+		if (hardTimeoutTimer) {
+			clearTimeout(hardTimeoutTimer);
+		}
+		if (signal) {
+			signal.removeEventListener("abort", abortHandler);
+		}
 		await pendingChunks;
+		if (resetSession) {
+			shellSessions.delete(sessionKey);
+		}
 	}
 }
 

package/src/ipy/executor.ts

@@ -1,8 +1,7 @@
 import * as path from "node:path";
-import { $env, isEnoent, logger } from "@oh-my-pi/pi-utils";
+import { isEnoent, logger } from "@oh-my-pi/pi-utils";
 import { getAgentDir, getProjectDir } from "@oh-my-pi/pi-utils/dirs";
 import { OutputSink } from "../session/streaming-output";
-import { time } from "../utils/timings";
 import { shutdownSharedGateway } from "./gateway-coordinator";
 import {
 	checkPythonKernelAvailability,
@@ -15,8 +14,6 @@ import {
 import { discoverPythonModules } from "./modules";
 import { PYTHON_PRELUDE } from "./prelude";
 
-const debugStartup = $env.PI_DEBUG_STARTUP ? (stage: string) => process.stderr.write(`[startup] ${stage}\n`) : () => {};
-
 const IDLE_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
 const MAX_KERNEL_SESSIONS = 4;
 const CLEANUP_INTERVAL_MS = 30 * 1000; // 30 seconds
@@ -228,10 +225,7 @@ export async function warmPythonEnvironment(
 	const isTestEnv = Bun.env.BUN_ENV === "test" || Bun.env.NODE_ENV === "test";
 	let cacheState: PreludeCacheState | null = null;
 	try {
-		debugStartup("warmPython:ensureKernel:start");
-		await ensureKernelAvailable(cwd);
-		debugStartup("warmPython:ensureKernel:done");
-		time("warmPython:ensureKernelAvailable");
+		await logger.timeAsync("warmPython:ensureKernelAvailable", () => ensureKernelAvailable(cwd));
 	} catch (err: unknown) {
 		const reason = err instanceof Error ? err.message : String(err);
 		cachedPreludeDocs = [];
@@ -255,16 +249,15 @@ export async function warmPythonEnvironment(
 	}
 	const resolvedSessionId = sessionId ?? `session:${cwd}`;
 	try {
-		debugStartup("warmPython:withKernelSession:start");
-		const docs = await withKernelSession(
-			resolvedSessionId,
-			cwd,
-			async kernel => kernel.introspectPrelude(),
-			useSharedGateway,
-			sessionFile,
+		const docs = await logger.timeAsync("warmPython:withKernelSession", () =>
+			withKernelSession(
+				resolvedSessionId,
+				cwd,
+				async kernel => kernel.introspectPrelude(),
+				useSharedGateway,
+				sessionFile,
+			),
 		);
-		debugStartup("warmPython:withKernelSession:done");
-		time("warmPython:withKernelSession");
 		cachedPreludeDocs = docs;
 		if (!isTestEnv && docs.length > 0) {
 			const state = cacheState ?? (await buildPreludeCacheState(cwd));
@@ -321,7 +314,6 @@ async function createKernelSession(
 	artifactsDir?: string,
 	isRetry?: boolean,
 ): Promise<KernelSession> {
-	debugStartup("kernel:createSession:entry");
 	const env: Record<string, string> | undefined =
 		sessionFile || artifactsDir
 			? {
@@ -332,10 +324,9 @@ async function createKernelSession(
 
 	let kernel: PythonKernel;
 	try {
-		debugStartup("kernel:PythonKernel.start:start");
-		kernel = await PythonKernel.start({ cwd, useSharedGateway, env });
-		debugStartup("kernel:PythonKernel.start:done");
-		time("createKernelSession:PythonKernel.start");
+		kernel = await logger.timeAsync("createKernelSession:PythonKernel.start", () =>
+			PythonKernel.start({ cwd, useSharedGateway, env }),
+		);
 	} catch (err) {
 		if (!isRetry && isResourceExhaustionError(err)) {
 			await recoverFromResourceExhaustion();
@@ -412,37 +403,56 @@ async function withKernelSession<T>(
 	sessionFile?: string,
 	artifactsDir?: string,
 ): Promise<T> {
-	debugStartup("kernel:withSession:entry");
 	let session = kernelSessions.get(sessionId);
 	if (!session) {
 		// Evict oldest session if at capacity
 		if (kernelSessions.size >= MAX_KERNEL_SESSIONS) {
 			await evictOldestSession();
 		}
-		session = await createKernelSession(sessionId, cwd, useSharedGateway, sessionFile, artifactsDir);
-		debugStartup("kernel:withSession:created");
+		session = await logger.timeAsync(
+			"kernel:createKernelSession",
+			createKernelSession,
+			sessionId,
+			cwd,
+			useSharedGateway,
+			sessionFile,
+			artifactsDir,
+		);
 		kernelSessions.set(sessionId, session);
 		startCleanupTimer();
 	}
 
 	const run = async (): Promise<T> => {
-		debugStartup("kernel:withSession:run");
 		session!.lastUsedAt = Date.now();
 		if (session!.dead || !session!.kernel.isAlive()) {
-			await restartKernelSession(session!, cwd, useSharedGateway, sessionFile, artifactsDir);
+			await logger.timeAsync(
+				"kernel:restartKernelSession",
+				restartKernelSession,
+				session!,
+				cwd,
+				useSharedGateway,
+				sessionFile,
+				artifactsDir,
+			);
 		}
 		try {
-			debugStartup("kernel:withSession:handler:start");
-			const result = await handler(session!.kernel);
-			debugStartup("kernel:withSession:handler:done");
+			const result = await logger.timeAsync("kernel:withSession:handler", handler, session!.kernel);
 			session!.restartCount = 0;
 			return result;
 		} catch (err) {
 			if (!session!.dead && session!.kernel.isAlive()) {
 				throw err;
 			}
-			await restartKernelSession(session!, cwd, useSharedGateway, sessionFile, artifactsDir);
-			const result = await handler(session!.kernel);
+			await logger.timeAsync(
+				"kernel:restartKernelSession",
+				restartKernelSession,
+				session!,
+				cwd,
+				useSharedGateway,
+				sessionFile,
+				artifactsDir,
+			);
+			const result = await logger.timeAsync("kernel:postRestart:handler", handler, session!.kernel);
 			session!.restartCount = 0;
 			return result;
 		}

package/src/ipy/gateway-coordinator.ts

@@ -6,7 +6,6 @@ import { getAgentDir } from "@oh-my-pi/pi-utils/dirs";
 import type { Subprocess } from "bun";
 import { Settings } from "../config/settings";
 import { getOrCreateSnapshot } from "../utils/shell-snapshot";
-import { time } from "../utils/timings";
 import { filterEnv, resolvePythonRuntime } from "./runtime";
 
 const GATEWAY_DIR_NAME = "python-gateway";
@@ -315,12 +314,9 @@ async function killGateway(pid: number, context: string): Promise<void> {
 export async function acquireSharedGateway(cwd: string): Promise<AcquireResult | null> {
 	try {
 		return await withGatewayLock(async () => {
-			time("acquireSharedGateway:lockAcquired");
-			const existingInfo = await readGatewayInfo();
-			time("acquireSharedGateway:readInfo");
+			const existingInfo = await logger.timeAsync("acquireSharedGateway:readInfo", () => readGatewayInfo());
 			if (existingInfo) {
-				if (await isGatewayAlive(existingInfo)) {
-					time("acquireSharedGateway:isAlive");
+				if (await logger.timeAsync("acquireSharedGateway:isAlive", () => isGatewayAlive(existingInfo))) {
 					localGatewayUrl = existingInfo.url;
 					isCoordinatorInitialized = true;
 					logger.debug("Reusing global Python gateway", { url: existingInfo.url });
@@ -334,8 +330,9 @@ export async function acquireSharedGateway(cwd: string): Promise<AcquireResult |
 				await clearGatewayInfo();
 			}
 
-			const { url, pid, pythonPath, venvPath } = await startGatewayProcess(cwd);
-			time("acquireSharedGateway:startGateway");
+			const { url, pid, pythonPath, venvPath } = await logger.timeAsync("acquireSharedGateway:startGateway", () =>
+				startGatewayProcess(cwd),
+			);
 			const info: GatewayInfo = {
 				url,
 				pid,

package/src/ipy/kernel.ts

@@ -1,7 +1,6 @@
 import { $env, logger, Snowflake } from "@oh-my-pi/pi-utils";
 import { $ } from "bun";
 import { Settings } from "../config/settings";
-import { time } from "../utils/timings";
 import { htmlToBasicMarkdown } from "../web/scrapers/types";
 import { acquireSharedGateway, releaseSharedGateway, shutdownSharedGateway } from "./gateway-coordinator";
 import { loadPythonModules } from "./modules";
@@ -13,8 +12,6 @@ const TEXT_DECODER = new TextDecoder();
 const TRACE_IPC = $env.PI_PYTHON_IPC_TRACE === "1";
 const PRELUDE_INTROSPECTION_SNIPPET = "import json\nprint(json.dumps(__omp_prelude_docs__()))";
 
-const debugStartup = $env.PI_DEBUG_STARTUP ? (stage: string) => process.stderr.write(`[startup] ${stage}\n`) : () => {};
-
 class SharedGatewayCreateError extends Error {
 	constructor(
 		readonly status: number,
@@ -334,10 +331,9 @@ export class PythonKernel {
 	}
 
 	static async start(options: KernelStartOptions): Promise<PythonKernel> {
-		debugStartup("PythonKernel.start:entry");
-		const availability = await checkPythonKernelAvailability(options.cwd);
-		debugStartup("PythonKernel.start:availCheck");
-		time("PythonKernel.start:availabilityCheck");
+		const availability = await logger.timeAsync("PythonKernel.start:availabilityCheck", () =>
+			checkPythonKernelAvailability(options.cwd),
+		);
 		if (!availability.ok) {
 			throw new Error(availability.reason ?? "Python kernel unavailable");
 		}
@@ -353,20 +349,18 @@ export class PythonKernel {
 
 		for (let attempt = 0; attempt < 2; attempt += 1) {
 			try {
-				debugStartup("PythonKernel.start:acquireShared:start");
-				const sharedResult = await acquireSharedGateway(options.cwd);
-				debugStartup("PythonKernel.start:acquireShared:done");
-				time("PythonKernel.start:acquireSharedGateway");
+				const sharedResult = await logger.timeAsync("PythonKernel.start:acquireSharedGateway", () =>
+					acquireSharedGateway(options.cwd),
+				);
 				if (!sharedResult) {
 					throw new Error("Shared Python gateway unavailable");
 				}
-				debugStartup("PythonKernel.start:startShared:start");
-				const kernel = await PythonKernel.#startWithSharedGateway(sharedResult.url, options.cwd, options.env);
-				debugStartup("PythonKernel.start:startShared:done");
-				time("PythonKernel.start:startWithSharedGateway");
+				const kernel = await logger.timeAsync("PythonKernel.start:startWithSharedGateway", () =>
+					PythonKernel.#startWithSharedGateway(sharedResult.url, options.cwd, options.env),
+				);
 				return kernel;
 			} catch (err) {
-				debugStartup("PythonKernel.start:sharedFailed");
+				logger.debug("PythonKernel.start:sharedFailed");
 				if (attempt === 0 && err instanceof SharedGatewayCreateError && err.status >= 500) {
 					logger.warn("Shared gateway kernel creation failed, retrying", {
 						status: err.status,
@@ -436,55 +430,42 @@ export class PythonKernel {
 		cwd: string,
 		env?: Record<string, string | undefined>,
 	): Promise<PythonKernel> {
-		debugStartup("sharedGateway:fetch:start");
-		const createResponse = await fetch(`${gatewayUrl}/api/kernels`, {
-			method: "POST",
-			headers: { "Content-Type": "application/json" },
-			body: JSON.stringify({ name: "python3" }),
-		});
-		debugStartup("sharedGateway:fetch:done");
-		time("startWithSharedGateway:createKernel");
+		const createResponse = await logger.timeAsync("startWithSharedGateway:createKernel", () =>
+			fetch(`${gatewayUrl}/api/kernels`, {
+				method: "POST",
+				headers: { "Content-Type": "application/json" },
+				body: JSON.stringify({ name: "python3" }),
+			}),
+		);
 
 		if (!createResponse.ok) {
-			debugStartup(`sharedGateway:fetch:notOk:${createResponse.status}`);
+			logger.debug(`sharedGateway:fetch:notOk:${createResponse.status}`);
 			await shutdownSharedGateway();
-			debugStartup("sharedGateway:fetch:shutdown");
 			const text = await createResponse.text();
-			debugStartup("sharedGateway:fetch:textRead");
 			throw new SharedGatewayCreateError(
 				createResponse.status,
 				`Failed to create kernel on shared gateway: ${text}`,
 			);
 		}
 
-		debugStartup("sharedGateway:json:start");
-		const kernelInfo = (await createResponse.json()) as { id: string };
-		debugStartup("sharedGateway:json:done");
+		const kernelInfo = await logger.timeAsync(
+			"startWithSharedGateway:parseJson",
+			() => createResponse.json() as Promise<{ id: string }>,
+		);
 		const kernelId = kernelInfo.id;
 
 		const kernel = new PythonKernel(Snowflake.next(), kernelId, gatewayUrl, Snowflake.next(), "omp", true);
-		debugStartup("sharedGateway:kernelCreated");
 
 		try {
-			debugStartup("sharedGateway:connectWS:start");
-			await kernel.#connectWebSocket();
-			debugStartup("sharedGateway:connectWS:done");
-			time("startWithSharedGateway:connectWS");
-			debugStartup("sharedGateway:initEnv:start");
-			await kernel.#initializeKernelEnvironment(cwd, env);
-			debugStartup("sharedGateway:initEnv:done");
-			time("startWithSharedGateway:initEnv");
-			debugStartup("sharedGateway:prelude:start");
-			const preludeResult = await kernel.execute(PYTHON_PRELUDE, { silent: true, storeHistory: false });
-			debugStartup("sharedGateway:prelude:done");
-			time("startWithSharedGateway:prelude");
+			await logger.timeAsync("startWithSharedGateway:connectWS", () => kernel.#connectWebSocket());
+			await logger.timeAsync("startWithSharedGateway:initEnv", () => kernel.#initializeKernelEnvironment(cwd, env));
+			const preludeResult = await logger.timeAsync("startWithSharedGateway:prelude", () =>
+				kernel.execute(PYTHON_PRELUDE, { silent: true, storeHistory: false }),
+			);
 			if (preludeResult.cancelled || preludeResult.status === "error") {
 				throw new Error("Failed to initialize Python kernel prelude");
 			}
-			debugStartup("sharedGateway:loadModules:start");
-			await loadPythonModules(kernel, { cwd });
-			debugStartup("sharedGateway:loadModules:done");
-			time("startWithSharedGateway:loadModules");
+			await logger.timeAsync("startWithSharedGateway:loadModules", () => loadPythonModules(kernel, { cwd }));
 			return kernel;
 		} catch (err: unknown) {
 			await kernel.shutdown();