frontend/src/editor/extensions/agent-highlight.ts

/**
 * Demo-only Tiptap extension that lets the Playwright showcase script mark a
 * paragraph as "the one an agent is currently working on" and have the
 * highlight survive any amount of concurrent Yjs syncing, remote typing,
 * or local selection changes.
 *
 * Why a ProseMirror decoration and not a DOM class?
 * -------------------------------------------------
 * In a Tiptap + Yjs editor, the DOM for a paragraph can be torn down and
 * re-created at any time: every remote update from Bob or Carol triggers
 * a sync transaction, and PM may decide to rebuild the `<p>` (especially
 * under the collaboration plugin). Any `className` we slap on the `<p>`
 * is lost on the next render, which is why previous versions of the
 * demo relied on a 150 ms polling timer to re-apply the class. That
 * worked, but it flickers and leaves a visible gap between "selection
 * happens" and "class comes back".
 *
 * ProseMirror decorations solve this at the right layer: they live in
 * editor state, are re-rendered on every view update, and are remapped
 * through doc changes automatically. Here we store a `{ from, to, phase }`
 * range in plugin state, decorate the paragraph containing `from` with
 * the appropriate `demo-agent-{pending|rewriting}` class, and let PM do
 * the rest.
 *
 * Usage (from the demo script or App.tsx):
 *   editor.commands.setAgentHighlight({ from, to, phase: "pending" });
 *   editor.commands.setAgentHighlight({ phase: "rewriting" });
 *   editor.commands.clearAgentHighlight();
 */
import { Extension } from "@tiptap/core";
import { Plugin, PluginKey } from "@tiptap/pm/state";
import { Decoration, DecorationSet } from "@tiptap/pm/view";

export type AgentHighlightPhase = "pending" | "rewriting";

interface AgentHighlightState {
  from: number | null;
  to: number | null;
  phase: AgentHighlightPhase | null;
}

interface AgentHighlightMeta {
  from?: number | null;
  to?: number | null;
  phase?: AgentHighlightPhase | null;
  clear?: boolean;
}

export const agentHighlightPluginKey =
  new PluginKey<AgentHighlightState>("agent-highlight");

declare module "@tiptap/core" {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface Commands<ReturnType> {
    agentHighlight: {
      /**
       * Set (or update) the highlighted range and phase. Omitting
       * `from`/`to` keeps the existing range; omitting `phase` keeps
       * the existing phase. Both can be updated independently so the
       * App.tsx rewrite handler can flip `pending -> rewriting` without
       * recomputing the range.
       */
      setAgentHighlight: (args: {
        from?: number;
        to?: number;
        phase?: AgentHighlightPhase;
      }) => ReturnType;
      /** Remove any active highlight. */
      clearAgentHighlight: () => ReturnType;
    };
  }
}

export const AgentHighlight = Extension.create({
  name: "agentHighlight",

  addProseMirrorPlugins() {
    return [
      new Plugin<AgentHighlightState>({
        key: agentHighlightPluginKey,
        state: {
          init: (): AgentHighlightState => ({
            from: null,
            to: null,
            phase: null,
          }),
          apply(tr, prev): AgentHighlightState {
            const meta = tr.getMeta(agentHighlightPluginKey) as
              | AgentHighlightMeta
              | undefined;
            if (meta) {
              if (meta.clear) {
                return { from: null, to: null, phase: null };
              }
              return {
                from: meta.from !== undefined ? meta.from : prev.from,
                to: meta.to !== undefined ? meta.to : prev.to,
                phase: meta.phase !== undefined ? meta.phase : prev.phase,
              };
            }

            if (prev.from !== null && prev.to !== null) {
              const from = tr.mapping.map(prev.from, 1);
              const to = tr.mapping.map(prev.to, -1);
              if (to > from) {
                if (from === prev.from && to === prev.to) return prev;
                return { ...prev, from, to };
              }
              return { from: null, to: null, phase: null };
            }

            return prev;
          },
        },
        props: {
          decorations(state) {
            const s = agentHighlightPluginKey.getState(state);
            if (!s || s.from === null || s.to === null || !s.phase) {
              return null;
            }

            let $pos;
            try {
              $pos = state.doc.resolve(s.from);
            } catch {
              return null;
            }

            const depth = $pos.depth;
            if (depth <= 0) return null;

            const nodeStart = $pos.before(depth);
            const nodeEnd = $pos.after(depth);
            const cls = `demo-agent-${s.phase}`;

            return DecorationSet.create(state.doc, [
              Decoration.node(nodeStart, nodeEnd, { class: cls }),
            ]);
          },
        },
      }),
    ];
  },

  addCommands() {
    return {
      setAgentHighlight:
        (args) =>
        ({ tr, dispatch }) => {
          if (dispatch) {
            tr.setMeta(agentHighlightPluginKey, args);
          }
          return true;
        },
      clearAgentHighlight:
        () =>
        ({ tr, dispatch }) => {
          if (dispatch) {
            tr.setMeta(agentHighlightPluginKey, { clear: true });
          }
          return true;
        },
    };
  },
});