mason: fix pi todowrite stale args rendering

Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/pi-plugin/src/index.ts

@@ -144,8 +144,10 @@ import {
 import { withTimeout } from "./timeout";
 import { registerMagicContextTools } from "./tools";
 import {
+	parseTodos,
 	registerTodoOverlay,
 	registerTodoStateLifecycle,
+	rememberTodowriteToolCallTodos,
 	setTodoSnapshot,
 } from "./tools/todo-view-pi";
 
@@ -199,7 +201,7 @@ type TodoOverlayUpdater = { update: (sessionId?: string) => void };
 
 type CompatiblePiTodoCapture = {
 	normalized: string;
-	todos: unknown[];
+	todos: Exclude<ReturnType<typeof parseTodos>, null>;
 };
 
 function getCompatiblePiTodoCapture(
@@ -208,7 +210,9 @@ function getCompatiblePiTodoCapture(
 	if (!Array.isArray(todos)) return null;
 	const normalized = normalizeTodoStateJson(todos);
 	if (normalized === null) return null;
-	return { normalized, todos };
+	const parsed = parseTodos(todos);
+	if (parsed === null) return null;
+	return { normalized, todos: parsed };
 }
 
 function applyCompatiblePiTodoCapture(args: {
@@ -217,8 +221,10 @@ function applyCompatiblePiTodoCapture(args: {
 	todowriteEnabled: boolean;
 	todoOverlay?: TodoOverlayUpdater;
 	persist: boolean;
+	toolCallId?: string;
 	capture: CompatiblePiTodoCapture;
 }): void {
+	rememberTodowriteToolCallTodos(args.toolCallId, args.capture.todos);
 	if (args.todowriteEnabled) {
 		setTodoSnapshot(args.sessionId, args.capture.todos);
 		args.todoOverlay?.update(args.sessionId);
@@ -233,7 +239,8 @@ function applyCompatiblePiTodoCapture(args: {
 /**
  * Capture a `todowrite` args.todos payload only when it matches Magic Context's
  * exact todo enum contract. Third-party Pi extensions can reuse the same tool
- * name, so incompatible shapes must leave `last_todo_state` untouched.
+ * name, so incompatible shapes must not update `last_todo_state` or the
+ * transcript render cache.
  */
 export function capturePiTodowriteArgsIfCompatible(args: {
 	db: ContextDatabase;
@@ -242,6 +249,7 @@ export function capturePiTodowriteArgsIfCompatible(args: {
 	todowriteEnabled: boolean;
 	todoOverlay?: TodoOverlayUpdater;
 	persist: boolean;
+	toolCallId?: string;
 }): boolean {
 	const capture = getCompatiblePiTodoCapture(args.todos);
 	if (capture === null) return false;
@@ -1597,6 +1605,10 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 				const todoArgs = event.args as
 					| { todos?: Array<{ status?: string }> }
 					| undefined;
+				const toolCallId =
+					typeof (event as { toolCallId?: unknown }).toolCallId === "string"
+						? (event as { toolCallId: string }).toolCallId
+						: undefined;
 				const todos = todoArgs?.todos;
 				const sessionMeta = Array.isArray(todos)
 					? getOrCreateSessionMeta(db, sessionId)
@@ -1607,8 +1619,9 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 				// state on EVERY todowrite call so the transform-time
 				// injection path in pi-pipeline.ts always has a current
 				// snapshot to replay on the next cache-busting pass.
-				// Cache-safe: this is a pure DB write with no message
-				// mutation. Subagents skip — they do not get synthetic
+				// Render-safe: this only stores validated todos in shared
+				// session state and the local tool-call cache; it does not
+				// mutate Pi messages. Subagents skip — they do not get synthetic
 				// todowrite injection. Foreign Pi extensions can share the
 				// `todowrite` name, so only the exact Magic Context todo
 				// shape updates the stored snapshot.
@@ -1619,6 +1632,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 					todowriteEnabled,
 					todoOverlay,
 					persist: Boolean(sessionMeta && !sessionMeta.isSubagent),
+					toolCallId,
 				});
 
 				if (

packages/pi-plugin/src/todo-capture-pi.test.ts

@@ -13,13 +13,42 @@ import { assistantToolCall, createTestDb } from "./test-utils.test";
 import {
 	__resetTodoSnapshotsForTests,
 	getTodoSnapshot,
+	renderTodowriteCall,
 } from "./tools/todo-view-pi";
 
+const identityTheme = {
+	fg: (_color: string, text: string) => text,
+	bold: (text: string) => text,
+} as never;
+
 beforeEach(() => {
 	__resetTodoSnapshotsForTests();
 });
 
 describe("Pi todowrite capture compatibility", () => {
+	it("tool_execution_start capture seeds the render cache for compatible payloads", () => {
+		const db = createTestDb();
+		try {
+			const captured = capturePiTodowriteArgsIfCompatible({
+				db,
+				sessionId: "ses-args-cache",
+				todos: [{ content: "Build", status: "in_progress" }],
+				todowriteEnabled: true,
+				persist: true,
+				toolCallId: "call-cache",
+			});
+
+			expect(captured).toBe(true);
+			expect(
+				renderTodowriteCall({}, identityTheme, {
+					toolCallId: "call-cache",
+				} as never).render(80)[0],
+			).toContain("Todos — 1 active");
+		} finally {
+			closeQuietly(db);
+		}
+	});
+
 	it("tool_execution_start capture ignores foreign status values", () => {
 		const db = createTestDb();
 		const previousState =
@@ -35,6 +64,7 @@ describe("Pi todowrite capture compatibility", () => {
 				todowriteEnabled: true,
 				todoOverlay: { update: () => overlayUpdates++ },
 				persist: true,
+				toolCallId: "call-foreign",
 			});
 
 			expect(captured).toBe(false);
@@ -43,6 +73,11 @@ describe("Pi todowrite capture compatibility", () => {
 			);
 			expect(getTodoSnapshot("ses-args").todos).toEqual([]);
 			expect(overlayUpdates).toBe(0);
+			expect(
+				renderTodowriteCall({}, identityTheme, {
+					toolCallId: "call-foreign",
+				} as never).render(80)[0],
+			).toContain("Todos — 0 active");
 		} finally {
 			closeQuietly(db);
 		}

packages/pi-plugin/src/tools/todo-view-pi.test.ts

@@ -8,6 +8,7 @@ import {
 	getTodoSnapshot,
 	registerTodoOverlay,
 	registerTodosCommand,
+	rememberTodowriteToolCallTodos,
 	setTodoSnapshot,
 	TodoOverlay,
 } from "./todo-view-pi";
@@ -89,6 +90,129 @@ describe("todowrite tool rendering", () => {
 			renderCall({ todos } as never, identityTheme, {} as never).render(80)[0],
 		).toContain("Todos — 3 active");
 	});
+	it("uses cached todos when the transcript component was created with empty args", () => {
+		const tool = createTodowriteTool();
+		const renderCall = tool.renderCall;
+		if (!renderCall) throw new Error("todowrite renderCall missing");
+		const todos = [
+			{ id: "1", content: "Work", status: "in_progress" as const },
+			...Array.from({ length: 9 }, (_, index) => ({
+				id: String(index + 2),
+				content: `Done ${index + 2}`,
+				status: "completed" as const,
+			})),
+		];
+		rememberTodowriteToolCallTodos("call-stale", todos);
+		expect(
+			renderCall({}, identityTheme, {
+				toolCallId: "call-stale",
+			} as never).render(80)[0],
+		).toContain("Todos — 1 active");
+	});
+	it("falls back to args when no cached tool-call todos exist", () => {
+		const tool = createTodowriteTool();
+		const renderCall = tool.renderCall;
+		if (!renderCall) throw new Error("todowrite renderCall missing");
+		const todos = [
+			{ id: "1", content: "Plan", status: "pending" as const },
+			{ id: "2", content: "Build", status: "in_progress" as const },
+			{ id: "3", content: "Done", status: "completed" as const },
+		];
+		expect(
+			renderCall({ todos } as never, identityTheme, {
+				toolCallId: "call-miss",
+			} as never).render(80)[0],
+		).toContain("Todos — 2 active");
+	});
+	it("renders identically when cache and args already agree", () => {
+		const tool = createTodowriteTool();
+		const renderCall = tool.renderCall;
+		if (!renderCall) throw new Error("todowrite renderCall missing");
+		const todos = [
+			{ id: "1", content: "Plan", status: "pending" as const },
+			{ id: "2", content: "Done", status: "completed" as const },
+		];
+		rememberTodowriteToolCallTodos("call-equal", todos);
+		expect(
+			renderCall({ todos } as never, identityTheme, {
+				toolCallId: "call-equal",
+			} as never).render(80),
+		).toEqual(
+			renderCall({ todos } as never, identityTheme, {
+				toolCallId: "call-equal-miss",
+			} as never).render(80),
+		);
+	});
+	it("uses cached todos for stale-empty results", () => {
+		const tool = createTodowriteTool();
+		const renderResult = tool.renderResult;
+		if (!renderResult) throw new Error("todowrite renderResult missing");
+		const todos = [
+			{ id: "1", content: "Work", status: "in_progress" as const },
+			{ id: "2", content: "Done", status: "completed" as const },
+		];
+		rememberTodowriteToolCallTodos("call-stale-result", todos);
+		const rendered = renderResult(
+			{
+				details: { todos: [] },
+				content: [{ type: "text", text: "[]" }],
+			},
+			{ expanded: false, isPartial: false },
+			identityTheme,
+			{ toolCallId: "call-stale-result" } as never,
+		)
+			.render(120)
+			.join("\n");
+		expect(rendered).toContain("◐ #1 Work");
+		expect(rendered).toContain("✓ #2 Done");
+	});
+	it("prefers parsed tool results over cached todos when both are present", () => {
+		const tool = createTodowriteTool();
+		const renderResult = tool.renderResult;
+		if (!renderResult) throw new Error("todowrite renderResult missing");
+		rememberTodowriteToolCallTodos("call-real-result", [
+			{ id: "stale", content: "Stale", status: "pending" },
+		]);
+		const rendered = renderResult(
+			{
+				details: {
+					todos: [{ id: "fresh", content: "Fresh", status: "pending" }],
+				},
+				content: [{ type: "text", text: "[]" }],
+			},
+			{ expanded: false, isPartial: false },
+			identityTheme,
+			{ toolCallId: "call-real-result" } as never,
+		)
+			.render(120)
+			.join("\n");
+		expect(rendered).toContain("○ #fresh Fresh");
+		expect(rendered).not.toContain("Stale");
+	});
+	it("evicts the oldest cached tool-call todos after the cap", () => {
+		const tool = createTodowriteTool();
+		const renderCall = tool.renderCall;
+		if (!renderCall) throw new Error("todowrite renderCall missing");
+		for (let index = 1; index <= 51; index++) {
+			rememberTodowriteToolCallTodos(`call-${index}`, [
+				{
+					id: String(index),
+					content: `Task ${index}`,
+					status: "pending",
+				},
+			]);
+		}
+		expect(
+			renderCall({}, identityTheme, { toolCallId: "call-1" } as never).render(
+				80,
+			)[0],
+		).toContain("Todos — 0 active");
+		expect(
+			renderCall({}, identityTheme, { toolCallId: "call-51" } as never).render(
+				80,
+			)[0],
+		).toContain("Todos — 1 active");
+	});
 	it("caps result rows at 12 with a +N more tail", async () => {
 		const tool = createTodowriteTool();
 		const todos = Array.from({ length: 14 }, (_, index) => ({
@@ -258,10 +382,16 @@ describe("TodoOverlay lifecycle", () => {
 		expect(widget.render(120)).toEqual([]);
 	});
 	it("seeds from session_meta.last_todo_state and clears on session switch", async () => {
+		const tool = createTodowriteTool();
+		const renderCall = tool.renderCall;
+		if (!renderCall) throw new Error("todowrite renderCall missing");
 		const handlers = new Map<
 			string,
 			(event: unknown, ctx: unknown) => Promise<void> | void
 		>();
+		rememberTodowriteToolCallTodos("call-session-switch", [
+			{ id: "cached", content: "Cached", status: "pending" },
+		]);
 		registerTodoOverlay(
 			{
 				on: (event, handler) => handlers.set(event, handler as never),
@@ -293,6 +423,11 @@ describe("TodoOverlay lifecycle", () => {
 			"magic-context-todos",
 			undefined,
 		]);
+		expect(
+			renderCall({}, identityTheme, {
+				toolCallId: "call-session-switch",
+			} as never).render(80)[0],
+		).toContain("Todos — 0 active");
 	});
 	it("caps content rows with a +N more tail", () => {
 		setTodoSnapshot(