From ec56794029655d5b6054d2bec5f94df23dea477f Mon Sep 17 00:00:00 2001 From: amosblomqvist Date: Fri, 11 Sep 2026 21:19:45 +0300 Subject: [PATCH] fix: plan block renders in flow right after the plan_init call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 23 ++++++++++++++---- tests/plan-blocks.test.ts | 50 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/src/core/plan-blocks.ts b/src/core/plan-blocks.ts index 65a6e13..32d60a4 100644 --- a/src/core/plan-blocks.ts +++ b/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 { - ...m, - parts: [...m.parts, { type: "data-plan-block", id: blockId, data: snapshot(tree, messages) } satisfies PlanBlockPart], - }; + // POSITION, not append: emission runs at persistTurn — after the whole + // turn is composed — so appending would bury the block under every + // 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; +} diff --git a/tests/plan-blocks.test.ts b/tests/plan-blocks.test.ts index 7093b4c..a7267fa 100644 --- a/tests/plan-blocks.test.ts +++ b/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",