import { chmodSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { cleanupSessionResources } from "@earendil-works/pi-ai";
import { Value } from "typebox/value";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import type { ExtensionContext } from "../src/core/extensions/types.js";
import type { KernelBootstrapProgressHandler } from "../src/core/kernel/bootstrap.js";
import {
	type ExecuteResult,
	KernelBusyAfterInterruptError,
	type KernelClient,
	ReplKernelManager,
} from "../src/core/kernel/index.js";
import { createIpythonToolDefinition, IpythonKernelProvisioner } from "../src/core/tools/ipython.js";

let tempDir = "";

function writeFakePython(opts: { sleepSeconds?: number } = {}): { python: string; countRuns: () => number } {
	const python = join(tempDir, "python");
	const countFile = join(tempDir, "runs");
	writeFileSync(
		python,
		[
			"#!/bin/sh",
			`echo run >> "${countFile}"`,
			...(opts.sleepSeconds ? [`sleep ${opts.sleepSeconds}`] : []),
			"exit 42",
			"",
		].join("\n"),
	);
	chmodSync(python, 0o755);
	const countRuns = () => {
		try {
			return readFileSync(countFile, "utf8").split("\n").filter(Boolean).length;
		} catch {
			return 0;
		}
	};
	return { python, countRuns };
}

/** The declared AgentToolResult shape omits `isError`; the tool still sets it. */
function isErrorResult(result: { content: unknown }): boolean {
	return (result as { isError?: boolean }).isError === true;
}

function okExecuteResult(): ExecuteResult {
	return { stdout: "ok", stderr: "", status: "ok", durationMs: 1 };
}

function createBusyKernelContext(
	select: (title: string, options: string[]) => Promise<string | undefined>,
	options: { throwWorkingMessage?: boolean } = {},
): {
	ctx: ExtensionContext;
	setWorkingMessage: ReturnType<typeof vi.fn>;
} {
	const setWorkingMessage = vi.fn(() => {
		if (options.throwWorkingMessage) {
			throw new Error("stale UI context");
		}
	});
	const ctx = {
		hasUI: true,
		ui: {
			select,
			setWorkingMessage,
		},
	} as unknown as ExtensionContext;
	return { ctx, setWorkingMessage };
}

function writeFakeReplRuntime(
	markerPath: string,
	options: { gatedExecute?: { startedPath: string; gatePath: string } } = {},
): string {
	const python = join(tempDir, "python-repl");
	const refuse = `emit({ event: "error", id: request.id, ename: "RuntimeError", evalue: "bootstrap refused", traceback: [] });
		emit({ event: "done", id: request.id, status: "error" });`;
	const executeBranch = options.gatedExecute
		? `fs.writeFileSync(${JSON.stringify(options.gatedExecute.startedPath)}, "1");
		const gate = setInterval(() => {
			if (!fs.existsSync(${JSON.stringify(options.gatedExecute.gatePath)})) return;
			clearInterval(gate);
			${refuse}
		}, 10);`
		: refuse;
	writeFileSync(
		python,
		`#!/usr/bin/env node
const fs = require("node:fs");
const readline = require("node:readline");
const emit = (event) => process.stdout.write(JSON.stringify(event) + "\\n");
emit({ event: "ready", protocol: 3, python: process.version });
const input = readline.createInterface({ input: process.stdin });
input.on("line", (line) => {
	const request = JSON.parse(line);
	if (request.type === "restore") {
		emit({ event: "done", id: request.id, status: "ok", restored: [], failed: [] });
		return;
	}
	if (request.type === "snapshot") {
		setTimeout(() => {
			fs.writeFileSync(${JSON.stringify(markerPath)}, "1");
			emit({ event: "done", id: request.id, status: "ok", saved: [], skipped: [], bytes: 0 });
		}, 150);
		return;
	}
	if (request.type === "execute") {
		${executeBranch}
		return;
	}
	if (request.type === "shutdown") {
		emit({ event: "done", id: request.id, status: "ok" });
		process.exit(0);
	}
});
`,
	);
	chmodSync(python, 0o755);
	return python;
}

/** One Python skill's runtime info, as the session passes it to the provisioner. */
function pythonSkillInfo(name: string): {
	name: string;
	importName: string;
	packagePath: string;
	pyprojectPath: string;
} {
	const importName = name.replaceAll("-", "_");
	return {
		name,
		importName,
		packagePath: join(tempDir, name),
		pyprojectPath: join(tempDir, name, "pyproject.toml"),
	};
}

/**
 * A runtime stand-in that accepts the bootstrap and answers `signatures`, recording every
 * request it receives. Without the capability it answers the way an older runtime does.
 */
function writeFakeDigestRuntime(options: { capabilities?: string[]; digest?: string } = {}): {
	python: string;
	requested: () => string[][];
} {
	const capable = options.capabilities?.includes("signatures") ?? false;
	const python = join(tempDir, "python-digest");
	const requestLog = join(tempDir, "digest-requests");
	const ready = { event: "ready", protocol: 3, python: "3.13.11", capabilities: options.capabilities ?? [] };
	const digest = options.digest ?? "# skill surface (generated)\n";
	const answer = capable
		? `emit({ event: "done", id: request.id, status: "ok", calls: 1, failed: [], digest: ${JSON.stringify(digest)} });`
		: `emit({ event: "error", id: null, ename: "ProtocolError", evalue: "unknown request type: signatures", traceback: [] });
			emit({ event: "done", id: request.id, status: "error", reason: "unknown request type: signatures" });`;
	writeFileSync(
		python,
		`#!/usr/bin/env node
const fs = require("node:fs");
const readline = require("node:readline");
const emit = (event) => process.stdout.write(JSON.stringify(event) + "\\n");
emit(${JSON.stringify(ready)});
const input = readline.createInterface({ input: process.stdin });
input.on("line", (line) => {
	const request = JSON.parse(line);
	if (request.type === "signatures") {
		fs.appendFileSync(${JSON.stringify(requestLog)}, JSON.stringify(request.skills) + "\\n");
		${answer}
		return;
	}
	if (request.type === "execute") {
		emit({ event: "done", id: request.id, status: "ok" });
		return;
	}
	if (request.type === "snapshot") {
		emit({ event: "done", id: request.id, status: "ok", saved: [], skipped: [], bytes: 0 });
		return;
	}
	if (request.type === "restore") {
		emit({ event: "done", id: request.id, status: "ok", restored: [], failed: [] });
		return;
	}
	if (request.type === "shutdown") {
		emit({ event: "done", id: request.id, status: "ok" });
		process.exit(0);
	}
});
`,
	);
	chmodSync(python, 0o755);
	const requested = () => {
		try {
			return readFileSync(requestLog, "utf8")
				.split("\n")
				.filter(Boolean)
				.map((line) => JSON.parse(line) as string[]);
		} catch {
			return []; // No request has been recorded yet.
		}
	};
	return { python, requested };
}

describe("IpythonKernelProvisioner", () => {
	beforeEach(() => {
		tempDir = mkdtempSync(join(tmpdir(), "prime-agent-provisioner-"));
	});

	afterEach(() => {
		if (tempDir) {
			rmSync(tempDir, { recursive: true, force: true });
			tempDir = "";
		}
	});

	it("does not surface a failed startup before the kernel's final snapshot flush finished", async () => {
		const marker = join(tempDir, "snapshot-flushed");
		const snapshotDir = join(tempDir, "snapshots");
		mkdirSync(snapshotDir, { recursive: true });
		const python = writeFakeReplRuntime(marker);
		const provisioner = new IpythonKernelProvisioner(tempDir, { python, snapshotDir });
		try {
			await expect(provisioner.ensure()).rejects.toThrow(/Failed to initialize rlm runtime/);
			// The failed kernel's teardown (final snapshot flush included) completed
			// before the failure surfaced, so a replacement provisioner gated on this
			// one cannot race the still-flushing kernel over the same snapshot files.
			expect(existsSync(marker)).toBe(true);
		} finally {
			await provisioner.dispose();
		}
	});

	it("memoizes concurrent ensure() calls into one startup", async () => {
		const { python, countRuns } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });

		const [a, b] = await Promise.allSettled([provisioner.ensure(), provisioner.ensure()]);
		expect(a.status).toBe("rejected");
		expect(b.status).toBe("rejected");
		expect(countRuns()).toBe(1);
	});

	it("retries after a failed startup instead of caching the rejection", async () => {
		const { python, countRuns } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });

		await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/);
		await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/);
		expect(countRuns()).toBe(2);
	});

	it("prewarm() swallows the failure and the next ensure() starts fresh", async () => {
		const { python, countRuns } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });

		provisioner.prewarm();
		expect(provisioner.manager).toBeUndefined();

		// Once the prewarm startup settles, ensure() must launch a second attempt.
		await vi.waitFor(async () => {
			await expect(provisioner.ensure()).rejects.toThrow();
			expect(countRuns()).toBeGreaterThanOrEqual(2);
		});
	});

	it("replays the current startup stage to listeners attaching mid-flight", async () => {
		const { python } = writeFakePython({ sleepSeconds: 1 });
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });

		provisioner.prewarm();
		const messages: string[] = [];
		const joined = provisioner.ensure((message) => messages.push(message));
		expect(messages).toContain("Starting Python kernel...");
		await expect(joined).rejects.toThrow();
	});

	it("dispose() settles a startup that is still in flight", async () => {
		const { python } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });

		provisioner.prewarm();
		await provisioner.dispose();
		expect(provisioner.manager).toBeUndefined();
	});

	it("skips the snapshot when dispose({ snapshot: false }) aborts a startup in flight", async () => {
		const marker = join(tempDir, "snapshot-flushed");
		const executeStarted = join(tempDir, "execute-started");
		const executeGate = join(tempDir, "execute-gate");
		const snapshotDir = join(tempDir, "snapshots");
		mkdirSync(snapshotDir, { recursive: true });
		const python = writeFakeReplRuntime(marker, {
			gatedExecute: { startedPath: executeStarted, gatePath: executeGate },
		});
		const provisioner = new IpythonKernelProvisioner(tempDir, { python, snapshotDir });

		const started = provisioner.ensure().catch(() => undefined);
		await vi.waitFor(() => expect(existsSync(executeStarted)).toBe(true));
		const disposed = provisioner.dispose({ snapshot: false });
		writeFileSync(executeGate, "1");
		await Promise.all([disposed, started]);
		expect(existsSync(marker)).toBe(false);
	});

	it("dispose({ snapshot: false }) skips the kernel's final snapshot flush", async () => {
		const { python } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });
		const shutdown = vi.fn(async () => {});
		Reflect.set(provisioner, "managerPromise", Promise.resolve({ shutdown }));

		await provisioner.dispose({ snapshot: false });
		expect(shutdown).toHaveBeenCalledWith({ snapshot: false, drainHostRequests: true });
	});

	it("dispose() before the boot slot prevents the kernel from spawning", async () => {
		const { python, countRuns } = writeFakePython();
		let release: () => void = () => {};
		const gate = new Promise<void>((r) => {
			release = r;
		});
		const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate });

		const started = provisioner.ensure().catch(() => {});
		const disposed = provisioner.dispose(); // aborts while the boot waits on readyGate
		release();
		await Promise.all([started, disposed]);
		expect(countRuns()).toBe(0); // disposed boot must never spawn a kernel
	});

	it("aborting the startup owner before the boot slot prevents the kernel from spawning", async () => {
		const { python, countRuns } = writeFakePython();
		let release: () => void = () => {};
		const gate = new Promise<void>((r) => {
			release = r;
		});
		const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate });
		const controller = new AbortController();

		const started = provisioner.ensure(undefined, controller.signal);
		controller.abort();
		await expect(started).rejects.toThrow("Python execution aborted");
		release();
		await new Promise((r) => setTimeout(r, 50));

		expect(countRuns()).toBe(0);
		expect(provisioner.manager).toBeUndefined();
	});

	it("waits for readyGate before starting the kernel", async () => {
		const { python, countRuns } = writeFakePython();
		let release: () => void = () => {};
		const gate = new Promise<void>((r) => {
			release = r;
		});
		const provisioner = new IpythonKernelProvisioner(tempDir, { python, readyGate: gate });

		const started = provisioner.ensure().catch(() => {});
		await new Promise((r) => setTimeout(r, 50));
		expect(countRuns()).toBe(0); // gated: must not spawn the kernel yet

		release();
		await started;
		expect(countRuns()).toBe(1);
	});

	it("namespace maintenance returns null when no kernel is running", async () => {
		const provisioner = new IpythonKernelProvisioner(tempDir, {});
		expect(await provisioner.listNamespaceNames()).toBeNull();
		expect(await provisioner.pruneOversizedVariables()).toBeNull();
	});

	it("does not dispose a running kernel when an ensure caller is aborted", async () => {
		const provisioner = new IpythonKernelProvisioner(tempDir, {});
		const dispose = vi.fn(async () => {});
		const manager = { dispose, isRunning: true } as unknown as KernelClient;
		Object.assign(
			provisioner as unknown as {
				managerPromise: Promise<KernelClient>;
				startedManager: KernelClient;
			},
			{
				managerPromise: Promise.resolve(manager),
				startedManager: manager,
			},
		);
		const controller = new AbortController();
		controller.abort();

		await expect(provisioner.ensure(undefined, controller.signal)).rejects.toThrow("Python execution aborted");
		expect(dispose).not.toHaveBeenCalled();
		expect(provisioner.manager).toBe(manager);
	});

	function primeKernelMemo(provisioner: IpythonKernelProvisioner, manager: KernelClient) {
		Object.assign(provisioner as unknown as { managerPromise: Promise<KernelClient>; startedManager: KernelClient }, {
			managerPromise: Promise.resolve(manager),
			startedManager: manager,
		});
	}

	it("drops a dead kernel memo so ensure() restarts instead of reusing it", async () => {
		const { python, countRuns } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, { python });
		primeKernelMemo(provisioner, { isRunning: false, isDefunct: true } as unknown as KernelClient);

		await expect(provisioner.ensure()).rejects.toThrow(/Kernel exited before ready/);
		expect(countRuns()).toBe(1);
	});

	it("keeps the memo for a kernel that is repairing itself, not defunct", async () => {
		const { countRuns } = writeFakePython();
		const provisioner = new IpythonKernelProvisioner(tempDir, {});
		const repairing = { isRunning: false, isDefunct: false } as unknown as KernelClient;
		primeKernelMemo(provisioner, repairing);

		// A second provisioner kernel during protocol repair would split the snapshot dir.
		await expect(provisioner.ensure()).resolves.toBe(repairing);
		expect(countRuns()).toBe(0);
	});

	it("removes startup progress listeners when an ensure caller is aborted", async () => {
		const provisioner = new IpythonKernelProvisioner(tempDir, {});
		Object.assign(
			provisioner as unknown as {
				managerPromise: Promise<KernelClient>;
			},
			{
				managerPromise: new Promise<KernelClient>(() => {}),
			},
		);
		const controller = new AbortController();
		const onProgress = vi.fn();

		const ensurePromise = provisioner.ensure(onProgress, controller.signal).catch(() => undefined);
		controller.abort();
		await ensurePromise;

		const internals = provisioner as unknown as {
			startupListeners: Set<KernelBootstrapProgressHandler>;
		};
		expect(internals.startupListeners.has(onProgress)).toBe(false);
	});

	it("surfaces backgroundOutput in details without changing model content", async () => {
		const execute = vi
			.fn<KernelClient["execute"]>()
			.mockResolvedValueOnce({ ...okExecuteResult(), backgroundOutput: "bg-line" });
		const manager = { execute } as unknown as KernelClient;
		const ensure = vi.fn(async () => manager);
		const kill = vi.fn(async () => {});
		const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner });

		const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, {} as ExtensionContext);

		expect(result.details.backgroundOutput).toBe("bg-line");
		expect(result.content).toEqual([{ type: "text", text: "ok\n[background output (unattributed)]\nbg-line" }]);
	});

	it("lets the user wait when an interrupted kernel is still busy", async () => {
		const execute = vi
			.fn<KernelClient["execute"]>()
			.mockRejectedValueOnce(new KernelBusyAfterInterruptError())
			.mockResolvedValueOnce(okExecuteResult());
		const manager = { execute } as unknown as KernelClient;
		const ensure = vi.fn(async () => manager);
		const kill = vi.fn(async () => {});
		const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner;
		const select = vi.fn(async () => "Wait and preserve state");
		const { ctx, setWorkingMessage } = createBusyKernelContext(select, { throwWorkingMessage: true });
		const tool = createIpythonToolDefinition(tempDir, { provisioner });

		const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, ctx);

		expect(result.details.status).toBe("ok");
		expect(result.details.kernelRestarted).toBe(false);
		expect(ensure).toHaveBeenCalledTimes(2);
		expect(kill).not.toHaveBeenCalled();
		expect(select).toHaveBeenCalledWith(
			expect.stringContaining("previous cell has not stopped"),
			["Wait and preserve state", "Kill kernel and restart"],
			{
				signal: undefined,
			},
		);
		expect(setWorkingMessage).toHaveBeenCalledWith("Waiting for Python kernel...");
		expect(setWorkingMessage).toHaveBeenLastCalledWith(undefined);
	});

	it("lets the user kill and restart a busy interrupted kernel", async () => {
		const busyManager = {
			execute: vi.fn<KernelClient["execute"]>().mockRejectedValueOnce(new KernelBusyAfterInterruptError()),
		} as unknown as KernelClient;
		const freshManager = {
			execute: vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult()),
		} as unknown as KernelClient;
		const ensure = vi.fn(async () => {
			return ensure.mock.calls.length === 1 ? busyManager : freshManager;
		});
		const kill = vi.fn(async () => {});
		const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner;
		const select = vi.fn(async () => "Kill kernel and restart");
		const { ctx, setWorkingMessage } = createBusyKernelContext(select, { throwWorkingMessage: true });
		const tool = createIpythonToolDefinition(tempDir, { provisioner });

		const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, ctx);
		const text = result.content[0]?.type === "text" ? result.content[0].text : "";

		expect(result.details.status).toBe("ok");
		expect(result.details.kernelRestarted).toBe(true);
		expect(text).toContain("<ipython_kernel_reset>");
		expect(text).toContain("Variables, imports, async tasks, and open resources");
		expect(text).toContain("ok");
		expect(ensure).toHaveBeenCalledTimes(2);
		expect(kill).toHaveBeenCalledTimes(1);
		expect(freshManager.execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ signal: undefined }));
		expect(setWorkingMessage).toHaveBeenCalledWith("Restarting Python kernel...");
		expect(setWorkingMessage).toHaveBeenLastCalledWith(undefined);
	});

	it("reports a timed-out cell's partial output and its deadline notice", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({
			...okExecuteResult(),
			stdout: "half-done\n",
			status: "aborted",
			timedOut: true,
			durationMs: 60_100,
		});
		const settle = vi.fn(async () => true);
		const manager = { execute, waitForInterruptedCellToSettle: settle } as unknown as KernelClient;
		const ensure = vi.fn(async () => manager);
		const kill = vi.fn(async () => {});
		const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });

		const result = await tool.execute(
			"tool-call",
			{ code: "print('half-done')" },
			undefined,
			undefined,
			{} as ExtensionContext,
		);
		const text = result.content[0]?.type === "text" ? result.content[0].text : "";

		expect(text).toContain("half-done");
		expect(text).toContain("<ipython_cell_timeout>");
		expect(text).toContain(
			"Cell deadline 60s hit at 60.1s; the cell was interrupted. Nothing from before it is lost: earlier cells' variables and imports survive, and everything this cell printed or already wrote is kept.",
		);
		expect(text).toContain("cell: for _ in range(60): await asyncio.sleep(5)");
		expect(text).toContain(
			'If you were waiting on work, do not poll inside a cell: start it detached and end your turn -- h = bash("nix build .#prime-agent") returns at once, and its completion notice wakes you with the result.',
		);
		expect(text).toContain(
			"If it was one genuinely long call (cargo/nix build, test suite, training), re-run it with a larger deadline: timeout=600 (seconds, max 3600).",
		);
		expect(text).toContain("</ipython_cell_timeout>");
		expect(text).not.toContain("The kernel was killed and restarted");
		expect(result.details.timedOut).toBe(true);
		expect(result.details.timeoutMs).toEqual({ requested: 60_000, effective: 60_000 });
		expect(result.details.status).toBe("aborted");
		expect(result.details.kernelRestarted).toBe(false);
		expect(isErrorResult(result)).toBe(true);
		expect(execute).toHaveBeenCalledWith("print('half-done')", expect.objectContaining({ timeoutMs: 60_000 }));
		expect(settle).toHaveBeenCalledTimes(1);
		expect(kill).not.toHaveBeenCalled();
	});

	it("kills the kernel when the interrupted cell never settles, and says so twice", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({
			...okExecuteResult(),
			status: "aborted",
			timedOut: true,
			durationMs: 60_100,
		});
		const settle = vi.fn(async () => false);
		const manager = { execute, waitForInterruptedCellToSettle: settle } as unknown as KernelClient;
		const ensure = vi.fn(async () => manager);
		const kill = vi.fn(async () => {});
		const provisioner = { ensure, kill } as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });

		const result = await tool.execute("tool-call", { code: "while True: pass" }, undefined, undefined, {
			hasUI: true,
			ui: { select: vi.fn(), setWorkingMessage: vi.fn() },
		} as unknown as ExtensionContext);
		const text = result.content[0]?.type === "text" ? result.content[0].text : "";

		expect(kill).toHaveBeenCalledTimes(1);
		expect(ensure).toHaveBeenCalledTimes(1);
		expect(execute).toHaveBeenCalledTimes(1);
		expect(result.details.kernelRestarted).toBe(true);
		expect(result.details.timedOut).toBe(true);
		expect(text).toContain("<ipython_kernel_reset>");
		expect(text).toContain("<ipython_cell_timeout>");
		expect(text).toContain("The kernel was killed and restarted because the cell did not stop.");
		expect(isErrorResult(result)).toBe(true);
	});

	it("never asks the user before escalating a deadline", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({
			...okExecuteResult(),
			status: "aborted",
			timedOut: true,
			durationMs: 60_100,
		});
		const manager = {
			execute,
			waitForInterruptedCellToSettle: vi.fn(async () => false),
		} as unknown as KernelClient;
		const kill = vi.fn(async () => {});
		const provisioner = { ensure: vi.fn(async () => manager), kill } as unknown as IpythonKernelProvisioner;
		const select = vi.fn(async () => "Wait and preserve state");
		const { ctx, setWorkingMessage } = createBusyKernelContext(select);
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });

		await tool.execute("tool-call", { code: "while True: pass" }, undefined, undefined, ctx);

		// A watchdog that asks a human is not a watchdog: a child session has no UI.
		expect(select).not.toHaveBeenCalled();
		expect(kill).toHaveBeenCalledTimes(1);
		expect(setWorkingMessage).toHaveBeenCalledWith("Restarting Python kernel...");
	});

	it("clamps an over-ceiling timeout, names both numbers, and suggests nothing larger", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({
			...okExecuteResult(),
			status: "aborted",
			timedOut: true,
			durationMs: 3_600_500,
		});
		const manager = {
			execute,
			waitForInterruptedCellToSettle: vi.fn(async () => true),
		} as unknown as KernelClient;
		const provisioner = {
			ensure: vi.fn(async () => manager),
			kill: vi.fn(async () => {}),
		} as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });

		const result = await tool.execute(
			"tool-call",
			{ code: "while True: pass", timeout: 7200 },
			undefined,
			undefined,
			{
				hasUI: true,
				ui: { select: vi.fn(), setWorkingMessage: vi.fn() },
			} as unknown as ExtensionContext,
		);
		const text = result.content[0]?.type === "text" ? result.content[0].text : "";

		expect(execute).toHaveBeenCalledWith("while True: pass", expect.objectContaining({ timeoutMs: 3_600_000 }));
		expect(result.details.timeoutMs).toEqual({ requested: 7_200_000, effective: 3_600_000 });
		expect(text).toContain("timeout=7200s exceeds the 3600s ceiling; this cell used 3600s.");
		expect(text).toContain("Cell deadline 3600s hit at 3600.5s");
		expect(text).not.toContain("re-run it with a larger deadline");
		expect(text).toContain("start it detached and end your turn");
	});

	it("reports the clamp even when the clamped cell finishes in time", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce({
			...okExecuteResult(),
			status: "ok",
		});
		const manager = {
			execute,
			waitForInterruptedCellToSettle: vi.fn(async () => true),
		} as unknown as KernelClient;
		const provisioner = {
			ensure: vi.fn(async () => manager),
			kill: vi.fn(async () => {}),
		} as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });

		const result = await tool.execute("tool-call", { code: "1 + 1", timeout: 7200 }, undefined, undefined, {
			hasUI: true,
			ui: { select: vi.fn(), setWorkingMessage: vi.fn() },
		} as unknown as ExtensionContext);
		const text = result.content[0]?.type === "text" ? result.content[0].text : "";

		expect(text).toContain("timeout=7200s exceeds the 3600s ceiling; this cell used 3600s.");
		expect(text).not.toContain("<ipython_cell_timeout>");
		expect(result.details.timedOut).toBeUndefined();
	});

	it("drops the larger-deadline suggestion when the same code times out again", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValue({
			...okExecuteResult(),
			status: "aborted",
			timedOut: true,
			durationMs: 60_100,
		});
		const manager = {
			execute,
			waitForInterruptedCellToSettle: vi.fn(async () => true),
		} as unknown as KernelClient;
		const provisioner = {
			ensure: vi.fn(async () => manager),
			kill: vi.fn(async () => {}),
		} as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 60_000 });
		const ctx = { hasUI: true, ui: { select: vi.fn(), setWorkingMessage: vi.fn() } } as unknown as ExtensionContext;

		const first = await tool.execute("tool-call-1", { code: "while True: pass" }, undefined, undefined, ctx);
		const second = await tool.execute("tool-call-2", { code: "while True: pass" }, undefined, undefined, ctx);
		const other = await tool.execute(
			"tool-call-3",
			{ code: "while True: sleep_forever()" },
			undefined,
			undefined,
			ctx,
		);
		const textOf = (result: { content: unknown }) =>
			(result.content as { type: string; text?: string }[])[0]?.text ?? "";

		expect(textOf(first)).toContain("timeout=600 (seconds, max 3600)");
		// The same code again: the suggestion already failed once, so teach detach only.
		expect(textOf(second)).not.toContain("timeout=600");
		expect(textOf(second)).not.toContain("re-run it with a larger deadline");
		expect(textOf(second)).toContain("start it detached and end your turn");
		// Different code: a fresh first timeout, so the suggestion is back.
		expect(textOf(other)).toContain("timeout=600 (seconds, max 3600)");
	});

	it("rejects a non-positive or over-ceiling timeout before any kernel call", async () => {
		const execute = vi.fn<KernelClient["execute"]>();
		const ensure = vi.fn(async () => ({ execute }) as unknown as KernelClient);
		const provisioner = { ensure, kill: vi.fn(async () => {}) } as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner });

		// The same validator the agent loop runs before it invokes the tool.
		expect(Value.Check(tool.parameters, { code: "1" })).toBe(true);
		expect(Value.Check(tool.parameters, { code: "1", timeout: 3600 })).toBe(true);
		expect(Value.Check(tool.parameters, { code: "1", timeout: 0 })).toBe(false);
		expect(Value.Check(tool.parameters, { code: "1", timeout: -5 })).toBe(false);
		expect(Value.Check(tool.parameters, { code: "1", timeout: 3601 })).toBe(false);
		expect(ensure).not.toHaveBeenCalled();
		expect(execute).not.toHaveBeenCalled();
	});

	it("passes the session cell deadline to the manager and publishes it in the description", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult());
		const manager = { execute } as unknown as KernelClient;
		const provisioner = {
			ensure: vi.fn(async () => manager),
			kill: vi.fn(async () => {}),
		} as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: 120_000 });

		expect(tool.description).toContain("Cells stop at 120s by default (max 3600s)");
		const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, {
			hasUI: true,
			ui: { select: vi.fn(), setWorkingMessage: vi.fn() },
		} as unknown as ExtensionContext);

		expect(execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ timeoutMs: 120_000 }));
		expect(result.details.timedOut).toBeUndefined();
		expect(result.content).toEqual([{ type: "text", text: "ok" }]);
	});

	it("leaves the cell unbounded when the session turns the deadline off", async () => {
		const execute = vi.fn<KernelClient["execute"]>().mockResolvedValueOnce(okExecuteResult());
		const manager = { execute } as unknown as KernelClient;
		const provisioner = {
			ensure: vi.fn(async () => manager),
			kill: vi.fn(async () => {}),
		} as unknown as IpythonKernelProvisioner;
		const tool = createIpythonToolDefinition(tempDir, { provisioner, cellTimeoutMs: "off" });

		expect(tool.description).toContain("This session sets no cell deadline");
		const result = await tool.execute("tool-call", { code: "x = 1" }, undefined, undefined, {
			hasUI: true,
			ui: { select: vi.fn(), setWorkingMessage: vi.fn() },
		} as unknown as ExtensionContext);

		expect(execute).toHaveBeenCalledWith("x = 1", expect.objectContaining({ timeoutMs: undefined }));
		expect(result.details.timedOut).toBeUndefined();
	});

	it("caches the generated skill digest after the runtime bootstrap", async () => {
		const { python, requested } = writeFakeDigestRuntime({
			capabilities: ["signatures"],
			digest: "# skill surface (generated)\ngit.status(cwd?) -> str\n",
		});
		const onSkillDigest = vi.fn();
		const provisioner = new IpythonKernelProvisioner(tempDir, {
			python,
			pythonSkills: [pythonSkillInfo("git")],
			onSkillDigest,
		});

		try {
			await provisioner.ensure();
			expect(provisioner.skillDigest).toContain("# skill surface");
			expect(onSkillDigest).toHaveBeenCalledWith(provisioner.skillDigest);
			// The digest indexes the modules the bootstrap publishes, so the request names them.
			expect(requested()).toEqual([["git"]]);
		} finally {
			await provisioner.dispose();
		}
	});

	it("publishes no digest and sends nothing to a runtime without the capability", async () => {
		const { python, requested } = writeFakeDigestRuntime();
		const onSkillDigest = vi.fn();
		const provisioner = new IpythonKernelProvisioner(tempDir, {
			python,
			pythonSkills: [pythonSkillInfo("git")],
			onSkillDigest,
		});

		try {
			// An older runtime still starts; the missing capability is not a startup failure.
			await provisioner.ensure();
			expect(provisioner.skillDigest).toBeUndefined();
			expect(onSkillDigest).not.toHaveBeenCalled();
			expect(requested()).toEqual([]);
		} finally {
			await provisioner.dispose();
		}
	});

	it("does not delete the on-disk snapshot (the kernel survives compaction)", async () => {
		const snapshotDir = join(tempDir, "artifacts");
		const provisioner = new IpythonKernelProvisioner(tempDir, { snapshotDir });
		const dill = join(snapshotDir, "kernel-state.dill");
		const manifest = join(snapshotDir, "kernel-state.json");
		mkdirSync(snapshotDir, { recursive: true });
		writeFileSync(dill, "payload");
		writeFileSync(manifest, "{}");

		// listing the namespace must never touch the on-disk snapshot
		await provisioner.listNamespaceNames();

		expect(existsSync(dill)).toBe(true);
		expect(existsSync(manifest)).toBe(true);
	});
});

describe("ReplKernelManager session cleanup during startup", () => {
	beforeEach(() => {
		tempDir = mkdtempSync(join(tmpdir(), "prime-agent-kernel-cleanup-"));
	});

	afterEach(() => {
		if (tempDir) {
			rmSync(tempDir, { recursive: true, force: true });
			tempDir = "";
		}
	});

	it("disposes a kernel that is still booting when its session is cleaned up", async () => {
		const python = join(tempDir, "python");
		// Never emits the ready line - stays in the booting phase until killed.
		writeFileSync(python, ["#!/bin/sh", "sleep 30", ""].join("\n"));
		chmodSync(python, 0o755);
		const sessionId = `provisioner-test-${Date.now()}`;
		const manager = new ReplKernelManager({ python, cwd: tempDir, sessionId });

		try {
			const startup = manager.start();
			cleanupSessionResources(sessionId);
			await expect(startup).rejects.toThrow(/Kernel exited before ready|disposed during startup/);
			expect(manager.isRunning).toBe(false);
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
		}
	});
});