pi-map/tests/prompt-injection.test.ts

import { describe, it, expect } from "vitest";
import { mkdtempSync, writeFileSync, mkdirSync } from "fs";
import { join } from "path";
import { tmpdir } from "os";
import {
	buildRootPairBlock,
	hasRootPairMarker,
	buildPreInitHint,
	computeInjectionBudget,
	modeAllowsPreInitHint,
	modeAllowsInjection,
	modeRequiresProtocolPath,
	discoverContextWindow,
	estimateTokens,
	findAllArtifactPairs,
	buildInjectionPayload,
	ROOT_PAIR_START_MARKER,
	ROOT_PAIR_END_MARKER,
	TRUST_BOUNDARY_TEXT,
} from "../src/prompt-injection.js";

describe("prompt-injection helpers", () => {
	describe("canonical markers", () => {
		it("builds a root-pair block with markers", () => {
			const block = buildRootPairBlock("index content", "map content");
			expect(block).toContain(ROOT_PAIR_START_MARKER);
			expect(block).toContain(ROOT_PAIR_END_MARKER);
			expect(block).toContain(TRUST_BOUNDARY_TEXT);
			expect(block).toContain("index content");
			expect(block).toContain("map content");
		});

		it("detects root-pair marker in content", () => {
			const block = buildRootPairBlock("i", "m");
			expect(hasRootPairMarker(block)).toBe(true);
			expect(hasRootPairMarker("plain text")).toBe(false);
		});
	});

	describe("pre-init hint", () => {
		it("returns a visible hint with no synthetic artifacts", () => {
			const hint = buildPreInitHint();
			expect(hint).toContain("project_map_init");
			expect(hint).toContain("📋");
			expect(hint).not.toContain(ROOT_PAIR_START_MARKER);
			expect(hint).not.toContain("### Root index");
			expect(hint).not.toContain("### Root map");
		});
	});

	describe("budget calculation", () => {
		it("uses relative percent when smaller than absolute cap", () => {
			const budget = computeInjectionBudget(
				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
				200_000,
			);
			expect(budget).toBe(30_000);
		});

		it("uses absolute cap when smaller than relative percent", () => {
			const budget = computeInjectionBudget(
				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
				10_000_000,
			);
			expect(budget).toBe(100_000);
		});

		it("falls back to absolute cap when context window is unknown", () => {
			const budget = computeInjectionBudget(
				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
				undefined,
			);
			expect(budget).toBe(100_000);
		});

		it("falls back to absolute cap when context window is zero", () => {
			const budget = computeInjectionBudget(
				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
				0,
			);
			expect(budget).toBe(100_000);
		});
	});

	describe("mode helpers", () => {
		it("allows pre-init hint for all modes except off", () => {
			expect(modeAllowsPreInitHint("off")).toBe(false);
			expect(modeAllowsPreInitHint("advisory")).toBe(true);
			expect(modeAllowsPreInitHint("strong")).toBe(true);
			expect(modeAllowsPreInitHint("strict")).toBe(true);
		});

		it("allows injection only for strong and strict", () => {
			expect(modeAllowsInjection("off")).toBe(false);
			expect(modeAllowsInjection("advisory")).toBe(false);
			expect(modeAllowsInjection("strong")).toBe(true);
			expect(modeAllowsInjection("strict")).toBe(true);
		});

		it("requires protocol path only for strict", () => {
			expect(modeRequiresProtocolPath("off")).toBe(false);
			expect(modeRequiresProtocolPath("advisory")).toBe(false);
			expect(modeRequiresProtocolPath("strong")).toBe(false);
			expect(modeRequiresProtocolPath("strict")).toBe(true);
		});
	});

	describe("context-window discovery", () => {
		it("reads contextWindow from model metadata", () => {
			const ctx = { model: { contextWindow: 256_000 } };
			expect(discoverContextWindow(ctx)).toBe(256_000);
		});

		it("reads maxContextTokens from model metadata", () => {
			const ctx = { model: { maxContextTokens: 128_000 } };
			expect(discoverContextWindow(ctx)).toBe(128_000);
		});

		it("falls back to known model lookup by id", () => {
			expect(discoverContextWindow({ model: { id: "gpt-4o-mini" } })).toBe(
				128_000,
			);
			expect(discoverContextWindow({ model: { id: "kimi-for-coding" } })).toBe(
				200_000,
			);
		});

		it("returns undefined for unknown models", () => {
			expect(
				discoverContextWindow({ model: { id: "unknown-model" } }),
			).toBeUndefined();
		});

		it("returns undefined when model is missing", () => {
			expect(discoverContextWindow({})).toBeUndefined();
		});
	});

	describe("token estimation", () => {
		it("estimates tokens from character count", () => {
			expect(estimateTokens("")).toBe(0);
			expect(estimateTokens("abcd")).toBe(1);
			expect(estimateTokens("abcde")).toBe(2);
		});
	});

	describe("artifact-pair discovery", () => {
		it("finds paired artifacts shallow-first", () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
			mkdirSync(join(dir, "src", "nested"), { recursive: true });

			writeFileSync(join(dir, ".pi-map.md"), "# root");
			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
			writeFileSync(join(dir, "src", ".pi-map.md"), "# src");
			writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");
			writeFileSync(join(dir, "src", "nested", ".pi-map.md"), "# nested");
			writeFileSync(
				join(dir, "src", "nested", ".pi-map.index.md"),
				"# nested index",
			);

			const pairs = findAllArtifactPairs(dir);
			expect(pairs.map((p) => p.dir)).toEqual([".", "src", "src/nested"]);
		});

		it("ignores directories missing one of the paired artifacts", () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-pair-test-"));
			writeFileSync(join(dir, ".pi-map.md"), "# root");
			// missing .pi-map.index.md intentionally

			expect(findAllArtifactPairs(dir)).toHaveLength(0);
		});
	});

	describe("injection payload", () => {
		it("always includes the root pair", () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
			writeFileSync(join(dir, ".pi-map.md"), "# root map");
			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");

			const payload = buildInjectionPayload(
				dir,
				{ contextBudgetPercent: 15, contextBudgetMaxTokens: 100_000 },
				undefined,
			);

			expect(payload.content).toContain(ROOT_PAIR_START_MARKER);
			expect(payload.content).toContain(ROOT_PAIR_END_MARKER);
			expect(payload.content).toContain("# root map");
			expect(payload.content).toContain("# root index");
			expect(payload.display).toBe(false);
		});

		it("expands additional pairs when budget allows", () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
			mkdirSync(join(dir, "src"), { recursive: true });
			writeFileSync(join(dir, ".pi-map.md"), "# root map");
			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
			writeFileSync(join(dir, "src", ".pi-map.md"), "# src map");
			writeFileSync(join(dir, "src", ".pi-map.index.md"), "# src index");

			const payload = buildInjectionPayload(
				dir,
				{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10_000 },
				10_000,
			);

			expect(payload.content).toContain("# src map");
			expect(payload.content).toContain("# src index");
		});

		it("stops expanding when budget is exhausted", () => {
			const dir = mkdtempSync(join(tmpdir(), "pi-payload-test-"));
			mkdirSync(join(dir, "src"), { recursive: true });
			writeFileSync(join(dir, ".pi-map.md"), "# root map");
			writeFileSync(join(dir, ".pi-map.index.md"), "# root index");
			writeFileSync(
				join(dir, "src", ".pi-map.md"),
				`# src map ${"x".repeat(400)}`,
			);
			writeFileSync(
				join(dir, "src", ".pi-map.index.md"),
				`# src index ${"x".repeat(400)}`,
			);

			const payload = buildInjectionPayload(
				dir,
				{ contextBudgetPercent: 100, contextBudgetMaxTokens: 10 },
				10_000,
			);

			expect(payload.content).toContain("# root map");
			expect(payload.content).not.toContain("# src map x");
		});
	});
});