vscode-cowriting-plugin/src/attributionController.ts

/**
 * AttributionController — the thin editor-facing layer for F3 (spec §6.2).
 * Wires the pure AttributionTracker + PendingEditRegistry + Store/Anchorer to
 * the editor: human typing → human spans, seam edits → agent spans (INV-9),
 * decorations (Claude tint / human gutter border / toggle), save-time
 * persistence, load/external-change resolve-or-orphan (INV-1/INV-6), and the
 * orphan status-bar count. Sidecar self-write suppression and the shared
 * FileSystemWatcher live in CoauthorStore / extension.ts (the sidecar is
 * co-owned with ThreadController).
 */
import * as fs from "node:fs";
import * as vscode from "vscode";
import { CoauthorStore } from "./store";
import { newId, type AttributionRecord, type Provenance } from "./model";
import { buildFingerprint, resolve, type OffsetRange } from "./anchorer";
import { applyChange, coalesce, type LiveSpan } from "./attributionTracker";
import { minimizeReplace, PendingEditRegistry } from "./pendingEdits";

/** Test-facing snapshot of live attribution state for a document. */
export interface RenderedSpan {
  id: string;
  authorKind: "human" | "agent";
  authorId: string;
  turnId?: string;
  range: { start: number; end: number };
}

interface DocAttribution {
  docPath: string;
  spans: LiveSpan[];
  /** persisted records whose fingerprints failed to resolve (PUC-4). */
  orphans: AttributionRecord[];
  /** record metadata per live span id (updatedAt bookkeeping). */
  records: Map<string, AttributionRecord>;
  /**
   * true once this doc has had any span/record this session — lets save persist
   * deliberate deletion to empty (so stale records don't come back as phantom
   * orphans on reload).
   */
  hadAttributions: boolean;
}

const AGENT_DECO: vscode.DecorationRenderOptions = {
  backgroundColor: "rgba(99, 102, 241, 0.18)",
  overviewRulerColor: "rgba(99, 102, 241, 0.8)",
  overviewRulerLane: vscode.OverviewRulerLane.Right,
};
const HUMAN_DECO: vscode.DecorationRenderOptions = {
  borderColor: "rgba(16, 185, 129, 0.8)",
  borderStyle: "solid",
  borderWidth: "0 0 0 2px",
};

export class AttributionController implements vscode.Disposable {
  private readonly disposables: vscode.Disposable[] = [];
  private readonly docs = new Map<string, DocAttribution>();
  private readonly pending = new PendingEditRegistry();
  private readonly agentType = vscode.window.createTextEditorDecorationType(AGENT_DECO);
  private readonly humanType = vscode.window.createTextEditorDecorationType(HUMAN_DECO);
  private readonly statusItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 90);
  private readonly output = vscode.window.createOutputChannel("Cowriting Attribution");
  private visible = true;

  constructor(private readonly store: CoauthorStore, private readonly rootDir: string) {
    this.disposables.push(this.agentType, this.humanType, this.statusItem, this.output);
    this.disposables.push(
      vscode.commands.registerCommand("cowriting.toggleAttribution", () => this.toggle()),
      vscode.workspace.onDidChangeTextDocument((e) => this.onDidChange(e)),
      vscode.workspace.onDidSaveTextDocument((d) => this.onDidSave(d)),
      vscode.window.onDidChangeActiveTextEditor(() => this.renderActive()),
    );
  }

  private isTracked(document: vscode.TextDocument): boolean {
    return document.uri.scheme === "file" && document.uri.fsPath.startsWith(this.rootDir);
  }
  private docPathOf(uri: vscode.Uri): string {
    return vscode.workspace.asRelativePath(uri, false);
  }
  private currentAuthor(): Provenance {
    const id = vscode.workspace.getConfiguration("git").get<string>("user.name") || process.env.USER || "human";
    return { kind: "human", id };
  }
  private state(docPath: string): DocAttribution {
    let s = this.docs.get(docPath);
    if (!s) {
      s = { docPath, spans: [], orphans: [], records: new Map(), hadAttributions: false };
      this.docs.set(docPath, s);
    }
    return s;
  }

  // ---- PUC-4: load / resolve-or-orphan ---------------------------------------------

  /** Load the sidecar and re-resolve every attribution (live span | orphan). */
  loadAll(document: vscode.TextDocument): void {
    if (!this.isTracked(document)) return;
    const docPath = this.docPathOf(document.uri);
    const s = this.state(docPath);
    const artifact = this.store.load(docPath);
    s.spans = [];
    s.orphans = [];
    s.records.clear();
    if (artifact) {
      const text = document.getText();
      for (const rec of artifact.attributions) {
        const fp = artifact.anchors[rec.anchorId]?.fingerprint;
        const resolved = fp ? resolve(text, fp) : "orphaned";
        if (resolved === "orphaned") {
          s.orphans.push(rec);
        } else {
          s.spans.push({
            id: rec.id, start: resolved.start, end: resolved.end,
            author: rec.author, turnId: rec.turnId, createdAt: rec.createdAt,
          });
          s.records.set(rec.id, rec);
        }
      }
      s.spans = coalesce(s.spans);
      if (artifact.attributions.length > 0) s.hadAttributions = true;
    }
    this.render(document);
  }

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

  // ---- PUC-1/PUC-3: live tracking ---------------------------------------------------

  private onDidChange(e: vscode.TextDocumentChangeEvent): void {
    if (!this.isTracked(e.document) || e.contentChanges.length === 0) return;
    const docPath = this.docPathOf(e.document.uri);
    if (!e.document.isDirty && this.matchesDisk(e.document)) {
      // Disk sync (revert / external reload): buffer now equals the file on
      // disk — re-resolve, never attribute (PUC-4). A real edit can also
      // arrive with isDirty still false (VS Code flips the flag after the
      // change event), but then the buffer no longer matches the disk.
      this.loadAll(e.document);
      return;
    }
    const s = this.state(docPath);
    // One applyEdit = one change event, but the host may deliver a seam edit
    // as SEVERAL minimal hunks (word-level diffing). Match the EVENT's net
    // effect against the registry; on a hit the agent owns its FULL intended
    // replacement (INV-9) — apply it as ONE algebra edit, not per hunk.
    const hit = this.pending.matchEvent(
      docPath,
      e.contentChanges.map((c) => ({
        start: c.rangeOffset,
        end: c.rangeOffset + c.rangeLength,
        newLength: c.text.length,
      })),
    );
    if (hit) {
      const full = hit.full ?? { start: hit.start, end: hit.end, newLength: hit.newText.length };
      s.spans = applyChange(s.spans, full, hit.provenance, {
        newId: () => newId("at"),
        now: () => new Date().toISOString(),
        turnId: hit.turnId,
      });
    } else {
      // Sort descending by offset so earlier changes don't invalidate later offsets
      // (VS Code's order is undocumented; defensive sort is the safe guarantee).
      for (const change of [...e.contentChanges].sort((a, b) => b.rangeOffset - a.rangeOffset)) {
        const edit = {
          start: change.rangeOffset,
          end: change.rangeOffset + change.rangeLength,
          newLength: change.text.length,
        };
        s.spans = applyChange(s.spans, edit, this.currentAuthor(), {
          newId: () => newId("at"),
          now: () => new Date().toISOString(),
        });
      }
    }
    if (s.spans.length > 0) s.hadAttributions = true;
    this.render(e.document);
  }

  /**
   * True when the document buffer is byte-identical to the file on disk.
   * Fires only on `!isDirty` change events — first change after every clean
   * state (frequent under `files.autoSave: afterDelay`) and reverts/reloads;
   * O(file size) sync read. A size pre-check avoids the full read in the
   * common case: if the byte lengths differ (accounting for a possible 3-byte
   * UTF-8 BOM that `document.getText()` never includes) we return early. When
   * a BOM is present the disk read is stripped before comparing so a genuine
   * revert does not clobber attribution.
   */
  private matchesDisk(document: vscode.TextDocument): boolean {
    try {
      const bufLen = Buffer.byteLength(document.getText(), "utf8");
      const stat = fs.statSync(document.uri.fsPath);
      // Allow size === bufLen (no BOM) OR size === bufLen + 3 (UTF-8 BOM).
      if (stat.size !== bufLen && stat.size !== bufLen + 3) return false;
      const disk = fs.readFileSync(document.uri.fsPath, "utf8").replace(/^/, "");
      return disk === document.getText();
    } catch {
      // Unreadable/missing file: treat as a real edit (attribute), per
      // fail-open honesty — misclassifying a sync as an edit is recoverable.
      return false;
    }
  }

  // ---- the seam (INV-9) ---------------------------------------------------------------

  /**
   * The ONLY machine-edit ingress: register the exact expected change, then
   * apply it as a WorkspaceEdit. Returns false (applying nothing) on a stale
   * document version or a rejected edit — never partial-applies (spec §6.9).
   * The pending edit is unregistered unconditionally afterwards (a no-op when
   * `match` already consumed it), so a no-op edit never leaks a registration.
   * The replace is self-minimized (common prefix/suffix trimmed) to mirror the
   * host's WorkspaceEdit diff-minimization, so registered == applied == delivered.
   * Callers issuing concurrent edits on the same document must serialize them or
   * pass `expectedVersion` (offsets are computed against the call-time snapshot).
   */
  async applyAgentEdit(
    document: vscode.TextDocument,
    range: vscode.Range,
    newText: string,
    provenance: Provenance,
    opts?: { expectedVersion?: number; turnId?: string },
  ): Promise<boolean> {
    if (!this.isTracked(document)) return false;
    if (opts?.expectedVersion !== undefined && document.version !== opts.expectedVersion) return false;
    const docPath = this.docPathOf(document.uri);
    const startOffset = document.offsetAt(range.start);
    const endOffset = document.offsetAt(range.end);
    const oldText = document.getText(range);
    const { prefix, suffix } = minimizeReplace(oldText, newText);
    const minStart = startOffset + prefix;
    const minEnd = endOffset - suffix;
    const minText = newText.slice(prefix, newText.length - suffix);
    if (minStart === minEnd && minText.length === 0) {
      // Replacement equals the existing text — a no-op; nothing to attribute.
      return true;
    }
    const pendingEdit = {
      docPath,
      start: minStart,
      end: minEnd,
      newText: minText,
      provenance,
      turnId: opts?.turnId,
      full: { start: startOffset, end: endOffset, newLength: newText.length },
    };
    this.pending.register(pendingEdit);
    const we = new vscode.WorkspaceEdit();
    we.replace(document.uri, new vscode.Range(document.positionAt(minStart), document.positionAt(minEnd)), minText);
    const ok = await vscode.workspace.applyEdit(we);
    const removed = this.pending.unregister(pendingEdit);
    if (ok && removed) {
      // workspace.applyEdit resolves AFTER onDidChangeTextDocument is dispatched
      // synchronously to listeners, so a matching change event should have already
      // consumed the registration before we reach here. If it is still present,
      // the host minimized the diff differently than we predicted — attribution
      // may be wrong for this edit (INV-9).
      this.output.appendLine(
        "WARN: seam edit applied but its change event never matched the registration " +
        "(host minimized differently?) — the edit may be mis-attributed (INV-9).",
      );
    }
    return ok;
  }

  // ---- PUC-4: persistence on save ----------------------------------------------------

  private onDidSave(document: vscode.TextDocument): void {
    if (!this.isTracked(document)) return;
    const docPath = this.docPathOf(document.uri);
    const s = this.docs.get(docPath);
    // Allow save when hadAttributions is true even if spans/orphans are now empty:
    // that means the user deliberately deleted all attributed text, and we must
    // persist a.attributions=[] so stale records don't return as phantom orphans.
    if (!s || (!s.hadAttributions && s.spans.length === 0 && s.orphans.length === 0)) return;
    const text = document.getText();
    const now = new Date().toISOString();
    const records: AttributionRecord[] = [];
    const anchorOf = new Map<string, OffsetRange>();
    for (const span of s.spans) {
      const prev = s.records.get(span.id);
      const rec: AttributionRecord = {
        id: span.id,
        anchorId: prev?.anchorId ?? newId("a"),
        author: span.author,
        createdAt: span.createdAt,
        updatedAt: prev ? prev.updatedAt : now,
        ...(span.turnId !== undefined ? { turnId: span.turnId } : {}),
      };
      anchorOf.set(rec.anchorId, { start: span.start, end: span.end });
      records.push(rec);
      s.records.set(span.id, rec);
    }
    this.store.update(docPath, (a) => {
      for (const rec of records) {
        const range = anchorOf.get(rec.anchorId)!;
        const fp = buildFingerprint(text, range);
        const prevFp = a.anchors[rec.anchorId]?.fingerprint;
        if (prevFp && JSON.stringify(prevFp) !== JSON.stringify(fp)) rec.updatedAt = now;
        a.anchors[rec.anchorId] = { fingerprint: fp };
      }
      // Orphans ride along unchanged (recoverable, spec §6.9): their records
      // stay in attributions[], so update()'s prune keeps their anchors too.
      a.attributions = [...records, ...s.orphans];
    });
  }

  // ---- PUC-5: rendering ----------------------------------------------------------------

  private toggle(): void {
    this.visible = !this.visible;
    this.renderActive();
  }

  private renderActive(): void {
    const editor = vscode.window.activeTextEditor;
    if (editor) this.render(editor.document);
  }

  private render(document: vscode.TextDocument): void {
    if (!this.isTracked(document)) return;
    const s = this.docs.get(this.docPathOf(document.uri));
    const spans = this.visible && s ? s.spans : [];
    const toRange = (sp: LiveSpan) =>
      new vscode.Range(document.positionAt(sp.start), document.positionAt(sp.end));
    const agentRanges = spans.filter((x) => x.author.kind === "agent").map(toRange);
    const humanRanges = spans.filter((x) => x.author.kind === "human").map(toRange);
    // Apply decorations to ALL visible split-editors showing this document, not
    // just the first match — each editor pane has its own decoration layer.
    for (const editor of vscode.window.visibleTextEditors) {
      if (editor.document === document) {
        editor.setDecorations(this.agentType, agentRanges);
        editor.setDecorations(this.humanType, humanRanges);
      }
    }
    if (document === vscode.window.activeTextEditor?.document) {
      this.renderStatus(s);
    }
  }

  private renderStatus(s: DocAttribution | undefined): void {
    const n = s?.orphans.length ?? 0;
    if (n === 0) {
      this.statusItem.hide();
      return;
    }
    this.statusItem.text = `$(warning) ${n} orphaned attribution${n === 1 ? "" : "s"}`;
    this.statusItem.tooltip = "Cowriting: attribution anchors that no longer resolve (see output channel)";
    this.statusItem.show();
    this.output.clear();
    for (const o of s!.orphans) this.output.appendLine(`orphaned ${o.id} (${o.author.kind}:${o.author.id})`);
  }

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

  getSpans(docPath: string): RenderedSpan[] {
    const s = this.docs.get(docPath);
    if (!s) return [];
    return s.spans.map((sp) => ({
      id: sp.id,
      authorKind: sp.author.kind,
      authorId: sp.author.id,
      turnId: sp.turnId,
      range: { start: sp.start, end: sp.end },
    }));
  }
  getOrphanCount(docPath: string): number {
    return this.docs.get(docPath)?.orphans.length ?? 0;
  }
  isVisible(): boolean {
    return this.visible;
  }

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