diff --git a/packages/coding-agent/src/ipy/cancellation.ts b/packages/coding-agent/src/ipy/cancellation.ts new file mode 100644 index 000000000..7e8852731 --- /dev/null +++ b/packages/coding-agent/src/ipy/cancellation.ts @@ -0,0 +1,28 @@ +export function getAbortReason(signal: AbortSignal | undefined, fallbackReason: string): Error { + if (signal?.reason instanceof Error) return signal.reason; + if (typeof signal?.reason === "string" && signal.reason.length > 0) { + return new Error(signal.reason); + } + + return new Error(fallbackReason); +} + +export function createCancellationError(name: "AbortError" | "TimeoutError", message: string): Error { + const error = new Error(message); + error.name = name; + return error; +} + +export function getExecutionCancellationError( + result: { timedOut?: boolean }, + signal: AbortSignal | undefined, + fallbackReason: string, +): Error { + if (signal?.aborted) { + return getAbortReason(signal, fallbackReason); + } + if (result.timedOut) { + return createCancellationError("TimeoutError", fallbackReason); + } + return createCancellationError("AbortError", fallbackReason); +} diff --git a/packages/coding-agent/src/ipy/executor.ts b/packages/coding-agent/src/ipy/executor.ts index 1a8b9a023..ad7abd1e0 100644 --- a/packages/coding-agent/src/ipy/executor.ts +++ b/packages/coding-agent/src/ipy/executor.ts @@ -24,6 +24,8 @@ export interface PythonExecutorOptions { cwd?: string; /** Timeout in milliseconds */ timeoutMs?: number; + /** Absolute wall-clock deadline in milliseconds since epoch */ + deadlineMs?: number; /** Callback for streaming output chunks (already sanitized) */ onChunk?: (chunk: string) => Promise | void; /** AbortSignal for cancellation */ @@ -86,6 +88,151 @@ const kernelSessions = new Map(); let cachedPreludeDocs: PreludeHelper[] | null = null; let cleanupTimer: NodeJS.Timeout | null = null; +interface KernelSessionExecutionOptions { + useSharedGateway?: boolean; + sessionFile?: string; + signal?: AbortSignal; + deadlineMs?: number; +} + +class PythonExecutionCancelledError extends Error { + readonly timedOut: boolean; + + constructor(timedOut: boolean) { + super(timedOut ? "Command timed out" : "Command aborted"); + this.name = timedOut ? "TimeoutError" : "AbortError"; + this.timedOut = timedOut; + } +} + +function getExecutionDeadlineMs(options?: Pick): number | undefined { + if (options?.deadlineMs !== undefined) return options.deadlineMs; + if (options?.timeoutMs === undefined) return undefined; + return Date.now() + options.timeoutMs; +} + +function getRemainingTimeoutMs(deadlineMs?: number): number | undefined { + if (deadlineMs === undefined) return undefined; + return deadlineMs - Date.now(); +} + +function requireRemainingTimeoutMs(deadlineMs?: number): number | undefined { + const remainingMs = getRemainingTimeoutMs(deadlineMs); + if (remainingMs === undefined) return undefined; + if (remainingMs <= 0) { + throw new PythonExecutionCancelledError(true); + } + return remainingMs; +} + +function isCancellationError(error: unknown): boolean { + return ( + error instanceof PythonExecutionCancelledError || + (error instanceof DOMException && (error.name === "AbortError" || error.name === "TimeoutError")) || + (error instanceof Error && (error.name === "AbortError" || error.name === "TimeoutError")) + ); +} + +function isTimedOutCancellation(error: unknown, signal?: AbortSignal): boolean { + if (error instanceof PythonExecutionCancelledError) return error.timedOut; + if (error instanceof DOMException) return error.name === "TimeoutError"; + if (error instanceof Error && error.name === "TimeoutError") return true; + const reason = signal?.reason; + if (reason instanceof DOMException) return reason.name === "TimeoutError"; + return reason instanceof Error ? reason.name === "TimeoutError" : false; +} + +async function waitForQueueTurn( + queue: Promise, + options: Pick, +): Promise { + if (options.signal?.aborted) { + throw new PythonExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); + } + + const remainingMs = getRemainingTimeoutMs(options.deadlineMs); + if (remainingMs !== undefined && remainingMs <= 0) { + throw new PythonExecutionCancelledError(true); + } + + if (!options.signal && remainingMs === undefined) { + await queue; + return; + } + + await new Promise((resolve, reject) => { + const cleanups: Array<() => void> = []; + const finish = (callback: () => void) => { + while (cleanups.length > 0) { + cleanups.pop()?.(); + } + callback(); + }; + + const onAbort = () => { + finish(() => + reject(new PythonExecutionCancelledError(isTimedOutCancellation(options.signal?.reason, options.signal))), + ); + }; + + if (options.signal) { + options.signal.addEventListener("abort", onAbort, { once: true }); + cleanups.push(() => options.signal?.removeEventListener("abort", onAbort)); + } + + if (remainingMs !== undefined) { + const timeout = setTimeout(() => { + finish(() => reject(new PythonExecutionCancelledError(true))); + }, remainingMs); + timeout.unref(); + cleanups.push(() => clearTimeout(timeout)); + } + + queue.then( + () => finish(resolve), + error => finish(() => reject(error)), + ); + }); +} + +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`; +} + +function createCancelledPythonResult(timedOut: boolean, timeoutMs?: number): PythonResult { + const output = timedOut ? (formatTimeoutAnnotation(timeoutMs) ?? "Command timed out") : ""; + const outputBytes = Buffer.byteLength(output, "utf-8"); + const outputLines = output.length > 0 ? 1 : 0; + return { + output, + exitCode: undefined, + cancelled: true, + truncated: false, + totalLines: outputLines, + totalBytes: outputBytes, + outputLines, + outputBytes, + displayOutputs: [], + stdinRequested: false, + }; +} + +function buildKernelStartOptions( + cwd: string, + env: Record | undefined, + options: KernelSessionExecutionOptions, +) { + return { + cwd, + env, + useSharedGateway: options.useSharedGateway, + signal: options.signal, + deadlineMs: options.deadlineMs, + }; +} + interface PreludeCacheSource { path: string; hash: string; @@ -247,13 +394,10 @@ export async function warmPythonEnvironment( const resolvedSessionId = sessionId ?? `session:${cwd}`; try { const docs = await logger.timeAsync("warmPython:withKernelSession", () => - withKernelSession( - resolvedSessionId, - cwd, - async kernel => kernel.introspectPrelude(), + withKernelSession(resolvedSessionId, cwd, async kernel => kernel.introspectPrelude(), { useSharedGateway, sessionFile, - ), + }), ); cachedPreludeDocs = docs; if (!isTestEnv && docs.length > 0) { @@ -306,21 +450,22 @@ async function recoverFromResourceExhaustion(): Promise { async function createKernelSession( sessionId: string, cwd: string, - useSharedGateway?: boolean, - sessionFile?: string, + options: KernelSessionExecutionOptions = {}, isRetry?: boolean, ): Promise { - const env: Record | undefined = sessionFile ? { PI_SESSION_FILE: sessionFile } : undefined; + requireRemainingTimeoutMs(options.deadlineMs); + const env: Record | undefined = options.sessionFile + ? { PI_SESSION_FILE: options.sessionFile } + : undefined; + const startOptions = buildKernelStartOptions(cwd, env, options); let kernel: PythonKernel; try { - kernel = await logger.timeAsync("createKernelSession:PythonKernel.start", () => - PythonKernel.start({ cwd, useSharedGateway, env }), - ); + kernel = await logger.timeAsync("createKernelSession:PythonKernel.start", () => PythonKernel.start(startOptions)); } catch (err) { if (!isRetry && isResourceExhaustionError(err)) { await recoverFromResourceExhaustion(); - return createKernelSession(sessionId, cwd, useSharedGateway, sessionFile, true); + return createKernelSession(sessionId, cwd, options, true); } throw err; } @@ -347,20 +492,23 @@ async function createKernelSession( async function restartKernelSession( session: KernelSession, cwd: string, - useSharedGateway?: boolean, - sessionFile?: string, + options: KernelSessionExecutionOptions = {}, ): Promise { session.restartCount += 1; if (session.restartCount > 1) { throw new Error("Python kernel restarted too many times in this session"); } + requireRemainingTimeoutMs(options.deadlineMs); try { await session.kernel.shutdown(); } catch (err) { logger.warn("Failed to shutdown crashed kernel", { error: err instanceof Error ? err.message : String(err) }); } - const env: Record | undefined = sessionFile ? { PI_SESSION_FILE: sessionFile } : undefined; - const kernel = await PythonKernel.start({ cwd, useSharedGateway, env }); + const env: Record | undefined = options.sessionFile + ? { PI_SESSION_FILE: options.sessionFile } + : undefined; + const startOptions = buildKernelStartOptions(cwd, env, options); + const kernel = await PythonKernel.start(startOptions); session.kernel = kernel; session.dead = false; session.lastUsedAt = Date.now(); @@ -382,23 +530,18 @@ async function withKernelSession( sessionId: string, cwd: string, handler: (kernel: PythonKernel) => Promise, - useSharedGateway?: boolean, - sessionFile?: string, + options: KernelSessionExecutionOptions = {}, ): Promise { let session = kernelSessions.get(sessionId); if (!session) { - // Evict oldest session if at capacity if (kernelSessions.size >= MAX_KERNEL_SESSIONS) { await evictOldestSession(); } - session = await logger.timeAsync( - "kernel:createKernelSession", - createKernelSession, - sessionId, - cwd, - useSharedGateway, - sessionFile, - ); + requireRemainingTimeoutMs(options.deadlineMs); + if (options.signal?.aborted) { + throw new PythonExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); + } + session = await logger.timeAsync("kernel:createKernelSession", createKernelSession, sessionId, cwd, options); kernelSessions.set(sessionId, session); startCleanupTimer(); } @@ -406,14 +549,7 @@ async function withKernelSession( const run = async (): Promise => { session!.lastUsedAt = Date.now(); if (session!.dead || !session!.kernel.isAlive()) { - await logger.timeAsync( - "kernel:restartKernelSession", - restartKernelSession, - session!, - cwd, - useSharedGateway, - sessionFile, - ); + await logger.timeAsync("kernel:restartKernelSession", restartKernelSession, session!, cwd, options); } try { const result = await logger.timeAsync("kernel:withSession:handler", handler, session!.kernel); @@ -423,26 +559,34 @@ async function withKernelSession( if (!session!.dead && session!.kernel.isAlive()) { throw err; } - await logger.timeAsync( - "kernel:restartKernelSession", - restartKernelSession, - session!, - cwd, - useSharedGateway, - sessionFile, - ); + await logger.timeAsync("kernel:restartKernelSession", restartKernelSession, session!, cwd, options); const result = await logger.timeAsync("kernel:postRestart:handler", handler, session!.kernel); session!.restartCount = 0; return result; } }; - const task = session.queue.then(run, run); - session.queue = task.then( - () => undefined, - () => undefined, - ); - return task; + const queue = session.queue; + let releaseTurn: (() => void) | undefined; + const turn = new Promise(resolve => { + releaseTurn = resolve; + }); + session.queue = queue + .then( + () => turn, + () => turn, + ) + .then( + () => undefined, + () => undefined, + ); + + try { + await waitForQueueTurn(queue, options); + return await run(); + } finally { + releaseTurn?.(); + } } async function executeWithKernel( @@ -456,19 +600,20 @@ async function executeWithKernel( artifactId: options?.artifactId, }); const displayOutputs: KernelDisplayOutput[] = []; + const deadlineMs = getExecutionDeadlineMs(options); + let executionTimeoutMs: number | undefined; try { + executionTimeoutMs = requireRemainingTimeoutMs(deadlineMs); const result = await kernel.execute(code, { signal: options?.signal, - timeoutMs: options?.timeoutMs, + timeoutMs: executionTimeoutMs, onChunk: text => sink.push(text), onDisplay: output => void displayOutputs.push(output), }); if (result.cancelled) { - const secs = options?.timeoutMs ? Math.round(options.timeoutMs / 1000) : undefined; - const annotation = - result.timedOut && secs !== undefined ? `Command timed out after ${secs} seconds` : undefined; + const annotation = result.timedOut ? formatTimeoutAnnotation(executionTimeoutMs) : undefined; return { exitCode: undefined, cancelled: true, @@ -497,6 +642,16 @@ async function executeWithKernel( ...(await sink.dump()), }; } catch (err) { + if (isCancellationError(err) || options?.signal?.aborted) { + const timedOut = isTimedOutCancellation(err, options?.signal); + return { + exitCode: undefined, + cancelled: true, + displayOutputs, + stdinRequested: false, + ...(await sink.dump(timedOut ? formatTimeoutAnnotation(executionTimeoutMs) : undefined)), + }; + } const error = err instanceof Error ? err : new Error(String(err)); logger.error("Python execution failed", { error: error.message }); throw error; @@ -513,34 +668,54 @@ export async function executePythonWithKernel( export async function executePython(code: string, options?: PythonExecutorOptions): Promise { const cwd = options?.cwd ?? getProjectDir(); - await ensureKernelAvailable(cwd); + const deadlineMs = getExecutionDeadlineMs(options); + const executionOptions: PythonExecutorOptions = { + ...(options ?? {}), + deadlineMs, + }; - const kernelMode = options?.kernelMode ?? "session"; - const useSharedGateway = options?.useSharedGateway; - const sessionFile = options?.sessionFile; - - if (kernelMode === "per-call") { - const env: Record | undefined = sessionFile ? { PI_SESSION_FILE: sessionFile } : undefined; - const kernel = await PythonKernel.start({ cwd, useSharedGateway, env }); - try { - return await executeWithKernel(kernel, code, options); - } finally { - await kernel.shutdown(); + try { + requireRemainingTimeoutMs(deadlineMs); + if (executionOptions.signal?.aborted) { + throw new PythonExecutionCancelledError( + isTimedOutCancellation(executionOptions.signal.reason, executionOptions.signal), + ); } - } - const sessionId = options?.sessionId ?? `session:${cwd}`; - if (options?.reset) { - const existing = kernelSessions.get(sessionId); - if (existing) { - await disposeKernelSession(existing); + await ensureKernelAvailable(cwd); + + const kernelMode = executionOptions.kernelMode ?? "session"; + const sessionFile = executionOptions.sessionFile; + + if (kernelMode === "per-call") { + const env: Record | undefined = sessionFile ? { PI_SESSION_FILE: sessionFile } : undefined; + requireRemainingTimeoutMs(deadlineMs); + const startOptions = buildKernelStartOptions(cwd, env, executionOptions); + const kernel = await PythonKernel.start(startOptions); + try { + return await executeWithKernel(kernel, code, executionOptions); + } finally { + await kernel.shutdown(); + } } + + const sessionId = executionOptions.sessionId ?? `session:${cwd}`; + if (executionOptions.reset) { + const existing = kernelSessions.get(sessionId); + if (existing) { + await disposeKernelSession(existing); + } + } + return await withKernelSession( + sessionId, + cwd, + async kernel => executeWithKernel(kernel, code, executionOptions), + executionOptions, + ); + } catch (err) { + if (isCancellationError(err) || executionOptions.signal?.aborted) { + return createCancelledPythonResult(isTimedOutCancellation(err, executionOptions.signal)); + } + throw err; } - return await withKernelSession( - sessionId, - cwd, - async kernel => executeWithKernel(kernel, code, options), - useSharedGateway, - sessionFile, - ); } diff --git a/packages/coding-agent/src/ipy/kernel.ts b/packages/coding-agent/src/ipy/kernel.ts index f0f028ea2..941ef6507 100644 --- a/packages/coding-agent/src/ipy/kernel.ts +++ b/packages/coding-agent/src/ipy/kernel.ts @@ -2,6 +2,7 @@ import { $env, logger, Snowflake } from "@oh-my-pi/pi-utils"; import { $ } from "bun"; import { Settings } from "../config/settings"; import { htmlToBasicMarkdown } from "../web/scrapers/types"; +import { createCancellationError, getAbortReason, getExecutionCancellationError } from "./cancellation"; import { acquireSharedGateway, releaseSharedGateway, shutdownSharedGateway } from "./gateway-coordinator"; import { loadPythonModules } from "./modules"; import { PYTHON_PRELUDE } from "./prelude"; @@ -35,6 +36,94 @@ function getExternalGatewayConfig(): ExternalGatewayConfig | null { }; } +const STARTUP_CLEANUP_TIMEOUT_MS = 2_000; +const WEBSOCKET_CONNECT_TIMEOUT_MS = 10_000; + +interface KernelLifecycleOptions { + signal?: AbortSignal; + deadlineMs?: number; +} + +interface KernelShutdownOptions { + signal?: AbortSignal; + timeoutMs?: number; +} + +function getRemainingTimeMs(deadlineMs?: number): number | undefined { + if (deadlineMs === undefined) return undefined; + return Math.max(0, deadlineMs - Date.now()); +} + +function throwIfStartupExecutionFailed( + result: Pick, + signal: AbortSignal | undefined, + failureMessage: string, +): void { + if (result.cancelled) { + throw getExecutionCancellationError(result, signal, failureMessage); + } + if (result.status === "error") { + throw new Error(failureMessage); + } +} + +function createAbortedSignal(reason: Error): AbortSignal { + const controller = new AbortController(); + controller.abort(reason); + return controller.signal; +} + +function combineAbortSignal( + options: KernelLifecycleOptions, + timeoutCapMs?: number, + fallbackReason = "Operation aborted", +): AbortSignal | undefined { + if (options.signal?.aborted) { + return options.signal; + } + + const signals: AbortSignal[] = []; + if (options.signal) { + signals.push(options.signal); + } + + const remainingMs = getRemainingTimeMs(options.deadlineMs); + const timeoutMs = + remainingMs === undefined + ? timeoutCapMs + : timeoutCapMs === undefined + ? remainingMs + : Math.min(remainingMs, timeoutCapMs); + + if (timeoutMs !== undefined) { + if (timeoutMs <= 0) { + return createAbortedSignal(createCancellationError("TimeoutError", fallbackReason)); + } + signals.push(AbortSignal.timeout(timeoutMs)); + } + + if (signals.length === 0) return undefined; + return signals.length === 1 ? signals[0] : AbortSignal.any(signals); +} + +function throwIfAborted(signal: AbortSignal | undefined, fallbackReason: string): void { + if (!signal?.aborted) return; + throw getAbortReason(signal, fallbackReason); +} + +function getStartupExecuteOptions(options: KernelLifecycleOptions): Pick { + return { + signal: combineAbortSignal(options, undefined, "Python kernel startup aborted"), + timeoutMs: getRemainingTimeMs(options.deadlineMs), + }; +} + +function getStartupCleanupTimeoutMs(deadlineMs?: number): number { + const remainingMs = getRemainingTimeMs(deadlineMs); + if (remainingMs === undefined || remainingMs <= 0) return STARTUP_CLEANUP_TIMEOUT_MS; + return Math.min(STARTUP_CLEANUP_TIMEOUT_MS, remainingMs); +} + export interface JupyterHeader { msg_id: string; session: string; @@ -93,7 +182,7 @@ export interface PreludeHelper { category: string; } -interface KernelStartOptions { +interface KernelStartOptions extends KernelLifecycleOptions { cwd: string; env?: Record; useSharedGateway?: boolean; @@ -339,9 +428,12 @@ export class PythonKernel { throw new Error(availability.reason ?? "Python kernel unavailable"); } + const startup = { signal: options.signal, deadlineMs: options.deadlineMs }; + const startupSignal = combineAbortSignal(startup, undefined, "Python kernel startup aborted"); + const externalConfig = getExternalGatewayConfig(); if (externalConfig) { - return PythonKernel.#startWithExternalGateway(externalConfig, options.cwd, options.env); + return PythonKernel.#startWithExternalGateway(externalConfig, options.cwd, options.env, startup); } if (options.useSharedGateway === false) { @@ -349,6 +441,7 @@ export class PythonKernel { } for (let attempt = 0; attempt < 2; attempt += 1) { + throwIfAborted(startupSignal, "Python kernel startup aborted"); try { const sharedResult = await logger.timeAsync("PythonKernel.start:acquireSharedGateway", () => acquireSharedGateway(options.cwd), @@ -357,7 +450,7 @@ export class PythonKernel { throw new Error("Shared Python gateway unavailable"); } const kernel = await logger.timeAsync("PythonKernel.start:startWithSharedGateway", () => - PythonKernel.#startWithSharedGateway(sharedResult.url, options.cwd, options.env), + PythonKernel.#startWithSharedGateway(sharedResult.url, options.cwd, options.env, startup), ); return kernel; } catch (err) { @@ -382,16 +475,20 @@ export class PythonKernel { config: ExternalGatewayConfig, cwd: string, env?: Record, + startup: KernelLifecycleOptions = {}, ): Promise { const headers: Record = { "Content-Type": "application/json" }; if (config.token) { headers.Authorization = `token ${config.token}`; } + const startupSignal = combineAbortSignal(startup, undefined, "Python kernel startup aborted"); + throwIfAborted(startupSignal, "Python kernel startup aborted"); const createResponse = await fetch(`${config.url}/api/kernels`, { method: "POST", headers, body: JSON.stringify({ name: "python3" }), + signal: startupSignal, }); if (!createResponse.ok) { @@ -412,16 +509,23 @@ export class PythonKernel { ); try { - await kernel.#connectWebSocket(); - await kernel.#initializeKernelEnvironment(cwd, env); - const preludeResult = await kernel.execute(PYTHON_PRELUDE, { silent: true, storeHistory: false }); - if (preludeResult.cancelled || preludeResult.status === "error") { - throw new Error("Failed to initialize Python kernel prelude"); - } - await loadPythonModules(kernel, { cwd }); + await kernel.#connectWebSocket(startup); + await kernel.#initializeKernelEnvironment(cwd, env, startup); + const preludeOptions = getStartupExecuteOptions(startup); + const preludeResult = await kernel.execute(PYTHON_PRELUDE, { + ...preludeOptions, + silent: true, + storeHistory: false, + }); + throwIfStartupExecutionFailed( + preludeResult, + preludeOptions.signal, + "Failed to initialize Python kernel prelude", + ); + await loadPythonModules(kernel, { cwd, signal: startup.signal, deadlineMs: startup.deadlineMs }); return kernel; } catch (err: unknown) { - await kernel.shutdown(); + await kernel.shutdown({ timeoutMs: getStartupCleanupTimeoutMs(startup.deadlineMs) }); throw err; } } @@ -430,12 +534,16 @@ export class PythonKernel { gatewayUrl: string, cwd: string, env?: Record, + startup: KernelLifecycleOptions = {}, ): Promise { + const startupSignal = combineAbortSignal(startup, undefined, "Python kernel startup aborted"); + throwIfAborted(startupSignal, "Python kernel startup aborted"); const createResponse = await logger.timeAsync("startWithSharedGateway:createKernel", () => fetch(`${gatewayUrl}/api/kernels`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ name: "python3" }), + signal: startupSignal, }), ); @@ -458,46 +566,70 @@ export class PythonKernel { const kernel = new PythonKernel(Snowflake.next(), kernelId, gatewayUrl, Snowflake.next(), "omp", true); try { - 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 }), + await logger.timeAsync("startWithSharedGateway:connectWS", () => kernel.#connectWebSocket(startup)); + await logger.timeAsync("startWithSharedGateway:initEnv", () => + kernel.#initializeKernelEnvironment(cwd, env, startup), + ); + const preludeOptions = getStartupExecuteOptions(startup); + const preludeResult = await logger.timeAsync("startWithSharedGateway:prelude", () => + kernel.execute(PYTHON_PRELUDE, { + ...preludeOptions, + silent: true, + storeHistory: false, + }), + ); + throwIfStartupExecutionFailed( + preludeResult, + preludeOptions.signal, + "Failed to initialize Python kernel prelude", + ); + await logger.timeAsync("startWithSharedGateway:loadModules", () => + loadPythonModules(kernel, { cwd, signal: startup.signal, deadlineMs: startup.deadlineMs }), ); - if (preludeResult.cancelled || preludeResult.status === "error") { - throw new Error("Failed to initialize Python kernel prelude"); - } - await logger.timeAsync("startWithSharedGateway:loadModules", () => loadPythonModules(kernel, { cwd })); return kernel; } catch (err: unknown) { - await kernel.shutdown(); + await kernel.shutdown({ timeoutMs: getStartupCleanupTimeoutMs(startup.deadlineMs) }); throw err; } } - async #connectWebSocket(): Promise { + async #connectWebSocket(options: KernelLifecycleOptions = {}): Promise { const wsBase = this.gatewayUrl.replace(/^http/, "ws"); let wsUrl = `${wsBase}/api/kernels/${this.kernelId}/channels`; if (this.#authToken) { wsUrl += `?token=${encodeURIComponent(this.#authToken)}`; } + const connectSignal = combineAbortSignal(options, WEBSOCKET_CONNECT_TIMEOUT_MS, "WebSocket connection timeout"); + throwIfAborted(connectSignal, "WebSocket connection timeout"); + const { promise, resolve, reject } = Promise.withResolvers(); const ws = new WebSocket(wsUrl); ws.binaryType = "arraybuffer"; let settled = false; - const timeout = setTimeout(() => { - ws.close(); - if (!settled) { - settled = true; - reject(new Error("WebSocket connection timeout")); + const finalize = (): void => { + if (connectSignal) { + connectSignal.removeEventListener("abort", onAbort); } - }, 10000); + }; + + const onAbort = () => { + ws.close(); + if (settled) return; + settled = true; + finalize(); + reject(getAbortReason(connectSignal, "WebSocket connection timeout")); + }; + + if (connectSignal) { + connectSignal.addEventListener("abort", onAbort, { once: true }); + } ws.onopen = () => { if (settled) return; settled = true; - clearTimeout(timeout); + finalize(); this.#ws = ws; resolve(); }; @@ -506,7 +638,7 @@ export class PythonKernel { const error = new Error(`WebSocket error: ${event}`); if (!settled) { settled = true; - clearTimeout(timeout); + finalize(); reject(error); return; } @@ -520,7 +652,7 @@ export class PythonKernel { this.#ws = null; if (!settled) { settled = true; - clearTimeout(timeout); + finalize(); reject(new Error("WebSocket closed before connection")); return; } @@ -561,7 +693,11 @@ export class PythonKernel { return promise; } - async #initializeKernelEnvironment(cwd: string, env?: Record): Promise { + async #initializeKernelEnvironment( + cwd: string, + env?: Record, + options: KernelLifecycleOptions = {}, + ): Promise { const envEntries = Object.entries(env ?? {}).filter(([, value]) => value !== undefined); const envPayload = Object.fromEntries(envEntries); const initScript = [ @@ -572,10 +708,13 @@ export class PythonKernel { "for __omp_key, __omp_val in __omp_env.items():\n os.environ[__omp_key] = __omp_val", "if __omp_cwd not in sys.path:\n sys.path.insert(0, __omp_cwd)", ].join("\n"); - const result = await this.execute(initScript, { silent: true, storeHistory: false }); - if (result.cancelled || result.status === "error") { - throw new Error("Failed to initialize Python kernel environment"); - } + const executeOptions = getStartupExecuteOptions(options); + const result = await this.execute(initScript, { + ...executeOptions, + silent: true, + storeHistory: false, + }); + throwIfStartupExecutionFailed(result, executeOptions.signal, "Failed to initialize Python kernel environment"); } #abortPendingExecutions(reason: string): void { @@ -842,16 +981,23 @@ export class PythonKernel { } } - async shutdown(): Promise { + async shutdown(options?: KernelShutdownOptions): Promise { if (this.#disposed) return; this.#disposed = true; this.#alive = false; this.#abortPendingExecutions("Kernel shutdown"); + const shutdownSignal = combineAbortSignal( + { signal: options?.signal }, + options?.timeoutMs, + "Python kernel shutdown timed out", + ); + try { await fetch(`${this.gatewayUrl}/api/kernels/${this.kernelId}`, { method: "DELETE", headers: this.#authHeaders(), + signal: shutdownSignal, }); } catch (err: unknown) { logger.warn("Failed to delete kernel via API", { error: err instanceof Error ? err.message : String(err) }); diff --git a/packages/coding-agent/src/ipy/modules.ts b/packages/coding-agent/src/ipy/modules.ts index 55595a242..062cb6aa9 100644 --- a/packages/coding-agent/src/ipy/modules.ts +++ b/packages/coding-agent/src/ipy/modules.ts @@ -1,6 +1,7 @@ import * as fs from "node:fs/promises"; import * as path from "node:path"; import { getAgentModulesDir, getProjectDir, getProjectModulesDir } from "@oh-my-pi/pi-utils"; +import { getExecutionCancellationError } from "./cancellation"; export type PythonModuleSource = "user" | "project"; @@ -13,13 +14,14 @@ export interface PythonModuleEntry { export interface PythonModuleExecuteResult { status: "ok" | "error"; cancelled: boolean; + timedOut?: boolean; error?: { name: string; value: string; traceback: string[] }; } export interface PythonModuleExecutor { execute: ( code: string, - options?: { silent?: boolean; storeHistory?: boolean }, + options?: { signal?: AbortSignal; timeoutMs?: number; silent?: boolean; storeHistory?: boolean }, ) => Promise; } @@ -30,6 +32,12 @@ export interface DiscoverPythonModulesOptions { agentDir?: string; } +export interface LoadPythonModulesOptions extends DiscoverPythonModulesOptions { + signal?: AbortSignal; + timeoutMs?: number; + deadlineMs?: number; +} + interface ModuleCandidate { name: string; path: string; @@ -61,6 +69,25 @@ async function readModuleContent(candidate: ModuleCandidate): Promise { const modules = await discoverPythonModules(options); for (const module of modules) { - const result = await executor.execute(module.content, { silent: true, storeHistory: false }); - if (result.cancelled || result.status === "error") { + const result = await executor.execute(module.content, { + signal: options.signal, + timeoutMs: requireModuleExecutionTimeoutMs(options), + silent: true, + storeHistory: false, + }); + if (result.cancelled) { + throw getExecutionCancellationError(result, options.signal, `Failed to load Python module ${module.path}`); + } + if (result.status === "error") { const details = result.error ? `${result.error.name}: ${result.error.value}` : "unknown error"; throw new Error(`Failed to load Python module ${module.path}: ${details}`); } diff --git a/packages/coding-agent/src/tools/python.ts b/packages/coding-agent/src/tools/python.ts index 6130f7fe0..290a05dfe 100644 --- a/packages/coding-agent/src/tools/python.ts +++ b/packages/coding-agent/src/tools/python.ts @@ -180,7 +180,8 @@ export class PythonTool implements AgentTool { // Clamp to reasonable range: 1s - 600s (10 min) const timeoutSec = clampTimeout("python", rawTimeout); const timeoutMs = timeoutSec * 1000; - const timeoutSignal = AbortSignal.timeout(timeoutMs); + const deadlineMs = Date.now() + timeoutMs; + const timeoutSignal = AbortSignal.timeout(Math.max(0, deadlineMs - Date.now())); const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal; let outputSink: OutputSink | undefined; let outputSummary: OutputSummary | undefined; @@ -267,7 +268,7 @@ export class PythonTool implements AgentTool { const sessionId = sessionFile ? `session:${sessionFile}:cwd:${commandCwd}` : `cwd:${commandCwd}`; const baseExecutorOptions: Omit = { cwd: commandCwd, - timeoutMs, + deadlineMs, signal: combinedSignal, sessionId, kernelMode: this.session.settings.get("python.kernelMode"), diff --git a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts index 812fc2ef4..ac0d0b5fb 100644 --- a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts +++ b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts @@ -239,5 +239,6 @@ describe("AgentSession auto-compaction queue resume", () => { expect(getRuntimeSignals()).toContain("todo:1/3"); expect(continueSpy).toHaveBeenCalledTimes(1); + await session.waitForIdle(); }); }); diff --git a/packages/coding-agent/test/core/python-executor-per-call.test.ts b/packages/coding-agent/test/core/python-executor-per-call.test.ts index 9ab0843c1..879704b84 100644 --- a/packages/coding-agent/test/core/python-executor-per-call.test.ts +++ b/packages/coding-agent/test/core/python-executor-per-call.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "bun:test"; +import { afterEach, describe, expect, it } from "bun:test"; import { executePython } from "@oh-my-pi/pi-coding-agent/ipy/executor"; import type { KernelExecuteOptions, KernelExecuteResult } from "@oh-my-pi/pi-coding-agent/ipy/kernel"; import { PythonKernel } from "@oh-my-pi/pi-coding-agent/ipy/kernel"; @@ -9,7 +9,114 @@ interface KernelStub { shutdown: () => Promise; } +type KernelStartOptions = Parameters[0]; + +const originalDateNow = Date.now; + +const originalStart = PythonKernel.start; + +function createCancellationError(name: "AbortError" | "TimeoutError", message: string): Error { + const error = new Error(message); + error.name = name; + return error; +} + +function rejectOnStartupCancellation(options: KernelStartOptions): Promise { + const { promise, reject } = Promise.withResolvers(); + let settled = false; + let timeout: NodeJS.Timeout | undefined; + const finish = (error: unknown) => { + if (settled) return; + settled = true; + if (timeout) clearTimeout(timeout); + options.signal?.removeEventListener("abort", onAbort); + reject(error); + }; + const onAbort = () => { + finish(options.signal?.reason ?? createCancellationError("AbortError", "Python kernel startup aborted")); + }; + + options.signal?.addEventListener("abort", onAbort, { once: true }); + if (options.deadlineMs !== undefined) { + const remainingMs = Math.max(0, options.deadlineMs - Date.now()); + timeout = setTimeout(() => { + finish(createCancellationError("TimeoutError", "Python kernel startup timed out")); + }, remainingMs); + timeout.unref(); + } + + return promise; +} + describe("executePython (per-call)", () => { + afterEach(() => { + PythonKernel.start = originalStart; + Date.now = originalDateNow; + }); + + it("returns a cancelled timeout result when kernel startup exceeds the deadline", async () => { + Bun.env.PI_PYTHON_SKIP_CHECK = "1"; + using tempDir = TempDir.createSync("@omp-python-executor-per-call-"); + + PythonKernel.start = async options => await rejectOnStartupCancellation(options); + + const result = await executePython("sleep(10)", { + kernelMode: "per-call", + timeoutMs: 25, + cwd: tempDir.path(), + }); + + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(result.output).toContain("Command timed out"); + }); + + it("returns a cancelled timeout result when the startup budget expires before kernel creation", async () => { + Bun.env.PI_PYTHON_SKIP_CHECK = "1"; + using tempDir = TempDir.createSync("@omp-python-executor-per-call-"); + + let nowCalls = 0; + Date.now = () => { + nowCalls += 1; + return nowCalls <= 2 ? 1_000 : 2_000; + }; + + const result = await executePython("sleep(10)", { + kernelMode: "per-call", + timeoutMs: 10, + cwd: tempDir.path(), + }); + + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(result.output).toContain("Command timed out"); + }); + + it("returns a cancelled result when caller aborts during kernel startup", async () => { + Bun.env.PI_PYTHON_SKIP_CHECK = "1"; + using tempDir = TempDir.createSync("@omp-python-executor-per-call-"); + const startupStarted = Promise.withResolvers(); + + PythonKernel.start = async options => { + startupStarted.resolve(); + return await rejectOnStartupCancellation(options); + }; + + const abortController = new AbortController(); + const resultPromise = executePython("sleep(10)", { + kernelMode: "per-call", + signal: abortController.signal, + cwd: tempDir.path(), + }); + await startupStarted.promise; + abortController.abort(createCancellationError("AbortError", "caller aborted")); + + const result = await resultPromise; + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(result.output).toBe(""); + }); + it("shuts down kernel on timed-out cancellation", async () => { Bun.env.PI_PYTHON_SKIP_CHECK = "1"; using tempDir = TempDir.createSync("@omp-python-executor-per-call-"); @@ -27,25 +134,17 @@ describe("executePython (per-call)", () => { }, }; - const kernelClass = PythonKernel as unknown as { - start: (options: { cwd: string }) => Promise; - }; - const originalStart = kernelClass.start; - kernelClass.start = async () => kernel; + PythonKernel.start = async () => kernel as unknown as PythonKernel; - try { - const result = await executePython("sleep(10)", { - kernelMode: "per-call", - timeoutMs: 2000, - cwd: tempDir.path(), - }); + const result = await executePython("sleep(10)", { + kernelMode: "per-call", + timeoutMs: 2000, + cwd: tempDir.path(), + }); - expect(result.cancelled).toBe(true); - expect(result.exitCode).toBeUndefined(); - expect(result.output).toContain("Command timed out after 2 seconds"); - expect(shutdownCalls).toBe(1); - } finally { - kernelClass.start = originalStart; - } + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(result.output).toContain("Command timed out after 2 seconds"); + expect(shutdownCalls).toBe(1); }); }); diff --git a/packages/coding-agent/test/core/python-executor.lifecycle.test.ts b/packages/coding-agent/test/core/python-executor.lifecycle.test.ts index 82f45b7fb..e8b166f8a 100644 --- a/packages/coding-agent/test/core/python-executor.lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-executor.lifecycle.test.ts @@ -9,9 +9,9 @@ import { Bun.env.PI_PYTHON_SKIP_CHECK = "1"; class FakeKernel { - private result: KernelExecuteResult; - private onExecute?: (options?: KernelExecuteOptions) => void; - private alive: boolean; + #result: KernelExecuteResult; + #onExecute?: (options?: KernelExecuteOptions) => void; + #alive: boolean; readonly executeCalls: string[] = []; shutdownCalls = 0; @@ -19,28 +19,28 @@ class FakeKernel { result: KernelExecuteResult, options: { alive?: boolean; onExecute?: (options?: KernelExecuteOptions) => void } = {}, ) { - this.result = result; - this.onExecute = options.onExecute; - this.alive = options.alive ?? true; + this.#result = result; + this.#onExecute = options.onExecute; + this.#alive = options.alive ?? true; } isAlive(): boolean { - return this.alive; + return this.#alive; } async execute(code: string, options?: KernelExecuteOptions): Promise { this.executeCalls.push(code); - this.onExecute?.(options); - return this.result; + this.#onExecute?.(options); + return this.#result; } async shutdown(): Promise { this.shutdownCalls += 1; - this.alive = false; + this.#alive = false; } async ping(): Promise { - return this.alive; + return this.#alive; } } @@ -115,6 +115,51 @@ describe("executePython session lifecycle", () => { expect(secondKernel.executeCalls).toEqual(["print('two')"]); }); + it("cancels queued session execution before it reaches the kernel", async () => { + const firstStarted = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + const kernel = new FakeKernel(okResult); + kernel.execute = async (code, options) => { + kernel.executeCalls.push(code); + if (kernel.executeCalls.length === 1) { + options?.onChunk?.("first\n"); + firstStarted.resolve(); + await releaseFirst.promise; + } + return okResult; + }; + let startCount = 0; + + PythonKernel.start = async () => { + startCount += 1; + return kernel as unknown as PythonKernel; + }; + + const firstPromise = executePython("print('one')", { sessionId: "session-queue" }); + await firstStarted.promise; + + const abortController = new AbortController(); + const secondPromise = executePython("print('two')", { + sessionId: "session-queue", + signal: abortController.signal, + }); + abortController.abort(Object.assign(new Error("queue wait cancelled"), { name: "AbortError" })); + + const second = await secondPromise; + expect(second.cancelled).toBe(true); + expect(second.exitCode).toBeUndefined(); + expect(second.output).toBe(""); + expect(kernel.executeCalls).toEqual(["print('one')"]); + + releaseFirst.resolve(); + const first = await firstPromise; + + expect(first.cancelled).toBe(false); + expect(first.output).toContain("first"); + expect(startCount).toBe(1); + expect(kernel.executeCalls).toEqual(["print('one')"]); + }); + it("uses per-call kernels when configured", async () => { const kernelA = new FakeKernel(okResult); const kernelB = new FakeKernel(okResult); diff --git a/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts b/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts index 95f1dc5b3..b325cdbe2 100644 --- a/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts @@ -186,6 +186,138 @@ describe("PythonKernel gateway lifecycle", () => { expect(kernel.isAlive()).toBe(false); }); + it("aborts stalled startup after websocket connect and cleans up the kernel", async () => { + vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ + url: "http://127.0.0.1:9999", + isShared: true, + }); + + let executeCallCount = 0; + const preludeStarted = Promise.withResolvers(); + vi.spyOn(PythonKernel.prototype, "execute").mockImplementation(async (_code, options) => { + executeCallCount += 1; + if (executeCallCount === 1) { + return { status: "ok", cancelled: false, timedOut: false, stdinRequested: false }; + } + preludeStarted.resolve(); + return await new Promise((_, reject) => { + const onAbort = () => { + const reason = options?.signal?.reason; + reject(reason instanceof Error ? reason : new Error("Python kernel startup aborted")); + }; + if (options?.signal?.aborted) { + onAbort(); + return; + } + options?.signal?.addEventListener("abort", onAbort, { once: true }); + }); + }); + + using _hook = hookFetch((input, init) => { + const url = String(input); + env.fetchCalls.push({ url, init }); + if (url.endsWith("/api/kernels") && init?.method === "POST") { + return createResponse({ ok: true, json: { id: "kernel-stalled" } }) as unknown as Response; + } + return createResponse({ ok: true }) as unknown as Response; + }); + + const abortController = new AbortController(); + const startPromise = PythonKernel.start({ cwd: tempDir.path(), signal: abortController.signal }); + await preludeStarted.promise; + abortController.abort(new Error("cancel startup")); + + const pending = Symbol("pending"); + const settled = await Promise.race([ + startPromise.then( + () => "resolved", + error => error, + ), + Bun.sleep(50).then(() => pending), + ]); + + expect(settled).toBeInstanceOf(Error); + expect(settled).not.toBe(pending); + expect((settled as Error).message).toContain("cancel startup"); + await expect(startPromise).rejects.toThrow("cancel startup"); + expect( + env.fetchCalls.some( + call => call.url.endsWith("/api/kernels/kernel-stalled") && call.init?.method === "DELETE", + ), + ).toBe(true); + expect(FakeWebSocket.instances.at(-1)?.readyState).toBe(FakeWebSocket.CLOSED); + }); + + it("preserves timeout classification when startup environment initialization is cancelled", async () => { + vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ + url: "http://127.0.0.1:9999", + isShared: true, + }); + + vi.spyOn(PythonKernel.prototype, "execute").mockResolvedValue({ + status: "ok", + cancelled: true, + timedOut: true, + stdinRequested: false, + }); + + using _hook = hookFetch((input, init) => { + const url = String(input); + env.fetchCalls.push({ url, init }); + if (url.endsWith("/api/kernels") && init?.method === "POST") { + return createResponse({ ok: true, json: { id: "kernel-init-timeout" } }) as unknown as Response; + } + return createResponse({ ok: true }) as unknown as Response; + }); + + await expect(PythonKernel.start({ cwd: tempDir.path() })).rejects.toMatchObject({ + name: "TimeoutError", + message: "Failed to initialize Python kernel environment", + }); + expect( + env.fetchCalls.some( + call => call.url.endsWith("/api/kernels/kernel-init-timeout") && call.init?.method === "DELETE", + ), + ).toBe(true); + expect(FakeWebSocket.instances.at(-1)?.readyState).toBe(FakeWebSocket.CLOSED); + }); + + it("preserves timeout classification when startup prelude execution is cancelled", async () => { + vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ + url: "http://127.0.0.1:9999", + isShared: true, + }); + + let executeCallCount = 0; + vi.spyOn(PythonKernel.prototype, "execute").mockImplementation(async () => { + executeCallCount += 1; + if (executeCallCount === 1) { + return { status: "ok", cancelled: false, timedOut: false, stdinRequested: false }; + } + return { status: "ok", cancelled: true, timedOut: true, stdinRequested: false }; + }); + + using _hook = hookFetch((input, init) => { + const url = String(input); + env.fetchCalls.push({ url, init }); + if (url.endsWith("/api/kernels") && init?.method === "POST") { + return createResponse({ ok: true, json: { id: "kernel-prelude-timeout" } }) as unknown as Response; + } + return createResponse({ ok: true }) as unknown as Response; + }); + + await expect(PythonKernel.start({ cwd: tempDir.path() })).rejects.toMatchObject({ + name: "TimeoutError", + message: "Failed to initialize Python kernel prelude", + }); + expect( + env.fetchCalls.some( + call => call.url.endsWith("/api/kernels/kernel-prelude-timeout") && call.init?.method === "DELETE", + ), + ).toBe(true); + expect(FakeWebSocket.instances.at(-1)?.readyState).toBe(FakeWebSocket.CLOSED); + }); + it("throws when shared gateway kernel creation never succeeds", async () => { using _runtime = stubKernelRuntime(); vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ diff --git a/packages/coding-agent/test/core/python-modules.test.ts b/packages/coding-agent/test/core/python-modules.test.ts index 727ca25dd..1c20b1d7f 100644 --- a/packages/coding-agent/test/core/python-modules.test.ts +++ b/packages/coding-agent/test/core/python-modules.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it } from "bun:test"; +import { afterEach, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as path from "node:path"; import { loadPythonModules, type PythonModuleExecutor } from "@oh-my-pi/pi-coding-agent/ipy/modules"; @@ -22,32 +22,112 @@ describe("python modules", () => { tempRoot.removeSync(); } tempRoot = null; + vi.restoreAllMocks(); }); - it("loads modules in sorted order with silent execution", async () => { + it("loads modules in sorted order and forwards startup execute options", async () => { tempRoot = TempDir.createSync("@omp-python-modules-"); const agentDir = path.join(tempRoot.path(), "agent"); const cwd = path.join(tempRoot.path(), "project"); + const signal = new AbortController().signal; await writeModule(getAgentModulesDir(agentDir), "beta.py", "user-omp"); await writeModule(getAgentModulesDir(agentDir), "alpha.py", "user-omp"); - const calls: Array<{ name: string; options?: { silent?: boolean; storeHistory?: boolean } }> = []; + const calls: Array<{ + name: string; + options?: { signal?: AbortSignal; timeoutMs?: number; silent?: boolean; storeHistory?: boolean }; + }> = []; const executor: PythonModuleExecutor = { - execute: async (code: string, options?: { silent?: boolean; storeHistory?: boolean }) => { + execute: async (code: string, options) => { const name = code.includes("def alpha") ? "alpha" : "beta"; calls.push({ name, options }); return { status: "ok", cancelled: false }; }, }; - await loadPythonModules(executor, { cwd, agentDir }); + await loadPythonModules(executor, { cwd, agentDir, signal, timeoutMs: 987 }); expect(calls.map(call => call.name)).toEqual(["alpha", "beta"]); for (const call of calls) { - expect(call.options).toEqual({ silent: true, storeHistory: false }); + expect(call.options).toEqual({ signal, timeoutMs: 987, silent: true, storeHistory: false }); } }); + it("derives module execution timeout from the remaining deadline", async () => { + tempRoot = TempDir.createSync("@omp-python-modules-"); + const agentDir = path.join(tempRoot.path(), "agent"); + const cwd = path.join(tempRoot.path(), "project"); + const signal = new AbortController().signal; + + await writeModule(getProjectModulesDir(cwd), "alpha.py", "project-omp"); + + const execute = vi.fn( + async ( + _code: string, + _options?: { + signal?: AbortSignal; + timeoutMs?: number; + silent?: boolean; + storeHistory?: boolean; + }, + ) => ({ status: "ok" as const, cancelled: false }), + ); + const executor: PythonModuleExecutor = { execute }; + vi.spyOn(Date, "now").mockReturnValue(10_000); + + await loadPythonModules(executor, { cwd, agentDir, signal, deadlineMs: 10_250 }); + expect(execute).toHaveBeenCalledTimes(1); + expect(execute).toHaveBeenCalledWith(expect.any(String), { + signal, + timeoutMs: 250, + silent: true, + storeHistory: false, + }); + }); + + it("fails fast when the module deadline expires before the next execution starts", async () => { + tempRoot = TempDir.createSync("@omp-python-modules-"); + const agentDir = path.join(tempRoot.path(), "agent"); + const cwd = path.join(tempRoot.path(), "project"); + const signal = new AbortController().signal; + + await writeModule(getProjectModulesDir(cwd), "alpha.py", "project-omp"); + await writeModule(getProjectModulesDir(cwd), "beta.py", "project-omp"); + + const execute = vi.fn(async () => ({ status: "ok" as const, cancelled: false })); + const executor: PythonModuleExecutor = { execute }; + vi.spyOn(Date, "now").mockReturnValueOnce(10_000).mockReturnValueOnce(10_300); + + await expect(loadPythonModules(executor, { cwd, agentDir, signal, deadlineMs: 10_250 })).rejects.toMatchObject({ + name: "TimeoutError", + message: "Python module loading timed out", + }); + expect(execute).toHaveBeenCalledTimes(1); + expect(execute).toHaveBeenCalledWith(expect.any(String), { + signal, + timeoutMs: 250, + silent: true, + storeHistory: false, + }); + }); + + it("preserves timeout classification when module execution is cancelled", async () => { + tempRoot = TempDir.createSync("@omp-python-modules-"); + const agentDir = path.join(tempRoot.path(), "agent"); + const cwd = path.join(tempRoot.path(), "project"); + + await writeModule(getProjectModulesDir(cwd), "alpha.py", "project-omp"); + + const execute = vi.fn(async () => ({ status: "ok" as const, cancelled: true, timedOut: true })); + const executor: PythonModuleExecutor = { execute }; + + await expect(loadPythonModules(executor, { cwd, agentDir })).rejects.toMatchObject({ + name: "TimeoutError", + message: expect.stringContaining("Failed to load Python module"), + }); + expect(execute).toHaveBeenCalledTimes(1); + }); + it("fails fast when a module fails to execute", async () => { tempRoot = TempDir.createSync("@omp-python-modules-"); const agentDir = path.join(tempRoot.path(), "agent"); diff --git a/packages/coding-agent/test/tools/python-execution.test.ts b/packages/coding-agent/test/tools/python-execution.test.ts index 7d0e00d91..2fdaee3f6 100644 --- a/packages/coding-agent/test/tools/python-execution.test.ts +++ b/packages/coding-agent/test/tools/python-execution.test.ts @@ -49,9 +49,12 @@ describe("python tool execution", () => { "print('hi')", expect.objectContaining({ cwd: tempDir.path(), - timeoutMs: 5000, + deadlineMs: expect.any(Number), + signal: expect.any(AbortSignal), + sessionFile: `${tempDir.path()}/session-file.jsonl`, sessionId: `session:${tempDir.path()}/session-file.jsonl:cwd:${tempDir.path()}`, kernelMode: "per-call", + useSharedGateway: true, reset: true, }), );