vscode-cowriting-plugin/test/e2e/suite/s40Provenance.test.ts

import * as assert from "assert";
import * as fs from "fs";
import * as path from "path";
import * as vscode from "vscode";
import type { CowritingApi } from "../../../src/extension";

const WS = process.env.E2E_WORKSPACE!;
const settle = () => new Promise((r) => setTimeout(r, 400));
const AGENT = { kind: "agent" as const, id: "claude", agent: { sdk: "@cline/sdk", model: "sonnet", sessionId: "e2e-s40" } };

async function getApi(): Promise<CowritingApi> {
  const ext = vscode.extensions.getExtension("benstull.vscode-cowriting-plugin")!;
  const api = (await ext.activate()) as CowritingApi;
  assert.ok(api?.attributionController, "exports attribution");
  return api;
}

async function freshDoc(rel: string, body: string): Promise<{ doc: vscode.TextDocument; key: string }> {
  const abs = path.join(WS, rel);
  fs.mkdirSync(path.dirname(abs), { recursive: true });
  fs.writeFileSync(abs, body, "utf8");
  const uri = vscode.Uri.file(abs);
  const doc = await vscode.workspace.openTextDocument(uri);
  await vscode.window.showTextDocument(doc);
  await settle();
  return { doc, key: uri.toString() };
}

// #40 (follow-up to #38): on undo/redo, restore each re-inserted char's EXACT
// prior author attribution rather than leaving it neutral (#38). Driven mid-edit
// (buffer dirty) so the attribution branch runs, not the disk-sync one.
suite("F3 #40 — undo/redo restores exact author attribution (host E2E, no LLM)", () => {
  test("undo of a deletion of AGENT text restores its agent span (not neutral, not human)", async () => {
    const { doc } = await freshDoc("docs/s40agent.md", "Human start. xxxx end.\n");
    const api = await getApi();
    const ctl = api.attributionController;

    // Make "xxxx" agent-authored via the seam → an agent span over "ROBOT".
    const t = "xxxx";
    const start = doc.getText().indexOf(t);
    const ok = await ctl.applyAgentEdit(
      doc,
      new vscode.Range(doc.positionAt(start), doc.positionAt(start + t.length)),
      "ROBOT",
      AGENT,
      { turnId: "turn-s40" },
    );
    assert.strictEqual(ok, true, "seam edit applies");
    await settle();
    const key = api.proposalController.keyFor(doc);
    assert.ok(
      ctl.getSpans(key).some((s) => s.authorKind === "agent" && doc.getText().slice(s.range.start, s.range.end).includes("ROBOT")),
      "ROBOT is agent-attributed after the seam edit",
    );

    // Human deletes "ROBOT" (buffer stays dirty: it already diverged from disk).
    const rs = doc.getText().indexOf("ROBOT");
    const del = new vscode.WorkspaceEdit();
    del.delete(doc.uri, new vscode.Range(doc.positionAt(rs), doc.positionAt(rs + "ROBOT".length)));
    assert.ok(await vscode.workspace.applyEdit(del), "delete applied");
    await settle();
    assert.ok(!doc.getText().includes("ROBOT"), "ROBOT deleted");

    // Undo the deletion → ROBOT re-inserted. #40: its AGENT span is restored.
    await vscode.window.showTextDocument(doc);
    await vscode.commands.executeCommand("undo");
    await settle();
    assert.ok(doc.getText().includes("ROBOT"), "undo restored ROBOT");
    assert.ok(doc.isDirty, "buffer still dirty → attribution branch ran");

    const r2 = doc.getText().indexOf("ROBOT");
    const spans = ctl.getSpans(key);
    const over = spans.filter((s) => s.range.start < r2 + 5 && s.range.end > r2);
    assert.ok(over.length > 0, "restored ROBOT carries a span");
    assert.ok(over.every((s) => s.authorKind === "agent"), `restored ROBOT is agent-attributed, got ${JSON.stringify(over)}`);
  });

  test("edit → undo → redo round-trips attribution to identical state", async () => {
    const { doc } = await freshDoc("docs/s40roundtrip.md", "Base alpha. yyyy omega.\n");
    const api = await getApi();
    const ctl = api.attributionController;
    const key = api.proposalController.keyFor(doc);

    const start = doc.getText().indexOf("yyyy");
    await ctl.applyAgentEdit(
      doc,
      new vscode.Range(doc.positionAt(start), doc.positionAt(start + 4)),
      "BLUEWORD",
      AGENT,
      { turnId: "turn-s40b" },
    );
    await settle();
    const norm = (k: string) =>
      ctl
        .getSpans(k)
        .map((s) => ({ a: s.authorKind, t: doc.getText().slice(s.range.start, s.range.end) }))
        .sort((x, y) => (x.t < y.t ? -1 : 1));
    const afterAgent = JSON.stringify(norm(key));
    assert.ok(afterAgent.includes("BLUEWORD") && afterAgent.includes("agent"), "agent span present after the edit");

    // Delete BLUEWORD, then undo (restore) then redo (re-delete) then undo again.
    const bs = doc.getText().indexOf("BLUEWORD");
    const del = new vscode.WorkspaceEdit();
    del.delete(doc.uri, new vscode.Range(doc.positionAt(bs), doc.positionAt(bs + "BLUEWORD".length)));
    await vscode.workspace.applyEdit(del);
    await settle();
    await vscode.window.showTextDocument(doc);
    await vscode.commands.executeCommand("undo"); // restore BLUEWORD
    await settle();
    assert.ok(doc.getText().includes("BLUEWORD"), "undo restored BLUEWORD");
    assert.strictEqual(JSON.stringify(norm(key)), afterAgent, "attribution after undo matches the pre-deletion state exactly");
  });
});