diff --git a/packages/plugin/src/hooks/magic-context/transform.test.ts b/packages/plugin/src/hooks/magic-context/transform.test.ts index b75d5691..864666a9 100644 --- a/packages/plugin/src/hooks/magic-context/transform.test.ts +++ b/packages/plugin/src/hooks/magic-context/transform.test.ts @@ -23,6 +23,7 @@ import { getLastNudgeUndropped, getOrCreateSessionMeta, getPendingOps, + getSourceContents, getTagById, getTagsBySession, incrementHistorianFailure, @@ -207,6 +208,142 @@ describe("createTransform", () => { expect(toolOutput(messages[1], 1)).toStartWith("§3§ "); }); + it("restores message content when tagging throws mid-pass (no partial §N§ prefixes)", async () => { + //#given + useTempDataHome("context-transform-tag-rollback-"); + const scheduler: Scheduler = { shouldExecute: mock(() => "defer" as const) }; + const db = openDatabase(); + // Real tagger whose assignTag throws on the SECOND call: the first text + // part is tagged + §1§-prefixed in place, then the second throws, so the + // pass aborts with message[0] already carrying a partial prefix. + const realTagger = createTagger(); + let assignCalls = 0; + const tagger: ReturnType = { + ...realTagger, + assignTag(...args: Parameters) { + assignCalls += 1; + if (assignCalls >= 2) { + throw new Error("forced tag failure after first assign"); + } + return realTagger.assignTag(...args); + }, + }; + const transform = createTransform({ + tagger, + scheduler, + contextUsageMap: new Map([ + [ + "ses-rollback", + { usage: { percentage: 46, inputTokens: 92_000 }, updatedAt: Date.now() }, + ], + ]), + db, + historyRefreshSessions: new Set(), + pendingMaterializationSessions: new Set(), + lastHeuristicsTurnId: new Map(), + clearReasoningAge: 50, + protectedTags: 0, + }); + const messages: TestMessage[] = [ + { + info: { id: "m-user", role: "user", sessionID: "ses-rollback" }, + parts: [{ type: "text", text: "Plan this change" }], + }, + { + info: { id: "m-assistant", role: "assistant" }, + parts: [{ type: "text", text: "Implemented the change" }], + }, + ]; + + //#when — tagging throws after the first part was already prefixed. + await transform({}, { messages }); + + //#then — the in-memory rollback leaves NO partial §N§ prefixes behind. + expect(assignCalls).toBeGreaterThanOrEqual(2); + for (const message of messages) { + for (const part of message.parts) { + if (part.type === "text") { + expect(part.text).not.toMatch(/§\d+§/); + } + } + } + expect(text(messages[0], 0)).toBe("Plan this change"); + expect(text(messages[1], 0)).toBe("Implemented the change"); + }); + + it("keeps persisting source_contents and replaying drops on the happy path", async () => { + //#given + useTempDataHome("context-transform-snapshot-happy-"); + const scheduler: Scheduler = { shouldExecute: mock(() => "defer" as const) }; + const db = openDatabase(); + const tagger = createTagger(); + const makeTransform = () => + createTransform({ + tagger, + scheduler, + contextUsageMap: new Map([ + [ + "ses-snapshot-happy", + { usage: { percentage: 30, inputTokens: 60_000 }, updatedAt: Date.now() }, + ], + ]), + db, + historyRefreshSessions: new Set(), + pendingMaterializationSessions: new Set(), + lastHeuristicsTurnId: new Map(), + clearReasoningAge: 50, + protectedTags: 0, + }); + const firstPass: TestMessage[] = [ + { + info: { id: "m-user", role: "user", sessionID: "ses-snapshot-happy" }, + parts: [{ type: "text", text: "Plan this change carefully" }], + }, + { + info: { id: "m-assistant", role: "assistant" }, + parts: [{ type: "text", text: "Implemented the whole thing in detail" }], + }, + ]; + + //#when — first (successful) pass tags both messages. + await makeTransform()({}, { messages: firstPass }); + + //#then — source_contents persisted the PRE-PREFIX original text per tag. + const tags = getTagsBySession(db, "ses-snapshot-happy"); + const userTag = tags.find((tag) => tag.messageId === "m-user:p0"); + const assistantTag = tags.find((tag) => tag.messageId === "m-assistant:p0"); + expect(userTag).toBeDefined(); + expect(assistantTag).toBeDefined(); + const sources = getSourceContents(db, "ses-snapshot-happy", [ + userTag!.tagNumber, + assistantTag!.tagNumber, + ]); + expect(sources.get(userTag!.tagNumber)).toBe("Plan this change carefully"); + expect(sources.get(assistantTag!.tagNumber)).toBe("Implemented the whole thing in detail"); + + //#given — the assistant tag is dropped out of band. + updateTagStatus(db, "ses-snapshot-happy", assistantTag!.tagNumber, "dropped"); + + //#when — a second successful pass replays the dropped status. + const secondPass: TestMessage[] = [ + { + info: { id: "m-user", role: "user", sessionID: "ses-snapshot-happy" }, + parts: [{ type: "text", text: "Plan this change carefully" }], + }, + { + info: { id: "m-assistant", role: "assistant" }, + parts: [{ type: "text", text: "Implemented the whole thing in detail" }], + }, + ]; + await makeTransform()({}, { messages: secondPass }); + + //#then — active user tag keeps prefixed content; dropped assistant tag is + // replayed to the [dropped §N§] sentinel (snapshot/restore is a no-op here). + expect(text(secondPass[0], 0)).toStartWith("§"); + expect(text(secondPass[0], 0)).toContain("Plan this change carefully"); + expect(text(secondPass[1], 0)).toBe(`[dropped §${assistantTag!.tagNumber}§]`); + }); + it("does not inject user messages for emergency nudges (handled by promptAsync)", async () => { //#given useTempDataHome("context-transform-no-user-nudge-"); diff --git a/packages/plugin/src/hooks/magic-context/transform.ts b/packages/plugin/src/hooks/magic-context/transform.ts index 2049b15c..8b159b14 100644 --- a/packages/plugin/src/hooks/magic-context/transform.ts +++ b/packages/plugin/src/hooks/magic-context/transform.ts @@ -158,6 +158,26 @@ function findLastAssistantModel( return null; } +/** + * Deep-clone the visible message window before tagMessages() mutates it. + * + * tagMessages walks the array and rewrites part text in place (assignTag → + * prependTag injects §N§ prefixes, and persisted source restores rewrite + * textPart.text). It intentionally runs WITHOUT an outer DB transaction (see + * tag-messages.ts) so a mid-walk throw can leave earlier messages carrying + * partial §N§ prefixes while the DB has no record of them — silently + * corrupting the prompt sent to the LLM. We snapshot {info, parts} so the + * catch path can roll the in-memory array back to its pre-tag state. + * + * Parts are JSON-like OpenCode payloads, so structuredClone is sufficient. + */ +function snapshotMessagesForTagging(messages: MessageLike[]): MessageLike[] { + return messages.map((message) => ({ + info: structuredClone(message.info), + parts: structuredClone(message.parts), + })); +} + export interface TransformDeps { tagger: Tagger; scheduler: Scheduler; @@ -1102,6 +1122,10 @@ export function createTransform(deps: TransformDeps) { logTransformTiming(sessionId, "injectTemporalMarkers", tTemporal); } + // Snapshot the pre-tag message window so a mid-walk tagging failure can be + // rolled back. tagMessages mutates parts in place and has no outer + // transaction, so without this a partial pass leaves stray §N§ prefixes. + const preTagSnapshot = snapshotMessagesForTagging(messages); let taggingSucceeded = false; try { const t0 = performance.now(); @@ -1145,6 +1169,13 @@ export function createTransform(deps: TransformDeps) { "transform tag persistence failed; continuing without tagging:", error, ); + // Roll back the in-memory mutations applied before the throw. + // tagMessages prefixes/rewrites part text in place, so a mid-walk + // failure can leave earlier messages with partial §N§ prefixes. + // Restore the ARRAY CONTENTS (not a local rebind) so the emitted + // prompt matches the pre-tag state; taggingSucceeded stays false so + // the downstream mutation stages remain gated. + messages.splice(0, messages.length, ...preTagSnapshot); // Drop in-memory tagger state for this session so the next pass // re-loads from the DB. Without this, a stale counter or stale // assignments map can keep producing the same UNIQUE collision