pi-ext

import { readdir, writeFile } from "node:fs/promises";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { getBuiltinModels } from "@earendil-works/pi-ai/providers/all";
import { Type } from "typebox";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";
import { isProjectableClaudeModel } from "../src/models";

const extensionPath = fileURLToPath(new URL("../index.ts", import.meta.url));

async function startSession(session: PiSession): Promise<void> {
	await session.session.extensionRunner.emit({
		type: "session_start",
		reason: "startup",
	});
}

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() {
	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 session: PiSession | undefined;

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

	afterEach(async () => {
		await session?.dispose();
		session = undefined;
	});

	it("loads standalone when global debug selects Klaus", async () => {
		session = await createPiSession({
			userSettings: { "pi-ext": { debug: "klaus" } },
			extensionPaths: [extensionPath],
		});
		await startSession(session);

		expect(session.extensions).toHaveLength(1);
		expect(await tracePaths(session.cwd)).toHaveLength(1);
	});

	it("loads one extension and registers every projectable Claude model", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		const runtime = session.session.modelRuntime;

		expect(session.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("follows a models.json reload without re-registering on turns", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		const runtime = session.session.modelRuntime;
		const before = runtime.getModel("klaus", "claude-opus-5");
		expect(before?.name).toBe(
			builtinClaudeCatalog().find((m) => m.id === "claude-opus-5")?.name,
		);

		await writeFile(
			join(session.agentDir, "models.json"),
			JSON.stringify({
				providers: {
					anthropic: {
						modelOverrides: { "claude-opus-5": { name: "Opus (work)" } },
					},
				},
			}),
		);
		await runtime.refresh({ allowNetwork: false });

		expect(runtime.getModel("klaus", "claude-opus-5")).toMatchObject({
			provider: "klaus",
			api: "klaus-agent-sdk",
			baseUrl: "https://klaus.invalid",
			name: "Opus (work)",
		});
		expect(runtime.getModel("anthropic", "claude-opus-5")?.name).toBe(
			"Opus (work)",
		);
	});

	it("copies Anthropic model metadata and prices Fable safety fallbacks", async () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		for (const model of builtinClaudeCatalog()) {
			if (!isProjectableClaudeModel(model)) continue;
			const source = builtinClaudeCatalog().find((m) => m.id === model.id);
			const projected = session.session.modelRuntime.getModel(
				"klaus",
				model.id,
			);
			expect(source).toBeDefined();
			expect(projected).toMatchObject({
				id: source?.id,
				name: source?.name,
				reasoning: source?.reasoning,
				input: source?.input,
				cost: source?.cost,
				contextWindow: source?.contextWindow,
				maxTokens: source?.maxTokens,
			});
			expect(projected?.thinkingLevelMap).toEqual(source?.thinkingLevelMap);
			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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		const runtime = session.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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		const runtime = session.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 () => {
		session = await createPiSession({ extensionPaths: [extensionPath] });
		await startSession(session);
		const runtime = session.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) => {
			session = await createPiSession({ extensionPaths: [extensionPath] });
			await startSession(session);
			const runtime = session.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();
		},
	);
});