Skip to content

Commit b629a7a

Browse files
ualtinokalfonso-magic-context
ualtinok
and
alfonso-magic-context
committed
mason: harden todowrite state capture
Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>
1 parent 212c26d commit b629a7a

8 files changed

Lines changed: 479 additions & 126 deletions

File tree

‎packages/pi-plugin/src/index.ts‎

Lines changed: 112 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,99 @@ export function persistPiMessageEndModelMeta(args: {
195195
}
196196
}
197197

198+
type TodoOverlayUpdater = { update: (sessionId?: string) => void };
199+
200+
type CompatiblePiTodoCapture = {
201+
normalized: string;
202+
todos: unknown[];
203+
};
204+
205+
function getCompatiblePiTodoCapture(
206+
todos: unknown,
207+
): CompatiblePiTodoCapture | null {
208+
if (!Array.isArray(todos)) return null;
209+
const normalized = normalizeTodoStateJson(todos);
210+
if (normalized === null) return null;
211+
return { normalized, todos };
212+
}
213+
214+
function applyCompatiblePiTodoCapture(args: {
215+
db: ContextDatabase;
216+
sessionId: string;
217+
todowriteEnabled: boolean;
218+
todoOverlay?: TodoOverlayUpdater;
219+
persist: boolean;
220+
capture: CompatiblePiTodoCapture;
221+
}): void {
222+
if (args.todowriteEnabled) {
223+
setTodoSnapshot(args.sessionId, args.capture.todos);
224+
args.todoOverlay?.update(args.sessionId);
225+
}
226+
if (args.persist) {
227+
updateSessionMeta(args.db, args.sessionId, {
228+
lastTodoState: args.capture.normalized,
229+
});
230+
}
231+
}
232+
233+
/**
234+
* Capture a `todowrite` args.todos payload only when it matches Magic Context's
235+
* exact todo enum contract. Third-party Pi extensions can reuse the same tool
236+
* name, so incompatible shapes must leave `last_todo_state` untouched.
237+
*/
238+
export function capturePiTodowriteArgsIfCompatible(args: {
239+
db: ContextDatabase;
240+
sessionId: string;
241+
todos: unknown;
242+
todowriteEnabled: boolean;
243+
todoOverlay?: TodoOverlayUpdater;
244+
persist: boolean;
245+
}): boolean {
246+
const capture = getCompatiblePiTodoCapture(args.todos);
247+
if (capture === null) return false;
248+
applyCompatiblePiTodoCapture({ ...args, capture });
249+
return true;
250+
}
251+
252+
/**
253+
* Scan an assistant `message_end` payload for the first compatible `todowrite`
254+
* call. This keeps interop with third-party tools that share the name but only
255+
* captures state when their payload matches Magic Context's todo enums exactly.
256+
*/
257+
export function capturePiTodowriteMessageIfCompatible(args: {
258+
db: ContextDatabase;
259+
sessionId: string;
260+
message: unknown;
261+
todowriteEnabled: boolean;
262+
todoOverlay?: TodoOverlayUpdater;
263+
persist: boolean;
264+
}): boolean {
265+
const msg = args.message as { role?: unknown; content?: unknown } | undefined;
266+
if (msg?.role !== "assistant" || !Array.isArray(msg.content)) {
267+
return false;
268+
}
269+
270+
for (const block of msg.content) {
271+
if (!block || typeof block !== "object") continue;
272+
const b = block as {
273+
type?: unknown;
274+
name?: unknown;
275+
arguments?: unknown;
276+
};
277+
if (b.type !== "toolCall") continue;
278+
if (typeof b.name !== "string") continue;
279+
if (b.name !== "todowrite") continue;
280+
const capture = getCompatiblePiTodoCapture(
281+
(b.arguments as { todos?: unknown } | null | undefined)?.todos,
282+
);
283+
if (capture === null) continue;
284+
applyCompatiblePiTodoCapture({ ...args, capture });
285+
return true;
286+
}
287+
288+
return false;
289+
}
290+
198291
function info(message: string, data?: unknown): void {
199292
log(`${PREFIX} ${message}`, data);
200293
}
@@ -1470,10 +1563,6 @@ export default async function (pi: ExtensionAPI): Promise<void> {
14701563
const sessionMeta = Array.isArray(todos)
14711564
? getOrCreateSessionMeta(db, sessionId)
14721565
: null;
1473-
if (todowriteEnabled && Array.isArray(todos)) {
1474-
setTodoSnapshot(sessionId, todos);
1475-
todoOverlay?.update(sessionId);
1476-
}
14771566

14781567
// Synthetic-todowrite snapshot capture (Pi parity with
14791568
// OpenCode hook-handlers.ts:386-401). Persist normalized
@@ -1482,15 +1571,17 @@ export default async function (pi: ExtensionAPI): Promise<void> {
14821571
// snapshot to replay on the next cache-busting pass.
14831572
// Cache-safe: this is a pure DB write with no message
14841573
// mutation. Subagents skip — they do not get synthetic
1485-
// todowrite injection.
1486-
if (sessionMeta && !sessionMeta.isSubagent) {
1487-
const normalizedTodos = normalizeTodoStateJson(todos);
1488-
if (normalizedTodos !== null) {
1489-
updateSessionMeta(db, sessionId, {
1490-
lastTodoState: normalizedTodos,
1491-
});
1492-
}
1493-
}
1574+
// todowrite injection. Foreign Pi extensions can share the
1575+
// `todowrite` name, so only the exact Magic Context todo
1576+
// shape updates the stored snapshot.
1577+
capturePiTodowriteArgsIfCompatible({
1578+
db,
1579+
sessionId,
1580+
todos,
1581+
todowriteEnabled,
1582+
todoOverlay,
1583+
persist: Boolean(sessionMeta && !sessionMeta.isSubagent),
1584+
});
14941585

14951586
if (
14961587
Array.isArray(todos) &&
@@ -1725,43 +1816,14 @@ export default async function (pi: ExtensionAPI): Promise<void> {
17251816
try {
17261817
const sessionMetaForTodo = getOrCreateSessionMeta(db, sessionId);
17271818
if (!sessionMetaForTodo.isSubagent) {
1728-
const msg = event.message as
1729-
| { role?: string; content?: unknown }
1730-
| undefined;
1731-
if (msg && msg.role === "assistant" && Array.isArray(msg.content)) {
1732-
for (const block of msg.content) {
1733-
if (!block || typeof block !== "object") continue;
1734-
const b = block as {
1735-
type?: unknown;
1736-
name?: unknown;
1737-
arguments?: unknown;
1738-
};
1739-
if (b.type !== "toolCall") continue;
1740-
if (typeof b.name !== "string") continue;
1741-
if (b.name !== "todowrite") {
1742-
continue;
1743-
}
1744-
const args = b.arguments as
1745-
| { todos?: unknown }
1746-
| null
1747-
| undefined;
1748-
const todos = args?.todos;
1749-
if (!Array.isArray(todos)) continue;
1750-
const normalized = normalizeTodoStateJson(todos);
1751-
if (normalized === null) continue;
1752-
if (todowriteEnabled) {
1753-
setTodoSnapshot(sessionId, todos);
1754-
todoOverlay?.update(sessionId);
1755-
}
1756-
updateSessionMeta(db, sessionId, {
1757-
lastTodoState: normalized,
1758-
});
1759-
// First valid todowrite block wins — mirrors OpenCode's
1760-
// `tool.execute.after` behavior of capturing one
1761-
// snapshot per tool invocation.
1762-
break;
1763-
}
1764-
}
1819+
capturePiTodowriteMessageIfCompatible({
1820+
db,
1821+
sessionId,
1822+
message: event.message,
1823+
todowriteEnabled,
1824+
todoOverlay,
1825+
persist: true,
1826+
});
17651827
}
17661828
} catch (err) {
17671829
warn("message_end: synthetic todowrite capture failed:", err);

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

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import { beforeEach, describe, expect, it } from "bun:test";
2+
import {
3+
getOrCreateSessionMeta,
4+
updateSessionMeta,
5+
} from "@magic-context/core/features/magic-context/storage";
6+
import { closeQuietly } from "@magic-context/core/shared/sqlite-helpers";
7+
8+
import {
9+
capturePiTodowriteArgsIfCompatible,
10+
capturePiTodowriteMessageIfCompatible,
11+
} from "./index";
12+
import { assistantToolCall, createTestDb } from "./test-utils.test";
13+
import {
14+
__resetTodoSnapshotsForTests,
15+
getTodoSnapshot,
16+
} from "./tools/todo-view-pi";
17+
18+
beforeEach(() => {
19+
__resetTodoSnapshotsForTests();
20+
});
21+
22+
describe("Pi todowrite capture compatibility", () => {
23+
it("tool_execution_start capture ignores foreign status values", () => {
24+
const db = createTestDb();
25+
const previousState =
26+
'[{"content":"Keep","status":"pending","priority":"medium"}]';
27+
let overlayUpdates = 0;
28+
try {
29+
updateSessionMeta(db, "ses-args", { lastTodoState: previousState });
30+
31+
const captured = capturePiTodowriteArgsIfCompatible({
32+
db,
33+
sessionId: "ses-args",
34+
todos: [{ content: "Foreign", status: "done" }],
35+
todowriteEnabled: true,
36+
todoOverlay: { update: () => overlayUpdates++ },
37+
persist: true,
38+
});
39+
40+
expect(captured).toBe(false);
41+
expect(getOrCreateSessionMeta(db, "ses-args").lastTodoState).toBe(
42+
previousState,
43+
);
44+
expect(getTodoSnapshot("ses-args").todos).toEqual([]);
45+
expect(overlayUpdates).toBe(0);
46+
} finally {
47+
closeQuietly(db);
48+
}
49+
});
50+
51+
it("message_end capture ignores foreign status values", () => {
52+
const db = createTestDb();
53+
const previousState =
54+
'[{"content":"Keep","status":"pending","priority":"medium"}]';
55+
let overlayUpdates = 0;
56+
try {
57+
updateSessionMeta(db, "ses-message-status", {
58+
lastTodoState: previousState,
59+
});
60+
61+
const captured = capturePiTodowriteMessageIfCompatible({
62+
db,
63+
sessionId: "ses-message-status",
64+
message: assistantToolCall("call-1", "todowrite", {
65+
todos: [{ content: "Foreign", status: "done" }],
66+
}),
67+
todowriteEnabled: true,
68+
todoOverlay: { update: () => overlayUpdates++ },
69+
persist: true,
70+
});
71+
72+
expect(captured).toBe(false);
73+
expect(
74+
getOrCreateSessionMeta(db, "ses-message-status").lastTodoState,
75+
).toBe(previousState);
76+
expect(getTodoSnapshot("ses-message-status").todos).toEqual([]);
77+
expect(overlayUpdates).toBe(0);
78+
} finally {
79+
closeQuietly(db);
80+
}
81+
});
82+
83+
it("message_end capture ignores absent or non-array todos", () => {
84+
const db = createTestDb();
85+
const previousState =
86+
'[{"content":"Keep","status":"pending","priority":"medium"}]';
87+
let overlayUpdates = 0;
88+
try {
89+
updateSessionMeta(db, "ses-message-shape", {
90+
lastTodoState: previousState,
91+
});
92+
93+
expect(
94+
capturePiTodowriteMessageIfCompatible({
95+
db,
96+
sessionId: "ses-message-shape",
97+
message: assistantToolCall("call-1", "todowrite", {}),
98+
todowriteEnabled: true,
99+
todoOverlay: { update: () => overlayUpdates++ },
100+
persist: true,
101+
}),
102+
).toBe(false);
103+
expect(
104+
capturePiTodowriteMessageIfCompatible({
105+
db,
106+
sessionId: "ses-message-shape",
107+
message: assistantToolCall("call-2", "todowrite", {
108+
todos: { content: "Not an array", status: "pending" },
109+
}),
110+
todowriteEnabled: true,
111+
todoOverlay: { update: () => overlayUpdates++ },
112+
persist: true,
113+
}),
114+
).toBe(false);
115+
116+
expect(
117+
getOrCreateSessionMeta(db, "ses-message-shape").lastTodoState,
118+
).toBe(previousState);
119+
expect(getTodoSnapshot("ses-message-shape").todos).toEqual([]);
120+
expect(overlayUpdates).toBe(0);
121+
} finally {
122+
closeQuietly(db);
123+
}
124+
});
125+
126+
it("message_end capture preserves interop for compatible foreign todowrite payloads", () => {
127+
const db = createTestDb();
128+
let overlayUpdates = 0;
129+
const todos = [{ content: "Interop", status: "pending", priority: "high" }];
130+
try {
131+
const captured = capturePiTodowriteMessageIfCompatible({
132+
db,
133+
sessionId: "ses-message-valid",
134+
message: assistantToolCall("call-1", "todowrite", { todos }),
135+
todowriteEnabled: true,
136+
todoOverlay: { update: () => overlayUpdates++ },
137+
persist: true,
138+
});
139+
140+
expect(captured).toBe(true);
141+
expect(
142+
getOrCreateSessionMeta(db, "ses-message-valid").lastTodoState,
143+
).toBe('[{"content":"Interop","status":"pending","priority":"high"}]');
144+
expect(getTodoSnapshot("ses-message-valid").todos).toEqual(todos);
145+
expect(overlayUpdates).toBe(1);
146+
} finally {
147+
closeQuietly(db);
148+
}
149+
});
150+
});

0 commit comments

Comments
 (0)