fix: plan block renders in flow right after the plan_init call

Emission runs at persistTurn — after the whole turn is composed — so
appending the block buried it under the same turn's teaching prose and
quiz calls (user-observed in chat_mtwgb9lj: the DAG appeared only after
the first check). Splice it in immediately after the (last) successful
plan_init part instead. Two positional tests added.

src/core/plan-blocks.ts

@@ -66,9 +66,24 @@ export function withPlanBlocks(messages: UIMessage[], tree: PlanTree | null | un
 	return messages.map((m) => {
 		if (m.role !== "assistant" || m.parts.some((p) => p.type === "data-plan-block")) return m;
 		if (!hasSuccessfulPlanInit(m)) return m;
-		return {
+		// POSITION, not append: emission runs at persistTurn — after the whole
-			...m,
+		// turn is composed — so appending would bury the block under every
-			parts: [...m.parts, { type: "data-plan-block", id: blockId, data: snapshot(tree, messages) } satisfies PlanBlockPart],
+		// teaching note and quiz of the same turn. The block renders IN FLOW
-		};
+		// immediately after the plan_init call it presents (after the LAST
+		// one if a turn carried multiple — a rewrite).
+		const idx = lastIndexOfSuccessfulPlanInit(m.parts);
+		const parts = m.parts.slice();
+		parts.splice(idx + 1, 0, { type: "data-plan-block", id: blockId, data: snapshot(tree, messages) } satisfies PlanBlockPart);
+		return { ...m, parts };
 	});
 }
+
+/** Index of the last tool-plan_init part with an ok output in this message. */
+function lastIndexOfSuccessfulPlanInit(parts: ReadonlyArray<{ type: string }>): number {
+	let idx = -1;
+	parts.forEach((p, i) => {
+		if (p.type !== "tool-plan_init") return;
+		if ((p as { output?: { ok?: unknown } }).output?.ok === true) idx = i;
+	});
+	return idx;
+}

tests/plan-blocks.test.ts

@@ -29,6 +29,17 @@ const TREE: PlanTree = {
 	version: 1,
 };
 
+/** A rewrite fixture: 3 fresh nodes (x, y, z), different from TREE. */
+const TREE2: PlanTree = {
+	nodes: [
+		{ id: "x", title: "X" },
+		{ id: "y", title: "Y" },
+		{ id: "z", title: "Z" },
+	],
+	edges: [{ from: "x", to: "y" }, { from: "y", to: "z" }],
+	version: 2,
+};
+
 /** An assistant message carrying a successful plan_init with the fixture tree. */
 function planInitMessage(id: string, tree: PlanTree = TREE): UIMessage {
 	return {
@@ -41,6 +52,11 @@ function planInitMessage(id: string, tree: PlanTree = TREE): UIMessage {
 	} as unknown as UIMessage;
 }
 
+/** A message built from raw parts (positional tests need exact part arrays). */
+function msg(role: string, parts: unknown[]): UIMessage {
+	return { id: `m_${Math.random().toString(36).slice(2, 8)}`, role, parts } as unknown as UIMessage;
+}
+
 function blocksOf(messages: UIMessage[]): PlanBlockPart[] {
 	return messages.flatMap((m) => m.parts.filter((p) => p.type === "data-plan-block") as PlanBlockPart[]);
 }
@@ -85,6 +101,40 @@ test("withPlanBlocks: a plan REWRITE (second successful plan_init) gets its own
 	assert.equal(blocks[1]!.data.nodes.length, 3);
 });
 
+test("withPlanBlocks: the block renders IN FLOW immediately after the plan_init part — not appended under the turn's later parts", () => {
+	// Real-world shape (user session chat_mtwgb9lj): the authoring turn
+	// carries plan_init, then presentation prose, then quiz calls — emission
+	// runs at persistTurn AFTER all of that exists.
+	const turn = msg("assistant", [
+		{ type: "step-start" },
+		{ type: "tool-plan_init", state: "output-available", output: { ok: true, tree: TREE } },
+		{ type: "step-start" },
+		{ type: "text", state: "done", text: "Here's the plan." },
+		{ type: "tool-quiz", state: "input-available", input: { question: "q?" } },
+	]);
+	const out = withPlanBlocks([turn], TREE);
+	const types = out[0]!.parts.map((p) => p.type);
+	assert.equal(types[2], "data-plan-block", "the block directly follows the plan_init part");
+	assert.equal(types.indexOf("data-plan-block"), 2);
+	assert.ok(types.indexOf("data-plan-block") < types.indexOf("text"), "block before the prose");
+	assert.ok(types.indexOf("data-plan-block") < types.indexOf("tool-quiz"), "block before the first check");
+	assert.equal(types.at(-1), "tool-quiz", "the quiz is still the turn's last part");
+});
+
+test("withPlanBlocks: with multiple plan_inits in one turn (a rewrite), the block follows the LAST one", () => {
+	const turn = msg("assistant", [
+		{ type: "tool-plan_init", state: "output-available", output: { ok: true, tree: TREE } },
+		{ type: "text", state: "done", text: "v1" },
+		{ type: "tool-plan_init", state: "output-available", output: { ok: true, tree: TREE2 } },
+		{ type: "text", state: "done", text: "v2" },
+	]);
+	const out = withPlanBlocks([turn], TREE2);
+	const types = out[0]!.parts.map((p) => p.type);
+	assert.equal(types[3], "data-plan-block", "block directly after the last successful init");
+	assert.equal(types[4], "text", "the trailing prose stays after the block");
+	assert.equal(out[0]!.parts.filter((p) => p.type === "data-plan-block").length, 1, "one block per message even on a same-turn rewrite");
+});
+
 test("withPlanBlocks: quizzes alone emit nothing (no neighborhood blocks anymore)", () => {
 	const quizTurn: UIMessage = {
 		id: "m1",