frontend/tests/agent-undo-batch.test.ts

/**
 * 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<unknown>;
  undoManager: Y.UndoManager;
  batch: ReturnType<typeof createAgentBatch>;
}

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.",
    ]);
  });
});