fix(coding-agent): scoped python kernel cleanup to sessions
scoped retained-kernel ownership to agent sessions and cleaned it up on session disposal cleaned up warmed python owners on session startup failure and rejected new direct and tool-based python starts during disposal, including async hook, preflight, and warmup races
This commit is contained in:
Carl
@@ -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
|
||||
|
||||
@@ -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<void>;
|
||||
restartCount: number;
|
||||
dead: boolean;
|
||||
needsRestart: boolean;
|
||||
kernelInvalidatedByRecovery: boolean;
|
||||
disposing: boolean;
|
||||
disposeCapacityPromise?: Promise<void>;
|
||||
resolveDisposeCapacity?: () => void;
|
||||
disposeAttemptPromise?: Promise<void>;
|
||||
resolveDisposeAttempt?: () => void;
|
||||
disposeResultPromise?: Promise<KernelDisposalResult>;
|
||||
disposeResultTimeoutMs?: number;
|
||||
nextDisposalRetryAt?: number;
|
||||
lastUsedAt: number;
|
||||
ownerIds: Set<string>;
|
||||
hasFallbackOwner: boolean;
|
||||
heartbeatTimer?: NodeJS.Timeout;
|
||||
}
|
||||
|
||||
const kernelSessions = new Map<string, KernelSession>();
|
||||
const disposingKernelSessions = new Set<KernelSession>();
|
||||
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<void>,
|
||||
async function waitForPromiseWithCancellation<T>(
|
||||
promise: Promise<T>,
|
||||
options: Pick<KernelSessionExecutionOptions, "signal" | "deadlineMs">,
|
||||
): Promise<void> {
|
||||
): Promise<T> {
|
||||
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<void>((resolve, reject) => {
|
||||
return await new Promise<T>((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<void>,
|
||||
options: Pick<KernelSessionExecutionOptions, "signal" | "deadlineMs">,
|
||||
): Promise<void> {
|
||||
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<KernelSessionExecutionOptions, "signal" | "deadlineMs">,
|
||||
): Promise<void> {
|
||||
retryPendingKernelSessionDisposals();
|
||||
|
||||
const disposalPromises: Promise<void>[] = [];
|
||||
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<KernelSessionExecutionOptions, "signal" | "deadlineMs">,
|
||||
): Promise<void> {
|
||||
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<void> {
|
||||
const now = Date.now();
|
||||
const toDispose: KernelSession[] = [];
|
||||
@@ -367,9 +515,8 @@ async function cleanupIdleSessions(): Promise<void> {
|
||||
await Promise.allSettled(toDispose.map(session => disposeKernelSession(session)));
|
||||
}
|
||||
|
||||
if (kernelSessions.size === 0) {
|
||||
stopCleanupTimer();
|
||||
}
|
||||
retryPendingKernelSessionDisposals(now);
|
||||
syncCleanupTimer();
|
||||
}
|
||||
|
||||
async function evictOldestSession(): Promise<void> {
|
||||
@@ -387,12 +534,29 @@ async function evictOldestSession(): Promise<void> {
|
||||
|
||||
export async function disposeAllKernelSessions(): Promise<void> {
|
||||
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<void> {
|
||||
const availability = await checkPythonKernelAvailability(cwd);
|
||||
export async function disposeKernelSessionsByOwner(ownerId: string): Promise<void> {
|
||||
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<KernelSessionExecutionOptions, "signal" | "deadlineMs"> = {},
|
||||
): Promise<void> {
|
||||
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<void> {
|
||||
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<string, string> | 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<string, string> | 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<void> {
|
||||
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<KernelDisposalResult> {
|
||||
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<KernelDisposalResult> {
|
||||
if (!session.disposeResultPromise) {
|
||||
session.disposeResultTimeoutMs = timeoutMs;
|
||||
const releaseDisposalAttempt = Promise.withResolvers<void>();
|
||||
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<KernelDisposalWaitResult | undefined> {
|
||||
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<void> {
|
||||
if (!session.disposing) {
|
||||
if (!beginDisposingKernelSession(session)) return;
|
||||
const releaseDisposalCapacity = Promise.withResolvers<void>();
|
||||
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<T>(
|
||||
@@ -570,10 +945,11 @@ async function withKernelSession<T>(
|
||||
options: KernelSessionExecutionOptions = {},
|
||||
): Promise<T> {
|
||||
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<T>(
|
||||
kernelSessions.set(sessionId, session);
|
||||
startCleanupTimer();
|
||||
}
|
||||
attachKernelOwner(sessionId, options.kernelOwnerId);
|
||||
|
||||
if (session.disposing) {
|
||||
return await withKernelSession(sessionId, cwd, handler, options);
|
||||
}
|
||||
|
||||
const run = async (): Promise<T> => {
|
||||
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<T>(
|
||||
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<T>(
|
||||
|
||||
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(
|
||||
|
||||
@@ -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<string, (msg: JupyterMessage) => void>();
|
||||
#channelHandlers = new Map<string, Set<(msg: JupyterMessage) => void>>();
|
||||
#pendingExecutions = new Map<string, (reason: string) => void>();
|
||||
|
||||
private constructor(
|
||||
readonly id: string,
|
||||
readonly kernelId: string,
|
||||
@@ -1004,11 +1009,18 @@ export class PythonKernel {
|
||||
}
|
||||
}
|
||||
|
||||
async shutdown(options?: KernelShutdownOptions): Promise<void> {
|
||||
if (this.#disposed) return;
|
||||
this.#disposed = true;
|
||||
this.#alive = false;
|
||||
this.#abortPendingExecutions("Kernel shutdown");
|
||||
async shutdown(options?: KernelShutdownOptions): Promise<KernelShutdownResult> {
|
||||
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 {
|
||||
|
||||
+789
-758
File diff suppressed because it is too large
Load Diff
@@ -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<AbortController>();
|
||||
#pythonKernelOwnerId: string;
|
||||
#pendingPythonMessages: PythonExecutionMessage[] = [];
|
||||
|
||||
#activePythonExecutions = new Set<Promise<unknown>>();
|
||||
#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<void> {
|
||||
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<PythonResult> {
|
||||
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<PythonResult> => {
|
||||
// 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<T>(execution: Promise<T>, abortController: AbortController): Promise<T> {
|
||||
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<boolean> {
|
||||
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<boolean> {
|
||||
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 */
|
||||
|
||||
@@ -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?<T>(execution: Promise<T>, abortController: AbortController): Promise<T>;
|
||||
/** 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", {
|
||||
|
||||
@@ -175,6 +175,7 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
|
||||
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<typeof pythonSchema> {
|
||||
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<typeof pythonSchema> {
|
||||
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<AgentToolResult<PythonToolDetails | undefined>> => {
|
||||
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<PythonExecutorOptions, "reset"> = {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<typeof OK_EXECUTION> | undefined;
|
||||
blockedExecutionStarted: (() => void) | undefined;
|
||||
blockedExecutionReject: ((error: Error) => void) | undefined;
|
||||
abortBlockedExecution = true;
|
||||
|
||||
isAlive(): boolean {
|
||||
return this.alive;
|
||||
}
|
||||
|
||||
async execute(code: string, options?: { signal?: AbortSignal }): Promise<typeof OK_EXECUTION> {
|
||||
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<typeof OK_EXECUTION>((_, 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<boolean> {
|
||||
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<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
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<void>();
|
||||
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<void>();
|
||||
const releaseWarmup = Promise.withResolvers<void>();
|
||||
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<void>();
|
||||
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<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
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<void>();
|
||||
const releaseHook = Promise.withResolvers<void>();
|
||||
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<void>();
|
||||
const releaseHook = Promise.withResolvers<void>();
|
||||
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<void>();
|
||||
const releaseArtifact = Promise.withResolvers<void>();
|
||||
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<typeof OK_EXECUTION>();
|
||||
const blockedExecutionStarted = Promise.withResolvers<void>();
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<KernelShutdownResult> => ({ 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<KernelShutdownResult>();
|
||||
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<KernelShutdownResult>());
|
||||
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<KernelShutdownResult> => {
|
||||
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<KernelShutdownResult>());
|
||||
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<KernelShutdownResult>());
|
||||
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<KernelShutdownResult>();
|
||||
let shutdownCallCount = 0;
|
||||
ownerKernel.shutdown = vi.fn(async (options?: FakeKernelShutdownOptions): Promise<KernelShutdownResult> => {
|
||||
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<KernelShutdownResult>();
|
||||
let shutdownCallCount = 0;
|
||||
ownerKernel.shutdown = vi.fn(async (_options?: FakeKernelShutdownOptions): Promise<KernelShutdownResult> => {
|
||||
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<KernelShutdownResult>());
|
||||
for (const [index, kernel] of retainedKernels.entries()) {
|
||||
let shutdownCallCount = 0;
|
||||
kernel.shutdown = vi.fn(async (_options?: FakeKernelShutdownOptions): Promise<KernelShutdownResult> => {
|
||||
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<void>();
|
||||
const globalExecutionStarted = Promise.withResolvers<void>();
|
||||
ownerKernel.execute = vi.fn(async () => {
|
||||
ownerExecutionStarted.resolve();
|
||||
return await new Promise<KernelExecuteResult>(() => {});
|
||||
});
|
||||
globalKernel.execute = vi.fn(async () => {
|
||||
globalExecutionStarted.resolve();
|
||||
return await new Promise<KernelExecuteResult>(() => {});
|
||||
});
|
||||
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<KernelShutdownResult> => ({ 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<KernelShutdownResult> => ({ 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);
|
||||
});
|
||||
});
|
||||
@@ -25,9 +25,10 @@ class FakeKernel {
|
||||
return this.alive;
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
async shutdown(): Promise<{ confirmed: boolean }> {
|
||||
this.shutdownCalls += 1;
|
||||
this.alive = false;
|
||||
return { confirmed: true };
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -34,9 +34,10 @@ class FakeKernel {
|
||||
return this.#result;
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
async shutdown(): Promise<{ confirmed: boolean }> {
|
||||
this.shutdownCalls += 1;
|
||||
this.#alive = false;
|
||||
return { confirmed: true };
|
||||
}
|
||||
|
||||
async ping(): Promise<boolean> {
|
||||
@@ -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" });
|
||||
|
||||
@@ -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<KernelShutdownResult> => ({ confirmed: true })),
|
||||
};
|
||||
const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernel);
|
||||
|
||||
|
||||
@@ -20,9 +20,10 @@ class FakeKernel {
|
||||
return { status: "ok", cancelled: false, timedOut: false, stdinRequested: false };
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
async shutdown(): Promise<{ confirmed: boolean }> {
|
||||
this.shutdownCalls += 1;
|
||||
this.alive = false;
|
||||
return { confirmed: true };
|
||||
}
|
||||
|
||||
isAlive(): boolean {
|
||||
|
||||
@@ -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<void>();
|
||||
const firstDeleteAborted = Promise.withResolvers<void>();
|
||||
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<Response>((_, 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<typeof pending>(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 });
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -13,12 +13,16 @@ function createSession(
|
||||
cwd: string,
|
||||
sessionFile: string,
|
||||
overrides?: Partial<Record<SettingPath, unknown>>,
|
||||
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;
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user