fix: stale editor view in cell action buttons (marimo-team#1029)

* fix: stale editor view in cell action buttons * fix
Benni-Math · Apr 16, 2024 · 8d71b9c · 8d71b9c
1 parent e493bf1
commit 8d71b9c
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 8 deletions.
diff --git a/frontend/src/components/editor/Cell.tsx b/frontend/src/components/editor/Cell.tsx
@@ -190,6 +190,9 @@ const CellComponent = (
     [editorView, prepareToRunEffects],
   );
 
+  // Callback to get the editor view.
+  const getEditorView = useCallback(() => editorView.current, [editorView]);
+
   const handleRun = useEvent(async () => {
     if (loading) {
       return;
@@ -385,7 +388,7 @@ const CellComponent = (
       cellId={cellId}
       config={cellConfig}
       status={status}
-      editorView={editorView.current}
+      getEditorView={getEditorView}
       hasOutput={hasOutput}
       name={name}
     >
@@ -459,7 +462,7 @@ const CellComponent = (
                 ref={cellActionDropdownRef}
                 cellId={cellId}
                 status={status}
-                editorView={editorView.current}
+                getEditorView={getEditorView}
                 name={name}
                 config={cellConfig}
                 hasOutput={hasOutput}

diff --git a/frontend/src/components/editor/actions/useCellActionButton.tsx b/frontend/src/components/editor/actions/useCellActionButton.tsx
@@ -44,8 +44,8 @@ export interface CellActionButtonProps
   extends Pick<CellData, "name" | "config"> {
   cellId: CellId;
   status: CellStatus;
-  editorView: EditorView | null;
   hasOutput: boolean;
+  getEditorView: () => EditorView | null;
 }
 
 interface Props {
@@ -66,11 +66,11 @@ export function useCellActionButtons({ cell }: Props) {
   const runCell = useRunCell(cell?.cellId);
   const { openModal } = useImperativeModal();
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
-
   if (!cell) {
     return [];
   }
-  const { cellId, config, editorView, name, hasOutput, status } = cell;
+  const { cellId, config, getEditorView, name, hasOutput, status } = cell;
+  const editorView = getEditorView();
 
   const toggleDisabled = async () => {
     const newConfig = { disabled: !config.disabled };

diff --git a/frontend/src/components/editor/cell/cell-context-menu.tsx b/frontend/src/components/editor/cell/cell-context-menu.tsx
@@ -48,7 +48,8 @@ export const CellActionsContextMenu = ({ children, ...props }: Props) => {
       label: "Paste",
       icon: <ClipboardPasteIcon size={13} strokeWidth={1.5} />,
       handle: async () => {
-        const { editorView } = props;
+        const { getEditorView } = props;
+        const editorView = getEditorView();
         if (!editorView) {
           return;
         }

diff --git a/frontend/src/core/cells/focus.ts b/frontend/src/core/cells/focus.ts
@@ -15,7 +15,7 @@ export const lastFocusedCellAtom = atom<{
   config: CellConfig;
   cellId: CellId;
   status: CellStatus;
-  editorView: EditorView | null;
+  getEditorView: () => EditorView | null;
   hasOutput: boolean;
 } | null>((get) => {
   const cellId = get(lastFocusedCellIdAtom);
@@ -29,13 +29,14 @@ export const lastFocusedCellAtom = atom<{
   if (!data || !runtime || !handle) {
     return null;
   }
+  const getEditorView = () => handle.editorView;
 
   return {
     cellId,
     name: data.name,
     config: data.config,
     status: runtime.status,
-    editorView: handle.editorView,
+    getEditorView: getEditorView,
     hasOutput: runtime.output !== null,
   };
 });