/** * Agent Undo Batch Tests * * Covers `docs/TESTS.md` 5.3 - undo batching for AI-agent edits. * * The agent runs multiple tool calls per turn (replaceSelection, applyDiff, * updateFrontmatter, ...). Each tool call triggers one or more Yjs * transactions. Without batching, the user would have to press Cmd+Z three * times to revert a single agent turn. We bracket the turn with * `startAgentBatch()` / `endAgentBatch()` so it collapses into one undo step. * * These tests target `createAgentBatch` against a real `Y.Doc` + * `UndoManager`, mirroring what `useAgentChat` wires at runtime. We don't * spin up TipTap or React because the contract being verified ("3 * collaborative edits inside the batch revert with a single undo()") is * purely a Yjs concern. */ import { describe, it, expect } from "vitest"; import * as Y from "yjs"; import { createAgentBatch } from "../src/editor/agent-undo-batch"; const AGENT_ORIGIN = Symbol("agent"); const USER_ORIGIN = Symbol("user"); interface Fixture { ydoc: Y.Doc; body: Y.XmlFragment; frontmatter: Y.Map; undoManager: Y.UndoManager; batch: ReturnType; } function makeFixture(): Fixture { const ydoc = new Y.Doc(); const body = ydoc.getXmlFragment("default"); const frontmatter = ydoc.getMap("frontmatter"); // Mirror `@tiptap/extension-collaboration`'s setup: one UndoManager that // watches both the doc body and the frontmatter map, so a single undo() // can revert text edits AND metadata mutations together. const undoManager = new Y.UndoManager([body, frontmatter], { captureTimeout: 500, trackedOrigins: new Set([AGENT_ORIGIN, USER_ORIGIN]), }); return { ydoc, body, frontmatter, undoManager, batch: createAgentBatch(undoManager) }; } /** * Helper that mimics a TipTap text insertion: appends a paragraph element * containing the given text to the document body. */ function appendParagraph(ydoc: Y.Doc, body: Y.XmlFragment, text: string, origin: symbol) { ydoc.transact(() => { const p = new Y.XmlElement("paragraph"); p.insert(0, [new Y.XmlText(text)]); body.push([p]); }, origin); } /** * Helper that mimics agent `applyDiff`: rewrites the text content of a * specific paragraph in place. */ function rewriteParagraph( ydoc: Y.Doc, body: Y.XmlFragment, index: number, newText: string, origin: symbol, ) { ydoc.transact(() => { const p = body.get(index) as Y.XmlElement; const t = p.get(0) as Y.XmlText; t.delete(0, t.length); t.insert(0, newText); }, origin); } function bodyAsText(body: Y.XmlFragment): string[] { const out: string[] = []; for (let i = 0; i < body.length; i++) { const node = body.get(i) as Y.XmlElement; const t = node.get(0) as Y.XmlText; out.push(t.toString()); } return out; } describe("agent-undo-batch — primitive contract", () => { it("startAgentBatch + endAgentBatch are idempotent", () => { const { batch } = makeFixture(); expect(batch.isActive()).toBe(false); batch.startAgentBatch(); expect(batch.isActive()).toBe(true); // Calling start again is a no-op. batch.startAgentBatch(); expect(batch.isActive()).toBe(true); batch.endAgentBatch(); expect(batch.isActive()).toBe(false); // Calling end again is a no-op. batch.endAgentBatch(); expect(batch.isActive()).toBe(false); }); it("is a noop when the UndoManager is null", () => { const batch = createAgentBatch(null); expect(() => batch.startAgentBatch()).not.toThrow(); expect(() => batch.endAgentBatch()).not.toThrow(); expect(batch.isActive()).toBe(false); }); }); describe("agent-undo-batch — 5.3.1 batched undo (P0)", () => { it("reverts 3 chained agent edits with a single undo()", () => { const { ydoc, body, frontmatter, undoManager, batch } = makeFixture(); // Simulate a prior user edit that should NOT be touched by the agent's // single undo step. We give it a separate origin so its capture window // is closed before the agent batch starts. appendParagraph(ydoc, body, "User wrote this first.", USER_ORIGIN); expect(bodyAsText(body)).toEqual(["User wrote this first."]); // --- Agent turn --------------------------------------------------------- batch.startAgentBatch(); // Tool 1: replaceSelection-style rewrite of the existing paragraph. rewriteParagraph(ydoc, body, 0, "User wrote this first. (refined by agent)", AGENT_ORIGIN); // Tool 2: insertAtCursor-style append of a new paragraph. appendParagraph(ydoc, body, "Agent added a follow-up.", AGENT_ORIGIN); // Tool 3: updateFrontmatter-style metadata change. ydoc.transact(() => { frontmatter.set("title", "New title from agent"); }, AGENT_ORIGIN); batch.endAgentBatch(); // Sanity check: all three changes landed. expect(bodyAsText(body)).toEqual([ "User wrote this first. (refined by agent)", "Agent added a follow-up.", ]); expect(frontmatter.get("title")).toBe("New title from agent"); // --- Undo once ---------------------------------------------------------- undoManager.undo(); // All three agent edits revert together; the user paragraph stays. expect(bodyAsText(body)).toEqual(["User wrote this first."]); expect(frontmatter.get("title")).toBeUndefined(); }); it("does not lump the previous user edit into the agent step", () => { const { ydoc, body, undoManager, batch } = makeFixture(); // Two distinct user edits, then an agent batch. appendParagraph(ydoc, body, "First user paragraph.", USER_ORIGIN); appendParagraph(ydoc, body, "Second user paragraph.", USER_ORIGIN); batch.startAgentBatch(); appendParagraph(ydoc, body, "Agent paragraph.", AGENT_ORIGIN); batch.endAgentBatch(); expect(bodyAsText(body)).toEqual([ "First user paragraph.", "Second user paragraph.", "Agent paragraph.", ]); // First undo: only the agent edit reverts. undoManager.undo(); expect(bodyAsText(body)).toEqual([ "First user paragraph.", "Second user paragraph.", ]); }); });