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

import * as assert from "node:assert";
import * as fs from "node:fs";
import * as path from "node:path";
import MarkdownIt from "markdown-it";
import * as vscode from "vscode";
import { cowritingMarkdownItPlugin } from "../../../src/previewAnnotations";
import { activateApi, settle, settleUntil } from "./helpers";

const WS = process.env.E2E_WORKSPACE!;

async function api() {
  const ext = vscode.extensions.getExtension("benstull.vscode-cowriting-plugin")!;
  return (await ext.activate()) as import("../../../src/extension").CowritingApi;
}

suite("coediting registry (PUC-7)", () => {
  test("enter → isCoediting true; exit → false; persists in list()", async () => {
    const { coeditingRegistry } = await api();
    const doc = await vscode.workspace.openTextDocument({ language: "markdown", content: "# t\n\nbody\n" });
    await vscode.window.showTextDocument(doc);
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    assert.strictEqual(coeditingRegistry.isCoediting(doc.uri), true);
    assert.ok(coeditingRegistry.list().includes(doc.uri.toString()));
    await vscode.commands.executeCommand("cowriting.stopCoediting");
    assert.strictEqual(coeditingRegistry.isCoediting(doc.uri), false);
  });

  // Task 3 (INV-10): a document that was never entered into coediting gets no
  // native-diff surface at all. QuickDiff's `provideOriginalResource` isn't
  // directly queryable from a host-E2E test, so this asserts the gate INPUTS
  // structurally: `isCoediting` stays false and the SCM change count — which
  // `ScmSurfaceController` only ever populates for a coediting doc — stays 0.
  test("non-entered doc: no coediting, no change count (structural, INV-10)", async () => {
    const api = await activateApi();
    const doc = await vscode.workspace.openTextDocument({ language: "markdown", content: "untouched\n" });
    await vscode.window.showTextDocument(doc);
    await settle();
    assert.strictEqual(api.coeditingRegistry.isCoediting(doc.uri), false);
    assert.strictEqual(api.scmSurfaceController.changeCount(doc.uri.toString()), 0);
  });

  // Task 4 (INV-10, PUC-7): the no-hijack scenario — a non-entered doc gets no
  // thread-creation surface at all; entering unlocks it; exiting hides the
  // rendered thread (not the sidecar-backed data); re-entering restores it.
  test("non-entered doc gets no surfaces; exit detaches; re-enter restores threads", async () => {
    const api = await activateApi();
    // Task 6 (D10): createThreadOnSelection now also fires the respond-in-
    // thread turn once the doc is coediting — stub it off (not under test
    // here; this suite asserts render/hide/restore, not the reply loop).
    api.threadController.setTurnRunnerForTest(async () => {
      throw new Error("coediting suite: reply-loop turn stubbed off");
    });
    const doc = await vscode.workspace.openTextDocument({ language: "markdown", content: "# a\n\npara\n" });
    const ed = await vscode.window.showTextDocument(doc);
    // not entered → no thread creation
    ed.selection = new vscode.Selection(2, 0, 2, 4);
    const before = await api.threadController.createThreadOnSelection("hi");
    assert.strictEqual(before, undefined);
    // entered → works
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    const id = await api.threadController.createThreadOnSelection("hi");
    assert.ok(id);
    // exit → rendered thread set empty; re-enter → restored from sidecar
    await vscode.commands.executeCommand("cowriting.stopCoediting");
    assert.strictEqual(api.threadController.getRendered(api.proposalController.keyFor(doc)).length, 0);
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    await settle();
    assert.strictEqual(api.threadController.getRendered(api.proposalController.keyFor(doc)).length, 1);
  });

  // PUC-7 restore-on-enter gap: unlike ThreadController (Task 4, above),
  // ProposalController.renderAll and AttributionController.loadAll were never
  // wired to registry.onDidChange — so a document's FIRST enter into
  // coediting with pre-existing sidecar proposals/attribution showed threads
  // but not proposal decorations/CodeLens or committed author-coloring until
  // a later edit happened to fire renderAll/loadAll. A same-session
  // exit-then-re-enter doesn't reproduce this on its own (nothing clears the
  // in-memory render caches on exit), so this test wipes those caches after
  // exit — simulating "never rendered in this session yet" — then re-enters
  // and asserts BOTH surfaces restore via the registry.onDidChange path
  // alone, with no subsequent edit.
  test("re-enter after exit restores proposal + attribution surfaces without an edit (PUC-7)", async () => {
    const api = await activateApi();
    // A workspace FILE (not untitled) — attribution only persists to the
    // sidecar on save, and this test needs the sidecar (not just in-memory
    // live state) to hold the seeded data before re-entering.
    const rel = "docs/puc7-restore.md";
    const abs = path.join(WS, rel);
    fs.writeFileSync(abs, "# t\n\nfoo bar baz\n", "utf8");
    const doc = await vscode.workspace.openTextDocument(vscode.Uri.file(abs));
    await vscode.window.showTextDocument(doc);
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    await settle();

    // Seed a pending proposal (F4) + a committed agent-attributed edit (F3) —
    // the two surfaces the gap left un-restored.
    const target = "bar";
    const start = doc.getText().indexOf(target);
    const proposalId = await vscode.commands.executeCommand<string>("cowriting.proposeAgentEdit", {
      uri: doc.uri.toString(),
      start,
      end: start + target.length,
      newText: "BAR-PROPOSED",
      model: "sonnet",
      sessionId: "e2e-puc7",
      turnId: "turn-puc7",
    });
    assert.ok(proposalId, "propose returns an id");
    // A real text change (not a no-op replace) — "foo" precedes the "bar"
    // target above, so it is unaffected by F12's optimistic-apply of the
    // proposal into the buffer.
    const fooStart = doc.getText().indexOf("foo");
    const applied = await api.attributionController.applyAgentEdit(
      doc,
      new vscode.Range(doc.positionAt(fooStart), doc.positionAt(fooStart + "foo".length)),
      "FOO-CLAUDE",
      { kind: "agent", id: "claude", agent: { sdk: "@cline/sdk", model: "sonnet", sessionId: "e2e-puc7" } },
      { turnId: "turn-puc7" },
    );
    assert.strictEqual(applied, true, "seam edit applies");
    await settle();
    await doc.save(); // persists attribution records to the sidecar (F3 PUC-4)

    const key = api.proposalController.keyFor(doc);
    assert.ok(api.proposalController.getRendered(key).length >= 1, "sanity: proposal rendered pre-exit");
    assert.ok(api.attributionController.getSpans(key).length >= 1, "sanity: attribution rendered pre-exit");

    await vscode.commands.executeCommand("cowriting.stopCoediting");

    // Wipe the controllers' in-memory render caches for this doc — nothing in
    // production does this on exit (by design: exit hides, it does not wipe
    // sidecar-backed state), so this reflection reproduces the "never
    // rendered this session" precondition of a genuinely fresh first-enter.
    (api.proposalController as unknown as { docs: Map<string, unknown> }).docs.delete(key);
    (api.attributionController as unknown as { docs: Map<string, unknown> }).docs.delete(key);
    assert.strictEqual(api.proposalController.getRendered(key).length, 0, "cache cleared pre-re-enter");
    assert.strictEqual(api.attributionController.getSpans(key).length, 0, "cache cleared pre-re-enter");

    // Re-enter — the registry.onDidChange subscribers (PUC-7 fix) must
    // restore both surfaces from the sidecar with no subsequent edit.
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    await settleUntil(
      () => api.proposalController.getRendered(key).length >= 1 && api.attributionController.getSpans(key).length >= 1,
    );
    assert.strictEqual(api.proposalController.getRendered(key).length, 1, "proposal restored on re-enter");
    assert.ok(api.attributionController.getSpans(key).length >= 1, "attribution restored on re-enter");
  });

  // Finding 1 (final whole-branch review, session native-surfaces-exec): the
  // preview-annotation host's INV-10 gate must never leak one document's
  // baseline/spans into another document's preview. The built-in preview's
  // rendered DOM isn't queryable from a host E2E test (§6.8, mirrored from
  // Task 7's own previewAnnotations.test.ts), so this drives the ACTUAL
  // production `previewAnnotationHost` through a real markdown-it instance
  // and asserts on the rendered HTML.
  test("INV-10: a non-coedited doc's preview never inherits another doc's annotations", async () => {
    const api = await activateApi();
    const docA = await vscode.workspace.openTextDocument({ language: "markdown", content: "alpha original\n" });
    await vscode.window.showTextDocument(docA);
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    await settle();
    const start = docA.getText().indexOf("original");
    const okA = await vscode.commands.executeCommand<boolean>("cowriting.applyAgentEdit", {
      uri: docA.uri.toString(),
      start,
      end: start + "original".length,
      newText: "REPLACED",
      model: "sonnet",
      sessionId: "e2e-inv10-leak",
      turnId: "turn-e2e-inv10-leak",
    });
    assert.strictEqual(okA, true, "seam edit applies to docA");
    await settle();

    // docB is opened (becomes the active editor, so it would win any
    // `lastCoeditedUri` fallback if the gate were wrong) but NEVER entered
    // into coediting.
    const docB = await vscode.workspace.openTextDocument({ language: "markdown", content: "bravo untouched\n" });
    await vscode.window.showTextDocument(docB);
    await settle();

    const md = new MarkdownIt();
    cowritingMarkdownItPlugin(md, api.previewAnnotationHost);
    const html = md.render(docB.getText(), { currentDocument: docB.uri });
    assert.ok(!html.includes("cw-ins-claude"), `docB must render with no annotations, got:\n${html}`);
    assert.ok(!html.includes("original"), `docA's baseline text must never leak into docB's preview:\n${html}`);
    assert.ok(html.includes("bravo untouched"), "docB renders its own content unannotated");
  });

  // Finding 1, failure mode (b): `lastCoeditedUri` is never re-validated
  // against the registry, so a doc's OWN preview stayed annotated after
  // `stopCoediting` — an INV-10 violation for the doc whose gate was just
  // closed, not just for a different doc.
  test("INV-10: stopCoediting immediately clears a doc's own preview annotations", async () => {
    const api = await activateApi();
    const doc = await vscode.workspace.openTextDocument({ language: "markdown", content: "gamma original\n" });
    await vscode.window.showTextDocument(doc);
    await vscode.commands.executeCommand("cowriting.coeditDocument");
    await settle();
    const start = doc.getText().indexOf("original");
    const ok = await vscode.commands.executeCommand<boolean>("cowriting.applyAgentEdit", {
      uri: doc.uri.toString(),
      start,
      end: start + "original".length,
      newText: "REPLACED",
      model: "sonnet",
      sessionId: "e2e-inv10-stop",
      turnId: "turn-e2e-inv10-stop",
    });
    assert.strictEqual(ok, true, "seam edit applies");
    await settle();

    const md = new MarkdownIt();
    cowritingMarkdownItPlugin(md, api.previewAnnotationHost);
    const before = md.render(doc.getText(), { currentDocument: doc.uri });
    assert.ok(before.includes("cw-ins-claude"), `sanity: annotated while coediting, got:\n${before}`);

    await vscode.commands.executeCommand("cowriting.stopCoediting");
    await settle();
    const after = md.render(doc.getText(), { currentDocument: doc.uri });
    assert.ok(!after.includes("cw-ins-claude"), `expected a clean render after stopCoediting, got:\n${after}`);
  });
});