@beyondwork/docx-react-component 1.0.18 → 1.0.20

package/src/ui-tailwind/editor-surface/tw-prosemirror-surface.tsx

@@ -1,6 +1,7 @@
 import React, {
   forwardRef,
   type FocusEventHandler,
+  useCallback,
   useEffect,
   useImperativeHandle,
   useMemo,
@@ -9,12 +10,16 @@ import React, {
 import { EditorView } from "prosemirror-view";
 
 import type {
+  DocumentNavigationSnapshot,
   EditorUser,
   RuntimeRenderSnapshot,
   SearchOptions,
   SearchResultSnapshot,
   SelectionSnapshot,
+  WorkflowScope,
 } from "../../api/public-types";
+import type { CanonicalDocumentEnvelope } from "../../core/state/editor-state.ts";
+import { searchDocument } from "../../runtime/document-search.ts";
 import {
   getTableSelectionDescriptor,
   type TableSelectionDescriptor,
@@ -24,22 +29,37 @@ import {
   type MarkupDisplay,
 } from "../../ui/headless/comment-decoration-model";
 import { createRevisionDecorationModel } from "../../ui/headless/revision-decoration-model";
-import { createPMStateFromSnapshot } from "./pm-state-from-snapshot";
+import {
+  createPMSelectionFromSnapshot,
+  createPMStateFromSnapshot,
+} from "./pm-state-from-snapshot";
 import {
   createCommandBridgePlugins,
   type CommandBridgeCallbacks,
 } from "./pm-command-bridge";
 import { buildDecorations } from "./pm-decorations";
+import { createContextualInteractionPlugin } from "./pm-contextual-ui";
+import {
+  finishPerfProbe,
+  incrementInvalidationCounter,
+  recordPerfSample,
+  startPerfProbe,
+} from "./perf-probe";
 import type { PositionMap } from "./pm-position-map";
 import {
   clearSearch as clearSearchPlugin,
-  createSearchExcerpt,
   createSearchPlugin,
   DEFAULT_SEARCH_HIGHLIGHT_COLOR,
   performSearch,
   searchPluginKey,
 } from "./search-plugin";
+import {
+  createSurfaceDecorationKey,
+  createSurfaceDocumentBuildKey,
+} from "./surface-build-keys";
 import { tableNodeViews } from "./tw-table-node-view";
+import type { SelectionToolbarAnchor } from "../../ui/headless/selection-toolbar-model";
+import type { MediaPreviewDescriptor } from "./pm-state-from-snapshot";
 
 /**
  * Same props interface as the legacy TwEditorSurface — drop-in replacement.
@@ -47,10 +67,14 @@ import { tableNodeViews } from "./tw-table-node-view";
 export interface TwProseMirrorSurfaceProps {
   currentUser: EditorUser;
   snapshot: RuntimeRenderSnapshot;
+  canonicalDocument: CanonicalDocumentEnvelope;
+  documentNavigation: DocumentNavigationSnapshot;
   reviewMode: "editing" | "review";
   markupDisplay: MarkupDisplay;
   activeRevisionId?: string;
   showTrackedChanges?: boolean;
+  /** When true, the surface renders inside the page workspace (vs canvas). */
+  isPageWorkspace?: boolean;
   onFocus: FocusEventHandler<HTMLDivElement>;
   onBlur: FocusEventHandler<HTMLDivElement>;
   onSelectionChange?: (selection: SelectionSnapshot) => void;
@@ -58,10 +82,14 @@ export interface TwProseMirrorSurfaceProps {
   onDeleteBackward?: () => void;
   onDeleteForward?: () => void;
   onInsertTab?: () => void;
+  onOutdentTab?: () => void;
   onInsertHardBreak?: () => void;
   onSplitParagraph?: () => void;
   onCommentActivated?: (commentId: string) => void;
   onRevisionActivated?: (revisionId: string) => void;
+  onSelectionToolbarAnchorChange?: (anchor: SelectionToolbarAnchor | null) => void;
+  mediaPreviews?: Record<string, MediaPreviewDescriptor>;
+  workflowScopes?: readonly WorkflowScope[];
 }
 
 export interface TwProseMirrorSurfaceRef {
@@ -82,6 +110,17 @@ export const TwProseMirrorSurface = forwardRef<
     onBlur,
   } = props;
   const surface = snapshot.surface;
+  const mediaPreviewKey = useMemo(
+    () =>
+      Object.entries(props.mediaPreviews ?? {})
+        .sort(([leftId], [rightId]) => leftId.localeCompare(rightId))
+        .map(
+          ([mediaId, preview]) =>
+            `${mediaId}:${preview.widthEmu ?? ""}:${preview.heightEmu ?? ""}:${preview.src}`,
+        )
+        .join("|"),
+    [props.mediaPreviews],
+  );
 
   const canEdit = Boolean(
     surface && snapshot.isReady && !snapshot.readOnly && !snapshot.fatalError,
@@ -92,19 +131,44 @@ export const TwProseMirrorSurface = forwardRef<
   const positionMapRef = useRef<PositionMap | null>(null);
   const callbacksRef = useRef<CommandBridgeCallbacks | null>(null);
   const activeSearchRef = useRef<{ query: string; options: SearchOptions } | null>(null);
+  const pendingTypingProbeRef = useRef<string | null>(null);
+  const pendingSelectionProbeRef = useRef<string | null>(null);
+  const documentBuildKeyRef = useRef<string | null>(null);
+  const decorationBuildKeyRef = useRef<string | null>(null);
+  const suppressSelectionEchoRef = useRef(false);
+  const selectionToolbarFrameRef = useRef<number | null>(null);
+  const lastSelectionToolbarMeasurementRef = useRef<{
+    key: string | null;
+    anchor: SelectionToolbarAnchor | null;
+  }>({
+    key: null,
+    anchor: null,
+  });
 
   // Keep callbacks ref up to date (avoids stale closures in PM plugins)
   callbacksRef.current = {
-    onInsertText: (text) => props.onInsertText?.(text),
+    onInsertText: (text) => {
+      pendingTypingProbeRef.current = startPerfProbe("typing");
+      props.onInsertText?.(text);
+    },
     onDeleteBackward: () => props.onDeleteBackward?.(),
     onDeleteForward: () => props.onDeleteForward?.(),
     onSplitParagraph: () => props.onSplitParagraph?.(),
     onInsertHardBreak: () => props.onInsertHardBreak?.(),
     onInsertTab: () => props.onInsertTab?.(),
+    onOutdentTab: () => props.onOutdentTab?.(),
     onUndo: () => {}, // Handled by toolbar, not PM
     onRedo: () => {}, // Handled by toolbar, not PM
-    onSelectionChange: (sel) => props.onSelectionChange?.(sel),
+    onSelectionChange: (sel) => {
+      pendingSelectionProbeRef.current = startPerfProbe("selection");
+      props.onSelectionChange?.(
+        snapshot.activeStory.kind === "main"
+          ? sel
+          : { ...sel, storyTarget: snapshot.activeStory },
+      );
+    },
     getPositionMap: () => positionMapRef.current,
+    isSelectionSyncSuppressed: () => suppressSelectionEchoRef.current,
   };
 
   // Comment/revision decoration models
@@ -119,6 +183,34 @@ export const TwProseMirrorSurface = forwardRef<
     () => createRevisionDecorationModel(snapshot.trackedChanges, props.activeRevisionId),
     [snapshot.trackedChanges, props.activeRevisionId],
   );
+  const documentBuildKey = useMemo(
+    () =>
+      createSurfaceDocumentBuildKey({
+        surface,
+        activeStory: snapshot.activeStory,
+        mediaPreviewKey,
+      }),
+    [mediaPreviewKey, snapshot.activeStory, surface],
+  );
+  const decorationBuildKey = useMemo(
+    () =>
+      createSurfaceDecorationKey({
+        markupDisplay,
+        showTrackedChanges,
+        canEdit,
+        activeCommentId: snapshot.comments.activeCommentId,
+        activeRevisionId: props.activeRevisionId,
+        workflowScopeSignature: JSON.stringify(props.workflowScopes ?? []),
+      }),
+    [
+      canEdit,
+      markupDisplay,
+      props.activeRevisionId,
+      props.workflowScopes,
+      showTrackedChanges,
+      snapshot.comments.activeCommentId,
+    ],
+  );
 
   // Create PM plugins (stable across renders — callbacks accessed via ref)
   const plugins = useMemo(() => {
@@ -130,28 +222,68 @@ export const TwProseMirrorSurface = forwardRef<
         onSplitParagraph: () => callbacksRef.current?.onSplitParagraph(),
         onInsertHardBreak: () => callbacksRef.current?.onInsertHardBreak(),
         onInsertTab: () => callbacksRef.current?.onInsertTab(),
+        onOutdentTab: () => callbacksRef.current?.onOutdentTab?.(),
         onUndo: () => callbacksRef.current?.onUndo(),
         onRedo: () => callbacksRef.current?.onRedo(),
         onSelectionChange: (sel) => callbacksRef.current?.onSelectionChange(sel),
         getPositionMap: () => callbacksRef.current?.getPositionMap() ?? null,
+        isSelectionSyncSuppressed: () =>
+          callbacksRef.current?.isSelectionSyncSuppressed?.() ?? false,
+      }),
+      createContextualInteractionPlugin({
+        onCommentActivated: (commentId) => props.onCommentActivated?.(commentId),
+        onRevisionActivated: (revisionId) => props.onRevisionActivated?.(revisionId),
       }),
       createSearchPlugin(),
     ];
-  }, []);
+  }, [props.onCommentActivated, props.onRevisionActivated]);
 
-  // Create or update PM view whenever surface becomes available or changes.
-  // The view is created lazily — if surface is null on first render (loading),
-  // it will be created when the runtime provides a real snapshot.
+  const applyDecorationProps = useCallback(
+    (view: EditorView, positionMap: PositionMap): void => {
+      const decorations = buildDecorations(
+        view.state.doc,
+        positionMap,
+        commentModel,
+        revisionModel,
+        markupDisplay,
+        showTrackedChanges,
+        props.workflowScopes,
+        snapshot.activeStory,
+      );
+      view.setProps({
+        editable: () => canEdit,
+        decorations: () => decorations,
+      });
+      decorationBuildKeyRef.current = decorationBuildKey;
+      recordPerfSample("pm.decorations");
+      incrementInvalidationCounter("pm.laneB.decorationUpdates");
+    },
+    [
+      canEdit,
+      commentModel,
+      decorationBuildKey,
+      markupDisplay,
+      revisionModel,
+      showTrackedChanges,
+      props.workflowScopes,
+    ],
+  );
+
+  // Create or update the PM document only when the structural key changes.
   useEffect(() => {
     if (!mountRef.current || !surface) return;
 
+    if (viewRef.current && documentBuildKeyRef.current === documentBuildKey) {
+      return;
+    }
+
     const { state, positionMap } = createPMStateFromSnapshot(
       surface,
       snapshot.selection,
       plugins,
+      props.mediaPreviews,
     );
     positionMapRef.current = positionMap;
-
     const decorations = buildDecorations(
       state.doc,
       positionMap,
@@ -159,7 +291,11 @@ export const TwProseMirrorSurface = forwardRef<
       revisionModel,
       markupDisplay,
       showTrackedChanges,
+      props.workflowScopes,
+      snapshot.activeStory,
     );
+    recordPerfSample("pm.rebuild");
+    incrementInvalidationCounter("pm.laneA.rebuilds");
 
     if (!viewRef.current) {
       // First time surface is available — create the EditorView
@@ -174,14 +310,16 @@ export const TwProseMirrorSurface = forwardRef<
         },
       });
       viewRef.current = view;
+      recordPerfSample("pm.mount");
     } else {
-      // View exists — update state and decorations
-      viewRef.current.setProps({
-        editable: () => canEdit,
-        decorations: () => decorations,
-      });
+      suppressSelectionEchoRef.current = true;
       viewRef.current.updateState(state);
+      queueMicrotask(() => {
+        suppressSelectionEchoRef.current = false;
+      });
     }
+    documentBuildKeyRef.current = documentBuildKey;
+    applyDecorationProps(viewRef.current, positionMap);
 
     if (activeSearchRef.current) {
       applySearch(
@@ -189,11 +327,74 @@ export const TwProseMirrorSurface = forwardRef<
         activeSearchRef.current.options,
       );
     }
-  }, [snapshot.revisionToken, surface, commentModel, revisionModel, markupDisplay, canEdit]);
+    if (pendingTypingProbeRef.current) {
+      finishPerfProbe(pendingTypingProbeRef.current);
+      pendingTypingProbeRef.current = null;
+    }
+  }, [
+    applyDecorationProps,
+    documentBuildKey,
+    surface,
+    snapshot.selection,
+    plugins,
+    props.mediaPreviews,
+  ]);
+
+  // Update decorations and editability without rebuilding the PM document.
+  useEffect(() => {
+    const view = viewRef.current;
+    const positionMap = positionMapRef.current;
+    if (!view || !surface || !positionMap) {
+      return;
+    }
+
+    if (decorationBuildKeyRef.current === decorationBuildKey) {
+      return;
+    }
+
+    applyDecorationProps(view, positionMap);
+  }, [applyDecorationProps, decorationBuildKey, surface]);
+
+  useEffect(() => {
+    const view = viewRef.current;
+    const positionMap = positionMapRef.current;
+    if (!view || !surface || !positionMap) {
+      return;
+    }
+
+    const nextSelection = createPMSelectionFromSnapshot(
+      view.state.doc,
+      positionMap,
+      snapshot.selection,
+    );
+    if (view.state.selection.eq(nextSelection)) {
+      return;
+    }
+
+    suppressSelectionEchoRef.current = true;
+    view.dispatch(view.state.tr.setSelection(nextSelection));
+    recordPerfSample("selection.sync");
+    queueMicrotask(() => {
+      suppressSelectionEchoRef.current = false;
+    });
+  }, [snapshot.selection, surface]);
+
+  useEffect(() => {
+    if (!pendingSelectionProbeRef.current) {
+      return;
+    }
+    finishPerfProbe(pendingSelectionProbeRef.current);
+    pendingSelectionProbeRef.current = null;
+  }, [snapshot.selection]);
 
   // Cleanup on unmount
   useEffect(() => {
     return () => {
+      const win = mountRef.current?.ownerDocument.defaultView;
+      if (selectionToolbarFrameRef.current !== null && win) {
+        win.cancelAnimationFrame(selectionToolbarFrameRef.current);
+        selectionToolbarFrameRef.current = null;
+      }
       viewRef.current?.destroy();
       viewRef.current = null;
     };
@@ -230,45 +431,77 @@ export const TwProseMirrorSurface = forwardRef<
 
   function applySearch(query: string, options: SearchOptions): SearchResultSnapshot[] {
     const view = viewRef.current;
-    const positionMap = positionMapRef.current;
-    if (!view || !positionMap) {
-      return [];
+    const hiddenDeletionRanges =
+      markupDisplay === "clean"
+        ? snapshot.trackedChanges.revisions
+            .filter(
+              (
+                revision,
+              ): revision is typeof revision & {
+                anchor: Extract<typeof revision.anchor, { kind: "range" }>;
+              } =>
+                revision.kind === "deletion" &&
+                revision.status === "active" &&
+                revision.anchor.kind === "range",
+            )
+            .map((revision) => ({
+              from: revision.anchor.from,
+              to: revision.anchor.to,
+            }))
+        : [];
+    if (view) {
+      const rawResults = performSearch(view.state, query, options)
+        .filter((result) => {
+          if (hiddenDeletionRanges.length === 0) {
+            return true;
+          }
+          const positionMap = positionMapRef.current;
+          if (!positionMap) {
+            return true;
+          }
+          const runtimeFrom = positionMap.pmToRuntime(result.from);
+          const runtimeTo = positionMap.pmToRuntime(result.to);
+          return !hiddenDeletionRanges.some(
+            (range) => runtimeFrom < range.to && runtimeTo > range.from,
+          );
+        })
+        .slice(0, options.limit ?? Number.POSITIVE_INFINITY);
+      view.dispatch(
+        view.state.tr.setMeta(searchPluginKey, {
+          results: rawResults,
+          highlightColor: DEFAULT_SEARCH_HIGHLIGHT_COLOR,
+        }),
+      );
     }
 
-    const rawResults = performSearch(view.state, query, options).slice(
-      0,
-      options.limit ?? Number.POSITIVE_INFINITY,
-    );
-    view.dispatch(
-      view.state.tr.setMeta(searchPluginKey, {
-        results: rawResults,
-        highlightColor: DEFAULT_SEARCH_HIGHLIGHT_COLOR,
-      }),
+    return filterHiddenDeletionSearchResults(
+      searchDocument(
+        props.canonicalDocument,
+        snapshot.selection,
+        snapshot.activeStory,
+        props.documentNavigation,
+        query,
+        options,
+      ),
+      hiddenDeletionRanges,
     );
+  }
 
-    const activeResultIndex = getActiveSearchResultIndex(
-      rawResults,
-      (position) => positionMap.pmToRuntime(position),
-      snapshot.selection,
-    );
-    const plainText = snapshot.surface?.plainText ?? "";
-    return rawResults.map((result, index) => {
-      const runtimeFrom = positionMap.pmToRuntime(result.from);
-      const runtimeTo = positionMap.pmToRuntime(result.to);
-      return {
-        resultId: `search-result-${index}`,
-        anchor: {
-          kind: "range",
-          from: runtimeFrom,
-          to: runtimeTo,
-          assoc: {
-            start: -1,
-            end: 1,
-          },
-        },
-        excerpt: createSearchExcerpt(plainText, runtimeFrom, runtimeTo),
-        isActive: index === activeResultIndex,
-      };
+  function filterHiddenDeletionSearchResults(
+    results: SearchResultSnapshot[],
+    hiddenRanges: Array<{ from: number; to: number }>,
+  ): SearchResultSnapshot[] {
+    if (hiddenRanges.length === 0) {
+      return results;
+    }
+    return results.filter((result) => {
+      const anchor = result.anchor;
+      if (anchor.kind !== "range") {
+        return true;
+      }
+      return !hiddenRanges.some(
+        (range) => anchor.from < range.to && anchor.to > range.from,
+      );
     });
   }
 
@@ -288,37 +521,192 @@ export const TwProseMirrorSurface = forwardRef<
       ? "font-[family-name:var(--font-legal-sans)]"
       : "font-[family-name:var(--font-legal-serif)]";
 
+  // Story focus indicator — runtime-backed, not DOM-only
+  const storyKind = snapshot.activeStory.kind;
+  const storyFocusAttr = storyKind !== "main" ? storyKind : undefined;
+
+  // Table focus cue — add subtle ring when selection head is inside a table block
+  const tableFocusClass = (() => {
+    if (!surface || !snapshot.selection) return "";
+    const head = snapshot.selection.head;
+    const inTable = surface.blocks.some(
+      (b) => b.kind === "table" && head >= b.from && head <= b.to,
+    );
+    return inTable ? "prosemirror-table-focus" : "";
+  })();
+
+  const workspaceLabel = props.isPageWorkspace ? "Document page" : "Document canvas";
+
+  const selectionToolbarMeasurementKey = useMemo(
+    () => buildSelectionToolbarMeasurementKey(snapshot.selection, snapshot.activeStory),
+    [snapshot.activeStory, snapshot.selection],
+  );
+
+  const emitSelectionToolbarAnchor = useCallback((): void => {
+    const callback = props.onSelectionToolbarAnchorChange;
+    if (!callback) {
+      return;
+    }
+
+    const nextAnchor = measureSelectionToolbarAnchor();
+    const previous = lastSelectionToolbarMeasurementRef.current;
+    if (
+      previous.key === selectionToolbarMeasurementKey &&
+      selectionToolbarAnchorsEqual(previous.anchor, nextAnchor)
+    ) {
+      return;
+    }
+
+    lastSelectionToolbarMeasurementRef.current = {
+      key: selectionToolbarMeasurementKey,
+      anchor: nextAnchor,
+    };
+    callback(nextAnchor);
+  }, [
+    props.onSelectionToolbarAnchorChange,
+    selectionToolbarMeasurementKey,
+    snapshot.activeStory,
+    snapshot.selection,
+  ]);
+
+  const scheduleSelectionToolbarAnchorUpdate = useCallback((): void => {
+    const callback = props.onSelectionToolbarAnchorChange;
+    const mount = mountRef.current;
+    const win = mount?.ownerDocument.defaultView;
+    if (!callback || !win) {
+      emitSelectionToolbarAnchor();
+      return;
+    }
+
+    if (selectionToolbarFrameRef.current !== null) {
+      win.cancelAnimationFrame(selectionToolbarFrameRef.current);
+      selectionToolbarFrameRef.current = null;
+    }
+
+    selectionToolbarFrameRef.current = win.requestAnimationFrame(() => {
+      selectionToolbarFrameRef.current = null;
+      emitSelectionToolbarAnchor();
+    });
+  }, [emitSelectionToolbarAnchor, props.onSelectionToolbarAnchorChange]);
+
+  useEffect(() => {
+    scheduleSelectionToolbarAnchorUpdate();
+  }, [
+    scheduleSelectionToolbarAnchorUpdate,
+    snapshot.revisionToken,
+    snapshot.selection,
+    snapshot.surface,
+    props.isPageWorkspace,
+  ]);
+
+  useEffect(() => {
+    const mount = mountRef.current;
+    const callback = props.onSelectionToolbarAnchorChange;
+    if (!mount || !callback) {
+      return;
+    }
+
+    const updateAnchor = () => {
+      scheduleSelectionToolbarAnchorUpdate();
+    };
+    const scrollRoot = mount.closest<HTMLElement>("[data-wre-scroll-root='true']");
+    const win = mount.ownerDocument.defaultView;
+    const resizeObserver =
+      typeof ResizeObserver !== "undefined"
+        ? new ResizeObserver(() => {
+          updateAnchor();
+        })
+        : null;
+
+    updateAnchor();
+    scrollRoot?.addEventListener("scroll", updateAnchor, { passive: true });
+    win?.addEventListener("resize", updateAnchor);
+    resizeObserver?.observe(mount);
+    if (scrollRoot) {
+      resizeObserver?.observe(scrollRoot);
+    }
+
+    return () => {
+      scrollRoot?.removeEventListener("scroll", updateAnchor);
+      win?.removeEventListener("resize", updateAnchor);
+      resizeObserver?.disconnect();
+    };
+  }, [
+    props.onSelectionToolbarAnchorChange,
+    scheduleSelectionToolbarAnchorUpdate,
+    snapshot.revisionToken,
+    snapshot.selection,
+  ]);
+
+  useEffect(() => {
+    return () => {
+      lastSelectionToolbarMeasurementRef.current = {
+        key: null,
+        anchor: null,
+      };
+      props.onSelectionToolbarAnchorChange?.(null);
+    };
+  }, [props.onSelectionToolbarAnchorChange]);
+
   return (
-    <section aria-label="Document canvas" className="min-w-0">
+    <section
+      aria-label={workspaceLabel}
+      className="min-w-0"
+      data-active-story={storyFocusAttr}
+      data-workspace={props.isPageWorkspace ? "page" : "canvas"}
+    >
       {/* ProseMirror mount point — document content including headings is editable */}
       {surface ? (
         <div
           ref={mountRef}
           role="textbox"
+          tabIndex={0}
           aria-multiline="true"
-          className={`px-12 py-10 ${fontClass} text-[15px] text-primary leading-relaxed prosemirror-surface outline-none`}
-          onFocus={onFocus as unknown as React.FocusEventHandler<HTMLDivElement>}
+          className={`px-12 py-10 ${fontClass} text-[15px] text-primary leading-relaxed prosemirror-surface outline-none ${tableFocusClass}`}
+          onFocus={(event) => {
+            onFocus(event);
+            if (event.target === event.currentTarget) {
+              viewRef.current?.focus();
+            }
+          }}
           onBlur={onBlur as unknown as React.FocusEventHandler<HTMLDivElement>}
-          onClick={(e) => {
-            // Activate comment or revision when clicking on decorated text
-            const target = e.target as HTMLElement;
-            const commentEl = target.closest?.("[data-comment-id]");
-            if (commentEl) {
-              const commentId = commentEl.getAttribute("data-comment-id");
-              if (commentId) {
-                props.onCommentActivated?.(commentId);
-                return;
-              }
+          onKeyDown={(event) => {
+            if (event.target !== event.currentTarget) {
+              return;
             }
-            const revisionEl = target.closest?.("[data-revision-id]");
-            if (revisionEl) {
-              const revisionId = revisionEl.getAttribute("data-revision-id");
-              if (revisionId) {
-                props.onRevisionActivated?.(revisionId);
-              }
+
+            switch (event.key) {
+              case "Backspace":
+                event.preventDefault();
+                props.onDeleteBackward?.();
+                return;
+              case "Delete":
+                event.preventDefault();
+                props.onDeleteForward?.();
+                return;
+              case "Enter":
+                event.preventDefault();
+                if (event.shiftKey) {
+                  props.onInsertHardBreak?.();
+                } else {
+                  props.onSplitParagraph?.();
+                }
+                return;
+              case "Tab":
+                event.preventDefault();
+                if (event.shiftKey) {
+                  props.onOutdentTab?.();
+                } else {
+                  props.onInsertTab?.();
+                }
+                return;
+              default:
+                return;
             }
           }}
           aria-label="Document surface"
+          data-wre-document-surface="true"
+          data-story-focus={storyFocusAttr}
         />
       ) : (
         <div className="px-12 pb-10">
@@ -337,27 +725,82 @@ export const TwProseMirrorSurface = forwardRef<
       ) : null}
     </section>
   );
+
+  function measureSelectionToolbarAnchor(): SelectionToolbarAnchor | null {
+    const callback = props.onSelectionToolbarAnchorChange;
+    const view = viewRef.current;
+    const mount = mountRef.current;
+    const positionMap = positionMapRef.current;
+    const range = snapshot.selection.activeRange;
+
+    if (!callback || !view || !mount || !positionMap || snapshot.selection.isCollapsed || range.kind !== "range") {
+      return null;
+    }
+
+    const rootRect = mount.getBoundingClientRect();
+    if (rootRect.width <= 0 || rootRect.height <= 0) {
+      return null;
+    }
+
+    try {
+      const pmFrom = positionMap.runtimeToPm(range.from);
+      const pmTo = positionMap.runtimeToPm(range.to);
+      const startRect = view.coordsAtPos(pmFrom);
+      const endRect = view.coordsAtPos(pmTo);
+      const left = Math.max(rootRect.left, Math.min(startRect.left, endRect.left));
+      const right = Math.min(rootRect.right, Math.max(startRect.right, endRect.right));
+      const top = Math.max(rootRect.top, Math.min(startRect.top, endRect.top));
+      const bottom = Math.min(rootRect.bottom, Math.max(startRect.bottom, endRect.bottom));
+
+      if (
+        !Number.isFinite(left) ||
+        !Number.isFinite(right) ||
+        !Number.isFinite(top) ||
+        !Number.isFinite(bottom) ||
+        right <= left ||
+        bottom <= top ||
+        bottom < rootRect.top ||
+        top > rootRect.bottom
+      ) {
+        return null;
+      }
+
+      return { left, right, top, bottom };
+    } catch {
+      return null;
+    }
+  }
 });
 
-function getActiveSearchResultIndex(
-  results: Array<{ from: number; to: number }>,
-  toRuntimePosition: (position: number) => number,
+function buildSelectionToolbarMeasurementKey(
   selection: SelectionSnapshot,
-): number {
-  if (results.length === 0) {
-    return -1;
+  activeStory: RuntimeRenderSnapshot["activeStory"],
+): string | null {
+  if (selection.isCollapsed || selection.activeRange.kind !== "range") {
+    return null;
   }
 
-  const selectionFrom = Math.min(selection.anchor, selection.head);
-  const selectionTo = Math.max(selection.anchor, selection.head);
-  const activeIndex = results.findIndex((result) => {
-    const from = toRuntimePosition(result.from);
-    const to = toRuntimePosition(result.to);
-    if (selectionFrom === selectionTo) {
-      return selectionFrom >= from && selectionFrom <= to;
-    }
-    return selectionFrom < to && selectionTo > from;
+  return JSON.stringify({
+    story: activeStory,
+    from: selection.activeRange.from,
+    to: selection.activeRange.to,
   });
+}
 
-  return activeIndex >= 0 ? activeIndex : 0;
+function selectionToolbarAnchorsEqual(
+  left: SelectionToolbarAnchor | null,
+  right: SelectionToolbarAnchor | null,
+): boolean {
+  if (left === right) {
+    return true;
+  }
+  if (!left || !right) {
+    return false;
+  }
+  return (
+    left.left === right.left &&
+    left.right === right.right &&
+    left.top === right.top &&
+    left.bottom === right.bottom
+  );
 }