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

/**
 * Broadcasts "this user's AI agent is currently focused on this range" to
 * every collaborator via Yjs awareness, and renders it as a ProseMirror
 * decoration tinted with the owner's color.
 *
 * Design notes
 * ------------
 * - We combine two Tiptap patterns that already exist separately in the
 *   ecosystem:
 *     1. the paid `@tiptap-pro/ai-toolkit` `setActiveSelection` feature,
 *        which freezes a range locally while the agent is working on it;
 *     2. `@tiptap/extension-collaboration-caret`, which shares cursor
 *        positions across peers via Yjs awareness using relative
 *        positions.
 *
 * - Ranges are stored in awareness as **Y.js relative positions**
 *   (JSON form for safe transport) so they remain meaningful on every
 *   peer's doc even as Bob or Carol are typing concurrently. We resolve
 *   them back to absolute positions at decoration time.
 *
 * - Every peer (including self) gets both the inline tint AND the
 *   floating "<Name> agent" label. When the chat panel is open the
 *   editor often loses DOM focus to the chat textarea, which makes
 *   the native `::selection` fade or disappear: the PM decoration is
 *   then the ONLY thing keeping the range visible on the owner's
 *   screen. To avoid stacking two tints when the editor still has
 *   focus, `_ui.css` hides the native `::selection` inside
 *   `.ProseMirror` while `.editor-app--chat-open` is set.
 *
 * Usage (from App.tsx):
 *   editor.commands.setAgentFocus({ from, to }); // on chat open
 *   editor.commands.clearAgentFocus();           // on chat close
 */
import { Extension } from "@tiptap/core";
import { Plugin, PluginKey } from "@tiptap/pm/state";
import { Decoration, DecorationSet } from "@tiptap/pm/view";
import {
  absolutePositionToRelativePosition,
  relativePositionToAbsolutePosition,
  ySyncPluginKey,
} from "@tiptap/y-tiptap";
import * as Y from "yjs";

export interface AgentFocusUser {
  name: string;
  color: string;
  avatarUrl?: string;
}

export interface AgentFocusOptions {
  // Any Yjs provider with a `.awareness` API (Hocuspocus, y-websocket, ...)
  provider: any;
  user: AgentFocusUser;
}

interface AwarenessFocus {
  anchor: unknown;
  head: unknown;
  user: AgentFocusUser;
}

export const agentFocusPluginKey = new PluginKey("agent-focus");

declare module "@tiptap/core" {
  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  interface Commands<ReturnType> {
    agentFocus: {
      /** Mark `[from, to]` as "my agent is working on this". */
      setAgentFocus: (args: { from: number; to: number }) => ReturnType;
      /** Remove our own agent-focus highlight. */
      clearAgentFocus: () => ReturnType;
      /** Update the awareness user (name / color / avatar). */
      updateAgentFocusUser: (user: AgentFocusUser) => ReturnType;
    };
  }
}

export const AgentFocus = Extension.create<AgentFocusOptions>({
  name: "agentFocus",

  addOptions() {
    return {
      provider: null,
      user: { name: "Anonymous", color: "#5c7cfa" },
    };
  },

  onCreate() {
    if (!this.options.provider) {
      throw new Error('The "provider" option is required for AgentFocus');
    }
  },

  addCommands() {
    // None of these commands mutate the ProseMirror document - they only
    // touch Yjs awareness. That's what guarantees they stay out of the
    // undo stack: `Y.UndoManager` tracks CRDT ops, awareness is
    // ephemeral metadata. Tiptap will still dispatch an empty transaction
    // after each command returns `true`, but an empty tr has
    // `tr.docChanged === false`, so y-sync ignores it and undo is clean.
    return {
      setAgentFocus:
        ({ from, to }) =>
        ({ state }) => {
          const ystate = ySyncPluginKey.getState(state);
          if (!ystate?.binding) return false;

          const size = state.doc.content.size;
          const a = Math.max(0, Math.min(from, size));
          const b = Math.max(0, Math.min(to, size));
          if (a === b) return false;
          const [lo, hi] = a < b ? [a, b] : [b, a];

          const anchor = absolutePositionToRelativePosition(
            lo,
            ystate.type,
            ystate.binding.mapping,
          );
          const head = absolutePositionToRelativePosition(
            hi,
            ystate.type,
            ystate.binding.mapping,
          );

          this.options.provider.awareness.setLocalStateField("agentFocus", {
            anchor: Y.relativePositionToJSON(anchor),
            head: Y.relativePositionToJSON(head),
            user: this.options.user,
          });
          return true;
        },

      clearAgentFocus: () => () => {
        this.options.provider.awareness.setLocalStateField("agentFocus", null);
        return true;
      },

      updateAgentFocusUser: (user) => () => {
        this.options.user = user;
        // Re-emit current focus with the updated user payload, if any.
        const current = this.options.provider.awareness.getLocalState()
          ?.agentFocus as AwarenessFocus | null | undefined;
        if (current?.anchor && current?.head) {
          this.options.provider.awareness.setLocalStateField("agentFocus", {
            ...current,
            user,
          });
        }
        return true;
      },
    };
  },

  addProseMirrorPlugins() {
    const awareness = this.options.provider.awareness;

    return [
      new Plugin({
        key: agentFocusPluginKey,
        state: {
          init: () => ({ stamp: 0 }),
          apply(tr, prev) {
            if (tr.getMeta(agentFocusPluginKey)) {
              return { stamp: prev.stamp + 1 };
            }
            return prev;
          },
        },
        props: {
          decorations(state) {
            const ystate = ySyncPluginKey.getState(state);
            if (!ystate?.binding) return null;

            const decorations: Decoration[] = [];
            const localClientId = awareness.clientID;

            awareness.getStates().forEach((s: any, clientId: number) => {
              const focus = s?.agentFocus as
                | AwarenessFocus
                | null
                | undefined;
              if (!focus || !focus.anchor || !focus.head) return;
              const isSelf = clientId === localClientId;

              let from: number | null = null;
              let to: number | null = null;
              try {
                from = relativePositionToAbsolutePosition(
                  ystate.doc,
                  ystate.type,
                  Y.createRelativePositionFromJSON(focus.anchor),
                  ystate.binding.mapping,
                );
                to = relativePositionToAbsolutePosition(
                  ystate.doc,
                  ystate.type,
                  Y.createRelativePositionFromJSON(focus.head),
                  ystate.binding.mapping,
                );
              } catch {
                return;
              }

              if (from == null || to == null || from === to) return;
              const [lo, hi] = from < to ? [from, to] : [to, from];

              const color = focus.user?.color || "#5c7cfa";
              const name = focus.user?.name || "";
              // Match Tiptap's `defaultSelectionBuilder` exactly: a flat
              // `{color}70` tint (hex + alpha), no underline, no shadow -
              // so an agent focus is visually the same as a collaborator
              // selection, only the label name ("… agent") tells them apart.
              const style = `background-color: ${color}70`;

              decorations.push(
                Decoration.inline(
                  lo,
                  hi,
                  {
                    class: `agent-focus${isSelf ? " agent-focus--self" : ""}`,
                    style,
                    "data-agent-focus-user": name,
                    "data-agent-focus-color": color,
                  } as any,
                  // inclusiveStart:false / inclusiveEnd:false -> typing
                  // at either edge should not extend the highlight, the
                  // Y.js relative position already pins the true
                  // boundary.
                  { inclusiveStart: false, inclusiveEnd: false } as any,
                ),
              );

              // Floating label anchored at the END of the range, i.e.
              // on the last line of the selection at the cursor
              // position. This mirrors `.collaboration-cursor__label`,
              // which also renders on the right of the last line (at
              // the selection head). Using `hi` instead of `lo` keeps
              // the badge visually coupled to "where the agent's cursor
              // would sit if it were typing right now".
              decorations.push(
                Decoration.widget(
                  hi,
                  () => {
                    const anchor = document.createElement("span");
                    anchor.className = "agent-focus__anchor";

                    const label = document.createElement("span");
                    label.className = "agent-focus__label";
                    label.setAttribute(
                      "style",
                      `background-color: ${color}`,
                    );

                    const avatarUrl = focus.user?.avatarUrl;
                    if (avatarUrl) {
                      const avatar = document.createElement("img");
                      avatar.src = avatarUrl;
                      avatar.className = "agent-focus__avatar";
                      avatar.alt = "";
                      label.appendChild(avatar);
                    }

                    const text = document.createElement("span");
                    text.className = "agent-focus__text";
                    text.textContent = `${name || "Someone"} agent`;
                    label.appendChild(text);

                    anchor.appendChild(label);
                    return anchor;
                  },
                  {
                    // `side: 1` -> the widget stays on the RIGHT side of
                    // the `hi` position, so if the user types inside the
                    // range the label doesn't get pushed into the middle
                    // of new content. `key` lets PM reuse the DOM on
                    // every awareness tick instead of re-creating it.
                    side: 1,
                    key: `agent-focus-label-${clientId}-${name}`,
                  } as any,
                ),
              );
            });

            if (decorations.length === 0) return null;
            return DecorationSet.create(state.doc, decorations);
          },
        },
        view(view) {
          // `change` fires only when an awareness state is actually added,
          // updated, or removed - NOT on heartbeat ticks. Listening to
          // `update` instead (as some examples do) would trigger a
          // re-render every ~30s for nothing.
          //
          // The dispatched transaction is meta-only: no steps, no doc
          // change, no selection change. y-sync ignores it because
          // `tr.docChanged` is false, and the Yjs UndoManager only
          // tracks CRDT operations - so this never pollutes the undo
          // stack. Agent-focus is 100% ephemeral decoration state.
          const handler = () => {
            if (view.isDestroyed) return;
            view.dispatch(
              view.state.tr.setMeta(agentFocusPluginKey, { refresh: true }),
            );
          };
          awareness.on("change", handler);
          return {
            destroy() {
              awareness.off("change", handler);
            },
          };
        },
      }),
    ];
  },
});