vscode-cowriting-plugin/src/proposalController.ts

/**
 * ProposalController — the editor-facing layer for F4 (spec
 * coauthoring-propose-accept §6.2). Owns the proposal lifecycle: the propose
 * ingress (INV-10: NEVER mutates the document), persistence at propose time,
 * resolve-or-flag on load/external change (INV-11 — a proposal's anchor is
 * immutable for its life: no save-time re-fingerprint, unlike threads),
 * anchor bookkeeping into state.live/state.unresolved (no in-editor UI —
 * F10/INV-32 makes the rendered preview the single review surface),
 * and the human-only accept/reject gestures (INV-12). Accept drives the seam
 * (AttributionController.applyAgentEdit, INV-9) so accepted text lands
 * Claude-attributed with zero new attribution code.
 */
import * as vscode from "vscode";
import { SidecarRouter, docIdentity } from "./sidecarRouter";
import { emptyArtifact, type Artifact, type Fingerprint, type Proposal, type Provenance } from "./model";
import { resolve, shift, buildFingerprint, type OffsetRange } from "./anchorer";
import { addProposal, removeProposal, setProposalApplied } from "./proposalModel";
import type { AttributionController } from "./attributionController";
import type { VersionGuard } from "./versionGuard";
import { isAuthorable } from "./workspacePath";
import { wordEditHunks, type ProposalView } from "./trackChangesModel";

/** Test-facing snapshot of what is currently rendered for a document. */
export interface RenderedProposal {
  id: string;
  /** true → resolves exactly, decidable; false → stale/orphaned (INV-11). */
  pending: boolean;
  /** always false — proposals are decide-only, no reply input (INV-12). */
  canReply: boolean;
  turnId?: string;
  range: { start: number; end: number };
}

interface DocState {
  docPath: string;
  uri: vscode.Uri;
  artifact: Artifact;
  /** proposal id -> live offset range (within-session optimization, INV-3). */
  live: Map<string, OffsetRange>;
  /** proposal ids whose anchor did not resolve at last render (stale/orphaned). */
  unresolved: Set<string>;
  /** ids already optimistically applied to the buffer (so the trigger doesn't re-apply). */
  applied: Set<string>;
}

export class ProposalController implements vscode.Disposable {
  private readonly disposables: vscode.Disposable[] = [];
  private readonly docs = new Map<string, DocState>(); // keyed by docPath
  private readonly statusItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 89);
  private readonly onDidChangeProposalsEmitter = new vscode.EventEmitter<{ uri: string }>();
  /** Fires on propose / accept / reject / external sidecar change (F10). */
  readonly onDidChangeProposals = this.onDidChangeProposalsEmitter.event;

  constructor(
    private readonly store: SidecarRouter,
    private readonly attribution: AttributionController,
    private readonly rootDir: string | undefined,
    private readonly guard: VersionGuard,
  ) {
    this.disposables.push(this.statusItem);
    this.disposables.push(this.onDidChangeProposalsEmitter);
    this.disposables.push(
      vscode.workspace.onDidChangeTextDocument((e) => this.onDidChange(e)),
    );
  }

  private fireChanged(document: vscode.TextDocument): void {
    this.onDidChangeProposalsEmitter.fire({ uri: document.uri.toString() });
  }

  private isTracked(document: vscode.TextDocument): boolean {
    return isAuthorable(document.uri.scheme);
  }
  /** The single document key (F8): repo-relative path in-workspace, URI string otherwise. */
  private keyOf(document: vscode.TextDocument): string {
    return this.store.keyOf(docIdentity(document));
  }
  /** The doc key F4 uses (F8 routing) — exposed for F10's preview. */
  keyFor(document: vscode.TextDocument): string {
    return this.keyOf(document);
  }
  /** Resolved proposal views for the F10 preview (anchorStart=null when unresolved). */
  listProposals(document: vscode.TextDocument): ProposalView[] {
    const docPath = this.keyOf(document);
    const artifact = this.store.load(docPath) ?? emptyArtifact(docPath);
    const text = document.getText();
    return artifact.proposals.map((p) => {
      const fp = artifact.anchors[p.anchorId]?.fingerprint;
      const resolved = fp ? resolve(text, fp) : "orphaned";
      return {
        id: p.id,
        anchorStart: resolved === "orphaned" ? null : resolved.start,
        anchorEnd: resolved === "orphaned" ? null : resolved.end,
        replaced: p.original ?? fp?.text ?? "",
        replacement: p.replacement,
        original: p.original,
        author: p.author?.kind === "agent" ? "claude" : "human",
      };
    });
  }
  private ensureState(document: vscode.TextDocument): DocState {
    const docPath = this.keyOf(document);
    let state = this.docs.get(docPath);
    if (!state) {
      state = {
        docPath,
        uri: document.uri,
        artifact: this.store.load(docPath) ?? emptyArtifact(docPath),
        live: new Map(),
        unresolved: new Set(),
        applied: new Set(),
      };
      this.docs.set(docPath, state);
    }
    return state;
  }

  // ---- PUC-1: the propose ingress (INV-10) -----------------------------------------

  /**
   * Record + render a pending proposal. The fingerprint is built by the CALLER
   * at gesture time (editSelection captures it BEFORE the turn, so mid-turn
   * edits can't skew the anchor). Never touches the document.
   */
  async propose(
    document: vscode.TextDocument,
    fp: Fingerprint,
    replacement: string,
    author: Provenance,
    opts?: { turnId?: string; instruction?: string; granularity?: "block" | "single" },
  ): Promise<string | undefined> {
    if (!this.isTracked(document)) return undefined;
    if (this.guard.isReadOnly(this.keyOf(document))) return undefined;
    const docPath = this.keyOf(document);
    let proposalId: string | undefined;
    this.store.update(docPath, (a) => {
      proposalId = addProposal(a, fp, replacement, author, opts).proposalId;
    });
    this.renderAll(document);
    return proposalId;
  }

  // ---- PUC-2/PUC-3: accept / reject (INV-11/INV-12) ----------------------------------

  /** Accept by id — F12: finalize the already-applied text in place (INV-51). */
  async acceptById(docPath: string, proposalId: string, opts?: { silent?: boolean }): Promise<boolean> {
    if (this.isApplied(docPath, proposalId)) {
      // INV-11 (applied path): if an external write mangled the optimistically-applied
      // text so the fingerprint no longer resolves, refuse finalize — same guard as the
      // legacy accept path.  Direct finalizeInPlace calls (CodeLens Accept gesture where
      // the user may have edited inside the applied span) bypass this check intentionally.
      const hit = this.byId(docPath, proposalId);
      if (hit) {
        const document = this.openDoc(hit.state);
        if (document) {
          const fp = hit.state.artifact.anchors[hit.proposal.anchorId]?.fingerprint;
          const resolved = fp ? resolve(document.getText(), fp) : "orphaned";
          if (resolved === "orphaned") {
            if (!opts?.silent) {
              void vscode.window.showWarningMessage(
                "Cowriting: this proposal's target text changed or is missing — undo to restore it, or reject to discard (it is never applied by guess).",
              );
            }
            return false;
          }
        }
      }
      return this.finalizeInPlace(docPath, proposalId);
    }
    // Fallback: a proposal that was never optimistically applied (e.g. orphaned at
    // apply time) keeps the legacy seam-apply accept.
    const hit = this.byId(docPath, proposalId);
    return hit ? this.accept(hit.state, hit.proposal, opts) : false;
  }

  /**
   * #46 (INV-42): accept EVERY pending proposal on a document in one gesture — a
   * batched application of the existing `acceptById` seam, not a new mechanism.
   * Block proposals take the INV-40 word-precise path automatically. Applied in
   * DESCENDING anchor order so an earlier accept never invalidates a later one's
   * offsets; proposals whose anchor can't resolve are SKIPPED (never force-applied)
   * and counted. Returns the applied-vs-skipped tally for the caller to report.
   */
  async acceptAllProposals(document: vscode.TextDocument): Promise<{ applied: number; skipped: number }> {
    if (!this.isTracked(document)) return { applied: 0, skipped: 0 };
    const state = this.ensureState(document);
    state.artifact = this.store.load(state.docPath) ?? emptyArtifact(state.docPath);
    const text = document.getText();
    const items = state.artifact.proposals.map((p) => {
      const fp = state.artifact.anchors[p.anchorId]?.fingerprint;
      const resolved = fp ? resolve(text, fp) : "orphaned";
      return { id: p.id, start: resolved === "orphaned" ? null : resolved.start };
    });
    const resolvable = items
      .filter((i): i is { id: string; start: number } => i.start !== null)
      .sort((a, b) => b.start - a.start);
    let applied = 0;
    let skipped = items.length - resolvable.length; // orphans, skipped up front
    for (const it of resolvable) {
      // silent: one batch report stands in for N per-proposal warnings.
      if (await this.acceptById(state.docPath, it.id, { silent: true })) applied++;
      else skipped++;
    }
    return { applied, skipped };
  }
  /** Reject by proposal id (test-facing twin of the thread-menu gesture). */
  rejectById(docPath: string, proposalId: string): boolean {
    const hit = this.byId(docPath, proposalId);
    if (!hit) return false;
    this.reject(hit.state, hit.proposal);
    return true;
  }

  /** Reject by id — F12: revert the applied text in place (INV-51). */
  async rejectByIdInPlace(docPath: string, proposalId: string): Promise<boolean> {
    if (this.isApplied(docPath, proposalId)) return this.revertInPlace(docPath, proposalId);
    return this.rejectById(docPath, proposalId);
  }

  private async accept(state: DocState, proposal: Proposal, opts?: { silent?: boolean }): Promise<boolean> {
    if (this.guard.isReadOnly(state.docPath)) return false;
    const document = this.openDoc(state);
    if (!document) return false;
    // INV-11: the fingerprint-guard — exact re-resolve at decision time.
    const fp = state.artifact.anchors[proposal.anchorId]?.fingerprint;
    const resolved = fp ? resolve(document.getText(), fp) : "orphaned";
    if (resolved === "orphaned") {
      // #46: accept-all suppresses per-proposal warnings (one batch report instead).
      if (!opts?.silent)
        void vscode.window.showWarningMessage(
          "Cowriting: this proposal's target text changed or is missing — undo to restore it, or reject to discard (it is never applied by guess).",
        );
      this.renderAll(document);
      return false;
    }
    // #47 (INV-40): a BLOCK proposal applies the whole block but attributes only
    // the words Claude actually changed; a single proposal applies its whole range.
    const ok =
      proposal.granularity === "block"
        ? await this.acceptBlock(document, resolved, proposal)
        : await this.attribution.applyAgentEdit(
            document,
            new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
            proposal.replacement,
            proposal.author,
            // No awaits between resolve and the seam call: document.version is current.
            { expectedVersion: document.version, turnId: proposal.turnId },
          );
    if (!ok) {
      if (!opts?.silent)
        void vscode.window.showWarningMessage(
          "Cowriting: the editor rejected the accept — the proposal is still pending.",
        );
      return false;
    }
    this.store.update(state.docPath, (a) => removeProposal(a, proposal.id));
    this.renderAll(document);
    return true;
  }

  /**
   * #47 (INV-40): accept a BLOCK proposal — apply Claude's whole block, but
   * attribute ONLY the runs Claude actually changed. The block is the DECISION
   * unit; the word is the ATTRIBUTION unit. An intra-block word sub-diff
   * (`wordEditHunks`, the raw engine INV-37 used, repurposed — un-anchored so the
   * runs stay disjoint for a batch apply) yields the changed runs; each lands
   * through the F4 seam (Claude-attributed), applied last-position-first so an
   * earlier run's offsets stay valid under the later ones. Unchanged spans within
   * the block are never touched, so their prior authorship stands. (Each run is
   * one seam edit / one undo step — see the spec's deferred note on undo grouping.)
   */
  private async acceptBlock(
    document: vscode.TextDocument,
    resolved: OffsetRange,
    proposal: Proposal,
  ): Promise<boolean> {
    const blockText = document.getText(
      new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
    );
    const subHunks = wordEditHunks(blockText, proposal.replacement);
    if (subHunks.length === 0) return true; // block already equals the proposal — nothing to attribute
    for (const h of [...subHunks].sort((a, b) => b.start - a.start)) {
      const range = new vscode.Range(
        document.positionAt(resolved.start + h.start),
        document.positionAt(resolved.start + h.end),
      );
      const ok = await this.attribution.applyAgentEdit(document, range, h.replacement, proposal.author, {
        expectedVersion: document.version,
        turnId: proposal.turnId,
      });
      if (!ok) return false;
    }
    return true;
  }

  /** True once this proposal's text is in the buffer (optimistic apply ran). */
  isApplied(docPath: string, proposalId: string): boolean {
    return this.docs.get(docPath)?.applied.has(proposalId) ?? false;
  }

  /**
   * F12/#64 (INV-48): optimistically apply a pending proposal INTO the buffer so
   * the editor shows the would-be-accepted result (editable). Reuses the F4
   * word-precise seam (block → per-word hunks, INV-40; single → whole range) but
   * with `landBaseline:false` (the change stays pending). Then re-anchors the
   * proposal to the applied text and stores the original (`setProposalApplied`), so
   * `resolve()` finds it in the mutated buffer and revert/decorate key off it.
   * Idempotent: a no-op if already applied.
   */
  async optimisticApply(document: vscode.TextDocument, proposalId: string): Promise<boolean> {
    if (!this.isTracked(document) || this.guard.isReadOnly(this.keyOf(document))) return false;
    const docPath = this.keyOf(document);
    const state = this.ensureState(document);
    if (state.applied.has(proposalId)) return true;
    state.artifact = this.store.load(docPath) ?? emptyArtifact(docPath);
    const proposal = state.artifact.proposals.find((p) => p.id === proposalId);
    const fp = proposal ? state.artifact.anchors[proposal.anchorId]?.fingerprint : undefined;
    if (!proposal || !fp) return false;
    // Reload-safety (INV-51/54): a proposal that already carries `original` was
    // optimistically applied in a PRIOR session — the buffer holds the applied text
    // and `fp` points at it, but this (fresh) controller's in-memory `applied` set is
    // empty. Re-applying would recapture `original` from the already-applied buffer
    // (= the replacement) and CLOBBER the true revert target, breaking Reject. Mark it
    // applied in memory and stop — `original` is captured exactly once, on first apply.
    if (proposal.original !== undefined) {
      state.applied.add(proposalId);
      this.renderAll(document);
      return true;
    }
    const resolved = resolve(document.getText(), fp);
    if (resolved === "orphaned") return false;
    const original = document.getText(
      new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
    );
    const ok =
      proposal.granularity === "block"
        ? await this.applyBlockOptimistic(document, resolved, proposal)
        : await this.attribution.applyAgentEdit(
            document,
            new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
            proposal.replacement,
            proposal.author,
            { expectedVersion: document.version, turnId: proposal.turnId, landBaseline: false },
          );
    if (!ok) return false;
    state.applied.add(proposalId);
    // Re-anchor to the applied text now in the buffer (its start is unchanged; its
    // end shifts by the net length delta of the replacement).
    const appliedStart = resolved.start;
    const appliedEnd = appliedStart + proposal.replacement.length;
    const appliedFp = buildFingerprint(document.getText(), { start: appliedStart, end: appliedEnd });
    this.store.update(docPath, (a) => setProposalApplied(a, proposalId, appliedFp, original));
    this.renderAll(document);
    return true;
  }

  /** Block optimistic apply: the INV-40 per-word hunks, but landBaseline:false. */
  private async applyBlockOptimistic(
    document: vscode.TextDocument,
    resolved: OffsetRange,
    proposal: Proposal,
  ): Promise<boolean> {
    const blockText = document.getText(
      new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
    );
    const subHunks = wordEditHunks(blockText, proposal.replacement);
    if (subHunks.length === 0) return true;
    for (const h of [...subHunks].sort((a, b) => b.start - a.start)) {
      const range = new vscode.Range(
        document.positionAt(resolved.start + h.start),
        document.positionAt(resolved.start + h.end),
      );
      const ok = await this.attribution.applyAgentEdit(document, range, h.replacement, proposal.author, {
        expectedVersion: document.version, turnId: proposal.turnId, landBaseline: false,
      });
      if (!ok) return false;
    }
    return true;
  }

  /**
   * F12/#64 (INV-51): ACCEPT an optimistically-applied proposal — the text is
   * already in the buffer, so this only advances the F6 baseline (machine-landing,
   * via `attribution.signalLanded`) and clears the proposal. No re-application.
   */
  async finalizeInPlace(docPath: string, proposalId: string): Promise<boolean> {
    const hit = this.byId(docPath, proposalId);
    if (!hit) return false;
    const document = this.openDoc(hit.state);
    if (!document) return false;
    this.attribution.signalLanded(document);
    this.store.update(docPath, (a) => removeProposal(a, proposalId));
    hit.state.applied.delete(proposalId);
    this.renderAll(document);
    return true;
  }

  /**
   * F12/#64 (INV-51): REJECT an optimistically-applied proposal — replace its live
   * applied span with the stored `original`, then clear it. Reverts the whole block
   * regardless of any in-place edits the human made to the inserted text.
   */
  async revertInPlace(docPath: string, proposalId: string): Promise<boolean> {
    const hit = this.byId(docPath, proposalId);
    if (!hit) return false;
    const document = this.openDoc(hit.state);
    if (!document) return false;
    const fp = hit.state.artifact.anchors[hit.proposal.anchorId]?.fingerprint;
    const resolved = fp ? resolve(document.getText(), fp) : "orphaned";
    if (resolved !== "orphaned" && hit.proposal.original !== undefined) {
      const we = new vscode.WorkspaceEdit();
      we.replace(
        document.uri,
        new vscode.Range(document.positionAt(resolved.start), document.positionAt(resolved.end)),
        hit.proposal.original,
      );
      if (!(await vscode.workspace.applyEdit(we))) return false;
    }
    this.store.update(docPath, (a) => removeProposal(a, proposalId));
    hit.state.applied.delete(proposalId);
    this.renderAll(document);
    return true;
  }

  /**
   * F12/#64 (INV-53): reject EVERY pending proposal on a document — revert each in
   * DESCENDING anchor order (so an earlier revert never shifts a later one's
   * offsets), symmetric with #46's accept-all. Returns the reverted count.
   */
  async rejectAll(document: vscode.TextDocument): Promise<{ reverted: number }> {
    if (!this.isTracked(document)) return { reverted: 0 };
    const docPath = this.keyOf(document);
    const state = this.ensureState(document);
    state.artifact = this.store.load(docPath) ?? emptyArtifact(docPath);
    const text = document.getText();
    const ordered = state.artifact.proposals
      .map((p) => {
        const fp = state.artifact.anchors[p.anchorId]?.fingerprint;
        const r = fp ? resolve(text, fp) : "orphaned";
        return { id: p.id, start: r === "orphaned" ? -1 : r.start };
      })
      .sort((a, b) => b.start - a.start);
    let reverted = 0;
    for (const it of ordered) if (await this.revertInPlace(docPath, it.id)) reverted++;
    return { reverted };
  }

  private reject(state: DocState, proposal: Proposal): void {
    if (this.guard.isReadOnly(state.docPath)) return;
    this.store.update(state.docPath, (a) => removeProposal(a, proposal.id));
    const document = this.openDoc(state);
    if (document) this.renderAll(document);
  }

  // ---- PUC-4: load / external change / resolve-or-flag --------------------------------

  /** Load + (re)render every pending proposal at its resolved anchor (or flagged). */
  renderAll(document: vscode.TextDocument): void {
    if (!this.isTracked(document)) return;
    const docPath = this.keyOf(document);
    const state = this.ensureState(document);
    state.artifact = this.store.load(docPath) ?? emptyArtifact(docPath);
    state.live.clear();
    state.unresolved.clear();
    const text = document.getText();
    for (const proposal of state.artifact.proposals) {
      const fp = state.artifact.anchors[proposal.anchorId]?.fingerprint;
      const resolved = fp ? resolve(text, fp) : "orphaned";
      if (resolved === "orphaned") {
        const line = fp ? Math.min(fp.lineHint, Math.max(0, document.lineCount - 1)) : 0;
        const off = document.offsetAt(new vscode.Position(line, 0));
        this.recordProposal(state, proposal, { start: off, end: off }, false);
      } else {
        this.recordProposal(state, proposal, resolved, true);
      }
    }
    this.renderStatus(state);
    this.fireChanged(document);
  }

  /** Shared-watcher entry point (extension.ts): a sidecar changed externally. */
  handleExternalSidecarChange(uri: vscode.Uri): void {
    for (const state of this.docs.values()) {
      if (this.store.sidecarPath(state.docPath) === uri.fsPath) {
        const doc = vscode.workspace.textDocuments.find((d) => this.keyOf(d) === state.docPath);
        if (doc) this.renderAll(doc);
      }
    }
  }

  private onDidChange(e: vscode.TextDocumentChangeEvent): void {
    const state = this.docs.get(this.keyOf(e.document));
    if (!state || state.live.size === 0) return;
    for (const change of e.contentChanges) {
      const edit = { start: change.rangeOffset, end: change.rangeOffset + change.rangeLength, newLength: change.text.length };
      for (const [id, range] of state.live) state.live.set(id, shift(range, edit));
    }
  }

  // ---- rendering -----------------------------------------------------------------------

  /**
   * Record a proposal's resolved anchor in the live/unresolved bookkeeping.
   * No editor UI (F10/INV-32: the rendered preview is the single review
   * surface) — this keeps state.live/state.unresolved populated so the preview
   * (SLICE-3) and the getRendered test seam can read it.
   */
  private recordProposal(state: DocState, proposal: Proposal, offsets: OffsetRange, pending: boolean): void {
    state.live.set(proposal.id, offsets);
    if (!pending) state.unresolved.add(proposal.id);
  }

  private renderStatus(state: DocState): void {
    const n = state.unresolved.size;
    if (n === 0) {
      this.statusItem.hide();
      return;
    }
    this.statusItem.text = `$(warning) ${n} stale proposal${n === 1 ? "" : "s"}`;
    this.statusItem.tooltip =
      "Cowriting: proposals whose target text changed or is missing — undo to restore, or reject";
    this.statusItem.show();
  }

  // ---- lookups ----------------------------------------------------------------------------

  private openDoc(state: DocState): vscode.TextDocument | undefined {
    return vscode.workspace.textDocuments.find((d) => this.keyOf(d) === state.docPath);
  }
  private byId(docPath: string, proposalId: string): { state: DocState; proposal: Proposal } | undefined {
    const state = this.docs.get(docPath);
    const proposal = state?.artifact.proposals.find((p) => p.id === proposalId);
    return state && proposal ? { state, proposal } : undefined;
  }

  // ---- test-facing surface ------------------------------------------------------------------

  getRendered(docPath: string): RenderedProposal[] {
    const state = this.docs.get(docPath);
    if (!state) return [];
    const out: RenderedProposal[] = [];
    for (const [id, off] of state.live) {
      const p = state.artifact.proposals.find((x) => x.id === id)!;
      out.push({ id, pending: !state.unresolved.has(id), canReply: false, turnId: p.turnId, range: { start: off.start, end: off.end } });
    }
    return out;
  }
  getStaleCount(docPath: string): number {
    return this.docs.get(docPath)?.unresolved.size ?? 0;
  }

  dispose(): void {
    for (const d of this.disposables) d.dispose();
  }
}