diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 520c75848..d3eba5ffe 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed retained Python kernel ownership so `AgentSession.dispose()` only shuts down kernels owned by that session, including warmup-created kernels + ## [14.1.0] - 2026-04-11 ### Added @@ -105,6 +109,7 @@ - Fixed typo in system prompt: 'backwards compatibiltity' → 'backwards compatibility' + ## [14.0.3] - 2026-04-09 ### Fixed diff --git a/packages/coding-agent/src/ipy/executor.ts b/packages/coding-agent/src/ipy/executor.ts index 7ed5520db..ebc0b3964 100644 --- a/packages/coding-agent/src/ipy/executor.ts +++ b/packages/coding-agent/src/ipy/executor.ts @@ -16,6 +16,7 @@ import { PYTHON_PRELUDE } from "./prelude"; const IDLE_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes const MAX_KERNEL_SESSIONS = 4; const CLEANUP_INTERVAL_MS = 30 * 1000; // 30 seconds +const OWNER_CLEANUP_KERNEL_SHUTDOWN_TIMEOUT_MS = 2_000; export type PythonKernelMode = "session" | "per-call"; @@ -32,6 +33,8 @@ export interface PythonExecutorOptions { signal?: AbortSignal; /** Session identifier for kernel reuse */ sessionId?: string; + /** Logical owner identifier for retained kernel cleanup */ + kernelOwnerId?: string; /** Kernel mode (session reuse vs per-call) */ kernelMode?: PythonKernelMode; /** Restart the kernel before executing */ @@ -80,11 +83,24 @@ interface KernelSession { queue: Promise; restartCount: number; dead: boolean; + needsRestart: boolean; + kernelInvalidatedByRecovery: boolean; + disposing: boolean; + disposeCapacityPromise?: Promise; + resolveDisposeCapacity?: () => void; + disposeAttemptPromise?: Promise; + resolveDisposeAttempt?: () => void; + disposeResultPromise?: Promise; + disposeResultTimeoutMs?: number; + nextDisposalRetryAt?: number; lastUsedAt: number; + ownerIds: Set; + hasFallbackOwner: boolean; heartbeatTimer?: NodeJS.Timeout; } const kernelSessions = new Map(); +const disposingKernelSessions = new Set(); let cachedPreludeDocs: PreludeHelper[] | null = null; let cleanupTimer: NodeJS.Timeout | null = null; @@ -93,6 +109,7 @@ interface KernelSessionExecutionOptions { sessionFile?: string; signal?: AbortSignal; deadlineMs?: number; + kernelOwnerId?: string; } class PythonExecutionCancelledError extends Error { @@ -142,10 +159,10 @@ function isTimedOutCancellation(error: unknown, signal?: AbortSignal): boolean { return reason instanceof Error ? reason.name === "TimeoutError" : false; } -async function waitForQueueTurn( - queue: Promise, +async function waitForPromiseWithCancellation( + promise: Promise, options: Pick, -): Promise { +): Promise { if (options.signal?.aborted) { throw new PythonExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); } @@ -156,11 +173,10 @@ async function waitForQueueTurn( } if (!options.signal && remainingMs === undefined) { - await queue; - return; + return await promise; } - await new Promise((resolve, reject) => { + return await new Promise((resolve, reject) => { const cleanups: Array<() => void> = []; const finish = (callback: () => void) => { while (cleanups.length > 0) { @@ -188,13 +204,20 @@ async function waitForQueueTurn( cleanups.push(() => clearTimeout(timeout)); } - queue.then( - () => finish(resolve), + promise.then( + value => finish(() => resolve(value)), error => finish(() => reject(error)), ); }); } +async function waitForQueueTurn( + queue: Promise, + options: Pick, +): Promise { + await waitForPromiseWithCancellation(queue, options); +} + function formatTimeoutAnnotation(timeoutMs?: number): string | undefined { if (timeoutMs === undefined) return "Command timed out"; const secs = Math.max(1, Math.round(timeoutMs / 1000)); @@ -352,6 +375,131 @@ function stopCleanupTimer(): void { } } +function attachKernelOwner(sessionId: string, ownerId?: string): boolean { + const session = kernelSessions.get(sessionId); + if (!session || session.disposing) return false; + if (ownerId !== undefined) { + if (session.hasFallbackOwner) { + session.ownerIds.delete(sessionId); + session.hasFallbackOwner = false; + } + session.ownerIds.add(ownerId); + } else if (session.hasFallbackOwner || session.ownerIds.size === 0) { + session.ownerIds.add(sessionId); + session.hasFallbackOwner = true; + } + session.lastUsedAt = Date.now(); + return true; +} + +function getRetainedKernelSessionCount(): number { + return kernelSessions.size + disposingKernelSessions.size; +} + +function syncCleanupTimer(): void { + if (kernelSessions.size === 0 && disposingKernelSessions.size === 0) { + stopCleanupTimer(); + return; + } + startCleanupTimer(); +} + +function retryPendingKernelSessionDisposals(now: number = Date.now()): void { + for (const session of disposingKernelSessions.values()) { + if (session.disposeResultPromise) continue; + if (session.nextDisposalRetryAt !== undefined && session.nextDisposalRetryAt > now) continue; + session.nextDisposalRetryAt = undefined; + void disposeKernelSession(session); + } +} + +function beginDisposingKernelSession(session: KernelSession): boolean { + if (session.disposing) return false; + session.disposing = true; + disposingKernelSessions.add(session); + if (kernelSessions.get(session.id) === session) { + kernelSessions.delete(session.id); + } + if (session.heartbeatTimer) { + clearInterval(session.heartbeatTimer); + session.heartbeatTimer = undefined; + } + syncCleanupTimer(); + return true; +} + +function finishDisposingKernelSession(session: KernelSession): void { + disposingKernelSessions.delete(session); + session.resolveDisposeCapacity?.(); + session.resolveDisposeCapacity = undefined; + session.disposeCapacityPromise = undefined; + session.resolveDisposeAttempt = undefined; + session.disposeAttemptPromise = undefined; + session.disposeResultPromise = undefined; + session.disposeResultTimeoutMs = undefined; + session.nextDisposalRetryAt = undefined; + session.kernelInvalidatedByRecovery = false; + syncCleanupTimer(); +} + +async function waitForDisposalCapacity( + options: Pick, +): Promise { + retryPendingKernelSessionDisposals(); + + const disposalPromises: Promise[] = []; + let nextRetryAt: number | undefined; + for (const session of disposingKernelSessions.values()) { + if (session.disposeCapacityPromise) { + disposalPromises.push( + session.disposeCapacityPromise.then( + () => undefined, + () => undefined, + ), + ); + } + if (session.disposeAttemptPromise) { + disposalPromises.push( + session.disposeAttemptPromise.then( + () => undefined, + () => undefined, + ), + ); + } + if (session.nextDisposalRetryAt !== undefined) { + nextRetryAt = + nextRetryAt === undefined + ? session.nextDisposalRetryAt + : Math.min(nextRetryAt, session.nextDisposalRetryAt); + } + } + if (disposalPromises.length > 0) { + await waitForPromiseWithCancellation(Promise.race(disposalPromises), options); + return; + } + if (nextRetryAt === undefined) return; + await waitForPromiseWithCancellation( + Bun.sleep(Math.max(0, nextRetryAt - Date.now())).then(() => undefined), + options, + ); +} + +async function ensureKernelSessionCapacity( + options: Pick, +): Promise { + while (getRetainedKernelSessionCount() >= MAX_KERNEL_SESSIONS) { + if (disposingKernelSessions.size > 0) { + await waitForDisposalCapacity(options); + continue; + } + if (kernelSessions.size === 0) { + await waitForDisposalCapacity(options); + continue; + } + await evictOldestSession(); + } +} + async function cleanupIdleSessions(): Promise { const now = Date.now(); const toDispose: KernelSession[] = []; @@ -367,9 +515,8 @@ async function cleanupIdleSessions(): Promise { await Promise.allSettled(toDispose.map(session => disposeKernelSession(session))); } - if (kernelSessions.size === 0) { - stopCleanupTimer(); - } + retryPendingKernelSessionDisposals(now); + syncCleanupTimer(); } async function evictOldestSession(): Promise { @@ -387,12 +534,29 @@ async function evictOldestSession(): Promise { export async function disposeAllKernelSessions(): Promise { stopCleanupTimer(); - const sessions = Array.from(kernelSessions.values()); + const sessions = Array.from(new Set([...kernelSessions.values(), ...disposingKernelSessions.values()])); await Promise.allSettled(sessions.map(session => disposeKernelSession(session))); } -async function ensureKernelAvailable(cwd: string): Promise { - const availability = await checkPythonKernelAvailability(cwd); +export async function disposeKernelSessionsByOwner(ownerId: string): Promise { + const sessionsToDispose: KernelSession[] = []; + for (const session of Array.from(kernelSessions.values())) { + if (session.disposing || !session.ownerIds.delete(ownerId)) continue; + if (session.ownerIds.size === 0) { + sessionsToDispose.push(session); + } + } + await Promise.allSettled( + sessionsToDispose.map(session => disposeKernelSession(session, OWNER_CLEANUP_KERNEL_SHUTDOWN_TIMEOUT_MS)), + ); + syncCleanupTimer(); +} + +async function ensureKernelAvailable( + cwd: string, + options: Pick = {}, +): Promise { + const availability = await waitForPromiseWithCancellation(checkPythonKernelAvailability(cwd), options); if (!availability.ok) { throw new Error(availability.reason ?? "Python kernel unavailable"); } @@ -403,10 +567,13 @@ export async function warmPythonEnvironment( sessionId?: string, useSharedGateway?: boolean, sessionFile?: string, + kernelOwnerId?: string, + signal?: AbortSignal, ): Promise<{ ok: boolean; reason?: string; docs: PreludeHelper[] }> { let cacheState: PreludeCacheState | null = null; + const resolvedSessionId = sessionId ?? `session:${cwd}`; try { - await logger.time("warmPython:ensureKernelAvailable", ensureKernelAvailable, cwd); + await logger.time("warmPython:ensureKernelAvailable", ensureKernelAvailable, cwd, { signal }); } catch (err: unknown) { const reason = err instanceof Error ? err.message : String(err); cachedPreludeDocs = []; @@ -418,6 +585,7 @@ export async function warmPythonEnvironment( const cached = await readPreludeCache(cacheState); if (cached) { cachedPreludeDocs = cached; + attachKernelOwner(resolvedSessionId, kernelOwnerId); return { ok: true, docs: cached }; } } catch (err) { @@ -426,19 +594,21 @@ export async function warmPythonEnvironment( } } if (cachedPreludeDocs && cachedPreludeDocs.length > 0) { + attachKernelOwner(resolvedSessionId, kernelOwnerId); return { ok: true, docs: cachedPreludeDocs }; } - const resolvedSessionId = sessionId ?? `session:${cwd}`; try { const docs = await logger.time( "warmPython:withKernelSession", withKernelSession, resolvedSessionId, cwd, - kernel => ensurePreludeDocsLoaded(kernel, cwd, { useSharedGateway, sessionFile }, cacheState), + kernel => ensurePreludeDocsLoaded(kernel, cwd, { useSharedGateway, sessionFile, signal }, cacheState), { useSharedGateway, sessionFile, + kernelOwnerId, + signal, }, ); return { ok: true, docs }; @@ -471,17 +641,55 @@ function isResourceExhaustionError(error: unknown): boolean { ); } +function clearDisposingKernelSessionTracking(): void { + for (const session of disposingKernelSessions.values()) { + if (session.heartbeatTimer) { + clearInterval(session.heartbeatTimer); + session.heartbeatTimer = undefined; + } + session.resolveDisposeCapacity?.(); + session.resolveDisposeCapacity = undefined; + session.disposeCapacityPromise = undefined; + session.resolveDisposeAttempt?.(); + session.resolveDisposeAttempt = undefined; + session.disposeAttemptPromise = undefined; + session.disposeResultPromise = undefined; + session.disposeResultTimeoutMs = undefined; + session.nextDisposalRetryAt = undefined; + } + disposingKernelSessions.clear(); +} + +function markLiveKernelSessionsForRecovery(): void { + for (const session of kernelSessions.values()) { + if (session.heartbeatTimer) { + clearInterval(session.heartbeatTimer); + session.heartbeatTimer = undefined; + } + session.needsRestart = true; + session.kernelInvalidatedByRecovery = session.kernel.isSharedGateway; + session.restartCount = 0; + } +} + async function recoverFromResourceExhaustion(): Promise { logger.warn("Resource exhaustion detected, recovering by restarting shared gateway"); stopCleanupTimer(); - const sessions = Array.from(kernelSessions.values()); - for (const session of sessions) { - if (session.heartbeatTimer) { - clearInterval(session.heartbeatTimer); - } - kernelSessions.delete(session.id); - } + markLiveKernelSessionsForRecovery(); + clearDisposingKernelSessionTracking(); await shutdownSharedGateway(); + syncCleanupTimer(); +} + +function ensureKernelHeartbeat(session: KernelSession): void { + if (session.heartbeatTimer) return; + session.heartbeatTimer = setInterval(() => { + if (session.dead || session.needsRestart) return; + if (!session.kernel.isAlive()) { + session.dead = true; + } + }, 5000); + session.heartbeatTimer.unref(); } async function createKernelSession( @@ -507,21 +715,25 @@ async function createKernelSession( throw err; } + const hasFallbackOwner = options.kernelOwnerId === undefined; + const initialOwnerId = options.kernelOwnerId ?? sessionId; const session: KernelSession = { id: sessionId, kernel, queue: Promise.resolve(), restartCount: 0, dead: false, + needsRestart: false, + kernelInvalidatedByRecovery: false, + disposing: false, + disposeResultPromise: undefined, + nextDisposalRetryAt: undefined, lastUsedAt: Date.now(), + ownerIds: new Set([initialOwnerId]), + hasFallbackOwner, }; - session.heartbeatTimer = setInterval(() => { - if (session.dead) return; - if (!session.kernel.isAlive()) { - session.dead = true; - } - }, 5000); + ensureKernelHeartbeat(session); return session; } @@ -537,30 +749,193 @@ async function restartKernelSession( } requireRemainingTimeoutMs(options.deadlineMs); try { - await session.kernel.shutdown(); + if (!session.kernelInvalidatedByRecovery) { + const shutdownTimeoutMs = requireRemainingTimeoutMs(options.deadlineMs); + const shutdownResult = await session.kernel.shutdown({ signal: options.signal, timeoutMs: shutdownTimeoutMs }); + if (!shutdownResult.confirmed) { + throw new Error("Failed to confirm crashed kernel shutdown before restart"); + } + } + 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.needsRestart = false; + session.kernelInvalidatedByRecovery = false; + session.lastUsedAt = Date.now(); + ensureKernelHeartbeat(session); } catch (err) { - logger.warn("Failed to shutdown crashed kernel", { error: err instanceof Error ? err.message : String(err) }); + session.restartCount = 0; + logger.warn("Failed to restart kernel", { error: err instanceof Error ? err.message : String(err) }); + throw err; } - 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(); } -async function disposeKernelSession(session: KernelSession): Promise { - if (session.heartbeatTimer) { - clearInterval(session.heartbeatTimer); +type KernelDisposalResult = { status: "confirmed" } | { status: "unconfirmed" } | { status: "failed"; err: unknown }; +type KernelDisposalWaitResult = KernelDisposalResult | { status: "timedOut" }; + +function createKernelDisposalResultPromise(session: KernelSession, timeoutMs?: number): Promise { + if (session.kernelInvalidatedByRecovery) { + return Promise.resolve({ status: "confirmed" as const }); } - try { - await session.kernel.shutdown(); - } catch (err) { - logger.warn("Failed to shutdown kernel", { error: err instanceof Error ? err.message : String(err) }); + return Promise.resolve() + .then(() => session.kernel.shutdown(timeoutMs === undefined ? undefined : { timeoutMs })) + .then( + result => (result.confirmed ? { status: "confirmed" as const } : { status: "unconfirmed" as const }), + (err: unknown) => ({ status: "failed" as const, err }), + ); +} + +function getOrStartKernelDisposalResultPromise( + session: KernelSession, + timeoutMs?: number, +): Promise { + if (!session.disposeResultPromise) { + session.disposeResultTimeoutMs = timeoutMs; + const releaseDisposalAttempt = Promise.withResolvers(); + session.disposeAttemptPromise = releaseDisposalAttempt.promise; + session.resolveDisposeAttempt = releaseDisposalAttempt.resolve; + const disposeResultPromise = createKernelDisposalResultPromise(session, timeoutMs); + void disposeResultPromise.then(result => { + if (result.status === "confirmed") { + finishDisposingKernelSession(session); + return; + } + if (session.disposing) { + session.nextDisposalRetryAt = Date.now() + CLEANUP_INTERVAL_MS; + syncCleanupTimer(); + } + }); + const disposalAttemptPromise = disposeResultPromise.finally(() => { + releaseDisposalAttempt.resolve(); + if (session.disposeResultPromise === disposalAttemptPromise) { + session.disposeResultPromise = undefined; + session.disposeResultTimeoutMs = undefined; + } + if (session.disposeAttemptPromise === releaseDisposalAttempt.promise) { + session.disposeAttemptPromise = undefined; + session.resolveDisposeAttempt = undefined; + } + }); + session.disposeResultPromise = disposalAttemptPromise; } - kernelSessions.delete(session.id); + return session.disposeResultPromise; +} + +async function waitForKernelSessionDisposal( + session: KernelSession, + timeoutMs?: number, +): Promise { + const disposeResultPromise = session.disposeResultPromise; + if (!disposeResultPromise) { + return undefined; + } + if (timeoutMs === undefined) { + return await disposeResultPromise; + } + + let timeoutId: NodeJS.Timeout | undefined; + const result = await Promise.race([ + disposeResultPromise, + new Promise<{ status: "timedOut" }>(resolve => { + timeoutId = setTimeout(() => resolve({ status: "timedOut" }), timeoutMs); + timeoutId.unref(); + }), + ]); + + if (timeoutId) { + clearTimeout(timeoutId); + } + return result; +} + +function retryKernelSessionDisposalInBackground(session: KernelSession): void { + session.nextDisposalRetryAt = undefined; + void disposeKernelSession(session); +} + +async function disposeKernelSession(session: KernelSession, shutdownTimeoutMs?: number): Promise { + if (!session.disposing) { + if (!beginDisposingKernelSession(session)) return; + const releaseDisposalCapacity = Promise.withResolvers(); + session.disposeCapacityPromise = releaseDisposalCapacity.promise; + session.resolveDisposeCapacity = releaseDisposalCapacity.resolve; + } + + if ( + shutdownTimeoutMs === undefined && + session.disposeResultPromise && + session.disposeResultTimeoutMs !== undefined + ) { + const inheritedResult = await session.disposeResultPromise; + if (inheritedResult.status === "confirmed") { + finishDisposingKernelSession(session); + return; + } + session.disposeResultPromise = undefined; + session.disposeResultTimeoutMs = undefined; + logger.warn("Retained kernel shutdown was not confirmed during owner cleanup; retrying without timeout", { + sessionId: session.id, + }); + } + + const inheritedBackgroundRetryTimeoutMs = + shutdownTimeoutMs === undefined && session.disposeResultPromise && session.disposeResultTimeoutMs === undefined + ? OWNER_CLEANUP_KERNEL_SHUTDOWN_TIMEOUT_MS + : shutdownTimeoutMs; + + getOrStartKernelDisposalResultPromise(session, shutdownTimeoutMs); + const result = await waitForKernelSessionDisposal(session, inheritedBackgroundRetryTimeoutMs); + if (!result) { + return; + } + if (result.status === "timedOut") { + logger.warn( + shutdownTimeoutMs === undefined + ? "Timed out waiting for retained kernel shutdown during global cleanup; retained capacity remains reserved" + : "Timed out shutting down retained kernel during owner cleanup", + { + sessionId: session.id, + timeoutMs: inheritedBackgroundRetryTimeoutMs, + }, + ); + if (shutdownTimeoutMs !== undefined) { + retryKernelSessionDisposalInBackground(session); + } + return; + } + if (result.status === "confirmed") { + finishDisposingKernelSession(session); + return; + } + if (result.status === "unconfirmed") { + logger.warn( + shutdownTimeoutMs === undefined + ? "Kernel shutdown was not confirmed; retained capacity remains reserved" + : "Retained kernel shutdown was not confirmed during owner cleanup", + { sessionId: session.id }, + ); + if (shutdownTimeoutMs !== undefined) { + retryKernelSessionDisposalInBackground(session); + } + return; + } + logger.warn( + shutdownTimeoutMs === undefined + ? "Failed to shutdown kernel" + : "Failed to shutdown retained kernel during owner cleanup", + { + sessionId: session.id, + error: result.err instanceof Error ? result.err.message : String(result.err), + }, + ); + if (shutdownTimeoutMs !== undefined) { + retryKernelSessionDisposalInBackground(session); + } + return; } async function withKernelSession( @@ -570,10 +945,11 @@ async function withKernelSession( options: KernelSessionExecutionOptions = {}, ): Promise { let session = kernelSessions.get(sessionId); + if (session?.disposing) { + session = undefined; + } if (!session) { - if (kernelSessions.size >= MAX_KERNEL_SESSIONS) { - await evictOldestSession(); - } + await ensureKernelSessionCapacity(options); requireRemainingTimeoutMs(options.deadlineMs); if (options.signal?.aborted) { throw new PythonExecutionCancelledError(isTimedOutCancellation(options.signal.reason, options.signal)); @@ -582,10 +958,15 @@ async function withKernelSession( kernelSessions.set(sessionId, session); startCleanupTimer(); } + attachKernelOwner(sessionId, options.kernelOwnerId); + + if (session.disposing) { + return await withKernelSession(sessionId, cwd, handler, options); + } const run = async (): Promise => { session!.lastUsedAt = Date.now(); - if (session!.dead || !session!.kernel.isAlive()) { + if (session!.dead || session!.needsRestart || !session!.kernel.isAlive()) { await logger.time("kernel:restartKernelSession", restartKernelSession, session!, cwd, options); } try { @@ -593,7 +974,7 @@ async function withKernelSession( session!.restartCount = 0; return result; } catch (err) { - if (!session!.dead && session!.kernel.isAlive()) { + if (!session!.dead && !session!.needsRestart && session!.kernel.isAlive()) { throw err; } await logger.time("kernel:restartKernelSession", restartKernelSession, session!, cwd, options); @@ -620,6 +1001,9 @@ async function withKernelSession( try { await waitForQueueTurn(queue, options); + if (session.disposing) { + return await withKernelSession(sessionId, cwd, handler, options); + } return await run(); } finally { releaseTurn?.(); @@ -741,6 +1125,9 @@ export async function executePython(code: string, options?: PythonExecutorOption const existing = kernelSessions.get(sessionId); if (existing) { await disposeKernelSession(existing); + if (existing.disposing && existing.nextDisposalRetryAt !== undefined) { + retryKernelSessionDisposalInBackground(existing); + } } } return await withKernelSession( diff --git a/packages/coding-agent/src/ipy/kernel.ts b/packages/coding-agent/src/ipy/kernel.ts index 685971643..a1945ba41 100644 --- a/packages/coding-agent/src/ipy/kernel.ts +++ b/packages/coding-agent/src/ipy/kernel.ts @@ -49,6 +49,10 @@ interface KernelShutdownOptions { timeoutMs?: number; } +export interface KernelShutdownResult { + confirmed: boolean; +} + function getRemainingTimeMs(deadlineMs?: number): number | undefined { if (deadlineMs === undefined) return undefined; return Math.max(0, deadlineMs - Date.now()); @@ -399,10 +403,11 @@ export class PythonKernel { #ws: WebSocket | null = null; #disposed = false; #alive = true; + #shutdownStarted = false; + #shutdownConfirmed = false; #messageHandlers = new Map void>(); #channelHandlers = new Map void>>(); #pendingExecutions = new Map void>(); - private constructor( readonly id: string, readonly kernelId: string, @@ -1004,11 +1009,18 @@ export class PythonKernel { } } - async shutdown(options?: KernelShutdownOptions): Promise { - if (this.#disposed) return; - this.#disposed = true; - this.#alive = false; - this.#abortPendingExecutions("Kernel shutdown"); + async shutdown(options?: KernelShutdownOptions): Promise { + if (this.#shutdownConfirmed) return { confirmed: true }; + if (!this.#shutdownStarted) { + this.#shutdownStarted = true; + this.#alive = false; + this.#abortPendingExecutions("Kernel shutdown"); + + if (this.#ws) { + this.#ws.close(); + this.#ws = null; + } + } const shutdownSignal = combineAbortSignal( { signal: options?.signal }, @@ -1016,24 +1028,38 @@ export class PythonKernel { "Python kernel shutdown timed out", ); + let confirmed = false; try { - await fetch(`${this.gatewayUrl}/api/kernels/${this.kernelId}`, { + const response = await fetch(`${this.gatewayUrl}/api/kernels/${this.kernelId}`, { method: "DELETE", headers: this.#authHeaders(), signal: shutdownSignal, }); + const deleteConfirmed = response.status === 404 || response.status === 410; + confirmed = response.ok || deleteConfirmed; + if (!confirmed) { + logger.warn("Kernel delete request was not confirmed", { + status: response.status, + statusText: response.statusText, + }); + } } catch (err: unknown) { logger.warn("Failed to delete kernel via API", { error: err instanceof Error ? err.message : String(err) }); } - - if (this.#ws) { - this.#ws.close(); - this.#ws = null; - } + this.#shutdownConfirmed = confirmed; + this.#disposed = confirmed; if (this.isSharedGateway) { - await releaseSharedGateway(); + try { + await releaseSharedGateway(); + } catch (err: unknown) { + logger.warn("Failed to release shared gateway after kernel shutdown", { + error: err instanceof Error ? err.message : String(err), + }); + } } + + return { confirmed }; } #sendMessage(msg: JupyterMessage): void { diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 3a9c22896..1b4c838c1 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -23,6 +23,7 @@ import { logger, postmortem, prompt, + Snowflake, } from "@oh-my-pi/pi-utils"; import chalk from "chalk"; import { AsyncJobManager, isBackgroundJobSupportEnabled } from "./async"; @@ -73,7 +74,7 @@ import { RuleProtocolHandler, SkillProtocolHandler, } from "./internal-urls"; -import { disposeAllKernelSessions } from "./ipy/executor"; +import { disposeAllKernelSessions, disposeKernelSessionsByOwner } from "./ipy/executor"; import { LSP_STARTUP_EVENT_CHANNEL, type LspStartupEvent } from "./lsp/startup-events"; import { discoverAndLoadMCPTools, type MCPManager, type MCPToolsLoadResult } from "./mcp"; import { @@ -202,6 +203,8 @@ export interface CreateAgentSessionOptions { enableLsp?: boolean; /** Skip Python kernel availability check and prelude warmup */ skipPythonPreflight?: boolean; + /** Force Python prelude warmup even when test env would normally skip it */ + forcePythonWarmup?: boolean; /** Tool names explicitly requested (enables disabled-by-default tools) */ toolNames?: string[]; @@ -835,8 +838,8 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} ); let agent: Agent; - let session: AgentSession; - + let session!: AgentSession; + let hasSession = false; const enableLsp = options.enableLsp ?? true; const backgroundJobsEnabled = isBackgroundJobSupportEnabled(settings); const asyncMaxJobs = Math.min(100, Math.max(1, settings.get("async.maxJobs") ?? 100)); @@ -889,437 +892,422 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }) : undefined; - const searchDb = options.searchDb ?? new SearchDb(getSearchDbDir(agentDir)); - const toolSession: ToolSession = { - cwd, - hasUI: options.hasUI ?? false, - enableLsp, - get hasEditTool() { - return !options.toolNames || options.toolNames.includes("edit"); - }, - skipPythonPreflight: options.skipPythonPreflight, - contextFiles, - skills, - eventBus, - outputSchema: options.outputSchema, - requireSubmitResultTool: options.requireSubmitResultTool, - taskDepth: options.taskDepth ?? 0, - getSessionFile: () => sessionManager.getSessionFile() ?? null, - getSessionId: () => sessionManager.getSessionId?.() ?? null, - getSessionSpawns: () => options.spawns ?? "*", - getModelString: () => (hasExplicitModel && model ? formatModelString(model) : undefined), - getActiveModelString: () => { - const activeModel = agent?.state.model; - if (activeModel) return formatModelString(activeModel); - // Fall back to initial model during tool creation (before agent exists) - if (model) return formatModelString(model); - return undefined; - }, - getPlanModeState: () => session.getPlanModeState(), - getCompactContext: () => session.formatCompactContext(), - getTodoPhases: () => session.getTodoPhases(), - setTodoPhases: phases => session.setTodoPhases(phases), - isMCPDiscoveryEnabled: () => session.isMCPDiscoveryEnabled(), - getDiscoverableMCPTools: () => session.getDiscoverableMCPTools(), - getDiscoverableMCPSearchIndex: () => session.getDiscoverableMCPSearchIndex(), - getSelectedMCPToolNames: () => session.getSelectedMCPToolNames(), - activateDiscoveredMCPTools: toolNames => session.activateDiscoveredMCPTools(toolNames), - getCheckpointState: () => session.getCheckpointState(), - setCheckpointState: state => session.setCheckpointState(state ?? undefined), - getToolChoiceQueue: () => session.toolChoiceQueue, - buildToolChoice: name => { - const m = session.model; - return m ? buildNamedToolChoice(name, m) : undefined; - }, - steer: msg => - session.agent.steer({ - role: "custom", - customType: msg.customType, - content: msg.content, - display: false, - details: msg.details, - attribution: "agent", - timestamp: Date.now(), - }), - peekQueueInvoker: () => session.peekQueueInvoker(), - allocateOutputArtifact: async toolType => { - try { - return await sessionManager.allocateArtifactPath(toolType); - } catch { - return {}; - } - }, - settings, - authStorage, - modelRegistry, - asyncJobManager, - searchDb, - }; + const pythonKernelOwnerId = `agent-session:${Snowflake.next()}`; - // Initialize internal URL router for internal protocols (agent://, artifact://, memory://, skill://, rule://, mcp://, local://) - const internalRouter = new InternalUrlRouter(); - const getArtifactsDir = () => sessionManager.getArtifactsDir(); - internalRouter.register(new AgentProtocolHandler({ getArtifactsDir })); - internalRouter.register(new ArtifactProtocolHandler({ getArtifactsDir })); - internalRouter.register( - new MemoryProtocolHandler({ - getMemoryRoot: () => getMemoryRoot(agentDir, settings.getCwd()), - }), - ); - internalRouter.register( - new LocalProtocolHandler({ - getArtifactsDir, - getSessionId: () => sessionManager.getSessionId(), - }), - ); - internalRouter.register( - new SkillProtocolHandler({ - getSkills: () => skills, - }), - ); - internalRouter.register( - new RuleProtocolHandler({ - getRules: () => [...rulebookRules, ...alwaysApplyRules], - }), - ); - internalRouter.register(new PiProtocolHandler()); - internalRouter.register(new JobsProtocolHandler({ getAsyncJobManager: () => asyncJobManager })); - internalRouter.register(new McpProtocolHandler({ getMcpManager: () => mcpManager })); - toolSession.internalRouter = internalRouter; - toolSession.getArtifactsDir = getArtifactsDir; - toolSession.agentOutputManager = new AgentOutputManager( - getArtifactsDir, - options.parentTaskPrefix ? { parentPrefix: options.parentTaskPrefix } : undefined, - ); - - // Create built-in tools (already wrapped with meta notice formatting) - const builtinTools = await logger.time("createAllTools", createTools, toolSession, options.toolNames); - - // Discover MCP tools from .mcp.json files - let mcpManager: MCPManager | undefined; - const enableMCP = options.enableMCP ?? true; - const customTools: CustomTool[] = []; - if (enableMCP) { - const mcpResult = await logger.time("discoverAndLoadMCPTools", discoverAndLoadMCPTools, cwd, { - onConnecting: serverNames => { - if (options.hasUI && serverNames.length > 0) { - process.stderr.write(`${chalk.gray(`Connecting to MCP servers: ${serverNames.join(", ")}…`)}\n`); + try { + const searchDb = options.searchDb ?? new SearchDb(getSearchDbDir(agentDir)); + const toolSession: ToolSession = { + cwd, + hasUI: options.hasUI ?? false, + enableLsp, + get hasEditTool() { + return !options.toolNames || options.toolNames.includes("edit"); + }, + skipPythonPreflight: options.skipPythonPreflight, + forcePythonWarmup: options.forcePythonWarmup, + contextFiles, + skills, + eventBus, + outputSchema: options.outputSchema, + requireSubmitResultTool: options.requireSubmitResultTool, + taskDepth: options.taskDepth ?? 0, + getSessionFile: () => sessionManager.getSessionFile() ?? null, + getPythonKernelOwnerId: () => pythonKernelOwnerId, + assertPythonExecutionAllowed: () => session?.assertPythonExecutionAllowed(), + trackPythonExecution: (execution, abortController) => session.trackPythonExecution(execution, abortController), + getSessionId: () => sessionManager.getSessionId?.() ?? null, + getSessionSpawns: () => options.spawns ?? "*", + getModelString: () => (hasExplicitModel && model ? formatModelString(model) : undefined), + getActiveModelString: () => { + const activeModel = agent?.state.model; + if (activeModel) return formatModelString(activeModel); + // Fall back to initial model during tool creation (before agent exists) + if (model) return formatModelString(model); + return undefined; + }, + getPlanModeState: () => session.getPlanModeState(), + getCompactContext: () => session.formatCompactContext(), + getTodoPhases: () => session.getTodoPhases(), + setTodoPhases: phases => session.setTodoPhases(phases), + isMCPDiscoveryEnabled: () => session.isMCPDiscoveryEnabled(), + getDiscoverableMCPTools: () => session.getDiscoverableMCPTools(), + getDiscoverableMCPSearchIndex: () => session.getDiscoverableMCPSearchIndex(), + getSelectedMCPToolNames: () => session.getSelectedMCPToolNames(), + activateDiscoveredMCPTools: toolNames => session.activateDiscoveredMCPTools(toolNames), + getCheckpointState: () => session.getCheckpointState(), + setCheckpointState: state => session.setCheckpointState(state ?? undefined), + getToolChoiceQueue: () => session.toolChoiceQueue, + buildToolChoice: name => { + const m = session.model; + return m ? buildNamedToolChoice(name, m) : undefined; + }, + steer: msg => + session.agent.steer({ + role: "custom", + customType: msg.customType, + content: msg.content, + display: false, + details: msg.details, + attribution: "agent", + timestamp: Date.now(), + }), + peekQueueInvoker: () => session.peekQueueInvoker(), + allocateOutputArtifact: async toolType => { + try { + return await sessionManager.allocateArtifactPath(toolType); + } catch { + return {}; } }, - enableProjectConfig: settings.get("mcp.enableProjectConfig") ?? true, - // Always filter Exa - we have native integration - filterExa: true, - // Filter browser MCP servers when builtin browser tool is active - filterBrowser: settings.get("browser.enabled") ?? false, - cacheStorage: settings.getStorage(), + settings, authStorage, - }); - mcpManager = mcpResult.manager; - toolSession.mcpManager = mcpManager; + modelRegistry, + asyncJobManager, + searchDb, + }; - if (settings.get("mcp.notifications")) { - mcpManager.setNotificationsEnabled(true); - } - // If we extracted Exa API keys from MCP configs and EXA_API_KEY isn't set, use the first one - if (mcpResult.exaApiKeys.length > 0 && !$env.EXA_API_KEY) { - Bun.env.EXA_API_KEY = mcpResult.exaApiKeys[0]; - } - - // Log MCP errors - for (const { path, error } of mcpResult.errors) { - logger.error("MCP tool load failed", { path, error }); - } - - if (mcpResult.tools.length > 0) { - // MCP tools are LoadedCustomTool, extract the tool property - customTools.push(...mcpResult.tools.map(loaded => loaded.tool)); - } - } - - // Add Gemini image tools if GEMINI_API_KEY (or GOOGLE_API_KEY) is available - const geminiImageTools = await logger.time("getGeminiImageTools", getGeminiImageTools); - if (geminiImageTools.length > 0) { - customTools.push(...(geminiImageTools as unknown as CustomTool[])); - } - - // Add web search tools - if (options.toolNames?.includes("web_search")) { - customTools.push(...getSearchTools()); - } - - // Discover and load custom tools from .omp/tools/, .claude/tools/, etc. - const builtInToolNames = builtinTools.map(t => t.name); - const discoveredCustomTools = await logger.time( - "discoverAndLoadCustomTools", - discoverAndLoadCustomTools, - [], - cwd, - builtInToolNames, - action => queueResolveHandler(toolSession, action), - ); - for (const { path, error } of discoveredCustomTools.errors) { - logger.error("Custom tool load failed", { path, error }); - } - if (discoveredCustomTools.tools.length > 0) { - customTools.push(...discoveredCustomTools.tools.map(loaded => loaded.tool)); - } - - const inlineExtensions: ExtensionFactory[] = options.extensions ? [...options.extensions] : []; - inlineExtensions.push(createAutoresearchExtension); - if (customTools.length > 0) { - inlineExtensions.push(createCustomToolsExtension(customTools)); - } - - // Load extensions (discovers from standard locations + configured paths) - let extensionsResult: LoadExtensionsResult; - if (options.disableExtensionDiscovery) { - const configuredPaths = options.additionalExtensionPaths ?? []; - extensionsResult = await logger.time("loadExtensions", loadExtensions, configuredPaths, cwd, eventBus); - for (const { path, error } of extensionsResult.errors) { - logger.error("Failed to load extension", { path, error }); - } - } else if (options.preloadedExtensions) { - extensionsResult = options.preloadedExtensions; - } else { - // Merge CLI extension paths with settings extension paths - const configuredPaths = [...(options.additionalExtensionPaths ?? []), ...(settings.get("extensions") ?? [])]; - const disabledExtensionIds = settings.get("disabledExtensions") ?? []; - extensionsResult = await logger.time( - "discoverAndLoadExtensions", - discoverAndLoadExtensions, - configuredPaths, - cwd, - eventBus, - disabledExtensionIds, + // Initialize internal URL router for internal protocols (agent://, artifact://, memory://, skill://, rule://, mcp://, local://) + const internalRouter = new InternalUrlRouter(); + const getArtifactsDir = () => sessionManager.getArtifactsDir(); + internalRouter.register(new AgentProtocolHandler({ getArtifactsDir })); + internalRouter.register(new ArtifactProtocolHandler({ getArtifactsDir })); + internalRouter.register( + new MemoryProtocolHandler({ + getMemoryRoot: () => getMemoryRoot(agentDir, settings.getCwd()), + }), + ); + internalRouter.register( + new LocalProtocolHandler({ + getArtifactsDir, + getSessionId: () => sessionManager.getSessionId(), + }), + ); + internalRouter.register( + new SkillProtocolHandler({ + getSkills: () => skills, + }), + ); + internalRouter.register( + new RuleProtocolHandler({ + getRules: () => [...rulebookRules, ...alwaysApplyRules], + }), + ); + internalRouter.register(new PiProtocolHandler()); + internalRouter.register(new JobsProtocolHandler({ getAsyncJobManager: () => asyncJobManager })); + internalRouter.register(new McpProtocolHandler({ getMcpManager: () => mcpManager })); + toolSession.internalRouter = internalRouter; + toolSession.getArtifactsDir = getArtifactsDir; + toolSession.agentOutputManager = new AgentOutputManager( + getArtifactsDir, + options.parentTaskPrefix ? { parentPrefix: options.parentTaskPrefix } : undefined, ); - for (const { path, error } of extensionsResult.errors) { - logger.error("Failed to load extension", { path, error }); - } - } - // Load inline extensions from factories - if (inlineExtensions.length > 0) { - for (let i = 0; i < inlineExtensions.length; i++) { - const factory = inlineExtensions[i]; - const loaded = await loadExtensionFromFactory( - factory, + // Create built-in tools (already wrapped with meta notice formatting) + const builtinTools = await logger.time("createAllTools", createTools, toolSession, options.toolNames); + + // Discover MCP tools from .mcp.json files + let mcpManager: MCPManager | undefined; + const enableMCP = options.enableMCP ?? true; + const customTools: CustomTool[] = []; + if (enableMCP) { + const mcpResult = await logger.time("discoverAndLoadMCPTools", discoverAndLoadMCPTools, cwd, { + onConnecting: serverNames => { + if (options.hasUI && serverNames.length > 0) { + process.stderr.write(`${chalk.gray(`Connecting to MCP servers: ${serverNames.join(", ")}…`)}\n`); + } + }, + enableProjectConfig: settings.get("mcp.enableProjectConfig") ?? true, + // Always filter Exa - we have native integration + filterExa: true, + // Filter browser MCP servers when builtin browser tool is active + filterBrowser: settings.get("browser.enabled") ?? false, + cacheStorage: settings.getStorage(), + authStorage, + }); + mcpManager = mcpResult.manager; + toolSession.mcpManager = mcpManager; + + if (settings.get("mcp.notifications")) { + mcpManager.setNotificationsEnabled(true); + } + // If we extracted Exa API keys from MCP configs and EXA_API_KEY isn't set, use the first one + if (mcpResult.exaApiKeys.length > 0 && !$env.EXA_API_KEY) { + Bun.env.EXA_API_KEY = mcpResult.exaApiKeys[0]; + } + + // Log MCP errors + for (const { path, error } of mcpResult.errors) { + logger.error("MCP tool load failed", { path, error }); + } + + if (mcpResult.tools.length > 0) { + // MCP tools are LoadedCustomTool, extract the tool property + customTools.push(...mcpResult.tools.map(loaded => loaded.tool)); + } + } + + // Add Gemini image tools if GEMINI_API_KEY (or GOOGLE_API_KEY) is available + const geminiImageTools = await logger.time("getGeminiImageTools", getGeminiImageTools); + if (geminiImageTools.length > 0) { + customTools.push(...(geminiImageTools as unknown as CustomTool[])); + } + + // Add web search tools + if (options.toolNames?.includes("web_search")) { + customTools.push(...getSearchTools()); + } + + // Discover and load custom tools from .omp/tools/, .claude/tools/, etc. + const builtInToolNames = builtinTools.map(t => t.name); + const discoveredCustomTools = await logger.time( + "discoverAndLoadCustomTools", + discoverAndLoadCustomTools, + [], + cwd, + builtInToolNames, + action => queueResolveHandler(toolSession, action), + ); + for (const { path, error } of discoveredCustomTools.errors) { + logger.error("Custom tool load failed", { path, error }); + } + if (discoveredCustomTools.tools.length > 0) { + customTools.push(...discoveredCustomTools.tools.map(loaded => loaded.tool)); + } + + const inlineExtensions: ExtensionFactory[] = options.extensions ? [...options.extensions] : []; + inlineExtensions.push(createAutoresearchExtension); + if (customTools.length > 0) { + inlineExtensions.push(createCustomToolsExtension(customTools)); + } + + // Load extensions (discovers from standard locations + configured paths) + let extensionsResult: LoadExtensionsResult; + if (options.disableExtensionDiscovery) { + const configuredPaths = options.additionalExtensionPaths ?? []; + extensionsResult = await logger.time("loadExtensions", loadExtensions, configuredPaths, cwd, eventBus); + for (const { path, error } of extensionsResult.errors) { + logger.error("Failed to load extension", { path, error }); + } + } else if (options.preloadedExtensions) { + extensionsResult = options.preloadedExtensions; + } else { + // Merge CLI extension paths with settings extension paths + const configuredPaths = [...(options.additionalExtensionPaths ?? []), ...(settings.get("extensions") ?? [])]; + const disabledExtensionIds = settings.get("disabledExtensions") ?? []; + extensionsResult = await logger.time( + "discoverAndLoadExtensions", + discoverAndLoadExtensions, + configuredPaths, cwd, eventBus, - extensionsResult.runtime, - ``, + disabledExtensionIds, ); - extensionsResult.extensions.push(loaded); - } - } - - // Process provider registrations queued during extension loading. - // This must happen before the runner is created so that models registered by - // extensions are available for model selection on session resume / fallback. - const activeExtensionSources = extensionsResult.extensions.map(extension => extension.path); - modelRegistry.syncExtensionSources(activeExtensionSources); - for (const sourceId of new Set(activeExtensionSources)) { - modelRegistry.clearSourceRegistrations(sourceId); - } - if (extensionsResult.runtime.pendingProviderRegistrations.length > 0) { - for (const { name, config, sourceId } of extensionsResult.runtime.pendingProviderRegistrations) { - modelRegistry.registerProvider(name, config, sourceId); - } - extensionsResult.runtime.pendingProviderRegistrations = []; - } - - // Resolve deferred --model pattern now that extension models are registered. - if (!model && options.modelPattern) { - const availableModels = modelRegistry.getAll(); - const matchPreferences = { - usageOrder: settings.getStorage()?.getModelUsageOrder(), - }; - const { model: resolved } = parseModelPattern(options.modelPattern, availableModels, matchPreferences, { - modelRegistry, - }); - if (resolved) { - model = resolved; - modelFallbackMessage = undefined; - } else { - modelFallbackMessage = `Model "${options.modelPattern}" not found`; - } - } - - // Fall back to first available model with a valid API key. - // Skip fallback if the user explicitly requested a model via --model that wasn't found. - if (!model && !options.modelPattern) { - const allModels = modelRegistry.getAll(); - for (const candidate of allModels) { - if (await hasModelApiKey(candidate)) { - model = candidate; - break; + for (const { path, error } of extensionsResult.errors) { + logger.error("Failed to load extension", { path, error }); } } - if (model) { - if (modelFallbackMessage) { - modelFallbackMessage += `. Using ${model.provider}/${model.id}`; + + // Load inline extensions from factories + if (inlineExtensions.length > 0) { + for (let i = 0; i < inlineExtensions.length; i++) { + const factory = inlineExtensions[i]; + const loaded = await loadExtensionFromFactory( + factory, + cwd, + eventBus, + extensionsResult.runtime, + ``, + ); + extensionsResult.extensions.push(loaded); } - } else { - modelFallbackMessage = - "No models available. Use /login or set an API key environment variable. Then use /model to select a model."; } - } - // Discover custom commands (TypeScript slash commands) - const customCommandsResult: CustomCommandsLoadResult = options.disableExtensionDiscovery - ? { commands: [], errors: [] } - : await logger.time("discoverCustomCommands", loadCustomCommandsInternal, { cwd, agentDir }); - if (!options.disableExtensionDiscovery) { - for (const { path, error } of customCommandsResult.errors) { - logger.error("Failed to load custom command", { path, error }); + // Process provider registrations queued during extension loading. + // This must happen before the runner is created so that models registered by + // extensions are available for model selection on session resume / fallback. + const activeExtensionSources = extensionsResult.extensions.map(extension => extension.path); + modelRegistry.syncExtensionSources(activeExtensionSources); + for (const sourceId of new Set(activeExtensionSources)) { + modelRegistry.clearSourceRegistrations(sourceId); + } + if (extensionsResult.runtime.pendingProviderRegistrations.length > 0) { + for (const { name, config, sourceId } of extensionsResult.runtime.pendingProviderRegistrations) { + modelRegistry.registerProvider(name, config, sourceId); + } + extensionsResult.runtime.pendingProviderRegistrations = []; } - } - let extensionRunner: ExtensionRunner | undefined; - if (extensionsResult.extensions.length > 0) { - extensionRunner = new ExtensionRunner( - extensionsResult.extensions, - extensionsResult.runtime, - cwd, + // Resolve deferred --model pattern now that extension models are registered. + if (!model && options.modelPattern) { + const availableModels = modelRegistry.getAll(); + const matchPreferences = { + usageOrder: settings.getStorage()?.getModelUsageOrder(), + }; + const { model: resolved } = parseModelPattern(options.modelPattern, availableModels, matchPreferences, { + modelRegistry, + }); + if (resolved) { + model = resolved; + modelFallbackMessage = undefined; + } else { + modelFallbackMessage = `Model "${options.modelPattern}" not found`; + } + } + + // Fall back to first available model with a valid API key. + // Skip fallback if the user explicitly requested a model via --model that wasn't found. + if (!model && !options.modelPattern) { + const allModels = modelRegistry.getAll(); + for (const candidate of allModels) { + if (await hasModelApiKey(candidate)) { + model = candidate; + break; + } + } + if (model) { + if (modelFallbackMessage) { + modelFallbackMessage += `. Using ${model.provider}/${model.id}`; + } + } else { + modelFallbackMessage = + "No models available. Use /login or set an API key environment variable. Then use /model to select a model."; + } + } + + // Discover custom commands (TypeScript slash commands) + const customCommandsResult: CustomCommandsLoadResult = options.disableExtensionDiscovery + ? { commands: [], errors: [] } + : await logger.time("discoverCustomCommands", loadCustomCommandsInternal, { cwd, agentDir }); + if (!options.disableExtensionDiscovery) { + for (const { path, error } of customCommandsResult.errors) { + logger.error("Failed to load custom command", { path, error }); + } + } + + let extensionRunner: ExtensionRunner | undefined; + if (extensionsResult.extensions.length > 0) { + extensionRunner = new ExtensionRunner( + extensionsResult.extensions, + extensionsResult.runtime, + cwd, + sessionManager, + modelRegistry, + ); + } + + const getSessionContext = () => ({ sessionManager, modelRegistry, - ); - } - - const getSessionContext = () => ({ - sessionManager, - modelRegistry, - model: agent.state.model, - isIdle: () => !session.isStreaming, - hasQueuedMessages: () => session.queuedMessageCount > 0, - abort: () => { - session.abort(); - }, - settings, - }); - const toolContextStore = new ToolContextStore(getSessionContext); - - const registeredTools = extensionRunner?.getAllRegisteredTools() ?? []; - let wrappedExtensionTools: Tool[]; - - if (extensionRunner) { - // With extension runner: convert CustomTools to ToolDefinitions and wrap all together - const allCustomTools = [ - ...registeredTools, - ...(options.customTools?.map(tool => { - const definition = isCustomTool(tool) ? customToolToDefinition(tool) : tool; - return { definition, extensionPath: "" }; - }) ?? []), - ]; - wrappedExtensionTools = wrapRegisteredTools(allCustomTools, extensionRunner); - } else { - // Without extension runner: wrap CustomTools directly with CustomToolAdapter - // ToolDefinition items require ExtensionContext and cannot be used without a runner - const customToolContext = (): CustomToolContext => ({ - sessionManager, - modelRegistry, - model: agent?.state.model, - searchDb, - isIdle: () => !session?.isStreaming, - hasQueuedMessages: () => (session?.queuedMessageCount ?? 0) > 0, - abort: () => session?.abort(), + model: agent.state.model, + isIdle: () => !session.isStreaming, + hasQueuedMessages: () => session.queuedMessageCount > 0, + abort: () => { + session.abort(); + }, settings, }); - wrappedExtensionTools = (options.customTools ?? []) - .filter(isCustomTool) - .map(tool => CustomToolAdapter.wrap(tool, customToolContext)); - } + const toolContextStore = new ToolContextStore(getSessionContext); - // All built-in tools are active (conditional tools like git/ask return null from factory if disabled) - const toolRegistry = new Map(); - for (const tool of builtinTools) { - toolRegistry.set(tool.name, tool); - } - for (const tool of wrappedExtensionTools) { - toolRegistry.set(tool.name, tool); - } - if (extensionRunner) { - for (const tool of toolRegistry.values()) { - toolRegistry.set(tool.name, new ExtensionToolWrapper(tool, extensionRunner)); + const registeredTools = extensionRunner?.getAllRegisteredTools() ?? []; + let wrappedExtensionTools: Tool[]; + + if (extensionRunner) { + // With extension runner: convert CustomTools to ToolDefinitions and wrap all together + const allCustomTools = [ + ...registeredTools, + ...(options.customTools?.map(tool => { + const definition = isCustomTool(tool) ? customToolToDefinition(tool) : tool; + return { definition, extensionPath: "" }; + }) ?? []), + ]; + wrappedExtensionTools = wrapRegisteredTools(allCustomTools, extensionRunner); + } else { + // Without extension runner: wrap CustomTools directly with CustomToolAdapter + // ToolDefinition items require ExtensionContext and cannot be used without a runner + const customToolContext = (): CustomToolContext => ({ + sessionManager, + modelRegistry, + model: agent?.state.model, + searchDb, + isIdle: () => !session?.isStreaming, + hasQueuedMessages: () => (session?.queuedMessageCount ?? 0) > 0, + abort: () => session?.abort(), + settings, + }); + wrappedExtensionTools = (options.customTools ?? []) + .filter(isCustomTool) + .map(tool => CustomToolAdapter.wrap(tool, customToolContext)); } - } - if (model?.provider === "cursor") { - toolRegistry.delete("edit"); - } - const hasDeferrableTools = Array.from(toolRegistry.values()).some(tool => tool.deferrable === true); - if (!hasDeferrableTools) { - toolRegistry.delete("resolve"); - } else if (!toolRegistry.has("resolve")) { - const resolveTool = await logger.time("createTools:resolve:session", HIDDEN_TOOLS.resolve, toolSession); - if (resolveTool) { - toolRegistry.set(resolveTool.name, wrapToolWithMetaNotice(resolveTool)); + // All built-in tools are active (conditional tools like git/ask return null from factory if disabled) + const toolRegistry = new Map(); + for (const tool of builtinTools) { + toolRegistry.set(tool.name, tool); } - } - - let cursorEventEmitter: ((event: AgentEvent) => void) | undefined; - const cursorExecHandlers = new CursorExecHandlers({ - cwd, - tools: toolRegistry, - getToolContext: () => toolContextStore.getContext(), - emitEvent: event => cursorEventEmitter?.(event), - }); - - const repeatToolDescriptions = settings.get("repeatToolDescriptions"); - const eagerTasks = settings.get("task.eager"); - const intentField = settings.get("tools.intentTracing") || $flag("PI_INTENT_TRACING") ? INTENT_FIELD : undefined; - const rebuildSystemPrompt = async (toolNames: string[], tools: Map): Promise => { - toolContextStore.setToolNames(toolNames); - const discoverableMCPTools = mcpDiscoveryEnabled ? collectDiscoverableMCPTools(tools.values()) : []; - const discoverableMCPSummary = summarizeDiscoverableMCPTools(discoverableMCPTools); - const hasDiscoverableMCPTools = - mcpDiscoveryEnabled && toolNames.includes("search_tool_bm25") && discoverableMCPTools.length > 0; - const promptTools = buildSystemPromptToolMetadata(tools, { - search_tool_bm25: { description: renderSearchToolBm25Description(discoverableMCPTools) }, - }); - const memoryInstructions = await buildMemoryToolDeveloperInstructions(agentDir, settings); - - // Build combined append prompt: memory instructions + MCP server instructions - const serverInstructions = mcpManager?.getServerInstructions(); - let appendPrompt: string | undefined = memoryInstructions ?? undefined; - if (serverInstructions && serverInstructions.size > 0) { - const MAX_INSTRUCTIONS_LENGTH = 4000; - const parts: string[] = []; - if (appendPrompt) parts.push(appendPrompt); - parts.push( - "## MCP Server Instructions\n\nThe following instructions are provided by connected MCP servers. They are server-controlled and may not be verified.", - ); - for (const [srvName, srvInstructions] of serverInstructions) { - const truncated = - srvInstructions.length > MAX_INSTRUCTIONS_LENGTH - ? `${srvInstructions.slice(0, MAX_INSTRUCTIONS_LENGTH)}\n[truncated]` - : srvInstructions; - parts.push(`### ${srvName}\n${truncated}`); + for (const tool of wrappedExtensionTools) { + toolRegistry.set(tool.name, tool); + } + if (extensionRunner) { + for (const tool of toolRegistry.values()) { + toolRegistry.set(tool.name, new ExtensionToolWrapper(tool, extensionRunner)); } - appendPrompt = parts.join("\n\n"); } - const defaultPrompt = await buildSystemPromptInternal({ + if (model?.provider === "cursor") { + toolRegistry.delete("edit"); + } + + const hasDeferrableTools = Array.from(toolRegistry.values()).some(tool => tool.deferrable === true); + if (!hasDeferrableTools) { + toolRegistry.delete("resolve"); + } else if (!toolRegistry.has("resolve")) { + const resolveTool = await logger.time("createTools:resolve:session", HIDDEN_TOOLS.resolve, toolSession); + if (resolveTool) { + toolRegistry.set(resolveTool.name, wrapToolWithMetaNotice(resolveTool)); + } + } + + let cursorEventEmitter: ((event: AgentEvent) => void) | undefined; + const cursorExecHandlers = new CursorExecHandlers({ cwd, - skills, - contextFiles, - tools: promptTools, - toolNames, - rules: rulebookRules, - alwaysApplyRules, - skillsSettings: settings.getGroup("skills"), - appendSystemPrompt: appendPrompt, - repeatToolDescriptions, - intentField, - mcpDiscoveryMode: hasDiscoverableMCPTools, - mcpDiscoveryServerSummaries: discoverableMCPSummary.servers.map(formatDiscoverableMCPToolServerSummary), - eagerTasks, - secretsEnabled, + tools: toolRegistry, + getToolContext: () => toolContextStore.getContext(), + emitEvent: event => cursorEventEmitter?.(event), }); - if (options.systemPrompt === undefined) { - return defaultPrompt; - } - if (typeof options.systemPrompt === "string") { - return await buildSystemPromptInternal({ + const repeatToolDescriptions = settings.get("repeatToolDescriptions"); + const eagerTasks = settings.get("task.eager"); + const intentField = settings.get("tools.intentTracing") || $flag("PI_INTENT_TRACING") ? INTENT_FIELD : undefined; + const rebuildSystemPrompt = async (toolNames: string[], tools: Map): Promise => { + toolContextStore.setToolNames(toolNames); + const discoverableMCPTools = mcpDiscoveryEnabled ? collectDiscoverableMCPTools(tools.values()) : []; + const discoverableMCPSummary = summarizeDiscoverableMCPTools(discoverableMCPTools); + const hasDiscoverableMCPTools = + mcpDiscoveryEnabled && toolNames.includes("search_tool_bm25") && discoverableMCPTools.length > 0; + const promptTools = buildSystemPromptToolMetadata(tools, { + search_tool_bm25: { description: renderSearchToolBm25Description(discoverableMCPTools) }, + }); + const memoryInstructions = await buildMemoryToolDeveloperInstructions(agentDir, settings); + + // Build combined append prompt: memory instructions + MCP server instructions + const serverInstructions = mcpManager?.getServerInstructions(); + let appendPrompt: string | undefined = memoryInstructions ?? undefined; + if (serverInstructions && serverInstructions.size > 0) { + const MAX_INSTRUCTIONS_LENGTH = 4000; + const parts: string[] = []; + if (appendPrompt) parts.push(appendPrompt); + parts.push( + "## MCP Server Instructions\n\nThe following instructions are provided by connected MCP servers. They are server-controlled and may not be verified.", + ); + for (const [srvName, srvInstructions] of serverInstructions) { + const truncated = + srvInstructions.length > MAX_INSTRUCTIONS_LENGTH + ? `${srvInstructions.slice(0, MAX_INSTRUCTIONS_LENGTH)}\n[truncated]` + : srvInstructions; + parts.push(`### ${srvName}\n${truncated}`); + } + appendPrompt = parts.join("\n\n"); + } + const defaultPrompt = await buildSystemPromptInternal({ cwd, skills, contextFiles, @@ -1328,7 +1316,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} rules: rulebookRules, alwaysApplyRules, skillsSettings: settings.getGroup("skills"), - customPrompt: options.systemPrompt, appendSystemPrompt: appendPrompt, repeatToolDescriptions, intentField, @@ -1337,362 +1324,406 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} eagerTasks, secretsEnabled, }); - } - return options.systemPrompt(defaultPrompt); - }; - const toolNamesFromRegistry = Array.from(toolRegistry.keys()); - const requestedToolNames = options.toolNames?.map(name => name.toLowerCase()) ?? toolNamesFromRegistry; - const normalizedRequested = requestedToolNames.filter(name => toolRegistry.has(name)); - const includeExitPlanMode = requestedToolNames.includes("exit_plan_mode"); - const mcpDiscoveryEnabled = settings.get("mcp.discoveryMode") ?? false; - const defaultInactiveToolNames = new Set( - registeredTools.filter(tool => tool.definition.defaultInactive).map(tool => tool.definition.name), - ); - const requestedActiveToolNames = includeExitPlanMode - ? normalizedRequested - : normalizedRequested.filter(name => name !== "exit_plan_mode"); - const initialRequestedActiveToolNames = options.toolNames - ? requestedActiveToolNames - : requestedActiveToolNames.filter(name => !defaultInactiveToolNames.has(name)); - const explicitlyRequestedMCPToolNames = options.toolNames - ? requestedActiveToolNames.filter(name => name.startsWith("mcp_")) - : []; - const discoveryDefaultServers = new Set( - (settings.get("mcp.discoveryDefaultServers") ?? []).map(serverName => serverName.trim()).filter(Boolean), - ); - const discoveryDefaultServerToolNames = mcpDiscoveryEnabled - ? selectDiscoverableMCPToolNamesByServer( - collectDiscoverableMCPTools(toolRegistry.values()), - discoveryDefaultServers, - ) - : []; - let initialSelectedMCPToolNames: string[] = []; - let defaultSelectedMCPToolNames: string[] = []; - let initialToolNames = [...initialRequestedActiveToolNames]; - if (mcpDiscoveryEnabled) { - const restoredSelectedMCPToolNames = existingSession.selectedMCPToolNames.filter(name => toolRegistry.has(name)); - defaultSelectedMCPToolNames = [ - ...new Set([...discoveryDefaultServerToolNames, ...explicitlyRequestedMCPToolNames]), - ]; - initialSelectedMCPToolNames = existingSession.hasPersistedMCPToolSelection - ? restoredSelectedMCPToolNames - : [...new Set([...restoredSelectedMCPToolNames, ...defaultSelectedMCPToolNames])]; - initialToolNames = [ - ...new Set([ - ...initialRequestedActiveToolNames.filter(name => !name.startsWith("mcp_")), - ...initialSelectedMCPToolNames, - ]), - ]; - } - - // Custom tools and extension-registered tools are always included regardless of toolNames filter - const alwaysInclude: string[] = [ - ...(options.customTools?.map(t => (isCustomTool(t) ? t.name : t.name)) ?? []), - ...registeredTools.filter(t => !t.definition.defaultInactive).map(t => t.definition.name), - ]; - for (const name of alwaysInclude) { - if (mcpDiscoveryEnabled && name.startsWith("mcp_")) { - continue; - } - if (toolRegistry.has(name) && !initialToolNames.includes(name)) { - initialToolNames.push(name); - } - } - - const systemPrompt = await logger.time("buildSystemPrompt", rebuildSystemPrompt, initialToolNames, toolRegistry); - - const promptTemplates = - options.promptTemplates ?? (await logger.time("discoverPromptTemplates", discoverPromptTemplates, cwd, agentDir)); - toolSession.promptTemplates = promptTemplates; - - const slashCommands = - options.slashCommands ?? (await logger.time("discoverSlashCommands", discoverSlashCommands, cwd)); - - // Create convertToLlm wrapper that filters images if blockImages is enabled (defense-in-depth) - const convertToLlmWithBlockImages = (messages: AgentMessage[]): Message[] => { - const converted = convertToLlm(messages); - // Check setting dynamically so mid-session changes take effect - if (!settings.get("images.blockImages")) { - return converted; - } - // Filter out ImageContent from all messages, replacing with text placeholder - return converted.map(msg => { - if (msg.role === "user" || msg.role === "toolResult") { - const content = msg.content; - if (Array.isArray(content)) { - const hasImages = content.some(c => c.type === "image"); - if (hasImages) { - const filteredContent = content - .map(c => (c.type === "image" ? { type: "text" as const, text: "Image reading is disabled." } : c)) - .filter((c, i, arr) => { - // Dedupe consecutive "Image reading is disabled." texts - if (!(c.type === "text" && c.text === "Image reading is disabled." && i > 0)) return true; - const prev = arr[i - 1]; - return !(prev.type === "text" && prev.text === "Image reading is disabled."); - }); - return { ...msg, content: filteredContent }; - } - } + if (options.systemPrompt === undefined) { + return defaultPrompt; } - return msg; - }); - }; - - // Final convertToLlm: chain block-images filter with secret obfuscation - const convertToLlmFinal = (messages: AgentMessage[]): Message[] => { - const converted = convertToLlmWithBlockImages(messages); - if (!obfuscator?.hasSecrets()) return converted; - return obfuscateMessages(obfuscator, converted); - }; - const transformContext = extensionRunner - ? async (messages: AgentMessage[], _signal?: AbortSignal) => { - return await extensionRunner.emitContext(messages); + if (typeof options.systemPrompt === "string") { + return await buildSystemPromptInternal({ + cwd, + skills, + contextFiles, + tools: promptTools, + toolNames, + rules: rulebookRules, + alwaysApplyRules, + skillsSettings: settings.getGroup("skills"), + customPrompt: options.systemPrompt, + appendSystemPrompt: appendPrompt, + repeatToolDescriptions, + intentField, + mcpDiscoveryMode: hasDiscoverableMCPTools, + mcpDiscoveryServerSummaries: discoverableMCPSummary.servers.map(formatDiscoverableMCPToolServerSummary), + eagerTasks, + secretsEnabled, + }); } - : undefined; - const onPayload = extensionRunner - ? async (payload: unknown, _model?: Model) => { - return await extensionRunner.emitBeforeProviderRequest(payload); - } - : undefined; - - const setToolUIContext = (uiContext: ExtensionUIContext, hasUI: boolean) => { - toolContextStore.setUIContext(uiContext, hasUI); - }; - - const initialTools = initialToolNames - .map(name => toolRegistry.get(name)) - .filter((tool): tool is AgentTool => tool !== undefined); - - const openaiWebsocketSetting = settings.get("providers.openaiWebsockets") ?? "auto"; - const preferOpenAICodexWebsockets = - openaiWebsocketSetting === "on" ? true : openaiWebsocketSetting === "off" ? false : undefined; - const serviceTierSetting = settings.get("serviceTier"); - - const initialServiceTier = hasServiceTierEntry - ? existingSession.serviceTier - : serviceTierSetting === "none" - ? undefined - : serviceTierSetting; - - agent = new Agent({ - initialState: { - systemPrompt, - model, - thinkingLevel: toReasoningEffort(thinkingLevel), - tools: initialTools, - }, - convertToLlm: convertToLlmFinal, - onPayload, - sessionId: providerSessionId, - transformContext, - steeringMode: settings.get("steeringMode") ?? "one-at-a-time", - followUpMode: settings.get("followUpMode") ?? "one-at-a-time", - interruptMode: settings.get("interruptMode") ?? "immediate", - thinkingBudgets: settings.getGroup("thinkingBudgets"), - temperature: settings.get("temperature") >= 0 ? settings.get("temperature") : undefined, - topP: settings.get("topP") >= 0 ? settings.get("topP") : undefined, - topK: settings.get("topK") >= 0 ? settings.get("topK") : undefined, - minP: settings.get("minP") >= 0 ? settings.get("minP") : undefined, - presencePenalty: settings.get("presencePenalty") >= 0 ? settings.get("presencePenalty") : undefined, - repetitionPenalty: settings.get("repetitionPenalty") >= 0 ? settings.get("repetitionPenalty") : undefined, - serviceTier: initialServiceTier, - kimiApiFormat: settings.get("providers.kimiApiFormat") ?? "anthropic", - preferWebsockets: preferOpenAICodexWebsockets, - getToolContext: tc => toolContextStore.getContext(tc), - getApiKey: async provider => { - // Use the provider-facing session id for sticky credential selection so cache keys - // and provider auth affinity stay aligned across fresh benchmark sessions. - const key = await modelRegistry.getApiKeyForProvider(provider, providerSessionId); - if (!key) { - throw new Error(`No API key found for provider "${provider}"`); - } - return key; - }, - cursorExecHandlers, - transformToolCallArguments: (args, _toolName) => { - let result = args; - const maxTimeout = settings.get("tools.maxTimeout"); - if (maxTimeout > 0 && typeof result.timeout === "number") { - result = { ...result, timeout: Math.min(result.timeout, maxTimeout) }; - } - if (obfuscator?.hasSecrets()) { - result = obfuscator.deobfuscateObject(result); - } - return result; - }, - intentTracing: !!intentField, - getToolChoice: () => session?.nextToolChoice(), - }); - - cursorEventEmitter = event => agent.emitExternalEvent(event); - - // Restore messages if session has existing data - if (hasExistingSession) { - agent.replaceMessages(existingSession.messages); - } else { - // Save initial model and thinking level for new sessions so they can be restored on resume - if (model) { - sessionManager.appendModelChange(`${model.provider}/${model.id}`); - } - sessionManager.appendThinkingLevelChange(thinkingLevel); - } - - session = new AgentSession({ - agent, - thinkingLevel, - sessionManager, - settings, - scopedModels: options.scopedModels, - promptTemplates, - slashCommands, - extensionRunner, - customCommands: customCommandsResult.commands, - skills, - skillWarnings, - skillsSettings: settings.getGroup("skills"), - modelRegistry, - toolRegistry, - transformContext, - onPayload, - convertToLlm: convertToLlmFinal, - rebuildSystemPrompt, - mcpDiscoveryEnabled, - initialSelectedMCPToolNames, - defaultSelectedMCPToolNames, - persistInitialMCPToolSelection: !hasExistingSession, - defaultSelectedMCPServerNames: [...discoveryDefaultServers], - ttsrManager, - obfuscator, - asyncJobManager, - searchDb, - }); - - if (model?.api === "openai-codex-responses") { - const codexModel = model; - const codexTransport = getOpenAICodexTransportDetails(codexModel, { - sessionId: providerSessionId, - baseUrl: codexModel.baseUrl, - preferWebsockets: preferOpenAICodexWebsockets, - providerSessionState: session.providerSessionState, - }); - if (codexTransport.websocketPreferred) { - void (async () => { - try { - const codexPrewarmApiKey = await modelRegistry.getApiKey(codexModel, providerSessionId); - if (!codexPrewarmApiKey) return; - await logger.time("prewarmOpenAICodexResponses", prewarmOpenAICodexResponses, codexModel, { - apiKey: codexPrewarmApiKey, - sessionId: providerSessionId, - preferWebsockets: preferOpenAICodexWebsockets, - providerSessionState: session.providerSessionState, - }); - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error); - logger.debug("Codex websocket prewarm failed", { - error: errorMessage, - provider: codexModel.provider, - model: codexModel.id, - }); - } - })(); - } - } - - // Start LSP warmup in the background so startup does not block on language server initialization. - let lspServers: CreateAgentSessionResult["lspServers"]; - if (enableLsp && settings.get("lsp.diagnosticsOnWrite")) { - lspServers = discoverStartupLspServers(cwd); - if (lspServers.length > 0) { - void (async () => { - try { - const result = await logger.time("warmupLspServers", warmupLspServers, cwd); - const serversByName = new Map(result.servers.map(server => [server.name, server] as const)); - for (const server of lspServers ?? []) { - const next = serversByName.get(server.name); - if (!next) continue; - server.status = next.status; - server.fileTypes = next.fileTypes; - server.error = next.error; - } - const event: LspStartupEvent = { - type: "completed", - servers: result.servers, - }; - eventBus.emit(LSP_STARTUP_EVENT_CHANNEL, event); - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error); - logger.warn("LSP server warmup failed", { cwd, error: errorMessage }); - for (const server of lspServers ?? []) { - server.status = "error"; - server.error = errorMessage; - } - const event: LspStartupEvent = { - type: "failed", - error: errorMessage, - }; - eventBus.emit(LSP_STARTUP_EVENT_CHANNEL, event); - } - })(); - } - } - - logger.time("startMemoryStartupTask", () => - startMemoryStartupTask({ - session, - settings, - modelRegistry, - agentDir, - taskDepth, - }), - ); - - // Wire MCP manager callbacks to session for reactive tool updates - if (mcpManager) { - mcpManager.setOnToolsChanged(tools => { - void session.refreshMCPTools(tools); - }); - // Wire prompt refresh → rebuild MCP prompt slash commands - mcpManager.setOnPromptsChanged(serverName => { - const promptCommands = buildMCPPromptCommands(mcpManager); - session.setMCPPromptCommands(promptCommands); - logger.debug("MCP prompt commands refreshed", { path: `mcp:${serverName}` }); - }); - const notificationDebounceTimers = new Map(); - const clearDebounceTimers = () => { - for (const timer of notificationDebounceTimers.values()) clearTimeout(timer); - notificationDebounceTimers.clear(); + return options.systemPrompt(defaultPrompt); }; - postmortem.register("mcp-notification-cleanup", clearDebounceTimers); - mcpManager.setOnResourcesChanged((serverName, uri) => { - logger.debug("MCP resources changed", { path: `mcp:${serverName}`, uri }); - if (!settings.get("mcp.notifications")) return; - const debounceMs = settings.get("mcp.notificationDebounceMs"); - const key = `${serverName}:${uri}`; - const existing = notificationDebounceTimers.get(key); - if (existing) clearTimeout(existing); - notificationDebounceTimers.set( - key, - setTimeout(() => { - notificationDebounceTimers.delete(key); - // Re-check: user may have disabled notifications during the debounce window - if (!settings.get("mcp.notifications")) return; - void session.followUp( - `[MCP notification] Server "${serverName}" reports resource \`${uri}\` was updated. Use read(path="mcp://${uri}") to inspect if relevant.`, - ); - }, debounceMs), - ); - }); - } - logger.time("createAgentSession:return"); - return { - session, - extensionsResult, - setToolUIContext, - mcpManager, - modelFallbackMessage, - lspServers, - eventBus, - }; + const toolNamesFromRegistry = Array.from(toolRegistry.keys()); + const requestedToolNames = options.toolNames?.map(name => name.toLowerCase()) ?? toolNamesFromRegistry; + const normalizedRequested = requestedToolNames.filter(name => toolRegistry.has(name)); + const includeExitPlanMode = requestedToolNames.includes("exit_plan_mode"); + const mcpDiscoveryEnabled = settings.get("mcp.discoveryMode") ?? false; + const defaultInactiveToolNames = new Set( + registeredTools.filter(tool => tool.definition.defaultInactive).map(tool => tool.definition.name), + ); + const requestedActiveToolNames = includeExitPlanMode + ? normalizedRequested + : normalizedRequested.filter(name => name !== "exit_plan_mode"); + const initialRequestedActiveToolNames = options.toolNames + ? requestedActiveToolNames + : requestedActiveToolNames.filter(name => !defaultInactiveToolNames.has(name)); + const explicitlyRequestedMCPToolNames = options.toolNames + ? requestedActiveToolNames.filter(name => name.startsWith("mcp_")) + : []; + const discoveryDefaultServers = new Set( + (settings.get("mcp.discoveryDefaultServers") ?? []).map(serverName => serverName.trim()).filter(Boolean), + ); + const discoveryDefaultServerToolNames = mcpDiscoveryEnabled + ? selectDiscoverableMCPToolNamesByServer( + collectDiscoverableMCPTools(toolRegistry.values()), + discoveryDefaultServers, + ) + : []; + let initialSelectedMCPToolNames: string[] = []; + let defaultSelectedMCPToolNames: string[] = []; + let initialToolNames = [...initialRequestedActiveToolNames]; + if (mcpDiscoveryEnabled) { + const restoredSelectedMCPToolNames = existingSession.selectedMCPToolNames.filter(name => + toolRegistry.has(name), + ); + defaultSelectedMCPToolNames = [ + ...new Set([...discoveryDefaultServerToolNames, ...explicitlyRequestedMCPToolNames]), + ]; + initialSelectedMCPToolNames = existingSession.hasPersistedMCPToolSelection + ? restoredSelectedMCPToolNames + : [...new Set([...restoredSelectedMCPToolNames, ...defaultSelectedMCPToolNames])]; + initialToolNames = [ + ...new Set([ + ...initialRequestedActiveToolNames.filter(name => !name.startsWith("mcp_")), + ...initialSelectedMCPToolNames, + ]), + ]; + } + + // Custom tools and extension-registered tools are always included regardless of toolNames filter + const alwaysInclude: string[] = [ + ...(options.customTools?.map(t => (isCustomTool(t) ? t.name : t.name)) ?? []), + ...registeredTools.filter(t => !t.definition.defaultInactive).map(t => t.definition.name), + ]; + for (const name of alwaysInclude) { + if (mcpDiscoveryEnabled && name.startsWith("mcp_")) { + continue; + } + if (toolRegistry.has(name) && !initialToolNames.includes(name)) { + initialToolNames.push(name); + } + } + + const systemPrompt = await logger.time("buildSystemPrompt", rebuildSystemPrompt, initialToolNames, toolRegistry); + + const promptTemplates = + options.promptTemplates ?? + (await logger.time("discoverPromptTemplates", discoverPromptTemplates, cwd, agentDir)); + toolSession.promptTemplates = promptTemplates; + + const slashCommands = + options.slashCommands ?? (await logger.time("discoverSlashCommands", discoverSlashCommands, cwd)); + + // Create convertToLlm wrapper that filters images if blockImages is enabled (defense-in-depth) + const convertToLlmWithBlockImages = (messages: AgentMessage[]): Message[] => { + const converted = convertToLlm(messages); + // Check setting dynamically so mid-session changes take effect + if (!settings.get("images.blockImages")) { + return converted; + } + // Filter out ImageContent from all messages, replacing with text placeholder + return converted.map(msg => { + if (msg.role === "user" || msg.role === "toolResult") { + const content = msg.content; + if (Array.isArray(content)) { + const hasImages = content.some(c => c.type === "image"); + if (hasImages) { + const filteredContent = content + .map(c => + c.type === "image" ? { type: "text" as const, text: "Image reading is disabled." } : c, + ) + .filter((c, i, arr) => { + // Dedupe consecutive "Image reading is disabled." texts + if (!(c.type === "text" && c.text === "Image reading is disabled." && i > 0)) return true; + const prev = arr[i - 1]; + return !(prev.type === "text" && prev.text === "Image reading is disabled."); + }); + return { ...msg, content: filteredContent }; + } + } + } + return msg; + }); + }; + + // Final convertToLlm: chain block-images filter with secret obfuscation + const convertToLlmFinal = (messages: AgentMessage[]): Message[] => { + const converted = convertToLlmWithBlockImages(messages); + if (!obfuscator?.hasSecrets()) return converted; + return obfuscateMessages(obfuscator, converted); + }; + const transformContext = extensionRunner + ? async (messages: AgentMessage[], _signal?: AbortSignal) => { + return await extensionRunner.emitContext(messages); + } + : undefined; + const onPayload = extensionRunner + ? async (payload: unknown, _model?: Model) => { + return await extensionRunner.emitBeforeProviderRequest(payload); + } + : undefined; + + const setToolUIContext = (uiContext: ExtensionUIContext, hasUI: boolean) => { + toolContextStore.setUIContext(uiContext, hasUI); + }; + + const initialTools = initialToolNames + .map(name => toolRegistry.get(name)) + .filter((tool): tool is AgentTool => tool !== undefined); + + const openaiWebsocketSetting = settings.get("providers.openaiWebsockets") ?? "auto"; + const preferOpenAICodexWebsockets = + openaiWebsocketSetting === "on" ? true : openaiWebsocketSetting === "off" ? false : undefined; + const serviceTierSetting = settings.get("serviceTier"); + + const initialServiceTier = hasServiceTierEntry + ? existingSession.serviceTier + : serviceTierSetting === "none" + ? undefined + : serviceTierSetting; + + agent = new Agent({ + initialState: { + systemPrompt, + model, + thinkingLevel: toReasoningEffort(thinkingLevel), + tools: initialTools, + }, + convertToLlm: convertToLlmFinal, + onPayload, + sessionId: providerSessionId, + transformContext, + steeringMode: settings.get("steeringMode") ?? "one-at-a-time", + followUpMode: settings.get("followUpMode") ?? "one-at-a-time", + interruptMode: settings.get("interruptMode") ?? "immediate", + thinkingBudgets: settings.getGroup("thinkingBudgets"), + temperature: settings.get("temperature") >= 0 ? settings.get("temperature") : undefined, + topP: settings.get("topP") >= 0 ? settings.get("topP") : undefined, + topK: settings.get("topK") >= 0 ? settings.get("topK") : undefined, + minP: settings.get("minP") >= 0 ? settings.get("minP") : undefined, + presencePenalty: settings.get("presencePenalty") >= 0 ? settings.get("presencePenalty") : undefined, + repetitionPenalty: settings.get("repetitionPenalty") >= 0 ? settings.get("repetitionPenalty") : undefined, + serviceTier: initialServiceTier, + kimiApiFormat: settings.get("providers.kimiApiFormat") ?? "anthropic", + preferWebsockets: preferOpenAICodexWebsockets, + getToolContext: tc => toolContextStore.getContext(tc), + getApiKey: async provider => { + // Use the provider-facing session id for sticky credential selection so cache keys + // and provider auth affinity stay aligned across fresh benchmark sessions. + const key = await modelRegistry.getApiKeyForProvider(provider, providerSessionId); + if (!key) { + throw new Error(`No API key found for provider "${provider}"`); + } + return key; + }, + cursorExecHandlers, + transformToolCallArguments: (args, _toolName) => { + let result = args; + const maxTimeout = settings.get("tools.maxTimeout"); + if (maxTimeout > 0 && typeof result.timeout === "number") { + result = { ...result, timeout: Math.min(result.timeout, maxTimeout) }; + } + if (obfuscator?.hasSecrets()) { + result = obfuscator.deobfuscateObject(result); + } + return result; + }, + intentTracing: !!intentField, + getToolChoice: () => session?.nextToolChoice(), + }); + + cursorEventEmitter = event => agent.emitExternalEvent(event); + + // Restore messages if session has existing data + if (hasExistingSession) { + agent.replaceMessages(existingSession.messages); + } else { + // Save initial model and thinking level for new sessions so they can be restored on resume + if (model) { + sessionManager.appendModelChange(`${model.provider}/${model.id}`); + } + sessionManager.appendThinkingLevelChange(thinkingLevel); + } + + session = new AgentSession({ + agent, + thinkingLevel, + sessionManager, + settings, + pythonKernelOwnerId, + scopedModels: options.scopedModels, + promptTemplates, + slashCommands, + extensionRunner, + customCommands: customCommandsResult.commands, + skills, + skillWarnings, + skillsSettings: settings.getGroup("skills"), + modelRegistry, + toolRegistry, + transformContext, + onPayload, + convertToLlm: convertToLlmFinal, + rebuildSystemPrompt, + mcpDiscoveryEnabled, + initialSelectedMCPToolNames, + defaultSelectedMCPToolNames, + persistInitialMCPToolSelection: !hasExistingSession, + defaultSelectedMCPServerNames: [...discoveryDefaultServers], + ttsrManager, + obfuscator, + asyncJobManager, + searchDb, + }); + hasSession = true; + + if (model?.api === "openai-codex-responses") { + const codexModel = model; + const codexTransport = getOpenAICodexTransportDetails(codexModel, { + sessionId: providerSessionId, + baseUrl: codexModel.baseUrl, + preferWebsockets: preferOpenAICodexWebsockets, + providerSessionState: session.providerSessionState, + }); + if (codexTransport.websocketPreferred) { + void (async () => { + try { + const codexPrewarmApiKey = await modelRegistry.getApiKey(codexModel, providerSessionId); + if (!codexPrewarmApiKey) return; + await logger.time("prewarmOpenAICodexResponses", prewarmOpenAICodexResponses, codexModel, { + apiKey: codexPrewarmApiKey, + sessionId: providerSessionId, + preferWebsockets: preferOpenAICodexWebsockets, + providerSessionState: session.providerSessionState, + }); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + logger.debug("Codex websocket prewarm failed", { + error: errorMessage, + provider: codexModel.provider, + model: codexModel.id, + }); + } + })(); + } + } + + // Start LSP warmup in the background so startup does not block on language server initialization. + let lspServers: CreateAgentSessionResult["lspServers"]; + if (enableLsp && settings.get("lsp.diagnosticsOnWrite")) { + lspServers = discoverStartupLspServers(cwd); + if (lspServers.length > 0) { + void (async () => { + try { + const result = await logger.time("warmupLspServers", warmupLspServers, cwd); + const serversByName = new Map(result.servers.map(server => [server.name, server] as const)); + for (const server of lspServers ?? []) { + const next = serversByName.get(server.name); + if (!next) continue; + server.status = next.status; + server.fileTypes = next.fileTypes; + server.error = next.error; + } + const event: LspStartupEvent = { + type: "completed", + servers: result.servers, + }; + eventBus.emit(LSP_STARTUP_EVENT_CHANNEL, event); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + logger.warn("LSP server warmup failed", { cwd, error: errorMessage }); + for (const server of lspServers ?? []) { + server.status = "error"; + server.error = errorMessage; + } + const event: LspStartupEvent = { + type: "failed", + error: errorMessage, + }; + eventBus.emit(LSP_STARTUP_EVENT_CHANNEL, event); + } + })(); + } + } + + logger.time("startMemoryStartupTask", () => + startMemoryStartupTask({ + session, + settings, + modelRegistry, + agentDir, + taskDepth, + }), + ); + + // Wire MCP manager callbacks to session for reactive tool updates + if (mcpManager) { + mcpManager.setOnToolsChanged(tools => { + void session.refreshMCPTools(tools); + }); + // Wire prompt refresh → rebuild MCP prompt slash commands + mcpManager.setOnPromptsChanged(serverName => { + const promptCommands = buildMCPPromptCommands(mcpManager); + session.setMCPPromptCommands(promptCommands); + logger.debug("MCP prompt commands refreshed", { path: `mcp:${serverName}` }); + }); + const notificationDebounceTimers = new Map(); + const clearDebounceTimers = () => { + for (const timer of notificationDebounceTimers.values()) clearTimeout(timer); + notificationDebounceTimers.clear(); + }; + postmortem.register("mcp-notification-cleanup", clearDebounceTimers); + mcpManager.setOnResourcesChanged((serverName, uri) => { + logger.debug("MCP resources changed", { path: `mcp:${serverName}`, uri }); + if (!settings.get("mcp.notifications")) return; + const debounceMs = settings.get("mcp.notificationDebounceMs"); + const key = `${serverName}:${uri}`; + const existing = notificationDebounceTimers.get(key); + if (existing) clearTimeout(existing); + notificationDebounceTimers.set( + key, + setTimeout(() => { + notificationDebounceTimers.delete(key); + // Re-check: user may have disabled notifications during the debounce window + if (!settings.get("mcp.notifications")) return; + void session.followUp( + `[MCP notification] Server "${serverName}" reports resource \`${uri}\` was updated. Use read(path="mcp://${uri}") to inspect if relevant.`, + ); + }, debounceMs), + ); + }); + } + + logger.time("createAgentSession:return"); + return { + session, + extensionsResult, + setToolUIContext, + mcpManager, + modelFallbackMessage, + lspServers, + eventBus, + }; + } catch (error) { + try { + if (hasSession) { + await session.dispose(); + } else { + await disposeKernelSessionsByOwner(pythonKernelOwnerId); + } + } catch (cleanupError) { + logger.warn("Failed to clean up createAgentSession resources after startup error", { + error: cleanupError instanceof Error ? cleanupError.message : String(cleanupError), + }); + } + throw error; + } } diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 22b2ac473..e4abfdcff 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -50,7 +50,15 @@ import { parseRateLimitReason, } from "@oh-my-pi/pi-ai"; import { killTree, MacOSPowerAssertion, type SearchDb } from "@oh-my-pi/pi-natives"; -import { abortableSleep, getAgentDbPath, isEnoent, logger, prompt, setNativeKillTree } from "@oh-my-pi/pi-utils"; +import { + abortableSleep, + getAgentDbPath, + isEnoent, + logger, + prompt, + Snowflake, + setNativeKillTree, +} from "@oh-my-pi/pi-utils"; import type { AsyncJob, AsyncJobManager } from "../async"; import type { Rule } from "../capability/rule"; import { MODEL_ROLE_IDS, type ModelRegistry } from "../config/model-registry"; @@ -95,7 +103,11 @@ import type { HookCommandContext } from "../extensibility/hooks/types"; import type { Skill, SkillWarning } from "../extensibility/skills"; import { expandSlashCommand, type FileSlashCommand } from "../extensibility/slash-commands"; import { resolveLocalUrlToPath } from "../internal-urls"; -import { executePython as executePythonCommand, type PythonResult } from "../ipy/executor"; +import { + disposeKernelSessionsByOwner, + executePython as executePythonCommand, + type PythonResult, +} from "../ipy/executor"; import { buildDiscoverableMCPSearchIndex, collectDiscoverableMCPTools, @@ -246,6 +258,8 @@ export interface AgentSessionConfig { obfuscator?: SecretObfuscator; /** Shared native search DB for grep/glob/fuzzyFind-backed workflows. */ searchDb?: SearchDb; + /** Logical owner for retained Python kernels created by this session. */ + pythonKernelOwnerId?: string; } /** Options for AgentSession.prompt() */ @@ -452,9 +466,11 @@ export class AgentSession { #pendingBashMessages: BashExecutionMessage[] = []; // Python execution state - #pythonAbortController: AbortController | undefined = undefined; + #pythonAbortControllers = new Set(); + #pythonKernelOwnerId: string; #pendingPythonMessages: PythonExecutionMessage[] = []; - + #activePythonExecutions = new Set>(); + #pythonExecutionDisposing = false; // Extension system #extensionRunner: ExtensionRunner | undefined = undefined; #turnIndex = 0; @@ -547,6 +563,7 @@ export class AgentSession { this.searchDb = config.searchDb; this.#startPowerAssertion(); this.#asyncJobManager = config.asyncJobManager; + this.#pythonKernelOwnerId = config.pythonKernelOwnerId ?? `agent-session:${Snowflake.next()}`; this.#scopedModels = config.scopedModels ?? []; this.#thinkingLevel = config.thinkingLevel; this.#promptTemplates = config.promptTemplates ?? []; @@ -1798,6 +1815,7 @@ export class AgentSession { * Call this when completely done with the session. */ async dispose(): Promise { + this.#pythonExecutionDisposing = true; try { if (this.#extensionRunner?.hasHandlers("session_shutdown")) { await this.#extensionRunner.emit({ type: "session_shutdown" }); @@ -1812,6 +1830,13 @@ export class AgentSession { if (drained === false && deliveryState) { logger.warn("Async job completion deliveries still pending during dispose", { ...deliveryState }); } + const pythonExecutionsSettled = await this.#preparePythonExecutionsForDispose(); + if (!pythonExecutionsSettled) { + logger.warn( + "Detaching retained Python kernel ownership during dispose while Python execution is still active", + ); + } + await disposeKernelSessionsByOwner(this.#pythonKernelOwnerId); this.#stopPowerAssertion(); await this.sessionManager.close(); this.#closeAllProviderSessions("dispose"); @@ -5661,6 +5686,7 @@ export class AgentSession { ): Promise { const excludeFromContext = options?.excludeFromContext === true; const cwd = this.sessionManager.getCwd(); + this.assertPythonExecutionAllowed(); if (this.#extensionRunner?.hasHandlers("user_python")) { const hookResult = await this.#extensionRunner.emitUserPython({ @@ -5669,35 +5695,58 @@ export class AgentSession { excludeFromContext, cwd, }); + this.assertPythonExecutionAllowed(); if (hookResult?.result) { this.recordPythonResult(code, hookResult.result, options); return hookResult.result; } } - this.#pythonAbortController = new AbortController(); - - try { + const abortController = new AbortController(); + const execution = (async (): Promise => { // Use the same session ID as the Python tool for kernel sharing const sessionFile = this.sessionManager.getSessionFile(); const sessionId = sessionFile ? `session:${sessionFile}:cwd:${cwd}` : `cwd:${cwd}`; - const result = await executePythonCommand(code, { cwd, sessionId, + kernelOwnerId: this.#pythonKernelOwnerId, kernelMode: this.settings.get("python.kernelMode"), useSharedGateway: this.settings.get("python.sharedGateway"), onChunk, - signal: this.#pythonAbortController.signal, + signal: abortController.signal, }); - this.recordPythonResult(code, result, options); return result; - } finally { - this.#pythonAbortController = undefined; + })(); + return await this.trackPythonExecution(execution, abortController); + } + + assertPythonExecutionAllowed(): void { + if (this.#pythonExecutionDisposing) { + throw new Error("Python execution is unavailable while session disposal is in progress"); } } + /** + * Track Python work started outside AgentSession.executePython so dispose can await and abort it too. + */ + trackPythonExecution(execution: Promise, abortController: AbortController): Promise { + this.#pythonAbortControllers.add(abortController); + this.#activePythonExecutions.add(execution); + void execution.then( + () => { + this.#pythonAbortControllers.delete(abortController); + this.#activePythonExecutions.delete(execution); + }, + () => { + this.#pythonAbortControllers.delete(abortController); + this.#activePythonExecutions.delete(execution); + }, + ); + return execution; + } + /** * Record a Python execution result in session history. */ @@ -5728,12 +5777,46 @@ export class AgentSession { * Cancel running Python execution. */ abortPython(): void { - this.#pythonAbortController?.abort(); + for (const abortController of this.#pythonAbortControllers) { + abortController.abort(); + } + } + + async #waitForPythonExecutionsToSettle(timeoutMs: number): Promise { + const deadline = Date.now() + timeoutMs; + while (this.#activePythonExecutions.size > 0) { + const remainingMs = deadline - Date.now(); + if (remainingMs <= 0) { + return false; + } + const settled = await Promise.race([ + Promise.allSettled(Array.from(this.#activePythonExecutions)).then(() => true), + Bun.sleep(remainingMs).then(() => false), + ]); + if (!settled && this.#activePythonExecutions.size > 0) { + return false; + } + } + return true; + } + + async #preparePythonExecutionsForDispose(): Promise { + if (!(await this.#waitForPythonExecutionsToSettle(3_000))) { + logger.warn("Aborting active Python execution during dispose before retained kernel cleanup"); + this.abortPython(); + if (!(await this.#waitForPythonExecutionsToSettle(1_000))) { + logger.warn( + "Python execution is still active after dispose aborted all active runs; retained kernel ownership will still be detached", + ); + return false; + } + } + return true; } /** Whether a Python execution is currently running */ get isPythonRunning(): boolean { - return this.#pythonAbortController !== undefined; + return this.#pythonAbortControllers.size > 0; } /** Whether there are pending Python messages waiting to be flushed */ diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index f4dfe358c..b98ba7a17 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -116,6 +116,8 @@ export interface ToolSession { hasUI: boolean; /** Skip Python kernel availability check and warmup */ skipPythonPreflight?: boolean; + /** Force Python prelude warmup even when test env would normally skip it */ + forcePythonWarmup?: boolean; /** Pre-loaded context files (AGENTS.md, etc) */ contextFiles?: ContextFileEntry[]; /** Pre-loaded skills */ @@ -136,6 +138,12 @@ export interface ToolSession { taskDepth?: number; /** Get session file */ getSessionFile: () => string | null; + /** Get Python kernel owner ID for session-scoped retained-kernel cleanup */ + getPythonKernelOwnerId?: () => string | null; + /** Reject new Python work once session disposal has started. */ + assertPythonExecutionAllowed?: () => void; + /** Track tool-owned Python work so session disposal can await/abort it like direct session Python runs. */ + trackPythonExecution?(execution: Promise, abortController: AbortController): Promise; /** Get session ID */ getSessionId?: () => string | null; /** Get artifacts directory for artifact:// URLs */ @@ -297,7 +305,9 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P !skipPythonPreflight && pythonMode !== "bash-only" && (requestedTools === undefined || requestedTools.includes("python")); - const skipPythonWarm = isBunTestRuntime() || $flag("PI_PYTHON_SKIP_CHECK"); + const isTestEnv = isBunTestRuntime(); + const forcePythonWarmup = session.forcePythonWarmup === true; + const skipPythonWarm = (isTestEnv && !forcePythonWarmup) || $flag("PI_PYTHON_SKIP_CHECK"); if (shouldCheckPython) { const availability = await logger.time("createTools:pythonCheck", checkPythonKernelAvailability, session.cwd); pythonAvailable = availability.ok; @@ -307,6 +317,7 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P }); } else if (!skipPythonWarm && getPreludeDocs().length === 0) { const sessionFile = session.getSessionFile?.() ?? undefined; + const kernelOwnerId = session.getPythonKernelOwnerId?.() ?? undefined; const warmSessionId = sessionFile ? `session:${sessionFile}:cwd:${session.cwd}` : `cwd:${session.cwd}`; try { await logger.time( @@ -316,6 +327,7 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P warmSessionId, session.settings.get("python.sharedGateway"), sessionFile, + kernelOwnerId, ); } catch (err) { logger.warn("Failed to warm Python environment", { diff --git a/packages/coding-agent/src/tools/python.ts b/packages/coding-agent/src/tools/python.ts index 11b751dcd..c3dd645d5 100644 --- a/packages/coding-agent/src/tools/python.ts +++ b/packages/coding-agent/src/tools/python.ts @@ -175,6 +175,7 @@ export class PythonTool implements AgentTool { if (!this.session) { throw new ToolError("Python tool requires a session when not using proxy executor"); } + const session = this.session; const { cells, timeout: rawTimeout = 30, cwd, reset } = params; // Clamp to reasonable range: 1s - 600s (10 min) @@ -182,7 +183,10 @@ export class PythonTool implements AgentTool { const timeoutMs = timeoutSec * 1000; const deadlineMs = Date.now() + timeoutMs; const timeoutSignal = AbortSignal.timeout(Math.max(0, deadlineMs - Date.now())); - const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal; + const sessionAbortController = new AbortController(); + const combinedSignal = signal + ? AbortSignal.any([signal, timeoutSignal, sessionAbortController.signal]) + : AbortSignal.any([timeoutSignal, sessionAbortController.signal]); let outputSink: OutputSink | undefined; let outputSummary: OutputSummary | undefined; let outputDumped = false; @@ -193,311 +197,322 @@ export class PythonTool implements AgentTool { return outputSummary; }; - try { - if (signal?.aborted) { - throw new ToolAbortError(); - } - - const commandCwd = cwd ? resolveToCwd(cwd, this.session.cwd) : this.session.cwd; - let cwdStat: fs.Stats; + const execution = (async (): Promise> => { try { - cwdStat = await Bun.file(commandCwd).stat(); - } catch { - throw new ToolError(`Working directory does not exist: ${commandCwd}`); - } - if (!cwdStat.isDirectory()) { - throw new ToolError(`Working directory is not a directory: ${commandCwd}`); - } - - const tailBuffer = new TailBuffer(DEFAULT_MAX_BYTES * 2); - const jsonOutputs: unknown[] = []; - const images: ImageContent[] = []; - const statusEvents: PythonStatusEvent[] = []; - - const cellResults: PythonCellResult[] = cells.map((cell, index) => ({ - index, - title: cell.title, - code: cell.code, - output: "", - status: "pending", - })); - const cellOutputs: string[] = []; - - const appendTail = (text: string) => { - tailBuffer.append(text); - }; - - const buildUpdateDetails = (): PythonToolDetails => { - const details: PythonToolDetails = { - cells: cellResults.map(cell => ({ - ...cell, - statusEvents: cell.statusEvents ? [...cell.statusEvents] : undefined, - })), - }; - if (jsonOutputs.length > 0) { - details.jsonOutputs = jsonOutputs; + if (signal?.aborted) { + throw new ToolAbortError(); } - if (images.length > 0) { - details.images = images; - } - if (statusEvents.length > 0) { - details.statusEvents = statusEvents; - } - return details; - }; + session.assertPythonExecutionAllowed?.(); - const pushUpdate = () => { - if (!onUpdate) return; - const tailText = tailBuffer.text(); - onUpdate({ - content: [{ type: "text", text: tailText }], - details: buildUpdateDetails(), - }); - }; - - const sessionFile = this.session.getSessionFile?.() ?? undefined; - const { path: artifactPath, id: artifactId } = (await this.session.allocateOutputArtifact?.("python")) ?? {}; - outputSink = new OutputSink({ - artifactPath, - artifactId, - onChunk: chunk => { - appendTail(chunk); - pushUpdate(); - }, - }); - const sessionId = sessionFile ? `session:${sessionFile}:cwd:${commandCwd}` : `cwd:${commandCwd}`; - - if (getPreludeDocs().length === 0) { - const warmup = await warmPythonEnvironment( - commandCwd, - sessionId, - this.session.settings.get("python.sharedGateway"), - sessionFile ?? undefined, - ); - if (!warmup.ok) { - throw new ToolError(warmup.reason ?? "Python prelude helpers unavailable"); - } - } - - const baseExecutorOptions: Omit = { - cwd: commandCwd, - deadlineMs, - signal: combinedSignal, - sessionId, - kernelMode: this.session.settings.get("python.kernelMode"), - useSharedGateway: this.session.settings.get("python.sharedGateway"), - sessionFile: sessionFile ?? undefined, - }; - - for (let i = 0; i < cells.length; i++) { - const cell = cells[i]; - const isFirstCell = i === 0; - const cellResult = cellResults[i]; - cellResult.status = "running"; - cellResult.output = ""; - cellResult.statusEvents = undefined; - cellResult.exitCode = undefined; - cellResult.durationMs = undefined; - pushUpdate(); - - const executorOptions: PythonExecutorOptions = { - ...baseExecutorOptions, - reset: isFirstCell ? reset : false, - onChunk: chunk => { - outputSink!.push(chunk); - }, - }; - - const startTime = Date.now(); - const result = await executePython(cell.code, executorOptions); - const durationMs = Date.now() - startTime; - - const cellStatusEvents: PythonStatusEvent[] = []; - let cellHasMarkdown = false; - for (const output of result.displayOutputs) { - if (output.type === "json") { - jsonOutputs.push(output.data); - } - if (output.type === "image") { - images.push({ type: "image", data: output.data, mimeType: output.mimeType }); - } - if (output.type === "status") { - statusEvents.push(output.event); - cellStatusEvents.push(output.event); - } - if (output.type === "markdown") { - cellHasMarkdown = true; - } - } - - const cellOutput = result.output.trim(); - cellResult.output = cellOutput; - cellResult.exitCode = result.exitCode; - cellResult.durationMs = durationMs; - cellResult.statusEvents = cellStatusEvents.length > 0 ? cellStatusEvents : undefined; - cellResult.hasMarkdown = cellHasMarkdown || undefined; - - let combinedCellOutput = ""; - if (cells.length > 1) { - const cellHeader = `[${i + 1}/${cells.length}]`; - const cellTitle = cell.title ? ` ${cell.title}` : ""; - if (cellOutput) { - combinedCellOutput = `${cellHeader}${cellTitle}\n${cellOutput}`; - } else { - combinedCellOutput = `${cellHeader}${cellTitle} (ok)`; - } - cellOutputs.push(combinedCellOutput); - } else if (cellOutput) { - combinedCellOutput = cellOutput; - cellOutputs.push(combinedCellOutput); - } - - if (combinedCellOutput) { - const prefix = cellOutputs.length > 1 ? "\n\n" : ""; - appendTail(`${prefix}${combinedCellOutput}`); - } - - if (result.cancelled) { - cellResult.status = "error"; - pushUpdate(); - const errorMsg = result.output || "Command aborted"; - const combinedOutput = cellOutputs.join("\n\n"); - const outputText = - cells.length > 1 - ? `${combinedOutput}\n\nCell ${i + 1} aborted: ${errorMsg}` - : combinedOutput || errorMsg; - - const rawSummary = (await finalizeOutput()) ?? { - output: "", - truncated: false, - totalLines: 0, - totalBytes: 0, - outputLines: 0, - outputBytes: 0, - }; - const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; - const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); - const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); - const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); - const summaryForMeta: OutputSummary = { - output: combinedOutput, - truncated: rawSummary.truncated, - totalLines: outputLines + missingLines, - totalBytes: outputBytes + missingBytes, - outputLines, - outputBytes, - artifactId: rawSummary.artifactId, - }; - - const details: PythonToolDetails = { - cells: cellResults, - jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, - images: images.length > 0 ? images : undefined, - statusEvents: statusEvents.length > 0 ? statusEvents : undefined, - isError: true, - }; - - return toolResult(details) - .text(outputText) - .truncationFromSummary(summaryForMeta, { direction: "tail" }) - .done(); - } - - if (result.exitCode !== 0 && result.exitCode !== undefined) { - cellResult.status = "error"; - pushUpdate(); - const combinedOutput = cellOutputs.join("\n\n"); - const outputText = - cells.length > 1 - ? `${combinedOutput}\n\nCell ${i + 1} failed (exit code ${result.exitCode}). Earlier cells succeeded—their state persists. Fix only cell ${i + 1}.` - : combinedOutput - ? `${combinedOutput}\n\nCommand exited with code ${result.exitCode}` - : `Command exited with code ${result.exitCode}`; - - const rawSummary = (await finalizeOutput()) ?? { - output: "", - truncated: false, - totalLines: 0, - totalBytes: 0, - outputLines: 0, - outputBytes: 0, - }; - const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; - const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); - const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); - const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); - const summaryForMeta: OutputSummary = { - output: combinedOutput, - truncated: rawSummary.truncated, - totalLines: outputLines + missingLines, - totalBytes: outputBytes + missingBytes, - outputLines, - outputBytes, - artifactId: rawSummary.artifactId, - }; - - const details: PythonToolDetails = { - cells: cellResults, - jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, - images: images.length > 0 ? images : undefined, - statusEvents: statusEvents.length > 0 ? statusEvents : undefined, - isError: true, - }; - - return toolResult(details) - .text(outputText) - .truncationFromSummary(summaryForMeta, { direction: "tail" }) - .done(); - } - - cellResult.status = "complete"; - pushUpdate(); - } - - const combinedOutput = cellOutputs.join("\n\n"); - const outputText = - combinedOutput || (jsonOutputs.length > 0 || images.length > 0 ? "(no text output)" : "(no output)"); - const rawSummary = (await finalizeOutput()) ?? { - output: "", - truncated: false, - totalLines: 0, - totalBytes: 0, - outputLines: 0, - outputBytes: 0, - }; - const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; - const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); - const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); - const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); - const summaryForMeta: OutputSummary = { - output: combinedOutput, - truncated: rawSummary.truncated, - totalLines: outputLines + missingLines, - totalBytes: outputBytes + missingBytes, - outputLines, - outputBytes, - artifactId: rawSummary.artifactId, - }; - - const details: PythonToolDetails = { - cells: cellResults, - jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, - images: images.length > 0 ? images : undefined, - statusEvents: statusEvents.length > 0 ? statusEvents : undefined, - }; - - const resultBuilder = toolResult(details) - .text(outputText) - .truncationFromSummary(summaryForMeta, { direction: "tail" }); - - return resultBuilder.done(); - } finally { - if (!outputDumped) { + const commandCwd = cwd ? resolveToCwd(cwd, session.cwd) : session.cwd; + let cwdStat: fs.Stats; try { - await finalizeOutput(); - } catch {} + cwdStat = await Bun.file(commandCwd).stat(); + } catch { + throw new ToolError(`Working directory does not exist: ${commandCwd}`); + } + if (!cwdStat.isDirectory()) { + throw new ToolError(`Working directory is not a directory: ${commandCwd}`); + } + + const tailBuffer = new TailBuffer(DEFAULT_MAX_BYTES * 2); + const jsonOutputs: unknown[] = []; + const images: ImageContent[] = []; + const statusEvents: PythonStatusEvent[] = []; + + const cellResults: PythonCellResult[] = cells.map((cell, index) => ({ + index, + title: cell.title, + code: cell.code, + output: "", + status: "pending", + })); + const cellOutputs: string[] = []; + + const appendTail = (text: string) => { + tailBuffer.append(text); + }; + + const buildUpdateDetails = (): PythonToolDetails => { + const details: PythonToolDetails = { + cells: cellResults.map(cell => ({ + ...cell, + statusEvents: cell.statusEvents ? [...cell.statusEvents] : undefined, + })), + }; + if (jsonOutputs.length > 0) { + details.jsonOutputs = jsonOutputs; + } + if (images.length > 0) { + details.images = images; + } + if (statusEvents.length > 0) { + details.statusEvents = statusEvents; + } + return details; + }; + + const pushUpdate = () => { + if (!onUpdate) return; + const tailText = tailBuffer.text(); + onUpdate({ + content: [{ type: "text", text: tailText }], + details: buildUpdateDetails(), + }); + }; + + const sessionFile = session.getSessionFile?.() ?? undefined; + const kernelOwnerId = session.getPythonKernelOwnerId?.() ?? undefined; + const { path: artifactPath, id: artifactId } = (await session.allocateOutputArtifact?.("python")) ?? {}; + session.assertPythonExecutionAllowed?.(); + outputSink = new OutputSink({ + artifactPath, + artifactId, + onChunk: chunk => { + appendTail(chunk); + pushUpdate(); + }, + }); + const sessionId = sessionFile ? `session:${sessionFile}:cwd:${commandCwd}` : `cwd:${commandCwd}`; + + if (getPreludeDocs().length === 0) { + const warmup = await warmPythonEnvironment( + commandCwd, + sessionId, + session.settings.get("python.sharedGateway"), + sessionFile ?? undefined, + kernelOwnerId, + combinedSignal, + ); + if (!warmup.ok) { + if (combinedSignal.aborted) throw new ToolAbortError(); + throw new ToolError(warmup.reason ?? "Python prelude helpers unavailable"); + } + session.assertPythonExecutionAllowed?.(); + } + + const baseExecutorOptions = { + cwd: commandCwd, + deadlineMs, + signal: combinedSignal, + sessionId, + kernelMode: session.settings.get("python.kernelMode"), + useSharedGateway: session.settings.get("python.sharedGateway"), + sessionFile: sessionFile ?? undefined, + kernelOwnerId, + }; + + for (let i = 0; i < cells.length; i++) { + const cell = cells[i]; + const isFirstCell = i === 0; + const cellResult = cellResults[i]; + cellResult.status = "running"; + cellResult.output = ""; + cellResult.statusEvents = undefined; + cellResult.exitCode = undefined; + cellResult.durationMs = undefined; + pushUpdate(); + + const executorOptions: PythonExecutorOptions = { + ...baseExecutorOptions, + reset: isFirstCell ? reset : false, + onChunk: chunk => { + outputSink!.push(chunk); + }, + }; + + const startTime = Date.now(); + const result = await executePython(cell.code, executorOptions); + const durationMs = Date.now() - startTime; + + const cellStatusEvents: PythonStatusEvent[] = []; + let cellHasMarkdown = false; + for (const output of result.displayOutputs) { + if (output.type === "json") { + jsonOutputs.push(output.data); + } + if (output.type === "image") { + images.push({ type: "image", data: output.data, mimeType: output.mimeType }); + } + if (output.type === "status") { + statusEvents.push(output.event); + cellStatusEvents.push(output.event); + } + if (output.type === "markdown") { + cellHasMarkdown = true; + } + } + + const cellOutput = result.output.trim(); + cellResult.output = cellOutput; + cellResult.exitCode = result.exitCode; + cellResult.durationMs = durationMs; + cellResult.statusEvents = cellStatusEvents.length > 0 ? cellStatusEvents : undefined; + cellResult.hasMarkdown = cellHasMarkdown || undefined; + + let combinedCellOutput = ""; + if (cells.length > 1) { + const cellHeader = `[${i + 1}/${cells.length}]`; + const cellTitle = cell.title ? ` ${cell.title}` : ""; + if (cellOutput) { + combinedCellOutput = `${cellHeader}${cellTitle}\n${cellOutput}`; + } else { + combinedCellOutput = `${cellHeader}${cellTitle} (ok)`; + } + cellOutputs.push(combinedCellOutput); + } else if (cellOutput) { + combinedCellOutput = cellOutput; + cellOutputs.push(combinedCellOutput); + } + + if (combinedCellOutput) { + const prefix = cellOutputs.length > 1 ? "\n\n" : ""; + appendTail(`${prefix}${combinedCellOutput}`); + } + + if (result.cancelled) { + cellResult.status = "error"; + pushUpdate(); + const errorMsg = result.output || "Command aborted"; + const combinedOutput = cellOutputs.join("\n\n"); + const outputText = + cells.length > 1 + ? `${combinedOutput}\n\nCell ${i + 1} aborted: ${errorMsg}` + : combinedOutput || errorMsg; + + const rawSummary = (await finalizeOutput()) ?? { + output: "", + truncated: false, + totalLines: 0, + totalBytes: 0, + outputLines: 0, + outputBytes: 0, + }; + const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; + const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); + const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); + const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); + const summaryForMeta: OutputSummary = { + output: combinedOutput, + truncated: rawSummary.truncated, + totalLines: outputLines + missingLines, + totalBytes: outputBytes + missingBytes, + outputLines, + outputBytes, + artifactId: rawSummary.artifactId, + }; + + const details: PythonToolDetails = { + cells: cellResults, + jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, + images: images.length > 0 ? images : undefined, + statusEvents: statusEvents.length > 0 ? statusEvents : undefined, + isError: true, + }; + + return toolResult(details) + .text(outputText) + .truncationFromSummary(summaryForMeta, { direction: "tail" }) + .done(); + } + + if (result.exitCode !== 0 && result.exitCode !== undefined) { + cellResult.status = "error"; + pushUpdate(); + const combinedOutput = cellOutputs.join("\n\n"); + const outputText = + cells.length > 1 + ? `${combinedOutput}\n\nCell ${i + 1} failed (exit code ${result.exitCode}). Earlier cells succeeded—their state persists. Fix only cell ${i + 1}.` + : combinedOutput + ? `${combinedOutput}\n\nCommand exited with code ${result.exitCode}` + : `Command exited with code ${result.exitCode}`; + + const rawSummary = (await finalizeOutput()) ?? { + output: "", + truncated: false, + totalLines: 0, + totalBytes: 0, + outputLines: 0, + outputBytes: 0, + }; + const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; + const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); + const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); + const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); + const summaryForMeta: OutputSummary = { + output: combinedOutput, + truncated: rawSummary.truncated, + totalLines: outputLines + missingLines, + totalBytes: outputBytes + missingBytes, + outputLines, + outputBytes, + artifactId: rawSummary.artifactId, + }; + + const details: PythonToolDetails = { + cells: cellResults, + jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, + images: images.length > 0 ? images : undefined, + statusEvents: statusEvents.length > 0 ? statusEvents : undefined, + isError: true, + }; + + return toolResult(details) + .text(outputText) + .truncationFromSummary(summaryForMeta, { direction: "tail" }) + .done(); + } + + cellResult.status = "complete"; + pushUpdate(); + } + + const combinedOutput = cellOutputs.join("\n\n"); + const outputText = + combinedOutput || (jsonOutputs.length > 0 || images.length > 0 ? "(no text output)" : "(no output)"); + const rawSummary = (await finalizeOutput()) ?? { + output: "", + truncated: false, + totalLines: 0, + totalBytes: 0, + outputLines: 0, + outputBytes: 0, + }; + const outputLines = combinedOutput.length > 0 ? combinedOutput.split("\n").length : 0; + const outputBytes = Buffer.byteLength(combinedOutput, "utf-8"); + const missingLines = Math.max(0, rawSummary.totalLines - rawSummary.outputLines); + const missingBytes = Math.max(0, rawSummary.totalBytes - rawSummary.outputBytes); + const summaryForMeta: OutputSummary = { + output: combinedOutput, + truncated: rawSummary.truncated, + totalLines: outputLines + missingLines, + totalBytes: outputBytes + missingBytes, + outputLines, + outputBytes, + artifactId: rawSummary.artifactId, + }; + + const details: PythonToolDetails = { + cells: cellResults, + jsonOutputs: jsonOutputs.length > 0 ? jsonOutputs : undefined, + images: images.length > 0 ? images : undefined, + statusEvents: statusEvents.length > 0 ? statusEvents : undefined, + }; + + const resultBuilder = toolResult(details) + .text(outputText) + .truncationFromSummary(summaryForMeta, { direction: "tail" }); + + return resultBuilder.done(); + } finally { + if (!outputDumped) { + try { + await finalizeOutput(); + } catch {} + } } - } + })(); + return await (session.trackPythonExecution?.(execution, sessionAbortController) ?? execution); } } diff --git a/packages/coding-agent/test/agent-session-python-cleanup.test.ts b/packages/coding-agent/test/agent-session-python-cleanup.test.ts new file mode 100644 index 000000000..c137d0d49 --- /dev/null +++ b/packages/coding-agent/test/agent-session-python-cleanup.test.ts @@ -0,0 +1,803 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { getBundledModel } from "@oh-my-pi/pi-ai"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import * as pythonExecutor from "@oh-my-pi/pi-coding-agent/ipy/executor"; +import type { PreludeHelper, PythonKernel as PythonKernelInstance } from "@oh-my-pi/pi-coding-agent/ipy/kernel"; +import * as pythonKernel from "@oh-my-pi/pi-coding-agent/ipy/kernel"; +import * as memories from "@oh-my-pi/pi-coding-agent/memories"; +import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { Snowflake } from "@oh-my-pi/pi-utils"; + +const OK_EXECUTION = { status: "ok", cancelled: false, timedOut: false, stdinRequested: false } as const; + +class FakeKernel { + executeCalls: string[] = []; + shutdownCalls = 0; + alive = true; + blockedCode: string | undefined; + blockedExecution: Promise | undefined; + blockedExecutionStarted: (() => void) | undefined; + blockedExecutionReject: ((error: Error) => void) | undefined; + abortBlockedExecution = true; + + isAlive(): boolean { + return this.alive; + } + + async execute(code: string, options?: { signal?: AbortSignal }): Promise { + this.executeCalls.push(code); + if (code === this.blockedCode && this.blockedExecution) { + this.blockedExecutionStarted?.(); + if (!this.abortBlockedExecution || !options?.signal) { + return await this.blockedExecution; + } + return await Promise.race([ + this.blockedExecution, + new Promise((_, reject) => { + const onAbort = () => reject(new DOMException("Aborted", "AbortError")); + if (options.signal?.aborted) { + onAbort(); + return; + } + options.signal?.addEventListener("abort", onAbort, { once: true }); + }), + ]); + } + return OK_EXECUTION; + } + + async ping(): Promise { + return this.alive; + } + + shutdown = vi.fn(async () => { + this.shutdownCalls += 1; + this.alive = false; + this.blockedExecutionReject?.(new Error("Kernel shut down during execution")); + return { confirmed: true }; + }); +} + +const getModel = () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!model) throw new Error("Expected bundled model"); + return model; +}; + +const createTempProject = () => { + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-agent-session-python-cleanup-${Snowflake.next()}-`)); + const cwd = path.join(tempDir, "project"); + fs.mkdirSync(cwd, { recursive: true }); + return { tempDir, cwd }; +}; + +const createSession = async ( + tempDir: string, + cwd: string, + options: { extensions?: ExtensionFactory[]; sessionManager?: SessionManager } = {}, +) => + ( + await createAgentSession({ + cwd, + agentDir: tempDir, + sessionManager: options.sessionManager ?? SessionManager.inMemory(cwd), + settings: Settings.isolated({ "python.kernelMode": "session" }), + model: getModel(), + disableExtensionDiscovery: true, + extensions: options.extensions, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + forcePythonWarmup: true, + toolNames: ["python"], + }) + ).session; + +const stubPythonWarmup = () => + vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); + +const createWarmupKernel = (docs: PreludeHelper[] = []) => ({ + introspectPrelude: vi.fn().mockResolvedValue(docs), + execute: vi.fn(async () => OK_EXECUTION), + ping: vi.fn(async () => true), + isAlive: () => true, + shutdown: vi.fn(async () => ({ confirmed: true })), +}); + +describe("AgentSession python cleanup", () => { + const tempDirs: string[] = []; + + afterEach(async () => { + pythonExecutor.resetPreludeDocsCache(); + vi.restoreAllMocks(); + await pythonExecutor.disposeAllKernelSessions(); + for (const tempDir of tempDirs.splice(0)) { + fs.rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it("cleans up warmed Python owners when createAgentSession fails before session construction", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + + const unrelatedKernel = createWarmupKernel(); + const warmedKernel = createWarmupKernel([ + { name: "helper", signature: "helper()", docstring: "doc", category: "general" }, + ]); + const unrelatedCwd = path.join(tempDir, "unrelated-before"); + const throwingExtension: ExtensionFactory = () => { + throw new Error("Extension init failed"); + }; + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(pythonKernel.PythonKernel, "start") + .mockResolvedValueOnce(unrelatedKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(warmedKernel as unknown as PythonKernelInstance); + + await pythonExecutor.executePython("print('unrelated before')", { + cwd: unrelatedCwd, + sessionId: "unrelated-before-session", + kernelMode: "session", + kernelOwnerId: "other-owner", + }); + + await expect( + createAgentSession({ + cwd, + agentDir: tempDir, + sessionManager: SessionManager.inMemory(cwd), + settings: Settings.isolated({ "python.kernelMode": "session" }), + model: getModel(), + disableExtensionDiscovery: true, + extensions: [throwingExtension], + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + forcePythonWarmup: true, + toolNames: ["python"], + }), + ).rejects.toThrow("Extension init failed"); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(warmedKernel.introspectPrelude).toHaveBeenCalledTimes(1); + expect(warmedKernel.shutdown).toHaveBeenCalledTimes(1); + expect(unrelatedKernel.shutdown).not.toHaveBeenCalled(); + + await pythonExecutor.executePython("print('still alive before')", { + cwd: unrelatedCwd, + sessionId: "unrelated-before-session", + kernelMode: "session", + kernelOwnerId: "other-owner", + }); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(unrelatedKernel.execute).toHaveBeenCalledTimes(2); + }); + + it("cleans up warmed Python owners when createAgentSession fails after session construction", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + + const unrelatedKernel = createWarmupKernel(); + const warmedKernel = createWarmupKernel([ + { name: "helper", signature: "helper()", docstring: "doc", category: "general" }, + ]); + const unrelatedCwd = path.join(tempDir, "unrelated-after"); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(pythonKernel.PythonKernel, "start") + .mockResolvedValueOnce(unrelatedKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(warmedKernel as unknown as PythonKernelInstance); + vi.spyOn(memories, "startMemoryStartupTask").mockImplementation(() => { + throw new Error("Memory startup failed"); + }); + + await pythonExecutor.executePython("print('unrelated after')", { + cwd: unrelatedCwd, + sessionId: "unrelated-after-session", + kernelMode: "session", + kernelOwnerId: "other-owner", + }); + + await expect( + createAgentSession({ + cwd, + agentDir: tempDir, + sessionManager: SessionManager.inMemory(cwd), + settings: Settings.isolated({ "python.kernelMode": "session" }), + model: getModel(), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + forcePythonWarmup: true, + toolNames: ["python"], + }), + ).rejects.toThrow("Memory startup failed"); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(warmedKernel.introspectPrelude).toHaveBeenCalledTimes(1); + expect(warmedKernel.shutdown).toHaveBeenCalledTimes(1); + expect(unrelatedKernel.shutdown).not.toHaveBeenCalled(); + + await pythonExecutor.executePython("print('still alive after')", { + cwd: unrelatedCwd, + sessionId: "unrelated-after-session", + kernelMode: "session", + kernelOwnerId: "other-owner", + }); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(unrelatedKernel.execute).toHaveBeenCalledTimes(2); + }); + + it("waits for active SDK session Python work before releasing a shared retained kernel", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + stubPythonWarmup(); + + const kernel = new FakeKernel(); + const blockedExecution = Promise.withResolvers(); + const blockedExecutionStarted = Promise.withResolvers(); + let blockedExecutionSettled = false; + blockedExecution.promise.then( + () => { + blockedExecutionSettled = true; + }, + () => { + blockedExecutionSettled = true; + }, + ); + kernel.blockedCode = "print('first')"; + kernel.blockedExecution = blockedExecution.promise; + kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve(); + kernel.blockedExecutionReject = error => blockedExecution.reject(error); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(pythonKernel.PythonKernel, "start") + .mockResolvedValue(kernel as unknown as PythonKernelInstance); + const firstSession = await createSession(tempDir, cwd); + const secondSession = await createSession(tempDir, cwd); + expect(startSpy).toHaveBeenCalledTimes(0); + let firstDisposed = false; + + try { + const firstExecution = firstSession.executePython("print('first')"); + let firstExecutionSettled = false; + const observedFirstExecution = firstExecution.finally(() => { + firstExecutionSettled = true; + }); + await blockedExecutionStarted.promise; + + const disposeFirst = firstSession.dispose().then(() => { + expect(blockedExecutionSettled).toBe(true); + expect(firstExecutionSettled).toBe(true); + firstDisposed = true; + }); + await Bun.sleep(0); + expect(firstDisposed).toBe(false); + expect(blockedExecutionSettled).toBe(false); + expect(firstExecutionSettled).toBe(false); + + const secondExecution = secondSession.executePython("print('second')"); + await Bun.sleep(0); + + expect(firstDisposed).toBe(false); + expect(blockedExecutionSettled).toBe(false); + expect(firstExecutionSettled).toBe(false); + expect(kernel.shutdownCalls).toBe(0); + + blockedExecution.resolve(OK_EXECUTION); + await Promise.all([observedFirstExecution, secondExecution, disposeFirst]); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.shutdownCalls).toBe(0); + expect(kernel.executeCalls).toEqual(["print('first')", "print('second')"]); + + await secondSession.executePython("print('third')"); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.executeCalls).toEqual(["print('first')", "print('second')", "print('third')"]); + } finally { + if (!firstDisposed) { + await firstSession.dispose(); + } + await secondSession.dispose(); + } + + expect(kernel.shutdownCalls).toBe(1); + }); + + it("aborts tracked Python tool warmup during session dispose before executePython starts", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + + const blockedWarmupStarted = Promise.withResolvers(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "tool ok", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 7, + outputLines: 1, + outputBytes: 7, + displayOutputs: [], + stdinRequested: false, + }); + let warmupCallCount = 0; + const warmupSpy = vi + .spyOn(pythonExecutor, "warmPythonEnvironment") + .mockImplementation(async (_cwd, _sessionId, _useSharedGateway, _sessionFile, _kernelOwnerId, signal) => { + warmupCallCount += 1; + if (warmupCallCount === 1) { + return { ok: true, docs: [] }; + } + blockedWarmupStarted.resolve(); + return await new Promise<{ ok: boolean; reason?: string; docs: [] }>(resolve => { + const onAbort = () => resolve({ ok: false, reason: "Warmup aborted", docs: [] }); + if (signal?.aborted) { + onAbort(); + return; + } + signal?.addEventListener("abort", onAbort, { once: true }); + }); + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + + const session = await createSession(tempDir, cwd); + const pythonTool = session.getToolByName("python"); + expect(pythonTool).toBeDefined(); + let toolExecutionSettled = false; + const toolExecution = pythonTool! + .execute("call-id", { cells: [{ code: "print('tool')" }] }, undefined, undefined, undefined) + .finally(() => { + toolExecutionSettled = true; + }); + await blockedWarmupStarted.promise; + + let disposed = false; + const disposeSession = session.dispose().then(() => { + disposed = true; + }); + await Bun.sleep(0); + + expect(disposed).toBe(false); + expect(toolExecutionSettled).toBe(false); + expect(warmupSpy).toHaveBeenCalledTimes(2); + expect(executeSpy).not.toHaveBeenCalled(); + + await expect(toolExecution).rejects.toThrow("Operation aborted"); + await disposeSession; + + expect(disposed).toBe(true); + expect(toolExecutionSettled).toBe(true); + expect(executeSpy).not.toHaveBeenCalled(); + }); + + it("rejects Python tool starts when warmup finishes after dispose begins", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + + const blockedWarmupStarted = Promise.withResolvers(); + const releaseWarmup = Promise.withResolvers(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + let warmupCallCount = 0; + const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockImplementation(async () => { + warmupCallCount += 1; + if (warmupCallCount === 1) { + return { ok: true, docs: [] }; + } + blockedWarmupStarted.resolve(); + await releaseWarmup.promise; + return { ok: true, docs: [] }; + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + + const session = await createSession(tempDir, cwd); + const pythonTool = session.getToolByName("python"); + expect(pythonTool).toBeDefined(); + const toolExecution = pythonTool!.execute( + "call-id", + { cells: [{ code: "print('tool after warmup')" }] }, + undefined, + undefined, + undefined, + ); + await blockedWarmupStarted.promise; + const disposeSession = session.dispose(); + releaseWarmup.resolve(); + await expect(toolExecution).rejects.toThrow( + "Python execution is unavailable while session disposal is in progress", + ); + await disposeSession; + expect(warmupSpy).toHaveBeenCalledTimes(2); + expect(executeSpy).not.toHaveBeenCalled(); + }); + it("aborts tracked Python tool execution during session dispose after warmup completes", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + + const helper: PreludeHelper = { + name: "helper", + signature: "helper()", + docstring: "doc", + category: "general", + }; + let preludeDocs: PreludeHelper[] = []; + vi.spyOn(pythonExecutor, "getPreludeDocs").mockImplementation(() => preludeDocs); + const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockImplementation(async () => { + preludeDocs = [helper]; + return { ok: true, docs: preludeDocs }; + }); + const blockedExecuteStarted = Promise.withResolvers(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockImplementation(async (_code, options) => { + const signal = options?.signal; + if (!signal) { + throw new Error("Expected abort signal"); + } + blockedExecuteStarted.resolve(); + return await new Promise(resolve => { + const onAbort = () => + resolve({ + output: "Command aborted", + exitCode: undefined, + cancelled: true, + truncated: false, + totalLines: 1, + totalBytes: 15, + outputLines: 1, + outputBytes: 15, + displayOutputs: [], + stdinRequested: false, + }); + if (signal.aborted) { + onAbort(); + return; + } + signal.addEventListener("abort", onAbort, { once: true }); + }); + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + + const session = await createSession(tempDir, cwd); + const pythonTool = session.getToolByName("python"); + expect(pythonTool).toBeDefined(); + let toolExecutionSettled = false; + const toolExecution = pythonTool! + .execute("call-id", { cells: [{ code: "print('tool')" }] }, undefined, undefined, undefined) + .finally(() => { + toolExecutionSettled = true; + }); + await blockedExecuteStarted.promise; + + let disposed = false; + const disposeSession = session.dispose().then(() => { + disposed = true; + }); + await Bun.sleep(0); + + expect(disposed).toBe(false); + expect(toolExecutionSettled).toBe(false); + expect(warmupSpy).toHaveBeenCalledTimes(1); + expect(executeSpy).toHaveBeenCalledTimes(1); + + const [toolResult] = await Promise.all([toolExecution, disposeSession]); + + expect(disposed).toBe(true); + expect(toolExecutionSettled).toBe(true); + expect(warmupSpy).toHaveBeenCalledTimes(1); + expect(executeSpy).toHaveBeenCalledTimes(1); + expect(toolResult.details?.isError).toBe(true); + expect(toolResult.content).toContainEqual( + expect.objectContaining({ type: "text", text: expect.stringContaining("Command aborted") }), + ); + }); + + it("detaches retained kernel ownership even when dispose times out waiting for Python work", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + stubPythonWarmup(); + + const kernel = new FakeKernel(); + const blockedExecution = Promise.withResolvers(); + const blockedExecutionStarted = Promise.withResolvers(); + kernel.blockedCode = "print('blocked')"; + kernel.blockedExecution = blockedExecution.promise; + kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve(); + kernel.abortBlockedExecution = false; + + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(pythonKernel.PythonKernel, "start") + .mockResolvedValue(kernel as unknown as PythonKernelInstance); + + const firstSession = await createSession(tempDir, cwd); + const secondSession = await createSession(tempDir, cwd); + + await secondSession.executePython("print('owner-b warmup')"); + const firstExecution = firstSession.executePython("print('blocked')"); + await blockedExecutionStarted.promise; + let firstExecutionSettled = false; + void firstExecution.finally(() => { + firstExecutionSettled = true; + }); + + let firstDisposed = false; + const disposeFirst = firstSession.dispose().then(() => { + firstDisposed = true; + }); + await disposeFirst; + + expect(firstDisposed).toBe(true); + expect(firstExecutionSettled).toBe(false); + expect(kernel.shutdownCalls).toBe(0); + expect(startSpy).toHaveBeenCalledTimes(1); + + blockedExecution.resolve(OK_EXECUTION); + await expect(firstExecution).resolves.toMatchObject({ + cancelled: false, + exitCode: 0, + stdinRequested: false, + }); + await secondSession.executePython("print('owner-b after detach')"); + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.executeCalls).toEqual([ + "print('owner-b warmup')", + "print('blocked')", + "print('owner-b after detach')", + ]); + await secondSession.dispose(); + + expect(kernel.shutdownCalls).toBe(1); + }, 10000); + + it("rejects direct session Python starts once dispose begins", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + stubPythonWarmup(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + + const session = await createSession(tempDir, cwd); + const disposeSession = session.dispose(); + await expect(session.executePython("print('late')")).rejects.toThrow( + "Python execution is unavailable while session disposal is in progress", + ); + await disposeSession; + expect(executeSpy).not.toHaveBeenCalled(); + }); + + it("rejects direct session Python starts after an async user_python hook yields during dispose", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + const hookStarted = Promise.withResolvers(); + const releaseHook = Promise.withResolvers(); + const hookExtension: ExtensionFactory = api => { + api.on("user_python", async () => { + hookStarted.resolve(); + await releaseHook.promise; + return undefined; + }); + }; + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + + const session = await createSession(tempDir, cwd, { extensions: [hookExtension] }); + const execution = session.executePython("print('late after hook')"); + await hookStarted.promise; + await session.dispose(); + releaseHook.resolve(); + await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress"); + expect(executeSpy).not.toHaveBeenCalled(); + }); + + it("rejects async user_python hook results after dispose begins", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + const hookStarted = Promise.withResolvers(); + const releaseHook = Promise.withResolvers(); + const hookExtension: ExtensionFactory = api => { + api.on("user_python", async () => { + hookStarted.resolve(); + await releaseHook.promise; + return { + result: { + output: "hooked late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 11, + outputLines: 1, + outputBytes: 11, + displayOutputs: [], + stdinRequested: false, + }, + }; + }); + }; + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + + const session = await createSession(tempDir, cwd, { extensions: [hookExtension] }); + const execution = session.executePython("print('late hook result')"); + await hookStarted.promise; + await session.dispose(); + releaseHook.resolve(); + await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress"); + expect(executeSpy).not.toHaveBeenCalled(); + expect(session.messages.some(message => message.role === "pythonExecution")).toBe(false); + }); + + it("rejects Python tool starts once dispose begins", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + const warmupSpy = stubPythonWarmup(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + + const session = await createSession(tempDir, cwd); + const pythonTool = session.getToolByName("python"); + expect(pythonTool).toBeDefined(); + const disposeSession = session.dispose(); + await expect( + pythonTool!.execute("call-id", { cells: [{ code: "print('late')" }] }, undefined, undefined, undefined), + ).rejects.toThrow("Python execution is unavailable while session disposal is in progress"); + await disposeSession; + expect(warmupSpy).toHaveBeenCalledTimes(1); + expect(executeSpy).not.toHaveBeenCalled(); + }); + + it("rejects Python tool starts that reach async preflight after dispose begins", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + const warmupSpy = stubPythonWarmup(); + const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ + output: "late", + exitCode: 0, + cancelled: false, + truncated: false, + totalLines: 1, + totalBytes: 4, + outputLines: 1, + outputBytes: 4, + displayOutputs: [], + stdinRequested: false, + }); + const artifactStarted = Promise.withResolvers(); + const releaseArtifact = Promise.withResolvers(); + const sessionManager = SessionManager.inMemory(cwd); + vi.spyOn(sessionManager, "allocateArtifactPath").mockImplementation(async () => { + artifactStarted.resolve(); + await releaseArtifact.promise; + return {}; + }); + + const session = await createSession(tempDir, cwd, { sessionManager }); + const pythonTool = session.getToolByName("python"); + expect(pythonTool).toBeDefined(); + const execution = pythonTool!.execute( + "call-id", + { cells: [{ code: "print('late after artifact')" }] }, + undefined, + undefined, + undefined, + ); + await artifactStarted.promise; + const disposeSession = session.dispose(); + releaseArtifact.resolve(); + await expect(execution).rejects.toThrow("Python execution is unavailable while session disposal is in progress"); + await disposeSession; + expect(warmupSpy).toHaveBeenCalledTimes(1); + expect(executeSpy).not.toHaveBeenCalled(); + }); + + it("aborts every active Python execution owned by the session during dispose", async () => { + const { tempDir, cwd } = createTempProject(); + tempDirs.push(tempDir); + pythonExecutor.resetPreludeDocsCache(); + stubPythonWarmup(); + + const kernel = new FakeKernel(); + const blockedExecution = Promise.withResolvers(); + const blockedExecutionStarted = Promise.withResolvers(); + kernel.blockedCode = "print('first')"; + kernel.blockedExecution = blockedExecution.promise; + kernel.blockedExecutionStarted = () => blockedExecutionStarted.resolve(); + + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + vi.spyOn(pythonKernel.PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance); + + const session = await createSession(tempDir, cwd); + + const firstExecution = session.executePython("print('first')"); + await blockedExecutionStarted.promise; + const secondExecution = session.executePython("print('second')"); + await Bun.sleep(0); + + await session.dispose(); + const [firstResult, secondResult] = await Promise.all([firstExecution, secondExecution]); + + expect(firstResult.cancelled).toBe(true); + expect(secondResult.cancelled).toBe(true); + expect(kernel.executeCalls).toEqual(["print('first')"]); + expect(kernel.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 ff3acbd7f..1897a7513 100644 --- a/packages/coding-agent/test/core/python-executor-lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-executor-lifecycle.test.ts @@ -6,7 +6,9 @@ import { getProjectDir } from "@oh-my-pi/pi-utils"; class FakeKernel { execute = vi.fn(async () => this.result); - shutdown = vi.fn(async () => {}); + shutdown = vi.fn(async () => { + return { confirmed: true }; + }); ping = vi.fn(async () => true); alive = true; @@ -97,4 +99,26 @@ describe("executePython lifecycle", () => { expect(kernel.execute).toHaveBeenCalledTimes(0); expect(kernelNext.execute).toHaveBeenCalledTimes(1); }); + + it("retries a dead session restart after an unconfirmed shutdown", async () => { + const kernel = new FakeKernel(OK_RESULT); + const kernelNext = new FakeKernel(OK_RESULT); + kernel.alive = false; + kernel.shutdown.mockResolvedValueOnce({ confirmed: false }).mockResolvedValueOnce({ confirmed: true }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(pythonKernel.PythonKernel, "start") + .mockResolvedValueOnce(kernel as unknown as pythonKernel.PythonKernel) + .mockResolvedValueOnce(kernelNext as unknown as pythonKernel.PythonKernel); + + await expect( + executePython("1 + 1", { kernelMode: "session", sessionId: "retry-dead-session", cwd: getProjectDir() }), + ).rejects.toThrow("Failed to confirm crashed kernel shutdown before restart"); + await executePython("2 + 2", { kernelMode: "session", sessionId: "retry-dead-session", cwd: getProjectDir() }); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(kernel.shutdown).toHaveBeenCalledTimes(2); + expect(kernel.execute).toHaveBeenCalledTimes(0); + expect(kernelNext.execute).toHaveBeenCalledTimes(1); + }); }); diff --git a/packages/coding-agent/test/core/python-executor-owner-cleanup.test.ts b/packages/coding-agent/test/core/python-executor-owner-cleanup.test.ts new file mode 100644 index 000000000..b844b81e4 --- /dev/null +++ b/packages/coding-agent/test/core/python-executor-owner-cleanup.test.ts @@ -0,0 +1,857 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import { + disposeAllKernelSessions, + disposeKernelSessionsByOwner, + executePython, + resetPreludeDocsCache, + warmPythonEnvironment, +} from "@oh-my-pi/pi-coding-agent/ipy/executor"; +import * as gatewayCoordinator from "@oh-my-pi/pi-coding-agent/ipy/gateway-coordinator"; +import type { + KernelExecuteResult, + KernelShutdownResult, + PreludeHelper, + PythonKernel as PythonKernelInstance, +} from "@oh-my-pi/pi-coding-agent/ipy/kernel"; +import * as pythonKernel from "@oh-my-pi/pi-coding-agent/ipy/kernel"; +import { PythonKernel } from "@oh-my-pi/pi-coding-agent/ipy/kernel"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +const OK_RESULT: KernelExecuteResult = { + status: "ok", + cancelled: false, + timedOut: false, + stdinRequested: false, +}; + +type FakeKernelShutdownOptions = { timeoutMs?: number }; + +class FakeKernel { + execute = vi.fn(async () => OK_RESULT); + shutdown = vi.fn( + async (_options?: FakeKernelShutdownOptions): Promise => ({ confirmed: true }), + ); + ping = vi.fn(async () => true); + alive = true; + + isAlive(): boolean { + return this.alive; + } +} + +afterEach(async () => { + await disposeAllKernelSessions(); + resetPreludeDocsCache(); + vi.restoreAllMocks(); +}); + +describe("python executor owner cleanup", () => { + it("keeps shared retained kernels alive until the last owner is disposed", async () => { + const kernel = new FakeKernel(); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance); + + await executePython("1 + 1", { + cwd: "/tmp/shared-owner-kernel", + sessionId: "shared-session", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + await executePython("2 + 2", { + cwd: "/tmp/shared-owner-kernel", + sessionId: "shared-session", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.execute).toHaveBeenCalledTimes(2); + + await disposeKernelSessionsByOwner("owner-a"); + + expect(kernel.shutdown).not.toHaveBeenCalled(); + + await executePython("3 + 3", { + cwd: "/tmp/shared-owner-kernel", + sessionId: "shared-session", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.execute).toHaveBeenCalledTimes(3); + + await disposeKernelSessionsByOwner("owner-b"); + + expect(kernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("disposes every retained kernel owned by one owner across session ids and cwd values", async () => { + const kernelOne = new FakeKernel(); + const kernelTwo = new FakeKernel(); + const unrelatedKernel = new FakeKernel(); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(PythonKernel, "start") + .mockResolvedValueOnce(kernelOne as unknown as PythonKernelInstance) + .mockResolvedValueOnce(kernelTwo as unknown as PythonKernelInstance) + .mockResolvedValueOnce(unrelatedKernel as unknown as PythonKernelInstance); + + await executePython("print('one')", { + cwd: "/tmp/owner-a-one", + sessionId: "session-one", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + await executePython("print('two')", { + cwd: "/tmp/owner-a-two", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + await executePython("print('other')", { + cwd: "/tmp/owner-b-one", + sessionId: "session-other", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(3); + + await disposeKernelSessionsByOwner("owner-a"); + + expect(kernelOne.shutdown).toHaveBeenCalledTimes(1); + expect(kernelTwo.shutdown).toHaveBeenCalledTimes(1); + expect(unrelatedKernel.shutdown).not.toHaveBeenCalled(); + + await executePython("print('still alive')", { + cwd: "/tmp/owner-b-one", + sessionId: "session-other", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(3); + expect(unrelatedKernel.execute).toHaveBeenCalledTimes(2); + }); + + it("falls back to the retained session id when no explicit owner id is provided during execution", async () => { + const kernel = new FakeKernel(); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance); + + await executePython("1 + 1", { + cwd: "/tmp/fallback-owner-session", + sessionId: "fallback-session", + kernelMode: "session", + }); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.execute).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("fallback-session"); + + expect(kernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("does not reattach a kernel after owner disposal has already claimed it", async () => { + const disposingKernel = new FakeKernel(); + const replacementKernel = new FakeKernel(); + const shutdownDeferred = Promise.withResolvers(); + disposingKernel.shutdown = vi.fn(() => shutdownDeferred.promise); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(PythonKernel, "start") + .mockResolvedValueOnce(disposingKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(replacementKernel as unknown as PythonKernelInstance); + + await executePython("1 + 1", { + cwd: "/tmp/disposal-race-kernel", + sessionId: "race-session", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + + const disposal = disposeKernelSessionsByOwner("owner-a"); + await executePython("2 + 2", { + cwd: "/tmp/disposal-race-kernel", + sessionId: "race-session", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(2); + expect(disposingKernel.execute).toHaveBeenCalledTimes(1); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + expect(disposingKernel.shutdown).toHaveBeenCalledTimes(1); + expect(replacementKernel.shutdown).not.toHaveBeenCalled(); + + shutdownDeferred.resolve({ confirmed: true }); + await disposal; + + await disposeKernelSessionsByOwner("owner-b"); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("keeps tracked disposals counted against retained kernel capacity until shutdown settles", async () => { + const retainedKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const replacementKernel = new FakeKernel(); + const shutdownDeferreds = retainedKernels.map(() => Promise.withResolvers()); + for (const [index, kernel] of retainedKernels.entries()) { + kernel.shutdown = vi.fn(() => shutdownDeferreds[index]!.promise); + } + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of [...retainedKernels, replacementKernel]) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + + for (const [index] of retainedKernels.entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/capacity-tracking-${index}`, + sessionId: `capacity-session-${index}`, + kernelMode: "session", + }); + } + expect(startSpy).toHaveBeenCalledTimes(4); + + const globalDisposal = disposeAllKernelSessions(); + await Promise.resolve(); + + const fifthExecution = executePython("print('replacement')", { + cwd: "/tmp/capacity-tracking-replacement", + sessionId: "capacity-session-replacement", + kernelMode: "session", + }); + await Promise.resolve(); + + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + + shutdownDeferreds[0]!.resolve({ confirmed: true }); + await fifthExecution; + expect(startSpy).toHaveBeenCalledTimes(5); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + + for (const deferred of shutdownDeferreds.slice(1)) { + deferred.resolve({ confirmed: true }); + } + await globalDisposal; + await disposeAllKernelSessions(); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("returns a cancelled result when a dead session restart shutdown times out", async () => { + const kernel = new FakeKernel(); + kernel.alive = false; + let shutdownCallCount = 0; + kernel.shutdown = vi.fn(async (options?: FakeKernelShutdownOptions): Promise => { + shutdownCallCount += 1; + if (shutdownCallCount > 1) { + return { confirmed: true }; + } + return await new Promise((_, reject) => { + const timer = setTimeout( + () => reject(new DOMException("Python kernel shutdown timed out", "TimeoutError")), + options?.timeoutMs ?? 0, + ); + timer.unref?.(); + }); + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + + const result = await executePython("1 + 1", { + cwd: "/tmp/restart-timeout-session", + sessionId: "restart-timeout-session", + kernelMode: "session", + timeoutMs: 25, + }); + + expect(result.cancelled).toBe(true); + expect(result.exitCode).toBeUndefined(); + expect(kernel.shutdown).toHaveBeenCalledWith(expect.objectContaining({ timeoutMs: expect.any(Number) })); + expect(startSpy).toHaveBeenCalledTimes(1); + }); + it("clears stuck tracked disposals during resource-exhaustion recovery", async () => { + vi.useFakeTimers(); + try { + const staleKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const recoveredKernel = new FakeKernel(); + const laterKernel = new FakeKernel(); + const staleShutdownDeferreds = staleKernels.map(() => Promise.withResolvers()); + for (const [index, kernel] of staleKernels.entries()) { + kernel.shutdown = vi.fn(() => staleShutdownDeferreds[index]!.promise); + } + const shutdownSharedGatewaySpy = vi.spyOn(gatewayCoordinator, "shutdownSharedGateway").mockResolvedValue(); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of staleKernels) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + startSpy + .mockRejectedValueOnce(new Error("EMFILE: too many open files")) + .mockResolvedValueOnce(recoveredKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(laterKernel as unknown as PythonKernelInstance); + + for (const [index] of staleKernels.entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/recovery-stale-${index}`, + sessionId: `recovery-stale-session-${index}`, + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + } + + const ownerCleanup = disposeKernelSessionsByOwner("owner-a"); + await Promise.resolve(); + for (const kernel of staleKernels) { + expect(kernel.shutdown).toHaveBeenCalledWith({ timeoutMs: 2_000 }); + } + vi.advanceTimersByTime(2_000); + await ownerCleanup; + + await executePython("print('recovered')", { + cwd: "/tmp/recovery-after-emfile", + sessionId: "recovery-session", + kernelMode: "session", + }); + expect(shutdownSharedGatewaySpy).toHaveBeenCalledTimes(1); + expect(startSpy).toHaveBeenCalledTimes(5); + expect(recoveredKernel.execute).toHaveBeenCalledTimes(1); + + await executePython("print('later')", { + cwd: "/tmp/recovery-after-emfile-later", + sessionId: "recovery-session-later", + kernelMode: "session", + deadlineMs: Date.now() + 50, + }); + expect(startSpy).toHaveBeenCalledTimes(6); + expect(recoveredKernel.shutdown).not.toHaveBeenCalled(); + expect(laterKernel.execute).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("returns owner cleanup promptly but keeps retained capacity reserved until shutdown is confirmed", async () => { + vi.useFakeTimers(); + try { + const retainedKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const replacementKernel = new FakeKernel(); + const shutdownDeferreds = retainedKernels.map(() => Promise.withResolvers()); + for (const [index, kernel] of retainedKernels.entries()) { + kernel.shutdown = vi.fn(() => shutdownDeferreds[index]!.promise); + } + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of [...retainedKernels, replacementKernel]) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + + for (const [index] of retainedKernels.entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/owner-timeout-kernel-${index}`, + sessionId: `timeout-session-${index}`, + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + } + + let ownerCleanupResolved = false; + const ownerCleanup = disposeKernelSessionsByOwner("owner-a").then(() => { + ownerCleanupResolved = true; + }); + await Promise.resolve(); + + for (const kernel of retainedKernels) { + expect(kernel.shutdown).toHaveBeenCalledWith({ timeoutMs: 2_000 }); + } + expect(ownerCleanupResolved).toBe(false); + + vi.advanceTimersByTime(2_000); + await ownerCleanup; + expect(ownerCleanupResolved).toBe(true); + + const blockedExecution = executePython("print('replacement')", { + cwd: "/tmp/owner-timeout-kernel-replacement", + sessionId: "timeout-session-replacement", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + await Promise.resolve(); + + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + + shutdownDeferreds[0]!.resolve({ confirmed: true }); + await blockedExecution; + expect(startSpy).toHaveBeenCalledTimes(5); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + + for (const deferred of shutdownDeferreds.slice(1)) { + deferred.resolve({ confirmed: true }); + } + await Promise.resolve(); + await disposeAllKernelSessions(); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps owner-cleanup retries on the timer path without evicting unrelated live sessions", async () => { + vi.useFakeTimers(); + try { + const ownerKernel = new FakeKernel(); + const unrelatedKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const replacementKernel = new FakeKernel(); + const retryConfirmation = Promise.withResolvers(); + let shutdownCallCount = 0; + ownerKernel.shutdown = vi.fn(async (options?: FakeKernelShutdownOptions): Promise => { + shutdownCallCount += 1; + if (shutdownCallCount === 1) { + expect(options).toEqual({ timeoutMs: 2_000 }); + return { confirmed: false }; + } + if (shutdownCallCount === 2) { + expect(options).toBeUndefined(); + return { confirmed: false }; + } + return await retryConfirmation.promise; + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of [ownerKernel, ...unrelatedKernels, replacementKernel]) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + + await executePython("print('owner-a')", { + cwd: "/tmp/timer-retry-owner-a", + sessionId: "timer-retry-owner-a", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + for (const [index, ownerId] of ["owner-b", "owner-c", "owner-d"].entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/timer-retry-${ownerId}`, + sessionId: `timer-retry-${ownerId}`, + kernelMode: "session", + kernelOwnerId: ownerId, + }); + } + + await disposeKernelSessionsByOwner("owner-a"); + expect(ownerKernel.shutdown).toHaveBeenCalledTimes(2); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + const blockedExecution = executePython("print('replacement')", { + cwd: "/tmp/timer-retry-replacement", + sessionId: "timer-retry-replacement", + kernelMode: "session", + kernelOwnerId: "owner-e", + }); + await Promise.resolve(); + await Promise.resolve(); + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + vi.advanceTimersByTime(29_999); + await Promise.resolve(); + await Promise.resolve(); + expect(ownerKernel.shutdown).toHaveBeenCalledTimes(2); + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + vi.advanceTimersByTime(1); + await Promise.resolve(); + await Promise.resolve(); + expect(ownerKernel.shutdown).toHaveBeenCalledTimes(3); + expect(ownerKernel.shutdown).toHaveBeenNthCalledWith(3, undefined); + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + retryConfirmation.resolve({ confirmed: true }); + await blockedExecution; + expect(startSpy).toHaveBeenCalledTimes(5); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + await disposeAllKernelSessions(); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("waits for confirmed disposal capacity before evicting unrelated retained sessions", async () => { + const ownerKernel = new FakeKernel(); + const unrelatedKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const replacementKernel = new FakeKernel(); + const retryConfirmation = Promise.withResolvers(); + let shutdownCallCount = 0; + ownerKernel.shutdown = vi.fn(async (_options?: FakeKernelShutdownOptions): Promise => { + shutdownCallCount += 1; + if (shutdownCallCount === 1) { + return { confirmed: false }; + } + return await retryConfirmation.promise; + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of [ownerKernel, ...unrelatedKernels, replacementKernel]) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + + await executePython("print('owner-a')", { + cwd: "/tmp/unconfirmed-owner-cleanup-a", + sessionId: "unconfirmed-owner-cleanup-a", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + for (const [index, ownerId] of ["owner-b", "owner-c", "owner-d"].entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/unconfirmed-owner-cleanup-${ownerId}`, + sessionId: `unconfirmed-owner-cleanup-${ownerId}`, + kernelMode: "session", + kernelOwnerId: ownerId, + }); + } + + await disposeKernelSessionsByOwner("owner-a"); + expect(ownerKernel.shutdown).toHaveBeenCalledTimes(2); + expect(ownerKernel.shutdown).toHaveBeenNthCalledWith(1, { timeoutMs: 2_000 }); + expect(ownerKernel.shutdown).toHaveBeenNthCalledWith(2, undefined); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + const blockedExecution = executePython("print('replacement')", { + cwd: "/tmp/unconfirmed-owner-cleanup-replacement", + sessionId: "unconfirmed-owner-cleanup-replacement", + kernelMode: "session", + kernelOwnerId: "owner-e", + }); + await Promise.resolve(); + await Promise.resolve(); + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + for (const [index, ownerId] of ["owner-b", "owner-c", "owner-d"].entries()) { + await executePython(`print('reuse-${ownerId}')`, { + cwd: `/tmp/unconfirmed-owner-cleanup-${ownerId}`, + sessionId: `unconfirmed-owner-cleanup-${ownerId}`, + kernelMode: "session", + kernelOwnerId: ownerId, + }); + expect(unrelatedKernels[index]!.execute).toHaveBeenCalledTimes(2); + expect(unrelatedKernels[index]!.shutdown).not.toHaveBeenCalled(); + } + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + + retryConfirmation.resolve({ confirmed: true }); + await blockedExecution; + expect(startSpy).toHaveBeenCalledTimes(5); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + for (const kernel of unrelatedKernels) { + expect(kernel.shutdown).not.toHaveBeenCalled(); + } + + await disposeAllKernelSessions(); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("owner cleanup retries every shutdown in background and frees retained capacity one confirmation at a time", async () => { + const retainedKernels = [new FakeKernel(), new FakeKernel(), new FakeKernel(), new FakeKernel()]; + const replacementKernel = new FakeKernel(); + const laterKernel = new FakeKernel(); + const retryConfirmations = retainedKernels.map(() => Promise.withResolvers()); + for (const [index, kernel] of retainedKernels.entries()) { + let shutdownCallCount = 0; + kernel.shutdown = vi.fn(async (_options?: FakeKernelShutdownOptions): Promise => { + shutdownCallCount += 1; + if (shutdownCallCount === 1) { + return { confirmed: false }; + } + return await retryConfirmations[index]!.promise; + }); + } + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start"); + for (const kernel of [...retainedKernels, replacementKernel, laterKernel]) { + startSpy.mockResolvedValueOnce(kernel as unknown as PythonKernelInstance); + } + + for (const [index] of retainedKernels.entries()) { + await executePython(`print(${index})`, { + cwd: `/tmp/unconfirmed-capacity-${index}`, + sessionId: `unconfirmed-capacity-session-${index}`, + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + } + + await disposeKernelSessionsByOwner("owner-a"); + for (const kernel of retainedKernels) { + expect(kernel.shutdown).toHaveBeenCalledTimes(2); + expect(kernel.shutdown).toHaveBeenNthCalledWith(1, { timeoutMs: 2_000 }); + expect(kernel.shutdown).toHaveBeenNthCalledWith(2, undefined); + } + + let globalCleanupResolved = false; + const globalCleanup = disposeAllKernelSessions().then(() => { + globalCleanupResolved = true; + }); + await Promise.resolve(); + await Promise.resolve(); + + for (const kernel of retainedKernels) { + expect(kernel.shutdown).toHaveBeenCalledTimes(2); + } + expect(globalCleanupResolved).toBe(false); + + const blockedExecution = executePython("print('replacement')", { + cwd: "/tmp/unconfirmed-capacity-replacement", + sessionId: "unconfirmed-capacity-replacement", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + await Promise.resolve(); + await Promise.resolve(); + expect(startSpy).toHaveBeenCalledTimes(4); + expect(replacementKernel.execute).not.toHaveBeenCalled(); + + retryConfirmations[0]!.resolve({ confirmed: true }); + await blockedExecution; + expect(globalCleanupResolved).toBe(false); + expect(startSpy).toHaveBeenCalledTimes(5); + expect(replacementKernel.execute).toHaveBeenCalledTimes(1); + + const secondBlockedExecution = executePython("print('later')", { + cwd: "/tmp/unconfirmed-capacity-later", + sessionId: "unconfirmed-capacity-later", + kernelMode: "session", + kernelOwnerId: "owner-c", + }); + await Promise.resolve(); + await Promise.resolve(); + expect(globalCleanupResolved).toBe(false); + expect(startSpy).toHaveBeenCalledTimes(5); + expect(laterKernel.execute).not.toHaveBeenCalled(); + + retryConfirmations[1]!.resolve({ confirmed: true }); + await secondBlockedExecution; + expect(globalCleanupResolved).toBe(false); + expect(startSpy).toHaveBeenCalledTimes(6); + expect(laterKernel.execute).toHaveBeenCalledTimes(1); + + for (const confirmation of retryConfirmations.slice(2)) { + confirmation.resolve({ confirmed: true }); + } + await globalCleanup; + expect(globalCleanupResolved).toBe(true); + + await disposeAllKernelSessions(); + expect(replacementKernel.shutdown).toHaveBeenCalledTimes(1); + expect(laterKernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("does not let stuck retained executions block owner or global cleanup", async () => { + const ownerKernel = new FakeKernel(); + const globalKernel = new FakeKernel(); + const ownerExecutionStarted = Promise.withResolvers(); + const globalExecutionStarted = Promise.withResolvers(); + ownerKernel.execute = vi.fn(async () => { + ownerExecutionStarted.resolve(); + return await new Promise(() => {}); + }); + globalKernel.execute = vi.fn(async () => { + globalExecutionStarted.resolve(); + return await new Promise(() => {}); + }); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + vi.spyOn(PythonKernel, "start") + .mockResolvedValueOnce(ownerKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(globalKernel as unknown as PythonKernelInstance); + + void executePython("print('owner hangs')", { + cwd: "/tmp/stuck-owner-cleanup", + sessionId: "stuck-owner-session", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + await ownerExecutionStarted.promise; + + void executePython("print('global hangs')", { + cwd: "/tmp/stuck-global-cleanup", + sessionId: "stuck-global-session", + kernelMode: "session", + }); + await globalExecutionStarted.promise; + + const ownerCleanup = Promise.race([ + disposeKernelSessionsByOwner("owner-a").then(() => "disposed-owner" as const), + new Promise<"timeout">(resolve => setTimeout(() => resolve("timeout"), 50)), + ]); + await expect(ownerCleanup).resolves.toBe("disposed-owner"); + expect(ownerKernel.shutdown).toHaveBeenCalledTimes(1); + expect(globalKernel.shutdown).not.toHaveBeenCalled(); + + const globalCleanup = Promise.race([ + disposeAllKernelSessions().then(() => "disposed-all" as const), + new Promise<"timeout">(resolve => setTimeout(() => resolve("timeout"), 50)), + ]); + await expect(globalCleanup).resolves.toBe("disposed-all"); + expect(globalKernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("attaches cached warmup sessions to newly provided owners", async () => { + using tempDir = TempDir.createSync("@python-owner-warmup-"); + const docs: PreludeHelper[] = [ + { + name: "read", + signature: "(path)", + docstring: "Read file contents.", + category: "File I/O", + }, + ]; + const kernel = { + introspectPrelude: vi.fn().mockResolvedValue(docs), + execute: vi.fn(async () => OK_RESULT), + ping: vi.fn(async () => true), + isAlive: () => true, + shutdown: vi.fn(async (): Promise => ({ confirmed: true })), + }; + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance); + + const firstWarmup = await warmPythonEnvironment(tempDir.path(), "warm-session", true, undefined, "owner-a"); + expect(firstWarmup.ok).toBe(true); + expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1); + + const cachedWarmup = await warmPythonEnvironment(tempDir.path(), "warm-session", true, undefined, "owner-b"); + expect(cachedWarmup.ok).toBe(true); + expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1); + expect(startSpy).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("owner-a"); + expect(kernel.shutdown).not.toHaveBeenCalled(); + + await executePython("1 + 1", { + cwd: tempDir.path(), + sessionId: "warm-session", + kernelMode: "session", + kernelOwnerId: "owner-b", + }); + + expect(startSpy).toHaveBeenCalledTimes(1); + expect(kernel.execute).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("owner-b"); + expect(kernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("keeps cache-hit ownerless warmups provisional until an explicit owner takes over", async () => { + using tempDir = TempDir.createSync("@python-owner-fallback-"); + const docs: PreludeHelper[] = [ + { + name: "read", + signature: "(path)", + docstring: "Read file contents.", + category: "File I/O", + }, + ]; + const kernel = { + introspectPrelude: vi.fn().mockResolvedValue(docs), + execute: vi.fn(async () => OK_RESULT), + ping: vi.fn(async () => true), + isAlive: () => true, + shutdown: vi.fn(async (): Promise => ({ confirmed: true })), + }; + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernelInstance); + + const firstWarmup = await warmPythonEnvironment(tempDir.path(), "warm-fallback-session", true); + expect(firstWarmup.ok).toBe(true); + expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1); + + const cachedWarmup = await warmPythonEnvironment(tempDir.path(), "warm-fallback-session", true); + expect(cachedWarmup.ok).toBe(true); + expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1); + expect(startSpy).toHaveBeenCalledTimes(1); + + await executePython("1 + 1", { + cwd: tempDir.path(), + sessionId: "warm-fallback-session", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + + expect(kernel.execute).toHaveBeenCalledTimes(1); + + const postTakeoverWarmup = await warmPythonEnvironment(tempDir.path(), "warm-fallback-session", true); + expect(postTakeoverWarmup.ok).toBe(true); + expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1); + expect(startSpy).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("owner-a"); + expect(kernel.shutdown).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("warm-fallback-session"); + expect(kernel.shutdown).toHaveBeenCalledTimes(1); + }); + + it("leaves per-call kernels out of owner-scoped retained cleanup and keeps global cleanup intact", async () => { + const perCallKernel = new FakeKernel(); + const retainedKernel = new FakeKernel(); + const unownedRetainedKernel = new FakeKernel(); + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + const startSpy = vi + .spyOn(PythonKernel, "start") + .mockResolvedValueOnce(perCallKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(retainedKernel as unknown as PythonKernelInstance) + .mockResolvedValueOnce(unownedRetainedKernel as unknown as PythonKernelInstance); + + await executePython("print('per-call')", { + cwd: "/tmp/per-call-owner", + kernelMode: "per-call", + kernelOwnerId: "owner-a", + }); + await executePython("print('retained')", { + cwd: "/tmp/retained-owner", + sessionId: "retained-session", + kernelMode: "session", + kernelOwnerId: "owner-a", + }); + await executePython("print('unowned')", { + cwd: "/tmp/unowned-retained", + sessionId: "unowned-session", + kernelMode: "session", + }); + + expect(startSpy).toHaveBeenCalledTimes(3); + expect(perCallKernel.shutdown).toHaveBeenCalledTimes(1); + + await disposeKernelSessionsByOwner("owner-a"); + + expect(perCallKernel.shutdown).toHaveBeenCalledTimes(1); + expect(retainedKernel.shutdown).toHaveBeenCalledTimes(1); + expect(unownedRetainedKernel.shutdown).not.toHaveBeenCalled(); + + await disposeAllKernelSessions(); + + expect(unownedRetainedKernel.shutdown).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/coding-agent/test/core/python-executor-session.test.ts b/packages/coding-agent/test/core/python-executor-session.test.ts index ba8bb6b1b..df34a0571 100644 --- a/packages/coding-agent/test/core/python-executor-session.test.ts +++ b/packages/coding-agent/test/core/python-executor-session.test.ts @@ -25,9 +25,10 @@ class FakeKernel { return this.alive; } - async shutdown(): Promise { + async shutdown(): Promise<{ confirmed: boolean }> { this.shutdownCalls += 1; this.alive = false; + return { confirmed: true }; } } 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 e8b166f8a..e472e0591 100644 --- a/packages/coding-agent/test/core/python-executor.lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-executor.lifecycle.test.ts @@ -34,9 +34,10 @@ class FakeKernel { return this.#result; } - async shutdown(): Promise { + async shutdown(): Promise<{ confirmed: boolean }> { this.shutdownCalls += 1; this.#alive = false; + return { confirmed: true }; } async ping(): Promise { @@ -174,9 +175,11 @@ describe("executePython session lifecycle", () => { kernelA.shutdown = async () => { shutdownCount += 1; + return { confirmed: true }; }; kernelB.shutdown = async () => { shutdownCount += 1; + return { confirmed: true }; }; await executePython("print('one')", { kernelMode: "per-call" }); diff --git a/packages/coding-agent/test/core/python-executor.test.ts b/packages/coding-agent/test/core/python-executor.test.ts index fc1dd76c2..b4fc9c94f 100644 --- a/packages/coding-agent/test/core/python-executor.test.ts +++ b/packages/coding-agent/test/core/python-executor.test.ts @@ -11,6 +11,7 @@ import { import { type KernelExecuteOptions, type KernelExecuteResult, + type KernelShutdownResult, type PreludeHelper, PythonKernel, } from "@oh-my-pi/pi-coding-agent/ipy/kernel"; @@ -175,7 +176,7 @@ describe("warmPythonEnvironment", () => { introspectPrelude: vi.fn().mockResolvedValue(docs), ping: vi.fn().mockResolvedValue(true), isAlive: () => true, - shutdown: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn(async (): Promise => ({ confirmed: true })), }; const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernel); diff --git a/packages/coding-agent/test/core/python-kernel-session.test.ts b/packages/coding-agent/test/core/python-kernel-session.test.ts index 710458a54..4f17d12b3 100644 --- a/packages/coding-agent/test/core/python-kernel-session.test.ts +++ b/packages/coding-agent/test/core/python-kernel-session.test.ts @@ -20,9 +20,10 @@ class FakeKernel { return { status: "ok", cancelled: false, timedOut: false, stdinRequested: false }; } - async shutdown(): Promise { + async shutdown(): Promise<{ confirmed: boolean }> { this.shutdownCalls += 1; this.alive = false; + return { confirmed: true }; } isAlive(): boolean { 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 b325cdbe2..45cf6493b 100644 --- a/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts +++ b/packages/coding-agent/test/core/python-kernel.lifecycle.test.ts @@ -339,6 +339,97 @@ describe("PythonKernel gateway lifecycle", () => { ); }); + it("treats a retry against an already-missing kernel as confirmed shutdown", async () => { + using _runtime = stubKernelRuntime(); + vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ + url: "http://127.0.0.1:9999", + isShared: true, + }); + + let deleteCalls = 0; + 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-retry-delete" } }) as unknown as Response; + } + if (url.endsWith("/api/kernels/kernel-retry-delete") && init?.method === "DELETE") { + deleteCalls += 1; + if (deleteCalls === 1) { + return createResponse({ ok: false, status: 503, text: "not yet" }) as unknown as Response; + } + return createResponse({ ok: false, status: 404, text: "gone" }) as unknown as Response; + } + return createResponse({ ok: true }) as unknown as Response; + }); + + const kernel = await PythonKernel.start({ cwd: tempDir.path() }); + + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: false }); + expect(kernel.isAlive()).toBe(false); + expect(FakeWebSocket.instances.at(-1)?.readyState).toBe(FakeWebSocket.CLOSED); + + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: true }); + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: true }); + expect(deleteCalls).toBe(2); + }); + + it("returns unconfirmed when shutdown times out and can confirm on retry", async () => { + using _runtime = stubKernelRuntime(); + vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ + url: "http://127.0.0.1:9999", + isShared: true, + }); + let deleteCalls = 0; + const firstDeleteStarted = Promise.withResolvers(); + const firstDeleteAborted = Promise.withResolvers(); + 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-shutdown-timeout" } }) as unknown as Response; + } + if (url.endsWith("/api/kernels/kernel-shutdown-timeout") && init?.method === "DELETE") { + deleteCalls += 1; + if (deleteCalls === 1) { + firstDeleteStarted.resolve(); + return new Promise((_, reject) => { + const waitForAbort = () => { + if (init.signal?.aborted) { + firstDeleteAborted.resolve(); + const reason = init.signal.reason; + reject(reason instanceof Error ? reason : new Error("Python kernel shutdown timed out")); + return; + } + const poll = setTimeout(waitForAbort, 5); + poll.unref?.(); + }; + waitForAbort(); + }); + } + return createResponse({ ok: false, status: 404, text: "gone" }) as unknown as Response; + } + return createResponse({ ok: true }) as unknown as Response; + }); + const kernel = await PythonKernel.start({ cwd: tempDir.path() }); + const shutdownPromise = kernel.shutdown({ timeoutMs: 25 }); + await firstDeleteStarted.promise; + const pending = Symbol("pending"); + const settled = await Promise.race([ + shutdownPromise, + new Promise(resolve => { + const timer = setTimeout(() => resolve(pending), 250); + timer.unref?.(); + }), + ]); + await firstDeleteAborted.promise; + expect(settled).not.toBe(pending); + expect(settled).toEqual({ confirmed: false }); + expect(kernel.isAlive()).toBe(false); + expect(FakeWebSocket.instances.at(-1)?.readyState).toBe(FakeWebSocket.CLOSED); + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: true }); + expect(deleteCalls).toBe(2); + }); it("does not throw when shutdown API fails", async () => { using _runtime = stubKernelRuntime(); vi.spyOn(gatewayCoordinator, "acquireSharedGateway").mockResolvedValue({ @@ -360,6 +451,6 @@ describe("PythonKernel gateway lifecycle", () => { const kernel = await PythonKernel.start({ cwd: tempDir.path() }); - await expect(kernel.shutdown()).resolves.toBeUndefined(); + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: false }); }); }); diff --git a/packages/coding-agent/test/core/python-kernel.test.ts b/packages/coding-agent/test/core/python-kernel.test.ts index 2e2d464a0..a51195b11 100644 --- a/packages/coding-agent/test/core/python-kernel.test.ts +++ b/packages/coding-agent/test/core/python-kernel.test.ts @@ -294,13 +294,49 @@ describe("PythonKernel (external gateway)", () => { expect(chunks.join("")).toContain("result"); expect(displays).toEqual([{ type: "json", data: { answer: 42 } }]); - await kernel.shutdown(); + const shutdown = await kernel.shutdown(); + expect(shutdown).toEqual({ confirmed: true }); expect(fetchMock).toHaveBeenCalledWith("http://gateway.test/api/kernels/kernel-1", { method: "DELETE", headers: {}, }); }); + it("returns an unconfirmed shutdown result when kernel deletion is not acknowledged", async () => { + let deleteAttempts = 0; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/api/kernels") && init?.method === "POST") { + return new Response(JSON.stringify({ id: "kernel-delete-failure" }), { status: 201 }); + } + if (url.includes("/api/kernels/") && init?.method === "DELETE") { + deleteAttempts += 1; + if (deleteAttempts === 1) { + return new Response("delete failed", { status: 500, statusText: "Server Error" }); + } + return new Response("already gone", { status: 404, statusText: "Not Found" }); + } + return new Response("", { status: 200 }); + }); + using _hook = hookFetch((input, init) => fetchMock(String(input), init)); + + const kernelPromise = PythonKernel.start({ cwd: "/" }); + await Bun.sleep(10); + const ws = FakeWebSocket.lastInstance; + if (!ws) throw new Error("WebSocket not initialized"); + ws.setSendHandler(data => { + const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); + sendOkExecution(ws, msg.header.msg_id); + }); + + const kernel = await kernelPromise; + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: false }); + expect(deleteAttempts).toBe(1); + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: true }); + expect(deleteAttempts).toBe(2); + await expect(kernel.shutdown()).resolves.toEqual({ confirmed: true }); + expect(deleteAttempts).toBe(2); + }); + it("initializes the IPython prelude", async () => { const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { if (url.endsWith("/api/kernels") && init?.method === "POST") { diff --git a/packages/coding-agent/test/python-tool-settings.test.ts b/packages/coding-agent/test/python-tool-settings.test.ts index 9ddb31dde..3b95c27eb 100644 --- a/packages/coding-agent/test/python-tool-settings.test.ts +++ b/packages/coding-agent/test/python-tool-settings.test.ts @@ -13,12 +13,16 @@ function createSession( cwd: string, sessionFile: string, overrides?: Partial>, + kernelOwnerId?: string, + forcePythonWarmup = false, ): ToolSession { return { cwd, hasUI: false, getSessionFile: () => sessionFile, getSessionSpawns: () => null, + getPythonKernelOwnerId: () => kernelOwnerId ?? null, + forcePythonWarmup, settings: Settings.isolated({ "python.toolMode": "ipy-only", ...overrides }), }; } @@ -32,6 +36,7 @@ describe("python tool settings", () => { }); afterEach(() => { + pythonExecutor.resetPreludeDocsCache(); vi.restoreAllMocks(); fs.rmSync(testDir, { recursive: true, force: true }); }); @@ -55,8 +60,9 @@ describe("python tool settings", () => { expect(tools.map(tool => tool.name).sort()).toEqual(["bash", "exit_plan_mode"]); }); - it("passes kernel mode from settings to executor", async () => { - vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); + it("passes kernel owner and kernel mode from settings to executor", async () => { + vi.spyOn(pythonExecutor, "getPreludeDocs").mockReturnValue([]); + const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ output: "ok", exitCode: 0, @@ -71,17 +77,56 @@ describe("python tool settings", () => { }); const sessionFile = path.join(testDir, "session.jsonl"); - const session = createSession(testDir, sessionFile, { "python.kernelMode": "per-call" }); + const kernelOwnerId = "owner-456"; + const session = createSession(testDir, sessionFile, { "python.kernelMode": "per-call" }, kernelOwnerId); const pythonTool = new PythonTool(session); await pythonTool.execute("tool-call", { cells: [{ code: "print(1)" }] }); + expect(warmupSpy).toHaveBeenCalledWith( + testDir, + `session:${sessionFile}:cwd:${testDir}`, + true, + sessionFile, + kernelOwnerId, + expect.any(AbortSignal), + ); expect(executeSpy).toHaveBeenCalledWith( "print(1)", expect.objectContaining({ kernelMode: "per-call", sessionId: `session:${sessionFile}:cwd:${testDir}`, + kernelOwnerId, }), ); }); + + it("passes kernel owner into createTools warmup without changing session ids", async () => { + vi.spyOn(pythonKernel, "checkPythonKernelAvailability").mockResolvedValue({ ok: true }); + vi.spyOn(pythonExecutor, "getPreludeDocs").mockReturnValue([]); + const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); + + const sessionFile = path.join(testDir, "session-create-tools.jsonl"); + const kernelOwnerId = "owner-create-tools"; + const previousSkipCheck = Bun.env.PI_PYTHON_SKIP_CHECK; + + delete Bun.env.PI_PYTHON_SKIP_CHECK; + try { + await createTools(createSession(testDir, sessionFile, undefined, kernelOwnerId, true), ["python"]); + + expect(warmupSpy).toHaveBeenCalledWith( + testDir, + `session:${sessionFile}:cwd:${testDir}`, + true, + sessionFile, + kernelOwnerId, + ); + } finally { + if (previousSkipCheck === undefined) { + delete Bun.env.PI_PYTHON_SKIP_CHECK; + } else { + Bun.env.PI_PYTHON_SKIP_CHECK = previousSkipCheck; + } + } + }); }); diff --git a/packages/coding-agent/test/tools/python-execution.test.ts b/packages/coding-agent/test/tools/python-execution.test.ts index 463cbc6be..83ea63579 100644 --- a/packages/coding-agent/test/tools/python-execution.test.ts +++ b/packages/coding-agent/test/tools/python-execution.test.ts @@ -1,16 +1,17 @@ -import { describe, expect, it, vi } from "bun:test"; +import { afterEach, describe, expect, it, vi } from "bun:test"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import * as pythonExecutor from "@oh-my-pi/pi-coding-agent/ipy/executor"; import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; import { PythonTool } from "@oh-my-pi/pi-coding-agent/tools/python"; import { TempDir } from "@oh-my-pi/pi-utils"; -function createSession(cwd: string): ToolSession { +function createSession(cwd: string, kernelOwnerId?: string): ToolSession { return { cwd, hasUI: false, getSessionFile: () => `${cwd}/session-file.jsonl`, getSessionSpawns: () => "*", + getPythonKernelOwnerId: () => kernelOwnerId ?? null, settings: Settings.isolated({ "lsp.formatOnWrite": true, "bashInterceptor.enabled": true, @@ -21,9 +22,15 @@ function createSession(cwd: string): ToolSession { } describe("python tool execution", () => { - it("passes kernel options from settings and args", async () => { + afterEach(() => { + pythonExecutor.resetPreludeDocsCache(); + vi.restoreAllMocks(); + }); + + it("passes kernel owner and kernel options from settings and args", async () => { const tempDir = TempDir.createSync("@python-tool-"); - vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); + vi.spyOn(pythonExecutor, "getPreludeDocs").mockReturnValue([]); + const warmupSpy = vi.spyOn(pythonExecutor, "warmPythonEnvironment").mockResolvedValue({ ok: true, docs: [] }); const executeSpy = vi.spyOn(pythonExecutor, "executePython").mockResolvedValue({ output: "ok", exitCode: 0, @@ -37,7 +44,8 @@ describe("python tool execution", () => { stdinRequested: false, }); - const tool = new PythonTool(createSession(tempDir.path())); + const kernelOwnerId = "owner-123"; + const tool = new PythonTool(createSession(tempDir.path(), kernelOwnerId)); const result = await tool.execute( "call-id", { cells: [{ code: "print('hi')" }], timeout: 5, cwd: tempDir.path(), reset: true }, @@ -46,6 +54,14 @@ describe("python tool execution", () => { undefined, ); + expect(warmupSpy).toHaveBeenCalledWith( + tempDir.path(), + `session:${tempDir.path()}/session-file.jsonl:cwd:${tempDir.path()}`, + true, + `${tempDir.path()}/session-file.jsonl`, + kernelOwnerId, + expect.any(AbortSignal), + ); expect(executeSpy).toHaveBeenCalledWith( "print('hi')", expect.objectContaining({ @@ -54,6 +70,7 @@ describe("python tool execution", () => { signal: expect.any(AbortSignal), sessionFile: `${tempDir.path()}/session-file.jsonl`, sessionId: `session:${tempDir.path()}/session-file.jsonl:cwd:${tempDir.path()}`, + kernelOwnerId, kernelMode: "per-call", useSharedGateway: true, reset: true, @@ -62,7 +79,6 @@ describe("python tool execution", () => { const text = result.content.find(item => item.type === "text")?.text; expect(text).toBe("ok"); - executeSpy.mockRestore(); tempDir.removeSync(); }); });