@finos/legend-cube-builder 0.0.2 → 0.0.4

package/src/components/CubeEditor.tsx

@@ -35,22 +35,25 @@ import {
 } from '../__lib__/LegendCubeLabels.js';
 import { LEGEND_CUBE_TEST_ID } from '../__lib__/LegendCubeTesting.js';
 import { CubeEditorState } from '../stores/CubeEditorState.js';
+import {
+  type CubeEntrySource,
+  CubeEntrySourceState,
+} from '../stores/CubeEntrySource.js';
 import type { CubeHost } from '../stores/CubeHost.js';
+import { CubeAddItemsMenu } from './CubeAddItems.js';
 import { CubeButton } from './CubeButton.js';
 import { CubeCanvas } from './canvas/CubeCanvas.js';
-import { CubeNodeEditorPanel } from './editors/CubeNodeEditorPanel.js';
+import { CubeEntrySourceBanner } from './CubeEntrySourceBanner.js';
 import { CubeGridRegion } from './grid/CubeGridRegion.js';
 import { CubePalette } from './palette/CubePalette.js';
 import { CubeShowPureDialog } from './show-pure/CubeShowPureDialog.js';
+import { CubeExamplesDialog } from './examples/CubeExamplesDialog.js';
 import { CubeSourcePicker } from './source-picker/CubeSourcePicker.js';
 import { CubeSpecTransferDialog } from './spec-transfer/CubeSpecTransferDialog.js';
 
 /** The graph never takes more of the window's height than this, so the results stay in view (spec §17.3) */
 const MAX_GRAPH_SHARE_OF_WINDOW = 0.6;
 
-/** How wide the node editor opens, in pixels */
-const NODE_EDITOR_WIDTH = 400;
-
 /** The tallest the graph region may be, following the window's height */
 const useMaxGraphHeight = (): number => {
   const [windowHeight, setWindowHeight] = useState(() => window.innerHeight);
@@ -85,12 +88,17 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
             {CUBE_PENDING_LABEL.RESOLVING_SOURCE}
           </span>
         )}
+        {editorState.isTypingExtends && (
+          <span className="shrink-0 text-base text-[var(--color-text-secondary)]">
+            {CUBE_PENDING_LABEL.TYPING_COLUMNS}
+          </span>
+        )}
+        <CubeAddItemsMenu editorState={editorState} />
         <CubeButton
-          title={readOnly ? READ_ONLY_CUBE_TITLE : undefined}
-          disabled={readOnly}
-          onClick={() => editorState.sourcePicker.open()}
+          title="Open an example cube, or start one on sample data"
+          onClick={() => editorState.examples.open()}
         >
-          Add table
+          Examples
         </CubeButton>
         <CubeButton
           title={
@@ -98,8 +106,8 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
               ? READ_ONLY_CUBE_TITLE
               : `Undo the last change (${UNDO_SHORTCUT_LABEL})`
           }
-          disabled={!editorState.canUndo}
-          onClick={() => editorState.undo()}
+          disabled={!editorState.canUndo && !editorState.hasEditsToApply}
+          onClick={() => editorState.undoEdited()}
         >
           Undo
         </CubeButton>
@@ -111,6 +119,10 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
           }
           disabled={!editorState.execution.canExecute}
           onClick={() => {
+            // what it shows includes the node editor's edits, as Execute runs them
+            if (!editorState.nodeEditor.finishApplied()) {
+              return;
+            }
             flowResult(editorState.showPure.open()).catch(
               editorState.host.applicationStore.alertUnhandledError,
             );
@@ -125,15 +137,19 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
               : "Show the cube's spec, to copy or download"
           }
           disabled={readOnly}
-          onClick={() => editorState.specTransfer.openExport()}
+          onClick={() => {
+            if (editorState.nodeEditor.finishApplied()) {
+              editorState.specTransfer.openExport();
+            }
+          }}
         >
-          Export (dev)
+          Export
         </CubeButton>
         <CubeButton
           title="Open a cube from its spec, in place of this one"
           onClick={() => editorState.specTransfer.openImport()}
         >
-          Import (dev)
+          Import
         </CubeButton>
         <CubeButton
           title={
@@ -141,12 +157,21 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
               ? 'Hide the graph, leaving more room for the results'
               : 'Show the graph'
           }
-          onClick={() => editorState.setShowGraph(!showGraph)}
+          onClick={() => {
+            // the floating editor goes with the graph, applying its edits
+            if (!showGraph || editorState.nodeEditor.finish()) {
+              editorState.setShowGraph(!showGraph);
+            }
+          }}
         >
           {showGraph ? 'Hide graph' : 'Show graph'}
         </CubeButton>
       </div>
-      <PanelLoadingIndicator isLoading={editorState.isResolvingSources} />
+      <PanelLoadingIndicator
+        isLoading={
+          editorState.isResolvingSources || editorState.isTypingExtends
+        }
+      />
       {readOnly && (
         <div
           className="shrink-0 border-b border-[var(--color-border-default)] bg-[var(--color-status-warn-bg)] px-2 py-1 text-base text-[var(--color-status-warn)]"
@@ -170,6 +195,7 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
         </div>
       )}
       <CubeSourcePicker editorState={editorState} />
+      <CubeExamplesDialog editorState={editorState} />
       <CubeSpecTransferDialog editorState={editorState} />
       <CubeShowPureDialog editorState={editorState} />
     </div>
@@ -178,7 +204,8 @@ const CubeGraphHeader = observer((props: { editorState: CubeEditorState }) => {
 
 /**
  * The Legend Cube page: the query above, its results below, both always
- * shown (PLAN §7.1); hiding the graph leaves its header. The host gives it
+ * shown (PLAN §7.1); hiding the graph leaves its header. A node's editor
+ * floats below the node (PLAN §11.8). The host gives it
  * the engine, the models and the application store; the page's state lives
  * as long as the page.
  */
@@ -187,17 +214,26 @@ export const CubeEditor = observer(
     host: CubeHost;
     /** The cube to open; an empty one by default */
     initialDocument?: CubeDocument | undefined;
+    /** A source a link asks the page to start with, on an empty cube (spec §17.15) */
+    initialSource?: CubeEntrySource | undefined;
   }) => {
     const [editorState] = useState(
       () => new CubeEditorState(props.host, props.initialDocument),
     );
     useEffect(() => () => editorState.dispose(), [editorState]);
+    const [entry] = useState(() => new CubeEntrySourceState(editorState));
+    const { initialSource } = props;
+    useEffect(() => {
+      if (initialSource) {
+        flowResult(entry.open(initialSource)).catch(
+          editorState.host.applicationStore.alertUnhandledError,
+        );
+      }
+      // only as the page opens
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
     useCommands(editorState);
     const maxGraphHeight = useMaxGraphHeight();
-    const isEditorOpen = editorState.nodeEditor.nodeId !== undefined;
-    const editorPanel = getCollapsiblePanelGroupProps(!isEditorOpen, {
-      size: NODE_EDITOR_WIDTH,
-    });
     const { showGraph } = editorState.document.meta.presentation;
     // hiding the graph collapses its panel: the header and the results stay
     // where they are, so the grid keeps its state
@@ -209,30 +245,23 @@ export const CubeEditor = observer(
         data-testid={LEGEND_CUBE_TEST_ID.EDITOR}
       >
         <CubePalette editorState={editorState} />
-        <ResizablePanelGroup orientation="vertical">
-          <ResizablePanel {...editorPanel.remainingPanel} minSize={320}>
-            <div className="flex h-full min-w-0 flex-1 flex-col">
-              <CubeGraphHeader editorState={editorState} />
-              <ResizablePanelGroup orientation="horizontal">
-                <ResizablePanel
-                  {...graphPanel.collapsiblePanel}
-                  minSize={showGraph ? 96 : 0}
-                  maxSize={maxGraphHeight}
-                >
-                  {showGraph && <CubeCanvas editorState={editorState} />}
-                </ResizablePanel>
-                <ResizablePanelSplitter className={showGraph ? '' : 'hidden'} />
-                <ResizablePanel {...graphPanel.remainingPanel} minSize={96}>
-                  <CubeGridRegion editorState={editorState} />
-                </ResizablePanel>
-              </ResizablePanelGroup>
-            </div>
-          </ResizablePanel>
-          <ResizablePanelSplitter className={isEditorOpen ? '' : 'hidden'} />
-          <ResizablePanel {...editorPanel.collapsiblePanel} direction={-1}>
-            <CubeNodeEditorPanel editorState={editorState} />
-          </ResizablePanel>
-        </ResizablePanelGroup>
+        <div className="flex h-full min-w-0 flex-1 flex-col">
+          <CubeGraphHeader editorState={editorState} />
+          <CubeEntrySourceBanner entry={entry} />
+          <ResizablePanelGroup orientation="horizontal">
+            <ResizablePanel
+              {...graphPanel.collapsiblePanel}
+              minSize={showGraph ? 96 : 0}
+              maxSize={maxGraphHeight}
+            >
+              {showGraph && <CubeCanvas editorState={editorState} />}
+            </ResizablePanel>
+            <ResizablePanelSplitter className={showGraph ? '' : 'hidden'} />
+            <ResizablePanel {...graphPanel.remainingPanel} minSize={96}>
+              <CubeGridRegion editorState={editorState} />
+            </ResizablePanel>
+          </ResizablePanelGroup>
+        </div>
       </div>
     );
   },

package/src/components/CubeEntrySourceBanner.tsx

@@ -0,0 +1,64 @@
+/**
+ * Copyright (c) 2026-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { observer } from 'mobx-react-lite';
+import { LEGEND_CUBE_TEST_ID } from '../__lib__/LegendCubeTesting.js';
+import {
+  CUBE_ENTRY_SOURCE_ERROR_TITLE,
+  type CubeEntrySourceState,
+} from '../stores/CubeEntrySource.js';
+import { CubeButton } from './CubeButton.js';
+
+/** Why the source a link named couldn't be added, in the graph region (spec §17.13, U9(b)) */
+export const CubeEntrySourceBanner = observer(
+  (props: { entry: CubeEntrySourceState }) => {
+    const { entry } = props;
+    if (entry.isOpening) {
+      return (
+        <div
+          className="shrink-0 border-b border-[var(--color-border-default)] px-2 py-1 text-base text-[var(--color-text-secondary)]"
+          role="status"
+        >
+          Opening the linked source…
+        </div>
+      );
+    }
+    if (entry.error === undefined) {
+      return null;
+    }
+    return (
+      <div
+        className="flex shrink-0 items-start gap-2 border-b border-[var(--color-border-default)] bg-[var(--color-status-error-bg)] px-2 py-1 text-base text-[var(--color-status-error)]"
+        role="alert"
+        data-testid={LEGEND_CUBE_TEST_ID.ENTRY_SOURCE_ERROR}
+      >
+        <span className="min-w-0 flex-1 break-words">
+          <span className="font-medium">{CUBE_ENTRY_SOURCE_ERROR_TITLE}</span>{' '}
+          {entry.error.message}
+          {entry.error.detail !== undefined && (
+            <details>
+              <summary className="cursor-pointer">Details</summary>
+              <pre className="whitespace-pre-wrap text-sm">
+                {entry.error.detail}
+              </pre>
+            </details>
+          )}
+        </span>
+        <CubeButton onClick={() => entry.dismissError()}>Dismiss</CubeButton>
+      </div>
+    );
+  },
+);

package/src/components/CubeNodeIcon.tsx

@@ -15,10 +15,23 @@
  */
 
 import {
+  AlignBottomIcon,
+  AlignMiddleIcon,
+  AlignTopIcon,
   ArrowsJoinIcon,
+  CalculatorIcon,
+  CompareIcon,
+  CompressIcon,
+  DataCubeIcon,
+  DatabaseImportIcon,
   clsx,
   FilterIcon,
+  LayerGroupIcon,
+  PackageIcon,
+  PencilIcon,
   QuestionSquareIcon,
+  SigmaIcon,
+  SortIcon,
   TableIcon,
 } from '@finos/legend-art';
 
@@ -27,8 +40,22 @@ const NODE_ICONS: Readonly<
   Record<string, React.FC<{ className?: string; title?: string }>>
 > = {
   table: TableIcon,
+  dataProduct: PackageIcon,
+  ingest: DatabaseImportIcon,
   filter: FilterIcon,
   join: ArrowsJoinIcon,
+  concat: LayerGroupIcon,
+  difference: CompareIcon,
+  extend: CalculatorIcon,
+  limit: AlignTopIcon,
+  drop: AlignBottomIcon,
+  slice: AlignMiddleIcon,
+  distinct: CompressIcon,
+  restrict: DataCubeIcon.TableColumns,
+  group: DataCubeIcon.TableGroupBy,
+  partition: SigmaIcon,
+  rename: PencilIcon,
+  sort: SortIcon,
 };
 
 /** Whether a node definition's icon name maps to an icon, rather than the question mark */
@@ -41,7 +68,10 @@ export const CubeNodeIcon: React.FC<{
   className?: string;
 }> = (props) => {
   const { icon, className } = props;
+  // an own key only: a name such as `constructor` must not reach Object's
   const Icon =
-    (icon === undefined ? undefined : NODE_ICONS[icon]) ?? QuestionSquareIcon;
+    icon !== undefined && hasCubeNodeIcon(icon)
+      ? (NODE_ICONS[icon] ?? QuestionSquareIcon)
+      : QuestionSquareIcon;
   return <Icon className={clsx(className)} />;
 };

package/src/components/CubeSchemaColumnsTable.tsx

@@ -0,0 +1,49 @@
+/**
+ * Copyright (c) 2026-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type { Schema } from '@finos/legend-cube';
+import { getColumnTypeLabel } from '../__lib__/LegendCubeLabels.js';
+import { CubeColumnTypeIcon } from './editors/CubeColumnPicker.js';
+
+/** A source's columns, with their precise types and whether they can be empty */
+export const CubeSchemaColumnsTable: React.FC<{ schema: Schema }> = ({
+  schema,
+}) => (
+  <table className="w-full text-base" aria-label="Columns">
+    <thead>
+      <tr className="text-left text-[var(--color-text-secondary)]">
+        <th className="font-normal">Column</th>
+        <th className="font-normal">Type</th>
+      </tr>
+    </thead>
+    <tbody>
+      {schema.columns.map((column) => (
+        <tr key={column.name}>
+          <td className="break-all pr-2">{column.name}</td>
+          <td title={column.type.fullName}>
+            <span className="flex items-center gap-1">
+              <CubeColumnTypeIcon
+                type={column.type}
+                className="text-[var(--color-text-secondary)]"
+              />
+              {getColumnTypeLabel(column)}
+            </span>
+          </td>
+        </tr>
+      ))}
+    </tbody>
+  </table>
+);

package/src/components/canvas/CubeCanvas.tsx

@@ -36,6 +36,7 @@ import { observer } from 'mobx-react-lite';
 import { useEffect, useMemo, useRef } from 'react';
 import { useDrop } from 'react-dnd';
 import { LEGEND_CUBE_TEST_ID } from '../../__lib__/LegendCubeTesting.js';
+import { canAddCubeNode } from '../../stores/CubeAddPlacement.js';
 import type { CubeEditorState } from '../../stores/CubeEditorState.js';
 import {
   buildCubeCanvasEdges,
@@ -56,6 +57,7 @@ import {
 import { CubeCanvasEdge } from './CubeCanvasEdge.js';
 import { layoutCubeQuery } from './CubeCanvasLayout.js';
 import { CubeCanvasNode } from './CubeCanvasNode.js';
+import { CubeNodeEditorPopper } from '../editors/CubeNodeEditorPopper.js';
 
 // React Flow re-renders every node when these change identity
 const NODE_TYPES = { [CUBE_CANVAS_NODE_TYPE]: CubeCanvasNode };
@@ -93,14 +95,15 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
   );
   const edges = useMemo(() => buildCubeCanvasEdges(query), [query]);
   // the view fits the graph again whenever the layout moves a node, or the
-  // canvas gets wider or narrower (e.g. the node editor opens beside it),
-  // once React Flow has measured every node: it fits only measured ones
+  // canvas changes size (e.g. the splitter above the grid moves), once React
+  // Flow has measured every node: it fits only measured ones
   const layoutSignature = useMemo(
     () => JSON.stringify([...positions]),
     [positions],
   );
   const nodesInitialized = useNodesInitialized();
   const canvasWidth = useStore((state) => state.width);
+  const canvasHeight = useStore((state) => state.height);
   useEffect(() => {
     if (!nodesInitialized) {
       return undefined;
@@ -109,7 +112,7 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
       fitView(FIT_VIEW_OPTIONS).catch(noop());
     }, 0);
     return () => window.clearTimeout(timer);
-  }, [layoutSignature, canvasWidth, nodesInitialized, fitView]);
+  }, [layoutSignature, canvasWidth, canvasHeight, nodesInitialized, fitView]);
 
   return (
     <ReactFlow<CubeCanvasFlowNode, CubeCanvasFlowEdge>
@@ -133,7 +136,10 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
           return;
         }
         if (event.ctrlKey || event.metaKey) {
-          editorState.select(flowNode.id);
+          // applies and closes the open editor first, opening none (spec §17.5)
+          if (editorState.nodeEditor.finish()) {
+            editorState.select(flowNode.id);
+          }
         } else {
           editorState.nodeEditor.open(flowNode.id);
         }
@@ -158,14 +164,31 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
         }
         event.preventDefault();
         if (event.ctrlKey || event.metaKey) {
-          editorState.select(nodeId);
+          if (editorState.nodeEditor.finish()) {
+            editorState.select(nodeId);
+          }
         } else {
           editorState.nodeEditor.open(nodeId);
+          // the keyboard goes on in the editor, once it is shown
+          window.setTimeout(() =>
+            document
+              .querySelector<HTMLElement>(
+                `[data-testid="${LEGEND_CUBE_TEST_ID.NODE_EDITOR}"]`,
+              )
+              ?.focus(),
+          );
         }
       }}
       isValidConnection={(connection) =>
         !readOnly && isCubeCanvasConnectionValid(query, connection)
       }
+      // a connection changes the query, so the open editor applies first,
+      // whether the handles are dragged or clicked
+      onConnectStart={() => editorState.nodeEditor.finish()}
+      // a click on the canvas's background closes the floating editor; a
+      // press there that pans leaves it open
+      onPaneClick={() => editorState.nodeEditor.finish()}
+      onClickConnectStart={() => editorState.nodeEditor.finish()}
       onConnect={(connection) => {
         if (connection.targetHandle) {
           editorState.connect(
@@ -190,6 +213,7 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
       <Background />
       <Controls showInteractive={false} />
       <MiniMap pannable={true} zoomable={true} style={MINI_MAP_STYLE} />
+      <CubeNodeEditorPopper editorState={editorState} positions={positions} />
     </ReactFlow>
   );
 });
@@ -197,9 +221,10 @@ const CubeCanvasFlow = observer((props: { editorState: CubeEditorState }) => {
 /**
  * The query as a graph, laid out left to right (spec §17.3). Click a node to
  * edit it, Ctrl or Cmd-click it to run the query up to it, drag from a node's
- * output to another node's input to connect them. A palette item dropped
- * around the nodes is added unconnected. Right-click it, or a node, for the
- * context menu.
+ * output to another node's input to connect them. A transform dropped
+ * around the nodes goes after the selected node; a source, dropped anywhere,
+ * opens the source dialog on its tab (PLAN §11.8). Right-click it, or a
+ * node, for the context menu.
  */
 export const CubeCanvas = observer(
   (props: { editorState: CubeEditorState }) => {
@@ -208,7 +233,8 @@ export const CubeCanvas = observer(
     const [, dropConnector] = useDrop<CubePaletteDragItem>(
       () => ({
         accept: [CUBE_DND_TYPE.PALETTE_ITEM],
-        canDrop: () => !editorState.readOnly,
+        canDrop: (item) =>
+          !editorState.readOnly && canAddCubeNode(editorState, item.nodeType),
         // a node under the pointer, whether it took the drop or refused it,
         // keeps it from the canvas
         drop: (item, monitor) =>
@@ -229,21 +255,32 @@ export const CubeCanvas = observer(
       >
         <ContextMenu
           className="h-full w-full"
+          // an empty query has no menu: there is nothing to add after (U3)
+          disabled={editorState.document.query.isEmpty}
           content={<CubeCanvasContextMenu editorState={editorState} />}
           menuProps={{ elevation: 7 }}
         >
           {editorState.document.query.isEmpty ? (
             <div className="flex h-full items-center justify-center p-4 text-base text-[var(--color-text-secondary)]">
               <span>
-                No tables yet:{' '}
                 <button
                   className="text-[var(--color-accent)] underline disabled:text-[var(--color-text-disabled)] disabled:no-underline"
-                  disabled={editorState.readOnly}
-                  onClick={() => editorState.sourcePicker.open()}
+                  title={editorState.sourcePicker.disabledReason}
+                  disabled={
+                    editorState.sourcePicker.disabledReason !== undefined
+                  }
+                  onClick={() => editorState.sourcePicker.openToChoose()}
                 >
-                  add a table
+                  Connect to a source
                 </button>{' '}
-                to start.
+                to start a new one, or{' '}
+                <button
+                  className="text-[var(--color-accent)] underline"
+                  onClick={() => editorState.examples.open()}
+                >
+                  open an example
+                </button>
+                .
               </span>
             </div>
           ) : (

package/src/components/canvas/CubeCanvasContextMenu.tsx

@@ -18,21 +18,22 @@ import {
   MenuContent,
   MenuContentDivider,
   MenuContentItem,
-  MenuContentItemIcon,
-  MenuContentItemLabel,
 } from '@finos/legend-art';
-import type { AnyNodeDefinition } from '@finos/legend-cube';
 import { observer } from 'mobx-react-lite';
 import { forwardRef } from 'react';
 import type { CubeEditorState } from '../../stores/CubeEditorState.js';
-import { CubeNodeIcon } from '../CubeNodeIcon.js';
+import { CubeAddMenuItems } from '../CubeAddItems.js';
 
 /**
  * The canvas's context menu (spec §17.4): the palette, then what can be done
- * to the node it was opened on. On a node, a palette item is added after it;
- * around the nodes, on its own. Every item shows, disabled when it can't be
- * done, so Swap Inputs can be found. The node comes through props: a menu
- * opened around the nodes has none.
+ * to the node it was opened on. On a node, a transform is added after it;
+ * around the nodes, after the selected node; a source opens the source
+ * dialog on its tab (PLAN §11.8, U3 and U4). Every item shows, disabled when
+ * it can't be done, so Swap Inputs can be found. The node comes through
+ * props: a menu opened around the nodes has none. Each item applies the node
+ * editor first, so its edits are never dropped; a palette item adds nothing
+ * when they had to be. Swap Inputs on the edited node is its editor's own,
+ * which applies and follows it.
  */
 export const CubeCanvasContextMenu = observer(
   forwardRef<
@@ -41,45 +42,45 @@ export const CubeCanvasContextMenu = observer(
   >(function CubeCanvasContextMenu(props, ref) {
     const { editorState, nodeId } = props;
     const { query } = editorState.document;
-    const { registry } = editorState;
-    const paletteItem = (definition: AnyNodeDefinition): React.ReactNode => (
-      <MenuContentItem
-        key={definition.type}
-        disabled={!editorState.canAddNode(definition.type, nodeId)}
-        onClick={() => editorState.addNode(definition.type, nodeId)}
-      >
-        <MenuContentItemIcon>
-          <CubeNodeIcon icon={definition.icon} />
-        </MenuContentItemIcon>
-        <MenuContentItemLabel>
-          {definition.label}
-          {definition.beta ? ' (BETA)' : ''}
-        </MenuContentItemLabel>
-      </MenuContentItem>
-    );
+    const finishing =
+      (run: () => void): (() => void) =>
+      () => {
+        if (editorState.nodeEditor.finish()) {
+          run();
+        }
+      };
     return (
       <MenuContent ref={ref}>
-        {registry.sources.map(paletteItem)}
-        <MenuContentDivider />
-        {registry.transforms.map(paletteItem)}
+        <CubeAddMenuItems editorState={editorState} nodeId={nodeId} />
         <MenuContentDivider />
         <MenuContentItem
           title="Execute runs the query up to this node"
           disabled={nodeId === undefined || !query.canSelect(nodeId)}
-          onClick={() => nodeId !== undefined && editorState.select(nodeId)}
+          onClick={finishing(
+            () => nodeId !== undefined && editorState.select(nodeId),
+          )}
         >
           Select
         </MenuContentItem>
         <MenuContentItem
           disabled={nodeId === undefined || !editorState.canRemoveNode(nodeId)}
-          onClick={() => nodeId !== undefined && editorState.removeNode(nodeId)}
+          onClick={finishing(
+            () => nodeId !== undefined && editorState.removeNode(nodeId),
+          )}
         >
           Remove
         </MenuContentItem>
         <MenuContentItem
           title="Swap the node's two inputs, e.g. a Join's Left and Right"
           disabled={nodeId === undefined || !editorState.canSwapInputs(nodeId)}
-          onClick={() => nodeId !== undefined && editorState.swapInputs(nodeId)}
+          onClick={
+            nodeId !== undefined && nodeId === editorState.nodeEditor.nodeId
+              ? // its own editor applies, swaps and stays open on it
+                () => editorState.nodeEditor.swapInputs()
+              : finishing(
+                  () => nodeId !== undefined && editorState.swapInputs(nodeId),
+                )
+          }
         >
           Swap Inputs
         </MenuContentItem>