pi-ext

import type { Context, Model, Tool } from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import {
	contextFingerprint,
	createAssistant,
	emptyUsage,
	isKlausModelId,
	klausSelector,
	MODEL_SELECTORS,
	modelFacingContent,
	requiredConstrainedTool,
	toKlausTools,
} from "../src/protocol";

const model: Model<string> = {
	id: "claude-sonnet-5",
	name: "Claude Sonnet 5",
	api: "klaus-agent-sdk",
	provider: "klaus",
	baseUrl: "https://klaus.invalid",
	reasoning: true,
	input: ["text", "image"],
	cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
	contextWindow: 200_000,
	maxTokens: 64_000,
};

const tool = (
	name: string,
	constrainedSampling?: Tool["constrainedSampling"],
): Tool => ({
	name,
	description: `Description for ${name}`,
	parameters: Type.Object(
		{ value: Type.String({ minLength: 1 }) },
		{ additionalProperties: false },
	),
	constrainedSampling,
});

const assistant = (overrides: Partial<Context["messages"][number]> = {}) => ({
	role: "assistant" as const,
	content: [
		{
			type: "thinking" as const,
			thinking: "private",
			thinkingSignature: "sig",
		},
		{ type: "text" as const, text: "visible" },
		{
			type: "toolCall" as const,
			id: "call-1",
			name: "echo",
			arguments: { value: "hello" },
		},
	],
	api: "test",
	provider: "test",
	model: "test",
	usage: emptyUsage(),
	stopReason: "stop" as const,
	timestamp: 1,
	...overrides,
});

describe("Klaus model protocol", () => {
	it("accepts every current Claude catalog id shape, aliases included", () => {
		for (const id of Object.keys(MODEL_SELECTORS)) {
			expect(isKlausModelId(id)).toBe(true);
		}
		for (const id of [
			"claude-opus-5-5",
			"claude-fable-5-1",
			"claude-sonnet-4",
			"claude-haiku-4-5",
		]) {
			expect(isKlausModelId(id)).toBe(true);
		}
		for (const id of [
			"",
			"default",
			"claude-sonnet-5-latest",
			"claude-opus-4-5-20251101",
			"CLAUDE-SONNET-5",
			"claude-sonnet-5_",
			"claude--5",
			"__proto__",
		]) {
			expect(isKlausModelId(id)).toBe(false);
		}
	});

	it("maps known ids to short selectors and passes unknown ids through", () => {
		expect(klausSelector("claude-opus-5")).toBe("opus");
		expect(klausSelector("claude-fable-5-1")).toBe("fable");
		expect(klausSelector("claude-opus-5-5")).toBe("claude-opus-5-5");
	});

	it("projects Pi tools without changing names, descriptions, order, or schemas", () => {
		const tools = [tool("z_tool"), tool("A tool")];
		const projected = toKlausTools(tools);

		expect(projected).toEqual([
			{
				name: "z_tool",
				description: "Description for z_tool",
				inputSchema: tools[0]?.parameters,
			},
			{
				name: "A tool",
				description: "Description for A tool",
				inputSchema: tools[1]?.parameters,
			},
		]);
		expect(toKlausTools(undefined)).toEqual([]);
	});

	it("rejects only the first required JSON-schema constrained tool", () => {
		expect(
			requiredConstrainedTool([
				tool("plain"),
				tool("preferred", { type: "json_schema", strict: "prefer" }),
				tool("grammar", {
					type: "grammar",
					variants: { openai_regex: ".*" },
				}),
				tool("required", { type: "json_schema", strict: "require" }),
				tool("later", { type: "json_schema", strict: "require" }),
			]),
		).toBe("required");
		expect(requiredConstrainedTool(undefined)).toBeUndefined();
	});

	it("creates independent zeroed usage objects", () => {
		const first = emptyUsage();
		const second = emptyUsage();
		first.input = 9;
		first.cost.total = 4;

		expect(second).toEqual({
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		});
	});

	it("creates a pending assistant owned by the selected Klaus model", () => {
		const before = Date.now();
		const message = createAssistant(model);

		expect(message).toMatchObject({
			role: "assistant",
			content: [],
			api: "klaus-agent-sdk",
			provider: "klaus",
			model: "claude-sonnet-5",
			stopReason: "pending",
			usage: emptyUsage(),
		});
		expect(message.timestamp).toBeGreaterThanOrEqual(before);
		expect(message.timestamp).toBeLessThanOrEqual(Date.now());
	});

	it("forwards only model-facing text and image fields", () => {
		expect(
			modelFacingContent([
				{ type: "text", text: "hello" },
				{ type: "image", data: "base64", mimeType: "image/webp" },
			]),
		).toEqual([
			{ type: "text", text: "hello" },
			{ type: "image", data: "base64", mimeType: "image/webp" },
		]);
	});
});

describe("Klaus compatibility fingerprints", () => {
	const base = (): Context => ({
		systemPrompt: "system",
		messages: [
			{ role: "user", content: "question", timestamp: 1 },
			assistant(),
			{
				role: "toolResult",
				toolCallId: "call-1",
				toolName: "echo",
				content: [{ type: "text", text: "hello" }],
				isError: false,
				timestamp: 2,
				details: { private: "ignored" },
			},
		],
		tools: [tool("echo")],
	});

	it("is stable and always emits a SHA-256 hex digest", () => {
		const context = base();
		const first = contextFingerprint(context, "sonnet\0cwd");

		expect(first).toMatch(/^[a-f0-9]{64}$/);
		expect(contextFingerprint(structuredClone(context), "sonnet\0cwd")).toBe(
			first,
		);
	});

	it("ignores volatile accounting, timestamps, private reasoning, and tool details", () => {
		const first = base();
		const second = structuredClone(first);
		const user = second.messages[0];
		const answer = second.messages[1];
		const result = second.messages[2];
		if (
			user?.role !== "user" ||
			answer?.role !== "assistant" ||
			result?.role !== "toolResult"
		) {
			throw new Error("Broken fingerprint fixture.");
		}
		user.timestamp = 999;
		answer.timestamp = 999;
		answer.usage.input = 42;
		answer.usage.cost.total = 99;
		const thinking = answer.content.find((item) => item.type === "thinking");
		if (thinking?.type === "thinking") {
			thinking.thinking = "different private thought";
			thinking.thinkingSignature = "different signature";
		}
		result.timestamp = 999;
		result.details = { private: "different" };

		expect(contextFingerprint(first, "sonnet\0cwd")).toBe(
			contextFingerprint(second, "sonnet\0cwd"),
		);
	});

	it.each([
		["identity", (context: Context) => context, "opus\0cwd"],
		[
			"system prompt",
			(context: Context) => ({ ...context, systemPrompt: "other" }),
			"sonnet\0cwd",
		],
		[
			"user content",
			(context: Context) => ({
				...context,
				messages: context.messages.map((message, index) =>
					index === 0 && message.role === "user"
						? { ...message, content: "other" }
						: message,
				),
			}),
			"sonnet\0cwd",
		],
		[
			"assistant text",
			(context: Context) => ({
				...context,
				messages: context.messages.map((message, index) =>
					index === 1 && message.role === "assistant"
						? assistant({
								content: message.content.map((item) =>
									item.type === "text" ? { ...item, text: "other" } : item,
								),
							})
						: message,
				),
			}),
			"sonnet\0cwd",
		],
		[
			"tool arguments",
			(context: Context) => ({
				...context,
				messages: context.messages.map((message, index) =>
					index === 1 && message.role === "assistant"
						? assistant({
								content: message.content.map((item) =>
									item.type === "toolCall"
										? { ...item, arguments: { value: "other" } }
										: item,
								),
							})
						: message,
				),
			}),
			"sonnet\0cwd",
		],
		[
			"tool result error state",
			(context: Context) => ({
				...context,
				messages: context.messages.map((message) =>
					message.role === "toolResult"
						? { ...message, isError: true }
						: message,
				),
			}),
			"sonnet\0cwd",
		],
		[
			"active schema",
			(context: Context) => ({ ...context, tools: [tool("different")] }),
			"sonnet\0cwd",
		],
	] as const)("changes when %s changes", (_label, mutate, identity) => {
		const context = base();
		expect(contextFingerprint(mutate(context), identity)).not.toBe(
			contextFingerprint(context, "sonnet\0cwd"),
		);
	});

	it("distinguishes missing and empty message content without mutating context", () => {
		const context = base();
		const original = structuredClone(context);
		contextFingerprint(context);

		expect(context).toEqual(original);
		expect(contextFingerprint({ messages: [] })).not.toBe(
			contextFingerprint({
				messages: [{ role: "user", content: "", timestamp: 1 }],
			}),
		);
	});
});