feat(ui): add ui.selection.restore for capture round-trip (SD-2936)

packages/super-editor/src/ui/create-super-doc-ui.ts

@@ -14,15 +14,18 @@ import type {
   ScrollIntoViewOutput,
   TrackChangesListResult,
 } from '@superdoc/document-api';
+import { collectEntityHitsFromChain } from './entity-at.js';
 import { shallowEqual } from './equality.js';
 import { scrollRangeIntoView } from './scroll-into-view.js';
 import { getSelectionAnchorRect, getSelectionRects } from './selection-rects.js';
+import { restoreSelection } from './selection-restore.js';
 import { createCustomCommandsRegistry } from './custom-commands.js';
 import { createScope } from './scope.js';
 import type {
   CommandHandle,
   CommandsHandle,
   CommentsHandle,
+  ContextMenuItem,
   DocumentExportInput,
   DocumentHandle,
   DocumentSlice,
@@ -44,6 +47,8 @@ import type {
   ToolbarHandle,
   ToolbarSnapshotSlice,
   UIToolbarCommandState,
+  ViewportEntityAtInput,
+  ViewportEntityHit,
   ViewportGetRectInput,
   ViewportHandle,
   ViewportRect,
@@ -1162,6 +1167,15 @@ export function createSuperDocUI(options: SuperDocUIOptions): SuperDocUI {
           return handle;
         };
       }
+      // Custom-UI consumers building their own context menu pull
+      // contributed items here. Computed against the current snapshot
+      // (so `selection` matches what observers just saw) and the
+      // caller-supplied entities from `ui.viewport.entityAt`.
+      if (prop === 'getContextMenuItems') {
+        return (input?: { entities?: ViewportEntityHit[] }): ContextMenuItem[] => {
+          return customCommandsRegistry.getContextMenuItems(computeState(), input?.entities ?? []);
+        };
+      }
       // Custom-registered ids surface a typed handle from the registry.
       // Built-in ids fall through to the existing per-id cache so they
       // keep the same observe/execute shape they had before SD-2802.
@@ -1560,6 +1574,35 @@ export function createSuperDocUI(options: SuperDocUIOptions): SuperDocUI {
     async scrollIntoView(input: ScrollIntoViewInput): Promise<ScrollIntoViewOutput> {
       return runScrollIntoView(input);
     },
+
+    // The painter stamps `data-track-change-id` and `data-comment-ids`
+    // on each painted run; reading them back is what consumers were
+    // doing imperatively from `event.target.closest(...)` in
+    // contextmenu handlers. Centralizing the lookup here keeps the
+    // attribute names an implementation detail of the painter and
+    // surfaces a typed `EntityHit[]` consumers can switch on.
+    entityAt(input: ViewportEntityAtInput): ViewportEntityHit[] {
+      if (!input || typeof input.x !== 'number' || typeof input.y !== 'number') return [];
+      // The DOM `document` is reached through `globalThis.document`
+      // because the local `document: DocumentHandle` declared below
+      // would otherwise shadow it for type-checking. Guard SSR /
+      // non-browser stubs explicitly so the call doesn't throw in
+      // test environments without a global `document`.
+      const dom = (globalThis as { document?: Document }).document;
+      if (!dom || typeof dom.elementFromPoint !== 'function') {
+        return [];
+      }
+      // Scope the lookup to this controller's editor: a page mounting
+      // two SuperDoc instances would otherwise have one's entityAt
+      // return ids from the other's painted DOM. A null host (no
+      // editor mounted, post-destroy, SSR test stub) returns [].
+      const editor = resolveHostEditor(superdoc);
+      const host = editor?.presentationEditor?.visibleHost;
+      if (!host) return [];
+      const startEl = dom.elementFromPoint(input.x, input.y);
+      if (!startEl || !host.contains(startEl)) return [];
+      return collectEntityHitsFromChain(startEl);
+    },
   };
 
   // ---- ui.selection ------------------------------------------------------
@@ -1636,6 +1679,16 @@ export function createSuperDocUI(options: SuperDocUIOptions): SuperDocUI {
         capture,
       );
     },
+    restore(capture) {
+      // Routed editor: same rationale as `getRects(capture)` — block
+      // ids in a non-body capture only resolve in their own story
+      // editor's PM doc. When focus has moved to the body by call
+      // time, the routed editor is body and resolution returns
+      // `'stale'` rather than placing the selection on the wrong
+      // surface.
+      const editor = resolveRoutedEditor(superdoc);
+      return restoreSelection(editor as unknown as Parameters<typeof restoreSelection>[0], capture);
+    },
   };
 
   // ---- ui.document -------------------------------------------------------

packages/super-editor/src/ui/custom-commands.test.ts

@@ -949,3 +949,200 @@ describe('ui.commands.get', () => {
     ui.destroy();
   });
 });
+
+describe('ui.commands.getContextMenuItems', () => {
+  it('returns [] when no registered command carries a contextMenu field', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    ui.commands.register({ id: 'company.plain', execute: () => true });
+
+    expect(ui.commands.getContextMenuItems()).toEqual([]);
+
+    ui.destroy();
+  });
+
+  it('surfaces contributions, filling defaults for group / order', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    ui.commands.register({
+      id: 'company.acceptChange',
+      execute: () => true,
+      contextMenu: { label: 'Accept suggestion' },
+    });
+
+    expect(ui.commands.getContextMenuItems()).toEqual([
+      { id: 'company.acceptChange', label: 'Accept suggestion', group: 'custom', order: 0 },
+    ]);
+
+    ui.destroy();
+  });
+
+  it('filters items by the when predicate using caller-supplied entities + current selection', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    const whenSpy = vi.fn(({ entities }) => entities.some((e) => e.type === 'trackedChange'));
+    ui.commands.register({
+      id: 'company.acceptChange',
+      execute: () => true,
+      contextMenu: { label: 'Accept suggestion', group: 'review', when: whenSpy },
+    });
+
+    expect(ui.commands.getContextMenuItems({ entities: [{ type: 'comment', id: 'c1' }] })).toEqual([]);
+    expect(ui.commands.getContextMenuItems({ entities: [{ type: 'trackedChange', id: 'tc1' }] })).toHaveLength(1);
+
+    expect(whenSpy).toHaveBeenCalledTimes(2);
+    expect(whenSpy.mock.calls[0]![0].selection).toBeDefined();
+
+    ui.destroy();
+  });
+
+  it('sorts by built-in group order, then order, then registration sequence', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    // Built-in group order: format(0), clipboard(1), review(2), comment(3), link(4)
+    ui.commands.register({
+      id: 'a.review-2',
+      execute: () => true,
+      contextMenu: { label: 'Review B', group: 'review', order: 20 },
+    });
+    ui.commands.register({
+      id: 'b.format-1',
+      execute: () => true,
+      contextMenu: { label: 'Format A', group: 'format', order: 0 },
+    });
+    ui.commands.register({
+      id: 'c.review-1',
+      execute: () => true,
+      contextMenu: { label: 'Review A', group: 'review', order: 10 },
+    });
+    ui.commands.register({
+      id: 'd.review-tie-second',
+      execute: () => true,
+      contextMenu: { label: 'Review C', group: 'review', order: 10 },
+    });
+    ui.commands.register({
+      id: 'e.custom',
+      execute: () => true,
+      contextMenu: { label: 'Z', group: 'company.workflow', order: 0 },
+    });
+
+    expect(ui.commands.getContextMenuItems().map((i) => i.id)).toEqual([
+      'b.format-1',
+      'c.review-1',
+      'd.review-tie-second',
+      'a.review-2',
+      'e.custom',
+    ]);
+
+    ui.destroy();
+  });
+
+  it('plain custom commands (no contextMenu) do not anchor a custom group rank', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    // Register a plain command first (seq 0) — it has no contextMenu
+    // and must not claim the 'custom' fallback group's rank anchor.
+    ui.commands.register({ id: 'a.plain', execute: () => true });
+    // Register a workflow contribution (seq 1).
+    ui.commands.register({
+      id: 'b.workflow',
+      execute: () => true,
+      contextMenu: { label: 'Workflow A', group: 'company.workflow' },
+    });
+    // Register a 'custom' fallback group contribution (seq 2).
+    ui.commands.register({
+      id: 'c.custom',
+      execute: () => true,
+      contextMenu: { label: 'Default A' },
+    });
+
+    // 'company.workflow' (seq 1) must rank before 'custom' (seq 2).
+    // If the plain seq=0 command anchored 'custom', the order would
+    // flip.
+    expect(ui.commands.getContextMenuItems().map((i) => i.id)).toEqual(['b.workflow', 'c.custom']);
+
+    ui.destroy();
+  });
+
+  it('preserves a group rank anchor when one contributor is replaced and another remains', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    // Group 'workflow' opens with two contributors at seq 0 and seq 1.
+    ui.commands.register({
+      id: 'wf.first',
+      execute: () => true,
+      contextMenu: { label: 'WF 1', group: 'company.workflow', order: 0 },
+    });
+    ui.commands.register({
+      id: 'wf.second',
+      execute: () => true,
+      contextMenu: { label: 'WF 2', group: 'company.workflow', order: 1 },
+    });
+    // A second custom group registers at seq 2.
+    ui.commands.register({
+      id: 'rev.first',
+      execute: () => true,
+      contextMenu: { label: 'Rev 1', group: 'company.review-extras', order: 0 },
+    });
+
+    // Now replace `wf.first` — the new seq becomes 3, but `wf.second`
+    // still carries the original seq 1, so the workflow group's
+    // anchor must stay at 1 and render before 'review-extras' (seq 2).
+    ui.commands.register({
+      id: 'wf.first',
+      execute: () => true,
+      contextMenu: { label: 'WF 1 (replaced)', group: 'company.workflow', order: 0 },
+    });
+
+    expect(ui.commands.getContextMenuItems().map((i) => i.id)).toEqual(['wf.first', 'wf.second', 'rev.first']);
+
+    ui.destroy();
+  });
+
+  it('hides items whose when predicate throws and logs the error once per distinct message', () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    ui.commands.register({
+      id: 'company.flaky',
+      execute: () => true,
+      contextMenu: {
+        label: 'Flaky',
+        when: () => {
+          throw new Error('boom');
+        },
+      },
+    });
+
+    expect(ui.commands.getContextMenuItems()).toEqual([]);
+    expect(ui.commands.getContextMenuItems()).toEqual([]);
+    expect(errorSpy).toHaveBeenCalledTimes(1);
+
+    ui.destroy();
+  });
+
+  it("refuses 'getContextMenuItems' as a custom command id (Proxy collision)", () => {
+    const { superdoc } = makeStubs();
+    const ui = createSuperDocUI({ superdoc });
+
+    const reg = ui.commands.register({
+      id: 'getContextMenuItems',
+      execute: () => true,
+      contextMenu: { label: 'Should not register' },
+    });
+
+    expect(warnSpy).toHaveBeenCalled();
+    // Refused registrations get a no-op handle, so the contribution
+    // never enters the registry.
+    expect(ui.commands.getContextMenuItems()).toEqual([]);
+    reg.unregister();
+
+    ui.destroy();
+  });
+});