learnharness/tests/plan-tree.test.ts

/**
 * Plan-tree tests (W1) — two layers:
 *
 *  1. Pure ops (src/core/plan-tree.ts): every validation error path and
 *     version semantics (structure only — progress derives from the
 *     transcript, docs/probe-teach-split.md §3).
 *  2. Tool execute (src/agent/tools/plan.ts): against the REAL chat store —
 *     plan_init happy path, the re-callable rewrite, error results.
 */

import assert from "node:assert/strict";
import * as fs from "node:fs";
import * as path from "node:path";
import { after, describe, it } from "node:test";
import { derivePlanStatuses, derivedNodeStatus, type PlanTree, type PlanToolResult } from "../src/core/contracts.ts";
import { createTree, emptyTree, PlanTreeError } from "../src/core/plan-tree.ts";
import { planTools } from "../src/agent/tools/plan.ts";
import { parseMermaidPlan } from "../src/agent/tools/mermaid-plan.ts";
import { chatStore } from "../src/store/chats.ts";
import { PATHS } from "../config.ts";

const N = (id: string, title = id, summary?: string) => ({ id, title, ...(summary ? { summary } : {}) });
const E = (from: string, to: string) => ({ from, to });

function tree(ids: string[], edges: Array<[string, string]> = []): PlanTree {
	return createTree(
		ids.map((id) => N(id)),
		edges.map(([from, to]) => E(from, to)),
	);
}

/** Assert the op throws a PlanTreeError matching the given regex. */
function assertReject(fn: () => unknown, needle: RegExp): void {
	try {
		fn();
	} catch (err) {
		assert.ok(err instanceof PlanTreeError, `expected PlanTreeError, got ${err}`);
		assert.match(err.message, needle);
		return;
	}
	assert.fail(`expected a PlanTreeError matching ${needle}, but the op succeeded`);
}

// ─── 1. pure ops ─────────────────────────────────────────────────────────────

describe("plan-tree pure ops", () => {
	it("createTree builds version 1 with no mutation of its inputs", () => {
		const nodes = [N("a", "A"), N("b", "B")];
		const edges = [E("a", "b")];
		const t = createTree(nodes, edges);
		assert.equal(t.version, 1);
		assert.ok(!("currentNodeId" in t), "no stored progress — structure only");
		assert.ok(!("doneNodeIds" in t), "no done ledger of any kind");
		assert.deepEqual(t.nodes.map((n) => n.id), ["a", "b"]);
		assert.deepEqual(t.edges, [E("a", "b")]);
		// Input arrays untouched (deep-copied into the tree).
		assert.deepEqual(nodes, [N("a", "A"), N("b", "B")]);
		assert.deepEqual(edges, [E("a", "b")]);
	});

	it("createTree rejects duplicate node ids", () => {
		assertReject(() => createTree([N("a"), N("a")], []), /Duplicate node id "a"/);
	});

	it("createTree rejects edges with missing endpoints", () => {
		assertReject(() => createTree([N("a")], [E("a", "ghost")]), /unknown node id "ghost"/);
		assertReject(() => createTree([N("a")], [E("ghost", "a")]), /unknown node id "ghost"/);
	});

	it("createTree rejects self-loops, duplicates, and cycles", () => {
		assertReject(() => createTree([N("a")], [E("a", "a")]), /self-loop/);
		assertReject(() => createTree([N("a"), N("b")], [E("a", "b"), E("a", "b")]), /already exists/);
		// a→b→c→a is a cycle.
		assertReject(
			() => createTree([N("a"), N("b"), N("c")], [E("a", "b"), E("b", "c"), E("c", "a")]),
			/cycle/,
		);
	});

	it("createTree rejects invalid node shapes", () => {
		assertReject(() => createTree([N("")], []), /non-empty `id`/);
		assertReject(() => createTree([{ id: "a", title: "" }], []), /non-empty `title`/);
	});

});

// ─── 2. tool execute against the real store ─────────────────────────────────

const createdChats: string[] = [];
after(() => {
	for (const id of createdChats) {
		try {
			fs.unlinkSync(path.join(PATHS.chats, `${id}.json`));
		} catch {
			/* already gone */
		}
	}
});

/** Call a plan tool's execute directly with a minimal execution-options object. */
async function exec(tools: ReturnType<typeof planTools>, name: string, input: unknown): Promise<PlanToolResult> {
	const t = tools[name]!;
	assert.ok(typeof t.execute === "function", `${name} must be a server-side tool with execute`);
	return t.execute(input as never, { toolCallId: "test", messages: [] } as never) as Promise<PlanToolResult>;
}

function newChat(title: string): string {
	const chat = chatStore.createChat(title);
	createdChats.push(chat.id);
	return chat.id;
}

describe("plan_init — the author-once mermaid tool (execute against the real store)", () => {
	it("parses a mermaid block, builds the tree, persists, returns the full tree", async () => {
		const chatId = newChat("Plan: init");
		const tools = planTools({ chatId, chatStore });

		const res = await exec(tools, "plan_init", {
			mermaid: "graph TD\n  primitives[Number primitives] --> operators[Arithmetic operators]\n  operators --> expressions[Evaluating expressions]\n  detour[Side quest]\n",
		});
		assert.equal(res.ok, true);
		assert.equal(res.error, undefined);
		assert.equal(res.tree.version, 1);
		assert.deepEqual(
			res.tree.nodes.map((n) => n.id),
			["primitives", "operators", "expressions", "detour"],
		);
		assert.deepEqual(res.tree.edges, [
			{ from: "primitives", to: "operators" },
			{ from: "operators", to: "expressions" },
		]);

		const saved = chatStore.getChat(chatId)!.tree;
		assert.ok(saved, "tree persisted to the chat record");
		assert.deepEqual(saved, res.tree);
	});

	it("RE-CALLABLE: a rewrite replaces the tree wholesale (the mutation tools are gone)", async () => {
		const chatId = newChat("Plan: rewrite");
		const tools = planTools({ chatId, chatStore });
		await exec(tools, "plan_init", {
			mermaid: "graph TD\n  a[First attempt] --> b[Early shape]\n",
		});

		const res = await exec(tools, "plan_init", {
			mermaid: "graph TD\n  truth[The core truth]\n  goal[The sink] --> truth\n",
		});
		assert.equal(res.ok, true, "re-init is the sanctioned rewrite path now");
		assert.deepEqual(
			res.tree.nodes.map((n) => n.id),
			["truth", "goal"],
		);
		assert.deepEqual(
			chatStore.getChat(chatId)!.tree!.nodes.map((n) => n.id),
			["truth", "goal"],
			"rewrite persisted",
		);
	});

	it("parse failures collect EVERY problem with line numbers and persist nothing", async () => {
		const chatId = newChat("Plan: bad parse");
		const tools = planTools({ chatId, chatStore });

		const res = await exec(tools, "plan_init", {
			mermaid: "graph TD\n  Foo Bar[x] --> y\n  a --> a\n  subgraph hi\n",
		});
		assert.equal(res.ok, false);
		const err = res.error!;
		assert.match(err, /line 2/);
		assert.match(err, /line 3/);
		assert.match(err, /line 4/);
		assert.deepEqual(res.tree, emptyTree());
		assert.equal(chatStore.getChat(chatId)!.tree, undefined, "nothing persisted on failure");
	});

	it("structural failures (cycle) become ok:false without persisting", async () => {
		const chatId = newChat("Plan: cycle");
		const tools = planTools({ chatId, chatStore });
		const res = await exec(tools, "plan_init", {
			mermaid: "graph TD\n  a[A] --> b[B]\n  b --> a\n",
		});
		assert.equal(res.ok, false);
		assert.match(res.error!, /cycle/);
		assert.deepEqual(res.tree, emptyTree());
	});

	it("only plan_init exists — the seven mutation/tracking tools are gone by design", () => {
		const tools = planTools({ chatId: "chat_x", chatStore });
		assert.deepEqual(Object.keys(tools), ["plan_init"]);
	});

	it("unknown chat returns ok:false with the empty tree", async () => {
		const tools = planTools({ chatId: "chat_nonexistent", chatStore });
		const res = await exec(tools, "plan_init", { mermaid: "graph TD\n  a[A]\n" });
		assert.equal(res.ok, false);
		assert.match(res.error!, /Chat not found/);
		assert.deepEqual(res.tree, emptyTree());
	});
});

// ─── the mermaid subset parser ────────────────────────────────────────────────

describe("parseMermaidPlan (the constrained mermaid subset)", () => {
	const parse = (m: string) => parseMermaidPlan(m);

	it("accepts graph TD and flowchart TB headers (case-insensitive), comments, and blank lines", () => {
		for (const header of ["graph TD", "graph td", "flowchart TB", "Flowchart td"]) {
			const res = parse(`${header}\n%% a comment\n\n  a[Alpha]\n`);
			assert.equal(res.ok, true, header);
			assert.deepEqual(res.plan!.nodes, [{ id: "a", title: "Alpha" }]);
		}
	});

	it("bare ids fall back to id-as-title; bracket and paren labels both work", () => {
		const res = parse("graph TD\n  bare\n  bracket[Bracket title]\n  paren(\"Paren title\")\n");
		assert.deepEqual(res.plan!.nodes, [
			{ id: "bare", title: "bare" },
			{ id: "bracket", title: "Bracket title" },
			{ id: "paren", title: "Paren title" },
		]);
	});

	it("edges register bare endpoints as nodes and dedupe duplicate edges", () => {
		const res = parse("graph TD\n  a --> b[Bee]\n  a --> b[Bee]\n");
		assert.deepEqual(res.plan!.nodes, [
			{ id: "a", title: "a" },
			{ id: "b", title: "Bee" },
		]);
		assert.deepEqual(res.plan!.edges, [{ from: "a", to: "b" }]);
	});

	it("labels on both edge endpoints survive; forward references are fine", () => {
		const res = parse("graph TD\n  downstream[Later thing] --> upstream[Earlier thing]\n");
		assert.deepEqual(res.plan!.edges, [{ from: "downstream", to: "upstream" }]);
		assert.deepEqual(res.plan!.nodes, [
			{ id: "downstream", title: "Later thing" },
			{ id: "upstream", title: "Earlier thing" },
		]);
	});

	it("rejects: non-slug ids, self-loops, chained edges, subgraphs, missing header, empty", () => {
		assert.ok(!parse("graph TD\n  Foo Bar[x]\n").ok, "id with a space");
		assert.match(parse("graph TD\n  Foo Bar[x]\n").problems[0]!, /line 2.*slug/);
		assert.ok(!parse("graph TD\n  a --> a\n").ok, "self-loop");
		assert.ok(!parse("graph TD\n  a --> b --> c\n").ok, "chained edge");
		assert.match(parse("graph TD\n  a --> b --> c\n").problems[0]!, /one edge per line/);
		assert.ok(!parse("graph TD\n  subgraph cluster\n").ok, "subgraph");
		assert.ok(!parse("digraph G\n  a\n").ok, "wrong header");
		assert.ok(!parse("%% only a comment\n").ok, "empty plan");
	});

	it("the first label wins when a node line repeats (relabeling mid-block is ignored)", () => {
		const res = parse("graph TD\n  a[First label]\n  a[Second label]\n");
		assert.deepEqual(res.plan!.nodes, [{ id: "a", title: "First label" }]);
	});
});

// ─── derived statuses: position from CALLS, never results (§3) ─────────────────

describe("derivePlanStatuses (the map tracks position, not completion)", () => {
	const tree: PlanTree = {
		nodes: [
			{ id: "a", title: "A" },
			{ id: "b", title: "B" },
			{ id: "c", title: "C" },
		],
		edges: [
			{ from: "a", to: "b" },
			{ from: "b", to: "c" },
		],
		version: 1,
	};
	const quizPart = (state: string, nodeId: string | undefined, correct?: boolean) => ({
		type: "tool-quiz",
		state,
		input: nodeId === undefined ? {} : { nodeId },
		...(correct === undefined ? {} : { output: { correct } }),
	});
	type QuizPart = { type: string; state: string; input: unknown; output?: unknown };
	const msg = (...parts: QuizPart[]) => ({ parts });

	it("empty transcript → no current, nothing visited (no fallback current)", () => {
		const st = derivePlanStatuses(tree, []);
		assert.equal(st.current, null);
		assert.equal(st.visited.size, 0);
	});

	it("current = the node of the most recent tagged CALL — regardless of state or outcome", () => {
		const st = derivePlanStatuses(tree, [msg(quizPart("output-available", "b", false))]);
		assert.equal(st.current, "b", "even a FAILED check positions the map: the tag is about, not completion");
		assert.equal(st.visited.size, 0);
		const open = derivePlanStatuses(tree, [msg(quizPart("input-available", "c"))]);
		assert.equal(open.current, "c", "a merely ASKED question is already the position");
	});

	it("visited = earlier tagged nodes, sticky; results feed nothing", () => {
		const st = derivePlanStatuses(tree, [
			msg(quizPart("output-available", "a", true)),
			msg(quizPart("output-available", "b", true)),
		]);
		assert.equal(st.current, "b");
		assert.deepEqual([...st.visited], ["a"]);
		// a passed check never marks anything finished — only moving on does
		const again = derivePlanStatuses(tree, [msg(quizPart("output-available", "a", true))]);
		assert.equal(again.current, "a");
		assert.equal(again.visited.size, 0, "one correct answer on a does not finish it");
	});

	it("revisits keep both: A→B→A gives current A, visited {A,B}", () => {
		const st = derivePlanStatuses(tree, [
			msg(quizPart("output-available", "a", true)),
			msg(quizPart("output-available", "b", true)),
			msg(quizPart("input-available", "a")),
		]);
		assert.equal(st.current, "a");
		assert.deepEqual([...st.visited].sort(), ["a", "b"]);
	});

	it("several checks on one node hold the map still", () => {
		const st = derivePlanStatuses(tree, [
			msg(quizPart("output-available", "a", true)),
			msg(quizPart("output-available", "a", false)),
			msg(quizPart("input-available", "a")),
		]);
		assert.equal(st.current, "a");
		assert.equal(st.visited.size, 0);
	});

	it("untagged checks and unknown nodeIds are ad-hoc — they move nothing", () => {
		const st = derivePlanStatuses(tree, [
			msg(quizPart("output-available", undefined, true)),
			msg(quizPart("output-available", "ghost", true)),
		]);
		assert.equal(st.current, null);
		assert.equal(st.visited.size, 0);
	});

	it("a rewritten map derives against the CURRENT tree — dropped slugs stop counting", () => {
		const rewritten: PlanTree = { ...tree, nodes: [tree.nodes[0]!, tree.nodes[1]!] };
		const st = derivePlanStatuses(rewritten, [
			msg(quizPart("output-available", "a", true)),
			msg(quizPart("output-available", "c", true)), // c removed by the rewrite
		]);
		assert.equal(st.current, "a", "the orphaned tag is ignored");
		assert.equal(st.visited.size, 0);
	});

	it("derivedNodeStatus: current wins over visited; upcoming otherwise", () => {
		const st = derivePlanStatuses(tree, [msg(quizPart("output-available", "a", true))]);
		assert.equal(derivedNodeStatus(st, "a"), "current");
		assert.equal(derivedNodeStatus(st, "b"), "upcoming");
		const later = derivePlanStatuses(tree, [
			msg(quizPart("output-available", "a", true)),
			msg(quizPart("output-available", "b", true)),
		]);
		assert.equal(derivedNodeStatus(later, "a"), "visited");
		assert.equal(derivedNodeStatus(later, "b"), "current");
		assert.equal(derivedNodeStatus(later, "c"), "upcoming");
	});
});