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

package/src/core/python-kernel-display.test.ts

@@ -0,0 +1,54 @@
+import { describe, expect, it } from "bun:test";
+import { type KernelDisplayOutput, PythonKernel } from "./python-kernel";
+
+const renderDisplay = (
+	PythonKernel as unknown as {
+		prototype: {
+			renderDisplay: (content: Record<string, unknown>) => {
+				text: string;
+				outputs: KernelDisplayOutput[];
+			};
+		};
+	}
+).prototype.renderDisplay;
+
+describe("PythonKernel display rendering", () => {
+	it("normalizes text/plain output and returns no display outputs", () => {
+		const { text, outputs } = renderDisplay.call({} as PythonKernel, {
+			data: { "text/plain": "hello" },
+		});
+
+		expect(text).toBe("hello\n");
+		expect(outputs).toHaveLength(0);
+	});
+
+	it("collects image and json display outputs without text", () => {
+		const { text, outputs } = renderDisplay.call({} as PythonKernel, {
+			data: { "image/png": "abc", "application/json": { foo: "bar" } },
+		});
+
+		expect(text).toBe("");
+		expect(outputs).toEqual([
+			{ type: "image", data: "abc", mimeType: "image/png" },
+			{ type: "json", data: { foo: "bar" } },
+		]);
+	});
+
+	it("converts text/html to markdown", () => {
+		const { text, outputs } = renderDisplay.call({} as PythonKernel, {
+			data: { "text/html": "<p><strong>Hello</strong></p>" },
+		});
+
+		expect(outputs).toHaveLength(0);
+		expect(text).toBe("**Hello**\n");
+	});
+
+	it("combines text/plain with json output", () => {
+		const { text, outputs } = renderDisplay.call({} as PythonKernel, {
+			data: { "text/plain": "value", "application/json": { ok: true } },
+		});
+
+		expect(text).toBe("value\n");
+		expect(outputs).toEqual([{ type: "json", data: { ok: true } }]);
+	});
+});

package/src/core/python-kernel-env.test.ts

@@ -0,0 +1,138 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
+import * as shell from "../utils/shell";
+import * as shellSnapshot from "../utils/shell-snapshot";
+import { PythonKernel } from "./python-kernel";
+import { PYTHON_PRELUDE } from "./python-prelude";
+
+class FakeWebSocket {
+	static OPEN = 1;
+	static CLOSED = 3;
+	readyState = FakeWebSocket.OPEN;
+	binaryType = "arraybuffer";
+	url: string;
+	onopen?: () => void;
+	onerror?: (event: unknown) => void;
+	onclose?: () => void;
+	onmessage?: (event: { data: ArrayBuffer }) => void;
+
+	constructor(url: string) {
+		this.url = url;
+		queueMicrotask(() => {
+			this.onopen?.();
+		});
+	}
+
+	send(_data: ArrayBuffer) {}
+
+	close() {
+		this.readyState = FakeWebSocket.CLOSED;
+		this.onclose?.();
+	}
+}
+
+describe("PythonKernel.start (local gateway)", () => {
+	const originalEnv = { ...process.env };
+	const originalFetch = globalThis.fetch;
+	const originalWebSocket = globalThis.WebSocket;
+
+	beforeEach(() => {
+		process.env.BUN_ENV = "test";
+		delete process.env.OMP_PYTHON_GATEWAY_URL;
+		delete process.env.OMP_PYTHON_GATEWAY_TOKEN;
+		globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
+	});
+
+	afterEach(() => {
+		for (const key of Object.keys(process.env)) {
+			if (!(key in originalEnv)) {
+				delete process.env[key];
+			}
+		}
+		for (const [key, value] of Object.entries(originalEnv)) {
+			process.env[key] = value;
+		}
+		globalThis.fetch = originalFetch;
+		globalThis.WebSocket = originalWebSocket;
+		vi.restoreAllMocks();
+	});
+
+	it("filters environment variables before spawning gateway", async () => {
+		const fetchSpy = vi.fn(async (input: string | URL, init?: RequestInit) => {
+			const url = typeof input === "string" ? input : input.toString();
+			if (url.endsWith("/api/kernelspecs")) {
+				return new Response(JSON.stringify({}), { status: 200 });
+			}
+			if (url.endsWith("/api/kernels") && init?.method === "POST") {
+				return new Response(JSON.stringify({ id: "kernel-1" }), { status: 201 });
+			}
+			return new Response("", { status: 200 });
+		});
+		globalThis.fetch = fetchSpy as unknown as typeof fetch;
+
+		const shellSpy = vi.spyOn(shell, "getShellConfig").mockResolvedValue({
+			shell: "/bin/bash",
+			args: ["-lc"],
+			env: {
+				PATH: "/bin",
+				HOME: "/home/test",
+				OPENAI_API_KEY: "secret",
+				UNSAFE_TOKEN: "nope",
+				OMP_CUSTOM: "1",
+				LC_ALL: "en_US.UTF-8",
+			},
+			prefix: undefined,
+		});
+		const snapshotSpy = vi.spyOn(shellSnapshot, "getOrCreateSnapshot").mockResolvedValue(null);
+		const whichSpy = vi.spyOn(Bun, "which").mockReturnValue("/usr/bin/python");
+
+		let spawnEnv: Record<string, string | undefined> | undefined;
+		let spawnArgs: string[] | undefined;
+		const spawnSpy = vi.spyOn(Bun, "spawn").mockImplementation(((...args: unknown[]) => {
+			const [cmd, options] = args as [string[] | { cmd: string[] }, { env?: Record<string, string | undefined> }?];
+			spawnArgs = Array.isArray(cmd) ? cmd : cmd.cmd;
+			spawnEnv = options?.env;
+			return { pid: 1234, exited: Promise.resolve(0) } as unknown as Bun.Subprocess;
+		}) as unknown as typeof Bun.spawn);
+
+		const executeSpy = vi
+			.spyOn(PythonKernel.prototype, "execute")
+			.mockResolvedValue({ status: "ok", cancelled: false, timedOut: false, stdinRequested: false });
+
+		const kernel = await PythonKernel.start({ cwd: "/tmp/project", env: { CUSTOM_VAR: "ok" } });
+
+		const createCall = fetchSpy.mock.calls.find(([input, init]) => {
+			const url = typeof input === "string" ? input : input.toString();
+			return url.endsWith("/api/kernels") && init?.method === "POST";
+		});
+		expect(createCall).toBeDefined();
+		if (createCall) {
+			expect(JSON.parse(String(createCall[1]?.body ?? "{}"))).toEqual({ name: "python3" });
+		}
+
+		expect(spawnArgs).toContain("kernel_gateway");
+		expect(spawnEnv?.PATH).toBe("/bin");
+		expect(spawnEnv?.HOME).toBe("/home/test");
+		expect(spawnEnv?.OMP_CUSTOM).toBe("1");
+		expect(spawnEnv?.LC_ALL).toBe("en_US.UTF-8");
+		expect(spawnEnv?.CUSTOM_VAR).toBe("ok");
+		expect(spawnEnv?.OPENAI_API_KEY).toBeUndefined();
+		expect(spawnEnv?.UNSAFE_TOKEN).toBeUndefined();
+		expect(spawnEnv?.PYTHONPATH).toBe("/tmp/project");
+
+		expect(executeSpy).toHaveBeenCalledWith(
+			PYTHON_PRELUDE,
+			expect.objectContaining({
+				silent: true,
+				storeHistory: false,
+			}),
+		);
+
+		await kernel.shutdown();
+
+		shellSpy.mockRestore();
+		snapshotSpy.mockRestore();
+		whichSpy.mockRestore();
+		spawnSpy.mockRestore();
+		executeSpy.mockRestore();
+	});
+});

package/src/core/python-kernel-session.test.ts

@@ -0,0 +1,87 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import { disposeAllKernelSessions, executePython } from "./python-executor";
+import type { KernelExecuteOptions, KernelExecuteResult } from "./python-kernel";
+import { PythonKernel } from "./python-kernel";
+
+class FakeKernel {
+	executeCalls = 0;
+	shutdownCalls = 0;
+	alive = true;
+	readonly id: string;
+
+	constructor(id: string) {
+		this.id = id;
+	}
+
+	async execute(_code: string, options?: KernelExecuteOptions): Promise<KernelExecuteResult> {
+		this.executeCalls += 1;
+		options?.onChunk?.("ok\n");
+		return { status: "ok", cancelled: false, timedOut: false, stdinRequested: false };
+	}
+
+	async shutdown(): Promise<void> {
+		this.shutdownCalls += 1;
+		this.alive = false;
+	}
+
+	isAlive(): boolean {
+		return this.alive;
+	}
+
+	async ping(): Promise<boolean> {
+		return this.alive;
+	}
+}
+
+describe("executePython kernel reuse", () => {
+	const originalStart = PythonKernel.start;
+	let startCalls = 0;
+	let kernels: FakeKernel[] = [];
+
+	beforeEach(() => {
+		process.env.OMP_PYTHON_SKIP_CHECK = "1";
+		startCalls = 0;
+		kernels = [];
+		PythonKernel.start = (async () => {
+			startCalls += 1;
+			const kernel = new FakeKernel(`kernel-${startCalls}`);
+			kernels.push(kernel);
+			return kernel as unknown as PythonKernel;
+		}) as typeof PythonKernel.start;
+	});
+
+	afterEach(async () => {
+		PythonKernel.start = originalStart;
+		await disposeAllKernelSessions();
+	});
+
+	it("reuses kernels for session mode", async () => {
+		await executePython("print('one')", { cwd: "/tmp", sessionId: "session-a", kernelMode: "session" });
+		await executePython("print('two')", { cwd: "/tmp", sessionId: "session-a", kernelMode: "session" });
+
+		expect(startCalls).toBe(1);
+		expect(kernels[0]?.executeCalls).toBe(2);
+	});
+
+	it("creates and disposes per-call kernels", async () => {
+		await executePython("print('one')", { cwd: "/tmp", kernelMode: "per-call" });
+		await executePython("print('two')", { cwd: "/tmp", kernelMode: "per-call" });
+
+		expect(startCalls).toBe(2);
+		expect(kernels[0]?.shutdownCalls).toBe(1);
+		expect(kernels[1]?.shutdownCalls).toBe(1);
+	});
+
+	it("resets the session kernel when requested", async () => {
+		await executePython("print('one')", { cwd: "/tmp", sessionId: "session-b", kernelMode: "session" });
+		await executePython("print('two')", {
+			cwd: "/tmp",
+			sessionId: "session-b",
+			kernelMode: "session",
+			reset: true,
+		});
+
+		expect(startCalls).toBe(2);
+		expect(kernels[0]?.shutdownCalls).toBe(1);
+	});
+});

package/src/core/python-kernel-ws.test.ts

@@ -0,0 +1,104 @@
+import { describe, expect, it } from "bun:test";
+import { deserializeWebSocketMessage, type JupyterMessage, serializeWebSocketMessage } from "./python-kernel";
+
+const encoder = new TextEncoder();
+
+function buildFrame(message: Omit<JupyterMessage, "buffers">, buffers: Uint8Array[] = []): ArrayBuffer {
+	const msgBytes = encoder.encode(JSON.stringify(message));
+	const offsetCount = 1 + buffers.length;
+	const headerSize = 4 + offsetCount * 4;
+
+	let totalSize = headerSize + msgBytes.length;
+	for (const buffer of buffers) {
+		totalSize += buffer.length;
+	}
+
+	const frame = new ArrayBuffer(totalSize);
+	const view = new DataView(frame);
+	const bytes = new Uint8Array(frame);
+
+	view.setUint32(0, offsetCount, true);
+	view.setUint32(4, headerSize, true);
+	bytes.set(msgBytes, headerSize);
+
+	let offset = headerSize + msgBytes.length;
+	for (let i = 0; i < buffers.length; i++) {
+		view.setUint32(4 + (i + 1) * 4, offset, true);
+		bytes.set(buffers[i], offset);
+		offset += buffers[i].length;
+	}
+
+	return frame;
+}
+
+describe("deserializeWebSocketMessage", () => {
+	it("parses offset tables and buffers", () => {
+		const message = {
+			channel: "iopub",
+			header: {
+				msg_id: "msg-1",
+				session: "session-1",
+				username: "omp",
+				date: "2024-01-01T00:00:00Z",
+				msg_type: "stream",
+				version: "5.5",
+			},
+			parent_header: {},
+			metadata: {},
+			content: { text: "hello" },
+		};
+		const buffer = new Uint8Array([1, 2, 3]);
+		const frame = buildFrame(message, [buffer]);
+
+		const parsed = deserializeWebSocketMessage(frame);
+
+		expect(parsed).not.toBeNull();
+		expect(parsed?.header.msg_id).toBe("msg-1");
+		expect(parsed?.content).toEqual({ text: "hello" });
+		expect(parsed?.buffers?.[0]).toEqual(buffer);
+	});
+
+	it("returns null for invalid frames", () => {
+		const headerSize = 8;
+		const bytes = encoder.encode("not-json");
+		const frame = new ArrayBuffer(headerSize + bytes.length);
+		const view = new DataView(frame);
+		const data = new Uint8Array(frame);
+		view.setUint32(0, 1, true);
+		view.setUint32(4, headerSize, true);
+		data.set(bytes, headerSize);
+
+		expect(deserializeWebSocketMessage(frame)).toBeNull();
+		const emptyFrame = new ArrayBuffer(4);
+		new DataView(emptyFrame).setUint32(0, 0, true);
+		expect(deserializeWebSocketMessage(emptyFrame)).toBeNull();
+	});
+});
+
+describe("serializeWebSocketMessage", () => {
+	it("round trips message payloads", () => {
+		const message: JupyterMessage = {
+			channel: "shell",
+			header: {
+				msg_id: "msg-2",
+				session: "session-2",
+				username: "omp",
+				date: "2024-02-01T00:00:00Z",
+				msg_type: "execute_request",
+				version: "5.5",
+			},
+			parent_header: { parent: "root" },
+			metadata: { tag: "meta" },
+			content: { code: "print('hi')" },
+			buffers: [new Uint8Array([9, 8, 7])],
+		};
+
+		const frame = serializeWebSocketMessage(message);
+		const parsed = deserializeWebSocketMessage(frame);
+
+		expect(parsed).not.toBeNull();
+		expect(parsed?.header.msg_type).toBe("execute_request");
+		expect(parsed?.content).toEqual({ code: "print('hi')" });
+		expect(parsed?.buffers?.[0]).toEqual(new Uint8Array([9, 8, 7]));
+	});
+});

package/src/core/python-kernel.lifecycle.test.ts

@@ -0,0 +1,249 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { Subprocess } from "bun";
+import { PythonKernel } from "./python-kernel";
+
+type SpawnOptions = Parameters<typeof Bun.spawn>[1];
+
+type FetchCall = { url: string; init?: RequestInit };
+
+type FetchResponse = {
+	ok: boolean;
+	status: number;
+	json: () => Promise<unknown>;
+	text: () => Promise<string>;
+};
+
+type MockEnvironment = {
+	fetchCalls: FetchCall[];
+	spawnCalls: { cmd: string[]; options: SpawnOptions }[];
+};
+
+type MessageEventPayload = { data: ArrayBuffer };
+
+type WebSocketHandler = (event: unknown) => void;
+
+type WebSocketMessageHandler = (event: MessageEventPayload) => void;
+
+class FakeWebSocket {
+	static OPEN = 1;
+	static CLOSED = 3;
+	static instances: FakeWebSocket[] = [];
+
+	readyState = FakeWebSocket.OPEN;
+	binaryType = "arraybuffer";
+	url: string;
+	sent: ArrayBuffer[] = [];
+
+	onopen: WebSocketHandler | null = null;
+	onerror: WebSocketHandler | null = null;
+	onclose: WebSocketHandler | null = null;
+	onmessage: WebSocketMessageHandler | null = null;
+
+	constructor(url: string) {
+		this.url = url;
+		FakeWebSocket.instances.push(this);
+		queueMicrotask(() => {
+			this.onopen?.(undefined);
+		});
+	}
+
+	send(data: ArrayBuffer): void {
+		this.sent.push(data);
+	}
+
+	close(): void {
+		this.readyState = FakeWebSocket.CLOSED;
+		this.onclose?.(undefined);
+	}
+}
+
+const createResponse = (options: { ok: boolean; status?: number; json?: unknown; text?: string }): FetchResponse => {
+	return {
+		ok: options.ok,
+		status: options.status ?? (options.ok ? 200 : 500),
+		json: async () => options.json ?? {},
+		text: async () => options.text ?? "",
+	};
+};
+
+const createTempDir = () => mkdtempSync(join(tmpdir(), "omp-python-kernel-"));
+
+const createFakeProcess = (): Subprocess => {
+	const exited = new Promise<number>(() => undefined);
+	return { pid: 999999, exited } as Subprocess;
+};
+
+describe("PythonKernel gateway lifecycle", () => {
+	const originalFetch = globalThis.fetch;
+	const originalWebSocket = globalThis.WebSocket;
+	const originalSpawn = Bun.spawn;
+	const originalSleep = Bun.sleep;
+	const originalWhich = Bun.which;
+	const originalExecute = PythonKernel.prototype.execute;
+	const originalGatewayUrl = process.env.OMP_PYTHON_GATEWAY_URL;
+	const originalGatewayToken = process.env.OMP_PYTHON_GATEWAY_TOKEN;
+	const originalBunEnv = process.env.BUN_ENV;
+
+	let tempDir: string;
+	let env: MockEnvironment;
+
+	beforeEach(() => {
+		tempDir = createTempDir();
+		env = { fetchCalls: [], spawnCalls: [] };
+
+		process.env.BUN_ENV = "test";
+		delete process.env.OMP_PYTHON_GATEWAY_URL;
+		delete process.env.OMP_PYTHON_GATEWAY_TOKEN;
+
+		FakeWebSocket.instances = [];
+		globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket;
+
+		Bun.spawn = ((cmd: string[] | string, options?: SpawnOptions) => {
+			const normalized = Array.isArray(cmd) ? cmd : [cmd];
+			env.spawnCalls.push({ cmd: normalized, options: options ?? {} });
+			return createFakeProcess();
+		}) as typeof Bun.spawn;
+
+		Bun.sleep = (async () => undefined) as typeof Bun.sleep;
+
+		Bun.which = (() => "/usr/bin/python") as typeof Bun.which;
+
+		Object.defineProperty(PythonKernel.prototype, "execute", {
+			value: (async () => ({
+				status: "ok",
+				cancelled: false,
+				timedOut: false,
+				stdinRequested: false,
+			})) as typeof PythonKernel.prototype.execute,
+			configurable: true,
+		});
+	});
+
+	afterEach(() => {
+		if (tempDir) {
+			rmSync(tempDir, { recursive: true, force: true });
+		}
+
+		if (originalBunEnv === undefined) {
+			delete process.env.BUN_ENV;
+		} else {
+			process.env.BUN_ENV = originalBunEnv;
+		}
+		if (originalGatewayUrl === undefined) {
+			delete process.env.OMP_PYTHON_GATEWAY_URL;
+		} else {
+			process.env.OMP_PYTHON_GATEWAY_URL = originalGatewayUrl;
+		}
+		if (originalGatewayToken === undefined) {
+			delete process.env.OMP_PYTHON_GATEWAY_TOKEN;
+		} else {
+			process.env.OMP_PYTHON_GATEWAY_TOKEN = originalGatewayToken;
+		}
+
+		globalThis.fetch = originalFetch;
+		globalThis.WebSocket = originalWebSocket;
+
+		Bun.spawn = originalSpawn;
+		Bun.sleep = originalSleep;
+		Bun.which = originalWhich;
+		Object.defineProperty(PythonKernel.prototype, "execute", { value: originalExecute, configurable: true });
+	});
+
+	it("starts local gateway, polls readiness, interrupts, and shuts down", async () => {
+		let kernelspecAttempts = 0;
+		globalThis.fetch = (async (input: string | URL, init?: RequestInit) => {
+			const url = String(input);
+			env.fetchCalls.push({ url, init });
+
+			if (url.endsWith("/api/kernelspecs")) {
+				kernelspecAttempts += 1;
+				const ok = kernelspecAttempts >= 2;
+				return createResponse({ ok }) as unknown as Response;
+			}
+
+			if (url.endsWith("/api/kernels") && init?.method === "POST") {
+				return createResponse({ ok: true, json: { id: "kernel-123" } }) as unknown as Response;
+			}
+
+			return createResponse({ ok: true }) as unknown as Response;
+		}) as typeof fetch;
+
+		const kernel = await PythonKernel.start({ cwd: tempDir, useSharedGateway: false });
+
+		expect(env.spawnCalls).toHaveLength(1);
+		expect(env.spawnCalls[0].cmd).toEqual(
+			expect.arrayContaining([
+				"-m",
+				"kernel_gateway",
+				"--KernelGatewayApp.allow_origin=*",
+				"--JupyterApp.answer_yes=true",
+			]),
+		);
+		expect(env.fetchCalls.filter((call) => call.url.endsWith("/api/kernelspecs"))).toHaveLength(2);
+		expect(env.fetchCalls.some((call) => call.url.endsWith("/api/kernels") && call.init?.method === "POST")).toBe(
+			true,
+		);
+
+		await kernel.interrupt();
+		expect(env.fetchCalls.some((call) => call.url.includes("/interrupt") && call.init?.method === "POST")).toBe(true);
+		expect(FakeWebSocket.instances[0]?.sent.length).toBe(1);
+
+		await kernel.shutdown();
+		expect(env.fetchCalls.some((call) => call.init?.method === "DELETE")).toBe(true);
+		expect(kernel.isAlive()).toBe(false);
+	});
+
+	it("throws when gateway readiness never succeeds", async () => {
+		const originalNow = Date.now;
+		let now = 0;
+		Date.now = () => {
+			now += 1000;
+			return now;
+		};
+
+		try {
+			globalThis.fetch = (async (input: string | URL, init?: RequestInit) => {
+				const url = String(input);
+				env.fetchCalls.push({ url, init });
+				if (url.endsWith("/api/kernelspecs")) {
+					return createResponse({ ok: false, status: 503 }) as unknown as Response;
+				}
+				return createResponse({ ok: true }) as unknown as Response;
+			}) as typeof fetch;
+
+			await expect(PythonKernel.start({ cwd: tempDir, useSharedGateway: false })).rejects.toThrow(
+				"Kernel gateway failed to start",
+			);
+			expect(env.spawnCalls).toHaveLength(3);
+		} finally {
+			Date.now = originalNow;
+		}
+	});
+
+	it("does not throw when shutdown API fails", async () => {
+		let kernelspecAttempts = 0;
+		globalThis.fetch = (async (input: string | URL, init?: RequestInit) => {
+			const url = String(input);
+			env.fetchCalls.push({ url, init });
+			if (url.endsWith("/api/kernelspecs")) {
+				kernelspecAttempts += 1;
+				const ok = kernelspecAttempts >= 1;
+				return createResponse({ ok }) as unknown as Response;
+			}
+			if (url.endsWith("/api/kernels") && init?.method === "POST") {
+				return createResponse({ ok: true, json: { id: "kernel-456" } }) as unknown as Response;
+			}
+			if (init?.method === "DELETE") {
+				throw new Error("delete failed");
+			}
+			return createResponse({ ok: true }) as unknown as Response;
+		}) as typeof fetch;
+
+		const kernel = await PythonKernel.start({ cwd: tempDir });
+
+		await expect(kernel.shutdown()).resolves.toBeUndefined();
+	});
+});