pi-ext

import {
	access,
	mkdtemp,
	readFile,
	rm,
	stat,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";
import { afterEach, describe, expect, it } from "vitest";
import { type ChildRuntime, createChildRuntime } from "../src/runtime/index";

const runtimes: ChildRuntime[] = [];
const roots: string[] = [];
const originalEnvironment = { ...process.env };

afterEach(async () => {
	await Promise.all(runtimes.splice(0).map((runtime) => runtime.cleanup()));
	await Promise.all(
		roots.splice(0).map((root) => rm(root, { recursive: true, force: true })),
	);
	for (const key of Object.keys(process.env)) {
		if (!(key in originalEnvironment)) delete process.env[key];
	}
	Object.assign(process.env, originalEnvironment);
});

describe("Klaus isolated Claude child runtime", () => {
	it("builds a query-private allowlisted environment and filters headers", async () => {
		process.env.ANTHROPIC_API_KEY = "blocked";
		process.env.ANTHROPIC_CUSTOM_HEADERS = "x-ambient: blocked";
		process.env.CLAUDE_CODE_OAUTH_TOKEN = "blocked";
		process.env.AWS_SECRET_ACCESS_KEY = "blocked";
		process.env.UNRELATED_KLAUS_TEST = "blocked";
		process.env.KLAUS_E2E_BASE_URL = "http://127.0.0.1:43210";
		const runtime = await createChildRuntime(
			"fake",
			{
				"X-Custom": "present",
				"X-Empty": "",
				"X-Null": null,
				Authorization: "Bearer hostile",
				"x-api-key": "hostile",
				Host: "hostile.example",
				"Content-Type": "application/hostile",
				"x-klaus-private": "must-not-pass",
				"X-KLAUS-UPPER": "must-not-pass",
				"bad header": "must-not-pass",
				"X-Bad-Value": "line one\nline two",
				"anthropic-beta": "unknown-beta, context-1m-2025-08-07, another-beta",
			},
			{
				PATH: process.env.PATH ?? "",
				LANG: "klaus-test-locale",
				UNRELATED_KLAUS_TEST: "blocked",
				ANTHROPIC_API_KEY: "blocked",
			},
		);
		runtimes.push(runtime);

		expect(runtime.executable).toBeTypeOf("string");
		if (!runtime.executable)
			throw new Error("Klaus runtime executable missing.");
		expect((await stat(runtime.executable)).isFile()).toBe(true);
		expect(runtime.env).toMatchObject({
			PATH: process.env.PATH ?? "",
			LANG: "klaus-test-locale",
			HOME: runtime.configDir,
			CLAUDE_CONFIG_DIR: runtime.configDir,
			CLAUDE_CODE_OAUTH_TOKEN: "fake",
			CLAUDE_AGENT_SDK_CLIENT_APP: "klaus/1",
			CLAUDE_AGENT_SDK_MCP_NO_PREFIX: "1",
			CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: "1",
			CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY: "0",
			ENABLE_TOOL_SEARCH: "false",
			CLAUDE_CODE_MCP_AUTO_BACKGROUND_MS: "0",
			CLAUDE_CODE_MCP_TOOL_IDLE_TIMEOUT: "0",
			CLAUDE_CODE_DISABLE_NONSTREAMING_FALLBACK: "1",
			MAX_MCP_OUTPUT_TOKENS: "1000000",
			CLAUDE_CODE_DISABLE_CLAUDE_MDS: "1",
			CLAUDE_CODE_DISABLE_AUTO_MEMORY: "1",
			CLAUDE_CODE_DISABLE_BUNDLED_SKILLS: "1",
			CLAUDE_CODE_DISABLE_POLICY_SKILLS: "1",
			CLAUDE_AGENT_SDK_DISABLE_BUILTIN_AGENTS: "1",
			CLAUDE_CODE_DISABLE_WORKFLOWS: "1",
			CLAUDE_CODE_DISABLE_CRON: "1",
			CLAUDE_CODE_DISABLE_BACKGROUND_TASKS: "1",
			CLAUDE_CODE_DISABLE_FILE_CHECKPOINTING: "1",
			CLAUDE_CODE_TMPDIR: runtime.configDir,
			ANTHROPIC_BASE_URL: "http://127.0.0.1:43210",
			ANTHROPIC_CUSTOM_HEADERS: "X-Custom: present\nX-Empty: ",
		});
		for (const forbidden of [
			"ANTHROPIC_API_KEY",
			"AWS_SECRET_ACCESS_KEY",
			"UNRELATED_KLAUS_TEST",
		]) {
			expect(runtime.env).not.toHaveProperty(forbidden);
		}
		expect(runtime.betas).toEqual(["context-1m-2025-08-07"]);
		expect(runtime.env.ANTHROPIC_CUSTOM_HEADERS).not.toMatch(
			/auth|hostile|klaus|bad|ambient/i,
		);
	}, 60_000);

	it("maps Pi maxTokens onto the child output cap", async () => {
		process.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS = "hostile-ambient";
		const runtime = await createChildRuntime("fake", {}, {}, 32768);
		runtimes.push(runtime);
		expect(runtime.env.CLAUDE_CODE_MAX_OUTPUT_TOKENS).toBe("32768");

		const uncapped = await createChildRuntime("fake", {}, {});
		runtimes.push(uncapped);
		expect(uncapped.env).not.toHaveProperty("CLAUDE_CODE_MAX_OUTPUT_TOKENS");
	}, 60_000);

	it("rejects invalid maxTokens values", async () => {
		await expect(createChildRuntime("fake", {}, {}, 0)).rejects.toThrow(
			/invalid maxTokens/,
		);
		await expect(
			createChildRuntime("fake", {}, {}, Number.NaN),
		).rejects.toThrow(/invalid maxTokens/);
	}, 60_000);

	it("does not depend on a shared predictable klaus directory", async () => {
		const root = await mkdtemp(join(tmpdir(), "klaus-runtime-test-"));
		roots.push(root);
		await writeFile(join(root, "klaus"), "foreign obstacle");
		process.env.TMPDIR = root;
		process.env.TMP = root;
		process.env.TEMP = root;

		const runtime = await createChildRuntime("token", {}, {});
		runtimes.push(runtime);

		expect(dirname(runtime.configDir)).toBe(root);
		expect(await readFile(join(root, "klaus"), "utf8")).toBe(
			"foreign obstacle",
		);
	}, 60_000);

	it("creates unique config directories and removes only its own directory", async () => {
		const first = await createChildRuntime("token-one", {}, {});
		const second = await createChildRuntime("token-two", {}, {});
		runtimes.push(second);

		expect(first.configDir).not.toBe(second.configDir);
		expect(first.env.HOME).toBe(first.configDir);
		expect(second.env.HOME).toBe(second.configDir);
		await first.cleanup();
		await expect(access(first.configDir)).rejects.toMatchObject({
			code: "ENOENT",
		});
		await expect(access(second.configDir)).resolves.toBeUndefined();
	}, 60_000);
});