kandev/apps/web/components/task/chat/use-tiptap-editor.ts

"use client";

import { useRef, useImperativeHandle, useEffect, useLayoutEffect, useMemo } from "react";
import { useEditor, ReactNodeViewRenderer } from "@tiptap/react";
import Document from "@tiptap/extension-document";
import Paragraph from "@tiptap/extension-paragraph";
import Text from "@tiptap/extension-text";
import HardBreak from "@tiptap/extension-hard-break";
import History from "@tiptap/extension-history";
import { Plugin, PluginKey } from "@tiptap/pm/state";
import { matchesShortcut } from "@/lib/keyboard/utils";
import { getShortcut, type StoredShortcutOverrides } from "@/lib/keyboard/shortcut-overrides";
import { Extension } from "@tiptap/core";
import { useHistoryKeymap } from "./tiptap-editor-history";
import Code from "@tiptap/extension-code";
import CodeBlockLowlight from "@tiptap/extension-code-block-lowlight";
import { common, createLowlight } from "lowlight";
import { cn } from "@/lib/utils";
import { useAppStore } from "@/components/state-provider";
import { getChatDraftContent, setChatDraftContent } from "@/lib/local-storage";
import {
  extractEntityReferences,
  getMarkdownText,
  textToHtml,
  handleEditorPaste,
} from "./tiptap-helpers";
import type { ImagePasteIssue } from "./clipboard-attachments";
import { CodeBlockView } from "./tiptap-code-block-view";
import { DynamicPlaceholder, updateDynamicPlaceholder } from "./tiptap-dynamic-placeholder";
import { EntityReferenceNode } from "./tiptap-entity-reference-extension";
import { ContextMention } from "./tiptap-mention-extension";
import { SlashCommandNode } from "./tiptap-slash-command-extension";
import type { ContextFile } from "@/lib/state/context-files-store";
import type { TaskMentionData } from "@/hooks/use-inline-mention";
import type { SlashCommand } from "./slash-command-types";
import type { EntityReference } from "@/lib/types/entity-reference";
import type { MessageHistoryEntry } from "./message-history";

export type TipTapInputHandle = {
  focus: () => void;
  blur: () => void;
  getSelectionStart: () => number;
  getSelectionEnd: () => number;
  getValue: () => string;
  /**
   * The single character immediately before the current selection, or "" at
   * the very start. Read from the ProseMirror document rather than derived
   * from `getValue()`: selection offsets are doc positions (1-based, counting
   * node boundaries) while `getValue()` returns markdown, so indexing one
   * with the other is off by at least one and drifts further with mentions
   * and code blocks.
   */
  getCharBefore: () => string;
  setValue: (value: string) => void;
  clear: () => void;
  getTextareaElement: () => HTMLElement | null;
  insertText: (text: string, from: number, to: number) => void;
  getMentions: () => ContextFile[];
  getTaskMentions: () => TaskMentionData[];
  getEntityReferences: () => EntityReference[];
};

const lowlightInstance = createLowlight(common);
/** Viewport-width-independent on purpose. The 16px floor that stops iOS Safari
 *  from auto-zooming a focused field is owned by the `@media (any-pointer:
 *  coarse)` rule in `app/globals.css`, which is keyed on the input device
 *  rather than the window size. A width breakpoint here (the old
 *  `text-base … lg:text-sm`) made the composer text jump from 14px to 16px when
 *  a desktop window was dragged narrower than `lg`, which is not a touch
 *  device and needs no zoom guard. */
export const TIPTAP_EDITOR_TEXT_SIZE_CLASS = "text-sm leading-relaxed";

type UseTipTapEditorOptions = {
  value: string;
  onChange: (value: string) => void;
  onSubmit?: () => void;
  placeholder: string;
  disabled: boolean;
  className?: string;
  planModeEnabled: boolean;
  onPlanModeChange?: (enabled: boolean) => void;
  submitKey: "enter" | "cmd_enter";
  onFocus?: () => void;
  onBlur?: () => void;
  sessionId: string | null;
  onImagePaste?: (files: File[], issue?: ImagePasteIssue) => void;
  onTextInput?: (from: number, to: number, text: string) => void;
  onBeforeInput?: (inputType: string) => void;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  mentionSuggestion: any;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  slashSuggestion: any;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  entityReferenceSuggestion?: any;
  slashCommands: readonly SlashCommand[];
  /** True when a slash/@ suggestion menu is open with selectable items. Enter
   *  must defer to the suggestion plugin so the highlighted item is inserted
   *  instead of submitting the message. */
  isSuggestionMenuOpen: boolean;
  /** Returns the user's previous messages for this session, newest-first. The
   *  caller maintains the actual list; the editor reads it on each keypress so
   *  ArrowUp/ArrowDown navigate the latest history without prop churn. */
  getHistory: () => readonly MessageHistoryEntry[];
  /** Open the Ctrl+R fuzzy search overlay. The overlay lives in the parent
   *  component (so it can position itself relative to the editor) — the editor
   *  only knows when to open it. */
  onOpenReverseSearch: () => void;
  /** True while the reverse-search overlay owns focus. The editor must ignore
   *  ArrowUp/ArrowDown in that state so the overlay's own list navigation
   *  isn't shadowed. */
  isReverseSearchOpen: boolean;
  ref: React.ForwardedRef<TipTapInputHandle>;
};

function useTipTapRefs(opts: UseTipTapEditorOptions) {
  const onSubmitRef = useRef(opts.onSubmit);
  const submitKeyRef = useRef(opts.submitKey);
  const disabledRef = useRef(opts.disabled);
  const onChangeRef = useRef(opts.onChange);
  const onImagePasteRef = useRef(opts.onImagePaste);
  const onTextInputRef = useRef(opts.onTextInput);
  const onBeforeInputRef = useRef(opts.onBeforeInput);
  const sessionIdRef = useRef(opts.sessionId);
  const planModeEnabledRef = useRef(opts.planModeEnabled);
  const onPlanModeChangeRef = useRef(opts.onPlanModeChange);
  const isSuggestionMenuOpenRef = useRef(opts.isSuggestionMenuOpen);
  const getHistoryRef = useRef(opts.getHistory);
  const slashCommandsRef = useRef(opts.slashCommands);
  const getSlashCommandsRef = useRef(() => slashCommandsRef.current);
  const onOpenReverseSearchRef = useRef(opts.onOpenReverseSearch);
  const isReverseSearchOpenRef = useRef(opts.isReverseSearchOpen);
  const keyboardShortcuts = useAppStore((s) => s.userSettings.keyboardShortcuts);
  const keyboardShortcutsRef = useRef(keyboardShortcuts);
  useLayoutEffect(() => {
    onSubmitRef.current = opts.onSubmit;
    submitKeyRef.current = opts.submitKey;
    disabledRef.current = opts.disabled;
    onChangeRef.current = opts.onChange;
    onImagePasteRef.current = opts.onImagePaste;
    onTextInputRef.current = opts.onTextInput;
    onBeforeInputRef.current = opts.onBeforeInput;
    sessionIdRef.current = opts.sessionId;
    planModeEnabledRef.current = opts.planModeEnabled;
    onPlanModeChangeRef.current = opts.onPlanModeChange;
    isSuggestionMenuOpenRef.current = opts.isSuggestionMenuOpen;
    getHistoryRef.current = opts.getHistory;
    slashCommandsRef.current = opts.slashCommands;
    onOpenReverseSearchRef.current = opts.onOpenReverseSearch;
    isReverseSearchOpenRef.current = opts.isReverseSearchOpen;
    keyboardShortcutsRef.current = keyboardShortcuts;
  });
  return {
    onSubmitRef,
    submitKeyRef,
    disabledRef,
    onChangeRef,
    onImagePasteRef,
    onTextInputRef,
    onBeforeInputRef,
    sessionIdRef,
    planModeEnabledRef,
    onPlanModeChangeRef,
    isSuggestionMenuOpenRef,
    getHistoryRef,
    getSlashCommandsRef,
    onOpenReverseSearchRef,
    isReverseSearchOpenRef,
    keyboardShortcutsRef,
  };
}

export function buildEditorExtensions(args: {
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  mentionSuggestion: any;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  slashSuggestion: any;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  entityReferenceSuggestion?: any;
  submitKeymap: Extension;
  historyKeymap: Extension;
}) {
  return [
    Document,
    Paragraph,
    Text,
    HardBreak,
    History,
    Code,
    CodeBlockLowlight.extend({
      addNodeView() {
        return ReactNodeViewRenderer(CodeBlockView);
      },
    }).configure({ lowlight: lowlightInstance }),
    DynamicPlaceholder,
    SlashCommandNode,
    EntityReferenceNode,
    ContextMention.configure({
      suggestions: [
        args.mentionSuggestion,
        args.slashSuggestion,
        ...(args.entityReferenceSuggestion ? [args.entityReferenceSuggestion] : []),
      ],
    }),
    args.submitKeymap,
    args.historyKeymap,
  ];
}

function buildEditorProps(args: {
  planModeEnabled: boolean;
  className: string | undefined;
  onFocus: (() => void) | undefined;
  onBlur: (() => void) | undefined;
  onImagePasteRef: React.RefObject<((files: File[], issue?: ImagePasteIssue) => void) | undefined>;
  onTextInputRef: React.RefObject<((from: number, to: number, text: string) => void) | undefined>;
  onBeforeInputRef: React.RefObject<((inputType: string) => void) | undefined>;
}) {
  return {
    attributes: {
      "data-testid": "chat-input-editor",
      class: cn(
        "w-full h-full resize-none bg-transparent px-2 py-2 overflow-y-auto",
        TIPTAP_EDITOR_TEXT_SIZE_CLASS,
        "placeholder:text-muted-foreground",
        "focus:outline-none",
        "disabled:cursor-not-allowed disabled:opacity-50",
        args.planModeEnabled && "border-primary/40",
        args.className,
      ),
    },
    handlePaste: (view: import("@tiptap/pm/view").EditorView, event: ClipboardEvent) =>
      handleEditorPaste(view, event, args.onImagePasteRef),
    handleTextInput: (
      _view: import("@tiptap/pm/view").EditorView,
      from: number,
      to: number,
      text: string,
    ) => {
      args.onTextInputRef.current?.(from, to, text);
      return false;
    },
    handleDOMEvents: {
      focus: () => {
        args.onFocus?.();
        return false;
      },
      blur: () => {
        args.onBlur?.();
        return false;
      },
      beforeinput: (_view: import("@tiptap/pm/view").EditorView, event: Event) => {
        const inputType = (event as InputEvent).inputType;
        if (inputType?.startsWith("delete")) args.onBeforeInputRef.current?.(inputType);
        return false;
      },
    },
  };
}

export function useTipTapEditor(opts: UseTipTapEditorOptions) {
  const refs = useTipTapRefs(opts);
  const SubmitKeymap = useSubmitKeymap({
    disabledRef: refs.disabledRef,
    submitKeyRef: refs.submitKeyRef,
    onSubmitRef: refs.onSubmitRef,
    planModeEnabledRef: refs.planModeEnabledRef,
    onPlanModeChangeRef: refs.onPlanModeChangeRef,
    isSuggestionMenuOpenRef: refs.isSuggestionMenuOpenRef,
    keyboardShortcutsRef: refs.keyboardShortcutsRef,
  });
  const historyController = useHistoryKeymap({
    disabledRef: refs.disabledRef,
    isSuggestionMenuOpenRef: refs.isSuggestionMenuOpenRef,
    isReverseSearchOpenRef: refs.isReverseSearchOpenRef,
    getHistoryRef: refs.getHistoryRef,
    getSlashCommandsRef: refs.getSlashCommandsRef,
    onOpenReverseSearchRef: refs.onOpenReverseSearchRef,
    onChangeRef: refs.onChangeRef,
    keyboardShortcutsRef: refs.keyboardShortcutsRef,
  });
  const isSyncingRef = useRef(false);
  const initialSyncDoneRef = useRef(false);
  const editor = useEditor({
    immediatelyRender: false,
    extensions: buildEditorExtensions({
      mentionSuggestion: opts.mentionSuggestion,
      slashSuggestion: opts.slashSuggestion,
      entityReferenceSuggestion: opts.entityReferenceSuggestion,
      submitKeymap: SubmitKeymap,
      historyKeymap: historyController.extension,
    }),
    editorProps: buildEditorProps({
      planModeEnabled: opts.planModeEnabled,
      className: opts.className,
      onFocus: opts.onFocus,
      onBlur: opts.onBlur,
      onImagePasteRef: refs.onImagePasteRef,
      onTextInputRef: refs.onTextInputRef,
      onBeforeInputRef: refs.onBeforeInputRef,
    }),
    onUpdate: ({ editor: e }) => {
      if (isSyncingRef.current || !initialSyncDoneRef.current) return;
      const text = getMarkdownText(e);
      refs.onChangeRef.current(text);
      const sid = refs.sessionIdRef.current;
      if (sid) setChatDraftContent(sid, e.getJSON());
    },
    editable: !opts.disabled,
  });
  useSyncEditor({
    editor,
    disabled: opts.disabled,
    placeholder: opts.placeholder,
    sessionId: opts.sessionId,
    value: opts.value,
    isSyncingRef,
    initialSyncDoneRef,
    onChangeRef: refs.onChangeRef,
  });
  useEditorImperativeHandle(opts.ref, editor, opts.onChange, isSyncingRef);
  const applyHistoryEntry = useMemo(
    () => (index: number) => historyController.applyHistoryIndex(editor, index),
    [editor, historyController],
  );
  return { editor, applyHistoryEntry };
}

// ── Sync hook ─────────────────────────────────────────────────────

type SyncEditorOptions = {
  editor: ReturnType<typeof useEditor> | null;
  disabled: boolean;
  placeholder: string;
  sessionId: string | null;
  value: string;
  isSyncingRef: React.RefObject<boolean>;
  initialSyncDoneRef: React.RefObject<boolean>;
  onChangeRef: React.RefObject<(value: string) => void>;
};

/** Editor methods required to synchronize the disabled state. */
type DisabledStateEditor = {
  view: { hasFocus: () => boolean };
  setEditable: (editable: boolean) => void;
  commands: { focus: () => void };
};

/** Bound focus restoration while browser focus settles after re-enabling. */
const FOCUS_RESTORE_MAX_ATTEMPTS = 5;

/** Preserve editor focus across temporary disabled states without stealing focus. */
export function useSyncDisabledState(editor: DisabledStateEditor | null, disabled: boolean) {
  const hadFocusBeforeDisableRef = useRef(false);
  useEffect(() => {
    if (!editor) return;
    if (disabled) {
      hadFocusBeforeDisableRef.current = editor.view.hasFocus();
      editor.setEditable(false);
      return;
    }
    editor.setEditable(true);
    const hadFocus = hadFocusBeforeDisableRef.current;
    hadFocusBeforeDisableRef.current = false;
    if (!hadFocus) return;

    let frame = 0;
    let attempts = 0;
    const tick = () => {
      attempts += 1;
      if (!editor.view.hasFocus()) {
        if (!shouldRestoreFocusOnEnable(hadFocus)) return;
        editor.commands.focus();
      }
      if (attempts < FOCUS_RESTORE_MAX_ATTEMPTS) frame = requestAnimationFrame(tick);
    };
    frame = requestAnimationFrame(tick);
    return () => cancelAnimationFrame(frame);
  }, [editor, disabled]);
}

function useSyncEditor({
  editor,
  disabled,
  placeholder,
  sessionId,
  value,
  isSyncingRef,
  initialSyncDoneRef,
  onChangeRef,
}: SyncEditorOptions) {
  useSyncDisabledState(editor, disabled);

  // Sync placeholder via editor.storage. The DynamicPlaceholder extension reads
  // from editor.storage.dynamicPlaceholder.text at decoration time.
  useEffect(() => {
    if (!editor) return;
    updateDynamicPlaceholder(editor, placeholder);
  }, [editor, placeholder]);

  // Reset sync flag when session changes
  const prevSyncSessionRef = useRef(sessionId);
  useEffect(() => {
    if (sessionId === prevSyncSessionRef.current) return;
    prevSyncSessionRef.current = sessionId;
    initialSyncDoneRef.current = false;
  }, [sessionId, initialSyncDoneRef]);

  // Sync value prop changes
  useEffect(() => {
    syncEditorValue({ editor, sessionId, value, isSyncingRef, initialSyncDoneRef, onChangeRef });
  }, [editor, value, sessionId, isSyncingRef, initialSyncDoneRef, onChangeRef]);
}

type SyncEditorValueOptions = {
  editor: ReturnType<typeof useEditor> | null;
  sessionId: string | null;
  value: string;
  isSyncingRef: React.RefObject<boolean>;
  initialSyncDoneRef: React.RefObject<boolean>;
  onChangeRef: React.RefObject<(value: string) => void>;
};

function syncEditorValue({
  editor,
  sessionId,
  value,
  isSyncingRef,
  initialSyncDoneRef,
  onChangeRef,
}: SyncEditorValueOptions) {
  if (!editor) return;

  if (!initialSyncDoneRef.current) {
    const sid = sessionId;
    if (sid) {
      const savedContent = getChatDraftContent(sid);
      if (savedContent) {
        isSyncingRef.current = true;
        editor.commands.setContent(savedContent as import("@tiptap/core").Content);
        isSyncingRef.current = false;
        initialSyncDoneRef.current = true;
        onChangeRef.current(getMarkdownText(editor));
        return;
      }
    }
  }

  if (value === "") {
    if (!editor.isEmpty) {
      isSyncingRef.current = true;
      editor.commands.clearContent();
      isSyncingRef.current = false;
    }
    initialSyncDoneRef.current = true;
    return;
  }

  const currentText = getMarkdownText(editor);
  if (currentText === value) {
    initialSyncDoneRef.current = true;
    return;
  }

  isSyncingRef.current = true;
  editor.commands.setContent(textToHtml(value));
  isSyncingRef.current = false;
  initialSyncDoneRef.current = true;
}

// ── Focus-restore decision ───────────────────────────────────────────

/** Return true only when the editor had focus and no other control owns focus. */
export function shouldRestoreFocusOnEnable(hadFocusBeforeDisable: boolean): boolean {
  if (!hadFocusBeforeDisable) return false;
  const active = document.activeElement;
  return active === null || active === document.body;
}

// ── Submit shortcut decision ────────────────────────────────────────

export type SubmitShortcutDecision = "consume-noop" | "submit" | "defer";

/** Pure decision for whether an Enter/Mod-Enter press in the chat input should
 *  submit the message, defer to the next ProseMirror handler (e.g. the slash/
 *  mention suggestion plugin or paragraph-split), or no-op while disabled.
 *  Kept pure so the keymap contract is unit-testable without mounting TipTap. */
export function decideSubmitShortcut(args: {
  pressed: "enter" | "mod-enter";
  disabled: boolean;
  submitKey: "enter" | "cmd_enter";
  isSuggestionMenuOpen: boolean;
}): SubmitShortcutDecision {
  if (args.disabled) return "consume-noop";
  if (args.pressed === "enter") {
    if (args.submitKey !== "enter") return "defer";
    if (args.isSuggestionMenuOpen) return "defer";
    return "submit";
  }
  if (args.submitKey !== "cmd_enter") return "defer";
  return "submit";
}

// ── Submit keymap hook ──────────────────────────────────────────────

function useSubmitKeymap(refs: {
  disabledRef: React.RefObject<boolean | undefined>;
  submitKeyRef: React.RefObject<"enter" | "cmd_enter">;
  onSubmitRef: React.RefObject<(() => void) | undefined>;
  planModeEnabledRef: React.RefObject<boolean>;
  onPlanModeChangeRef: React.RefObject<((enabled: boolean) => void) | undefined>;
  isSuggestionMenuOpenRef: React.RefObject<boolean>;
  keyboardShortcutsRef: React.RefObject<StoredShortcutOverrides | undefined>;
}) {
  const {
    disabledRef,
    submitKeyRef,
    onSubmitRef,
    planModeEnabledRef,
    onPlanModeChangeRef,
    isSuggestionMenuOpenRef,
    keyboardShortcutsRef,
  } = refs;
  return useMemo(() => {
    return Extension.create({
      name: "submitKeymap",
      addKeyboardShortcuts() {
        const run = (pressed: "enter" | "mod-enter") => {
          const decision = decideSubmitShortcut({
            pressed,
            disabled: !!disabledRef.current,
            submitKey: submitKeyRef.current ?? "cmd_enter",
            isSuggestionMenuOpen: isSuggestionMenuOpenRef.current,
          });
          if (decision === "consume-noop") return true;
          if (decision === "submit") {
            onSubmitRef.current?.();
            return true;
          }
          return false;
        };
        return {
          Enter: () => run("enter"),
          "Mod-Enter": () => run("mod-enter"),
        };
      },
      addProseMirrorPlugins() {
        return [
          new Plugin({
            key: new PluginKey("planModeToggle"),
            props: {
              handleKeyDown: (_view, event) => {
                const shortcut = getShortcut("TOGGLE_PLAN_MODE", keyboardShortcutsRef.current);
                if (matchesShortcut(event, shortcut)) {
                  onPlanModeChangeRef.current?.(!planModeEnabledRef.current);
                  return true;
                }
                return false;
              },
            },
          }),
        ];
      },
    });
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);
}

// ── Imperative handle hook ──────────────────────────────────────────

function useEditorImperativeHandle(
  ref: React.ForwardedRef<TipTapInputHandle>,
  editor: ReturnType<typeof useEditor> | null,
  onChange: (value: string) => void,
  isSyncingRef: React.RefObject<boolean>,
) {
  useImperativeHandle(
    ref,
    () => ({
      focus: () => editor?.commands.focus(),
      blur: () => editor?.commands.blur(),
      getSelectionStart: () => editor?.state.selection.from ?? 0,
      getSelectionEnd: () => editor?.state.selection.to ?? 0,
      getValue: () => (editor ? getMarkdownText(editor) : ""),
      getCharBefore: () => {
        if (!editor) return "";
        const { from } = editor.state.selection;
        // Position 1 is the first character of the first block, so nothing
        // precedes it. `\n` for block/leaf gaps keeps a line start counting
        // as whitespace, which is what the spacing rule wants.
        if (from <= 1) return "";
        return editor.state.doc.textBetween(from - 1, from, "\n", "\n");
      },
      setValue: (v: string) => {
        if (!editor) return;
        isSyncingRef.current = true;
        if (v === "") {
          editor.commands.clearContent();
        } else {
          editor.commands.setContent(textToHtml(v));
        }
        isSyncingRef.current = false;
        onChange(v);
      },
      clear: () => {
        if (!editor) return;
        isSyncingRef.current = true;
        editor.commands.clearContent();
        isSyncingRef.current = false;
        onChange("");
      },
      getTextareaElement: () => editor?.view.dom ?? null,
      insertText: (text: string, from: number, to: number) => {
        if (!editor) return;
        editor.chain().focus().insertContentAt({ from, to }, text).run();
      },
      getMentions: () => {
        if (!editor) return [];
        const mentions: ContextFile[] = [];
        editor.state.doc.descendants((node) => {
          if (node.type.name === "contextMention") {
            const { kind, path, label } = node.attrs;
            if (kind === "file") mentions.push({ path, name: label, pinned: false });
            else if (kind === "prompt") mentions.push({ path, name: label, pinned: false });
            else if (kind === "plan") mentions.push({ path, name: label, pinned: false });
          }
        });
        return mentions;
      },
      getTaskMentions: () => {
        if (!editor) return [];
        const seen = new Set<string>();
        const mentions: TaskMentionData[] = [];
        editor.state.doc.descendants((node) => {
          if (node.type.name !== "contextMention") return;
          const { kind, label, taskId, workflowId, workflowStepId, taskState } = node.attrs;
          if (kind !== "task" || !taskId || !workflowId || !workflowStepId || seen.has(taskId))
            return;
          seen.add(taskId);
          mentions.push({
            taskId,
            title: label ?? taskId,
            workflowId,
            workflowStepId,
            state: taskState ?? null,
          });
        });
        return mentions;
      },
      getEntityReferences: () => (editor ? extractEntityReferences(editor.getJSON()) : []),
    }),
    [editor, onChange, isSyncingRef],
  );
}