pi-ext

import { describe, expect, it } from "vitest";
import type { KlausRequest } from "../src/protocol.js";
import { rewriteSubscriptionPrompt } from "../src/subscription-prompt.js";

function request(): KlausRequest {
	return {
		modelId: "claude-sonnet-5",
		selector: "sonnet",
		systemPrompt:
			"Pi documentation, read when the user asks about pi itself. See pi packages.\n<available_skills>Pi skill at .pi.</available_skills>",
		prompt: "Ask Pi about pi-ext.",
		tools: [
			{
				name: "pi_tool",
				description: "Use Pi.",
				inputSchema: { description: "Pi input" },
			},
		],
		cwd: "/work/pi/project",
		headers: { "x-pi-header": "pi" },
	};
}

describe("rewriteSubscriptionPrompt", () => {
	it("rewrites only the two subscription-triggering phrases", () => {
		const original = request();
		const warnings: string[] = [];
		const rewritten = rewriteSubscriptionPrompt(original, (message) =>
			warnings.push(message),
		);

		expect(rewritten.systemPrompt).toBe(
			"Pi documentation, read when the user asks about this harness. See harness packages.\n<available_skills>Pi skill at .pi.</available_skills>",
		);
		expect(rewritten.prompt).toBe(original.prompt);
		expect(rewritten.tools).toBe(original.tools);
		expect(rewritten.cwd).toBe(original.cwd);
		expect(rewritten.headers).toBe(original.headers);
		expect(original.systemPrompt).toContain("asks about pi itself");
		expect(warnings).toEqual([]);
	});

	it("does not warn for an intentional custom prompt", () => {
		const custom = request();
		custom.systemPrompt = "Custom system prompt.";
		const warnings: string[] = [];

		expect(
			rewriteSubscriptionPrompt(custom, (message) => warnings.push(message))
				.systemPrompt,
		).toBe(custom.systemPrompt);
		expect(warnings).toEqual([]);
	});

	it("warns when the recognizable Pi prompt loses both phrases", () => {
		const changed = request();
		changed.systemPrompt = "Pi documentation changed upstream.";
		const warnings: string[] = [];

		rewriteSubscriptionPrompt(changed, (message) => warnings.push(message));

		expect(warnings).toEqual([
			'Klaus did not find expected Pi system-prompt phrases: "asks about pi itself", "pi packages". Subscription routing may fail after upstream prompt changes.',
		]);
	});

	it("names one missing upstream phrase", () => {
		const changed = request();
		changed.systemPrompt = changed.systemPrompt.replace(
			"pi packages",
			"host packages",
		);
		const warnings: string[] = [];

		rewriteSubscriptionPrompt(changed, (message) => warnings.push(message));

		expect(warnings).toEqual([
			'Klaus did not find expected Pi system-prompt phrase: "pi packages". Subscription routing may fail after upstream prompt changes.',
		]);
	});

	it("rejects repeated trigger phrases instead of broadening the rewrite", () => {
		const repeated = request();
		repeated.systemPrompt = "asks about pi itself; asks about pi itself";

		expect(() => rewriteSubscriptionPrompt(repeated)).toThrow(
			"Klaus found repeated asks about pi itself.",
		);
	});
});