@oh-my-pi/pi-coding-agent 5.4.2 → 5.6.7

package/src/core/python-executor.result.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, it } from "bun:test";
+import { executePythonWithKernel, type PythonKernelExecutor } from "./python-executor";
+import type { KernelExecuteOptions, KernelExecuteResult } from "./python-kernel";
+
+class FakeKernel implements PythonKernelExecutor {
+	private result: KernelExecuteResult;
+	private onExecute?: (options?: KernelExecuteOptions) => void;
+
+	constructor(result: KernelExecuteResult, onExecute?: (options?: KernelExecuteOptions) => void) {
+		this.result = result;
+		this.onExecute = onExecute;
+	}
+
+	async execute(_code: string, options?: KernelExecuteOptions): Promise<KernelExecuteResult> {
+		this.onExecute?.(options);
+		return this.result;
+	}
+}
+
+describe("executePythonWithKernel result mapping", () => {
+	it("adds timeout annotation when cancelled", async () => {
+		const kernel = new FakeKernel({
+			status: "ok",
+			cancelled: true,
+			timedOut: true,
+			stdinRequested: false,
+		});
+
+		const result = await executePythonWithKernel(kernel, "sleep()", { timeout: 5000 });
+
+		expect(result.exitCode).toBeUndefined();
+		expect(result.cancelled).toBe(true);
+		expect(result.output).toContain("Command timed out after 5 seconds");
+	});
+
+	it("maps kernel error status to exit code 1", async () => {
+		const kernel = new FakeKernel(
+			{ status: "error", cancelled: false, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("Traceback...\n");
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "raise ValueError('boom')");
+
+		expect(result.exitCode).toBe(1);
+		expect(result.output).toContain("Traceback");
+	});
+});

package/src/core/python-executor.test.ts

@@ -0,0 +1,180 @@
+import { afterEach, describe, expect, it, vi } from "bun:test";
+import { rmSync } from "node:fs";
+import {
+	disposeAllKernelSessions,
+	executePythonWithKernel,
+	getPreludeDocs,
+	type PythonKernelExecutor,
+	resetPreludeDocsCache,
+	warmPythonEnvironment,
+} from "./python-executor";
+import { type KernelExecuteOptions, type KernelExecuteResult, type PreludeHelper, PythonKernel } from "./python-kernel";
+import { DEFAULT_MAX_BYTES } from "./tools/truncate";
+
+class FakeKernel implements PythonKernelExecutor {
+	private result: KernelExecuteResult;
+	private onExecute: (options?: KernelExecuteOptions) => void;
+
+	constructor(result: KernelExecuteResult, onExecute: (options?: KernelExecuteOptions) => void) {
+		this.result = result;
+		this.onExecute = onExecute;
+	}
+
+	async execute(_code: string, options?: KernelExecuteOptions): Promise<KernelExecuteResult> {
+		this.onExecute(options);
+		return this.result;
+	}
+}
+
+describe("executePythonWithKernel", () => {
+	it("captures text and display outputs", async () => {
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: false, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("hello\n");
+				options?.onDisplay?.({ type: "json", data: { foo: "bar" } });
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "print('hello')");
+
+		expect(result.exitCode).toBe(0);
+		expect(result.output).toContain("hello");
+		expect(result.displayOutputs).toHaveLength(1);
+	});
+
+	it("marks stdin request as error", async () => {
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: false, timedOut: false, stdinRequested: true },
+			() => {},
+		);
+
+		const result = await executePythonWithKernel(kernel, "input('prompt')");
+
+		expect(result.exitCode).toBe(1);
+		expect(result.stdinRequested).toBe(true);
+		expect(result.output).toContain("Kernel requested stdin; interactive input is not supported.");
+	});
+
+	it("maps error status to exit code 1", async () => {
+		const kernel = new FakeKernel(
+			{ status: "error", cancelled: false, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("Traceback\n");
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "raise ValueError('nope')");
+
+		expect(result.exitCode).toBe(1);
+		expect(result.cancelled).toBe(false);
+		expect(result.output).toContain("Traceback");
+	});
+
+	it("sanitizes streamed chunks", async () => {
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: false, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("\u001b[31mred\r\n");
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "print('red')");
+
+		expect(result.output).toBe("red\n");
+	});
+
+	it("returns cancelled result with timeout annotation", async () => {
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: true, timedOut: true, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("partial output\n");
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "while True: pass", { timeout: 4100 });
+
+		expect(result.exitCode).toBeUndefined();
+		expect(result.cancelled).toBe(true);
+		expect(result.output).toContain("Command timed out after 4 seconds");
+	});
+
+	it("returns cancelled result without timeout annotation", async () => {
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: true, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.("cancelled output\n");
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "while True: pass");
+
+		expect(result.exitCode).toBeUndefined();
+		expect(result.cancelled).toBe(true);
+		expect(result.output).toContain("cancelled output");
+		expect(result.output).not.toContain("Command timed out");
+	});
+
+	it("truncates large output and stores full output file", async () => {
+		const largeOutput = `${"x".repeat(DEFAULT_MAX_BYTES + 1024)}TAIL`;
+		const kernel = new FakeKernel(
+			{ status: "ok", cancelled: false, timedOut: false, stdinRequested: false },
+			(options) => {
+				options?.onChunk?.(largeOutput);
+			},
+		);
+
+		const result = await executePythonWithKernel(kernel, "print('big')");
+
+		expect(result.truncated).toBe(true);
+		expect(result.fullOutputPath).toBeDefined();
+		expect(result.output).toContain("TAIL");
+
+		const fullText = await Bun.file(result.fullOutputPath as string).text();
+		expect(fullText).toBe(largeOutput);
+
+		rmSync(result.fullOutputPath as string, { force: true });
+	});
+});
+
+afterEach(async () => {
+	await disposeAllKernelSessions();
+	resetPreludeDocsCache();
+	vi.restoreAllMocks();
+});
+
+describe("warmPythonEnvironment", () => {
+	it("caches prelude docs on warmup", async () => {
+		const previousSkip = process.env.OMP_PYTHON_SKIP_CHECK;
+		process.env.OMP_PYTHON_SKIP_CHECK = "1";
+		const docs: PreludeHelper[] = [
+			{
+				name: "read",
+				signature: "(path)",
+				docstring: "Read file contents.",
+				category: "File I/O",
+			},
+		];
+		const kernel = {
+			introspectPrelude: vi.fn().mockResolvedValue(docs),
+			ping: vi.fn().mockResolvedValue(true),
+			isAlive: () => true,
+			shutdown: vi.fn().mockResolvedValue(undefined),
+		};
+		const startSpy = vi.spyOn(PythonKernel, "start").mockResolvedValue(kernel as unknown as PythonKernel);
+
+		const result = await warmPythonEnvironment("/tmp/test", "session-1");
+
+		expect(result.ok).toBe(true);
+		expect(result.docs).toEqual(docs);
+		expect(getPreludeDocs()).toEqual(docs);
+		expect(kernel.introspectPrelude).toHaveBeenCalledTimes(1);
+
+		startSpy.mockRestore();
+		if (previousSkip === undefined) {
+			delete process.env.OMP_PYTHON_SKIP_CHECK;
+		} else {
+			process.env.OMP_PYTHON_SKIP_CHECK = previousSkip;
+		}
+	});
+});

package/src/core/python-executor.ts

@@ -0,0 +1,313 @@
+import stripAnsi from "strip-ansi";
+import { sanitizeBinaryOutput } from "../utils/shell";
+import { logger } from "./logger";
+import {
+	checkPythonKernelAvailability,
+	type KernelDisplayOutput,
+	type KernelExecuteOptions,
+	type KernelExecuteResult,
+	type PreludeHelper,
+	PythonKernel,
+} from "./python-kernel";
+import { createOutputSink } from "./streaming-output";
+import { DEFAULT_MAX_BYTES } from "./tools/truncate";
+
+export type PythonKernelMode = "session" | "per-call";
+
+export interface PythonExecutorOptions {
+	/** Working directory for command execution */
+	cwd?: string;
+	/** Timeout in milliseconds */
+	timeout?: number;
+	/** Callback for streaming output chunks (already sanitized) */
+	onChunk?: (chunk: string) => void;
+	/** AbortSignal for cancellation */
+	signal?: AbortSignal;
+	/** Session identifier for kernel reuse */
+	sessionId?: string;
+	/** Kernel mode (session reuse vs per-call) */
+	kernelMode?: PythonKernelMode;
+	/** Restart the kernel before executing */
+	reset?: boolean;
+	/** Use shared gateway across pi instances (default: true) */
+	useSharedGateway?: boolean;
+}
+
+export interface PythonKernelExecutor {
+	execute: (code: string, options?: KernelExecuteOptions) => Promise<KernelExecuteResult>;
+}
+
+export interface PythonResult {
+	/** Combined stdout + stderr output (sanitized, possibly truncated) */
+	output: string;
+	/** Execution exit code (0 ok, 1 error, undefined if cancelled) */
+	exitCode: number | undefined;
+	/** Whether the execution was cancelled via signal */
+	cancelled: boolean;
+	/** Whether the output was truncated */
+	truncated: boolean;
+	/** Path to temp file containing full output (if output exceeded truncation threshold) */
+	fullOutputPath?: string;
+	/** Rich display outputs captured from display_data/execute_result */
+	displayOutputs: KernelDisplayOutput[];
+	/** Whether stdin was requested */
+	stdinRequested: boolean;
+}
+
+interface KernelSession {
+	id: string;
+	kernel: PythonKernel;
+	queue: Promise<void>;
+	restartCount: number;
+	dead: boolean;
+	lastUsedAt: number;
+	heartbeatTimer?: NodeJS.Timeout;
+}
+
+const kernelSessions = new Map<string, KernelSession>();
+let cachedPreludeDocs: PreludeHelper[] | null = null;
+
+export async function disposeAllKernelSessions(): Promise<void> {
+	const sessions = Array.from(kernelSessions.values());
+	await Promise.allSettled(sessions.map((session) => disposeKernelSession(session)));
+}
+
+function sanitizeChunk(text: string): string {
+	return sanitizeBinaryOutput(stripAnsi(text)).replace(/\r/g, "");
+}
+
+async function ensureKernelAvailable(cwd: string): Promise<void> {
+	const availability = await checkPythonKernelAvailability(cwd);
+	if (!availability.ok) {
+		throw new Error(availability.reason ?? "Python kernel unavailable");
+	}
+}
+
+export async function warmPythonEnvironment(
+	cwd: string,
+	sessionId?: string,
+	useSharedGateway?: boolean,
+): Promise<{ ok: boolean; reason?: string; docs: PreludeHelper[] }> {
+	try {
+		await ensureKernelAvailable(cwd);
+	} catch (err: unknown) {
+		const reason = err instanceof Error ? err.message : String(err);
+		cachedPreludeDocs = [];
+		return { ok: false, reason, docs: [] };
+	}
+	if (cachedPreludeDocs && cachedPreludeDocs.length > 0) {
+		return { ok: true, docs: cachedPreludeDocs };
+	}
+	const resolvedSessionId = sessionId ?? `session:${cwd}`;
+	try {
+		const docs = await withKernelSession(
+			resolvedSessionId,
+			cwd,
+			async (kernel) => kernel.introspectPrelude(),
+			useSharedGateway,
+		);
+		cachedPreludeDocs = docs;
+		return { ok: true, docs };
+	} catch (err: unknown) {
+		const reason = err instanceof Error ? err.message : String(err);
+		cachedPreludeDocs = [];
+		return { ok: false, reason, docs: [] };
+	}
+}
+
+export function getPreludeDocs(): PreludeHelper[] {
+	return cachedPreludeDocs ?? [];
+}
+
+export function resetPreludeDocsCache(): void {
+	cachedPreludeDocs = null;
+}
+
+async function createKernelSession(sessionId: string, cwd: string, useSharedGateway?: boolean): Promise<KernelSession> {
+	const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+	const session: KernelSession = {
+		id: sessionId,
+		kernel,
+		queue: Promise.resolve(),
+		restartCount: 0,
+		dead: false,
+		lastUsedAt: Date.now(),
+	};
+
+	session.heartbeatTimer = setInterval(async () => {
+		if (session.dead) return;
+		const ok = await session.kernel.ping().catch(() => false);
+		if (!ok) {
+			session.dead = true;
+		}
+	}, 5000);
+
+	return session;
+}
+
+async function restartKernelSession(session: KernelSession, cwd: string, useSharedGateway?: boolean): Promise<void> {
+	session.restartCount += 1;
+	if (session.restartCount > 1) {
+		throw new Error("Python kernel restarted too many times in this session");
+	}
+	try {
+		await session.kernel.shutdown();
+	} catch (err) {
+		logger.warn("Failed to shutdown crashed kernel", { error: err instanceof Error ? err.message : String(err) });
+	}
+	const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+	session.kernel = kernel;
+	session.dead = false;
+	session.lastUsedAt = Date.now();
+}
+
+async function disposeKernelSession(session: KernelSession): Promise<void> {
+	if (session.heartbeatTimer) {
+		clearInterval(session.heartbeatTimer);
+	}
+	try {
+		await session.kernel.shutdown();
+	} catch (err) {
+		logger.warn("Failed to shutdown kernel", { error: err instanceof Error ? err.message : String(err) });
+	}
+	kernelSessions.delete(session.id);
+}
+
+async function withKernelSession<T>(
+	sessionId: string,
+	cwd: string,
+	handler: (kernel: PythonKernel) => Promise<T>,
+	useSharedGateway?: boolean,
+): Promise<T> {
+	let session = kernelSessions.get(sessionId);
+	if (!session) {
+		session = await createKernelSession(sessionId, cwd, useSharedGateway);
+		kernelSessions.set(sessionId, session);
+	}
+
+	const run = async (): Promise<T> => {
+		session!.lastUsedAt = Date.now();
+		if (session!.dead || !session!.kernel.isAlive()) {
+			await restartKernelSession(session!, cwd, useSharedGateway);
+		}
+		try {
+			const result = await handler(session!.kernel);
+			session!.restartCount = 0;
+			return result;
+		} catch (err) {
+			if (!session!.dead && session!.kernel.isAlive()) {
+				throw err;
+			}
+			await restartKernelSession(session!, cwd, useSharedGateway);
+			const result = await handler(session!.kernel);
+			session!.restartCount = 0;
+			return result;
+		}
+	};
+
+	const task = session.queue.then(run, run);
+	session.queue = task.then(
+		() => undefined,
+		() => undefined,
+	);
+	return task;
+}
+
+async function executeWithKernel(
+	kernel: PythonKernelExecutor,
+	code: string,
+	options: PythonExecutorOptions | undefined,
+): Promise<PythonResult> {
+	const sink = createOutputSink(DEFAULT_MAX_BYTES, DEFAULT_MAX_BYTES * 2, options?.onChunk);
+	const writer = sink.getWriter();
+	const displayOutputs: KernelDisplayOutput[] = [];
+
+	try {
+		const result = await kernel.execute(code, {
+			signal: options?.signal,
+			timeoutMs: options?.timeout,
+			onChunk: async (text) => {
+				await writer.write(sanitizeChunk(text));
+			},
+			onDisplay: async (output) => {
+				displayOutputs.push(output);
+			},
+		});
+
+		if (result.cancelled) {
+			const secs = options?.timeout ? Math.round(options.timeout / 1000) : undefined;
+			const annotation =
+				result.timedOut && secs !== undefined ? `Command timed out after ${secs} seconds` : undefined;
+			return {
+				exitCode: undefined,
+				cancelled: true,
+				displayOutputs,
+				stdinRequested: result.stdinRequested,
+				...sink.dump(annotation),
+			};
+		}
+
+		if (result.stdinRequested) {
+			return {
+				exitCode: 1,
+				cancelled: false,
+				displayOutputs,
+				stdinRequested: true,
+				...sink.dump("Kernel requested stdin; interactive input is not supported."),
+			};
+		}
+
+		const exitCode = result.status === "ok" ? 0 : 1;
+		return {
+			exitCode,
+			cancelled: false,
+			displayOutputs,
+			stdinRequested: false,
+			...sink.dump(),
+		};
+	} catch (err) {
+		const error = err instanceof Error ? err : new Error(String(err));
+		logger.error("Python execution failed", { error: error.message });
+		throw error;
+	} finally {
+		await writer.close();
+	}
+}
+
+export async function executePythonWithKernel(
+	kernel: PythonKernelExecutor,
+	code: string,
+	options?: PythonExecutorOptions,
+): Promise<PythonResult> {
+	return await executeWithKernel(kernel, code, options);
+}
+
+export async function executePython(code: string, options?: PythonExecutorOptions): Promise<PythonResult> {
+	const cwd = options?.cwd ?? process.cwd();
+	await ensureKernelAvailable(cwd);
+
+	const kernelMode = options?.kernelMode ?? "session";
+	const useSharedGateway = options?.useSharedGateway;
+	if (kernelMode === "per-call") {
+		const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+		try {
+			return await executeWithKernel(kernel, code, options);
+		} finally {
+			await kernel.shutdown();
+		}
+	}
+
+	const sessionId = options?.sessionId ?? `session:${cwd}`;
+	if (options?.reset) {
+		const existing = kernelSessions.get(sessionId);
+		if (existing) {
+			await disposeKernelSession(existing);
+		}
+	}
+	return await withKernelSession(
+		sessionId,
+		cwd,
+		async (kernel) => executeWithKernel(kernel, code, options),
+		useSharedGateway,
+	);
+}