pi-ext

import { mkdir, mkdtemp, readdir, rm, writeFile } from "node:fs/promises";
import { join } from "node:path";
import {
	getBuiltinModel,
	getBuiltinModels,
} from "@earendil-works/pi-ai/providers/all";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { createTestSession, type TestSession } from "../../../test/harness";
import { isProjectableClaudeModel } from "../src/models";

async function loadKlaus(pi: ExtensionAPI): Promise<void> {
	const { default: klaus } = await import("../index.js");
	await klaus(pi);
}

async function tracePaths(directory: string): Promise<string[]> {
	const paths: string[] = [];
	for (const entry of await readdir(directory, { withFileTypes: true })) {
		const path = join(directory, entry.name);
		if (entry.isDirectory()) paths.push(...(await tracePaths(path)));
		else if (entry.isFile() && entry.name.endsWith(".jsonl")) paths.push(path);
	}
	return paths;
}

function builtinClaudeCatalog(): ReturnType<typeof getBuiltinModels> {
	return getBuiltinModels("anthropic");
}

function installDispatchAuth(
	runtime: object,
	anthropicAuth: () => unknown,
): void {
	Object.defineProperty(runtime, "getAuth", {
		configurable: true,
		value: async (providerOrModel: unknown) => {
			const provider =
				typeof providerOrModel === "string"
					? providerOrModel
					: typeof providerOrModel === "object" &&
							providerOrModel !== null &&
							"provider" in providerOrModel
						? (providerOrModel as { provider?: unknown }).provider
						: undefined;
			if (provider === "klaus") {
				return { auth: { apiKey: "klaus" }, source: "Klaus test dispatch" };
			}
			if (provider === "anthropic") return anthropicAuth();
			return undefined;
		},
	});
}

describe("Klaus real Pi harness", () => {
	let t: TestSession | undefined;
	let root: string | undefined;

	beforeEach(() => {
		vi.resetModules();
	});

	afterEach(async () => {
		t?.dispose();
		t = undefined;
		if (root) await rm(root, { recursive: true, force: true });
		root = undefined;
	});

	it("loads standalone when global debug selects Klaus", async () => {
		root = await mkdtemp(join(process.cwd(), ".klaus-harness-test-"));
		const agentDirectory = join(root, ".test-home", ".pi", "agent");
		await mkdir(agentDirectory, { recursive: true });
		await writeFile(
			join(agentDirectory, "settings.json"),
			JSON.stringify({ "pi-ext": { debug: "klaus" } }),
		);
		t = await createTestSession({
			cwd: root,
			env: { LOCALAPPDATA: join(root, ".test-local-app-data") },
			extensionFactories: [loadKlaus],
		});

		expect(t.session.extensionRunner.extensions).toHaveLength(1);
		expect(await tracePaths(root)).toHaveLength(1);
	});

	it("loads one extension and registers every projectable Claude model", async () => {
		t = await createTestSession({ extensionFactories: [loadKlaus] });
		const runtime = t.session.modelRuntime;

		expect(t.session.extensionRunner.extensions).toHaveLength(1);
		for (const model of builtinClaudeCatalog()) {
			const registered = runtime.getModel("klaus", model.id);
			if (isProjectableClaudeModel(model)) {
				expect(registered).toMatchObject({
					provider: "klaus",
					id: model.id,
					api: "klaus-agent-sdk",
					baseUrl: "https://klaus.invalid",
				});
			} else {
				expect(registered).toBeUndefined();
			}
		}
		expect(runtime.getModel("klaus", "default")).toBeUndefined();
	});

	it("copies Anthropic model metadata and prices Fable safety fallbacks", async () => {
		t = await createTestSession({ extensionFactories: [loadKlaus] });
		for (const model of builtinClaudeCatalog()) {
			if (!isProjectableClaudeModel(model)) continue;
			const source = getBuiltinModel("anthropic", model.id);
			const projected = t.session.modelRuntime.getModel("klaus", model.id);
			expect(source).toBeDefined();
			expect(projected).toMatchObject({
				id: source?.id,
				name: source?.name,
				reasoning: source?.reasoning,
				thinkingLevelMap: source?.thinkingLevelMap,
				input: source?.input,
				cost: source?.cost,
				contextWindow: source?.contextWindow,
				maxTokens: source?.maxTokens,
			});
			expect(projected?.compat).toEqual(
				model.id === "claude-fable-5-1"
					? {
							...source?.compat,
							allowedFallbackModels: [
								{
									provider: "anthropic",
									model: "claude-opus-4-8",
									cost: {
										input: 5,
										output: 25,
										cacheRead: 0.5,
										cacheWrite: 6.25,
									},
								},
								{
									provider: "anthropic",
									model: "claude-opus-5",
									cost: {
										input: 5,
										output: 25,
										cacheRead: 0.5,
										cacheWrite: 6.25,
									},
								},
							],
						}
					: source?.compat,
			);
		}
	});

	it("rejects non-OAuth Anthropic auth before starting Claude", async () => {
		t = await createTestSession({ extensionFactories: [loadKlaus] });
		const runtime = t.session.modelRuntime;
		Object.defineProperty(runtime, "isUsingOAuth", {
			configurable: true,
			value: () => false,
		});
		const anthropicAuth = vi.fn();
		installDispatchAuth(runtime, anthropicAuth);
		const model = runtime.getModel("klaus", "claude-sonnet-5");
		if (!model) throw new Error("Klaus model was not registered.");

		const result = await runtime.completeSimple(
			model,
			{ messages: [{ role: "user", content: "hello", timestamp: 1 }] },
			{ sessionId: "auth-test" },
		);

		expect(result).toMatchObject({
			stopReason: "error",
			errorMessage:
				"Klaus requires Anthropic OAuth. Use Pi /login for Anthropic first.",
		});
		expect(anthropicAuth).not.toHaveBeenCalled();
	});

	it("rejects an empty resolved OAuth credential actionably", async () => {
		t = await createTestSession({ extensionFactories: [loadKlaus] });
		const runtime = t.session.modelRuntime;
		Object.defineProperty(runtime, "isUsingOAuth", {
			configurable: true,
			value: () => true,
		});
		const anthropicAuth = vi.fn(() => undefined);
		installDispatchAuth(runtime, anthropicAuth);
		const model = runtime.getModel("klaus", "claude-sonnet-5");
		if (!model) throw new Error("Klaus model was not registered.");

		const result = await runtime.completeSimple(
			model,
			{ messages: [{ role: "user", content: "hello", timestamp: 1 }] },
			{ sessionId: "auth-test" },
		);

		expect(result).toMatchObject({
			stopReason: "error",
			errorMessage: "Pi could not resolve Anthropic OAuth for Klaus.",
		});
		expect(anthropicAuth).toHaveBeenCalledOnce();
	});

	it("rejects required constrained sampling before auth or inference", async () => {
		t = await createTestSession({ extensionFactories: [loadKlaus] });
		const runtime = t.session.modelRuntime;
		const anthropicAuth = vi.fn();
		installDispatchAuth(runtime, anthropicAuth);
		const model = runtime.getModel("klaus", "claude-sonnet-5");
		if (!model) throw new Error("Klaus model was not registered.");

		const result = await runtime.completeSimple(
			model,
			{
				messages: [{ role: "user", content: "strict", timestamp: 1 }],
				tools: [
					{
						name: "strict_tool",
						description: "Requires provider-constrained JSON.",
						parameters: Type.Object({ value: Type.String() }),
						constrainedSampling: {
							type: "json_schema",
							strict: "require",
						},
					},
				],
			},
			{ sessionId: "strict-test" },
		);

		expect(result).toMatchObject({
			stopReason: "error",
			errorMessage:
				"Tool strict_tool requires constrained sampling, which Klaus cannot provide.",
		});
		expect(anthropicAuth).not.toHaveBeenCalled();
	});

	it.each([
		null,
		{},
		{ modelId: "claude-sonnet-5" },
		{ modelId: "unknown", selector: "sonnet" },
	] as const)(
		"rejects invalid payload-hook result %j before auth",
		async (payload) => {
			t = await createTestSession({ extensionFactories: [loadKlaus] });
			const runtime = t.session.modelRuntime;
			const anthropicAuth = vi.fn();
			installDispatchAuth(runtime, anthropicAuth);
			const model = runtime.getModel("klaus", "claude-sonnet-5");
			if (!model) throw new Error("Klaus model was not registered.");

			const result = await runtime.completeSimple(
				model,
				{ messages: [{ role: "user", content: "payload", timestamp: 1 }] },
				{ sessionId: "payload-test", onPayload: () => payload },
			);

			expect(result).toMatchObject({
				stopReason: "error",
				errorMessage: "Klaus onPayload returned an invalid request.",
			});
			expect(anthropicAuth).not.toHaveBeenCalled();
		},
	);
});