fix(frontend): include unmounted cells in LSP and Copilot context (#10677)

Author: kirangadhave

frontend/src/core/codemirror/copilot/__tests__/getCodes.test.ts

@@ -5,13 +5,18 @@ import { EditorView } from "@codemirror/view";
 import { describe, expect, it } from "vitest";
 import { cellId, variableName } from "@/__tests__/branded";
 import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
+import { createCell } from "@/core/cells/types";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import { MultiColumn } from "@/utils/id-tree";
 import { cellConfigExtension } from "../../config/extension";
 import { adaptiveLanguageConfiguration } from "../../language/extension";
-import { getCodes, getTopologicalCellIds } from "../getCodes";
+import {
+  getCodes,
+  getTopologicalCellIds,
+  topologicalCodesAtom,
+} from "../getCodes";
 
 const Cells = {
   cell1: cellId("cell1"),
@@ -208,4 +213,64 @@ describe("getCodes", () => {
     const result = getCodes(otherCode);
     expect(result).toEqual("import os\nx = 1\ny = x + 1\nprint('Hello World')");
   });
+
+  it("should read code from the store when no editor is mounted", () => {
+    const otherCode = "print('Hello World')";
+    store.set(notebookAtom, {
+      ...initialNotebookState(),
+      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellData: {
+        [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
+        [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),
+      },
+    });
+    const result = getCodes(otherCode);
+    expect(result).toEqual("import os\nx = 1\nprint('Hello World')");
+  });
+
+  it("should prefer the mounted editor code over the store code", () => {
+    const otherCode = "print('Hello World')";
+    store.set(notebookAtom, {
+      ...initialNotebookState(),
+      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellData: {
+        [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
+        [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),
+      },
+      cellHandles: {
+        [Cells.cell2]: { current: createMockEditorView("x = 2") },
+      },
+    });
+    const result = getCodes(otherCode);
+    expect(result).toEqual("import os\nx = 2\nprint('Hello World')");
+  });
+});
+
+describe("topologicalCodesAtom", () => {
+  it("should include every cell when no editor is mounted", () => {
+    store.set(notebookAtom, {
+      ...initialNotebookState(),
+      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2, Cells.cell3]]),
+      cellData: {
+        [Cells.cell1]: createCell({ id: Cells.cell1, code: "import os" }),
+        [Cells.cell2]: createCell({ id: Cells.cell2, code: "x = 1" }),
+        [Cells.cell3]: createCell({ id: Cells.cell3, code: "y = x + 1" }),
+      },
+    });
+    store.set(variablesAtom, {
+      [Variables.var1]: {
+        name: Variables.var1,
+        declaredBy: [Cells.cell1],
+        usedBy: [Cells.cell2, Cells.cell3],
+      },
+    });
+
+    const { cellIds, codes } = store.get(topologicalCodesAtom);
+    expect(cellIds).toEqual([Cells.cell1, Cells.cell2, Cells.cell3]);
+    expect(codes).toEqual({
+      [Cells.cell1]: "import os",
+      [Cells.cell2]: "x = 1",
+      [Cells.cell3]: "y = x + 1",
+    });
+  });
 });

frontend/src/core/codemirror/copilot/getCodes.ts

@@ -1,13 +1,32 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { atom } from "jotai";
-import { getAllEditorViews, notebookAtom } from "@/core/cells/cells";
+import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
+import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { Variables } from "@/core/variables/types";
 import { Objects } from "@/utils/objects";
 import { getEditorCodeAsPython } from "../language/utils";
 
+const notebookCellCodes = atom((get) => {
+  const notebook = get(notebookAtom);
+  const codes = Objects.fromEntries(
+    notebook.cellIds.inOrderIds.map((id) => {
+      // A cell editor mounts progressively, so the notebook store is the only
+      // complete source of code until every editor is built.
+      const editorView = notebook.cellHandles[id]?.current?.editorViewOrNull;
+      return [
+        id,
+        editorView
+          ? getEditorCodeAsPython(editorView)
+          : (notebook.cellData[id]?.code ?? ""),
+      ];
+    }),
+  );
+  return codes;
+});
+
 export function getCodes(otherCode: string) {
   const codes = getOtherCellsCode(otherCode);
 
@@ -18,15 +37,8 @@ export function getOtherCellsCode(otherCode: string) {
   // Get all other cells' code
   // Put `import` statements at the top, as it can help copilot give better suggestions
   // TODO: we should sort this topologically
-  const codes = getAllEditorViews()
-    .map((editorView) => {
-      const code = getEditorCodeAsPython(editorView);
-      if (code === otherCode) {
-        return null;
-      }
-      return code;
-    })
-    .filter(Boolean)
+  const codes = Object.values(store.get(notebookCellCodes))
+    .filter((code) => code !== "" && code !== otherCode)
     .toSorted((a, b) => {
       if (a.startsWith("import") && !b.startsWith("import")) {
         return -1;
@@ -40,21 +52,6 @@ export function getOtherCellsCode(otherCode: string) {
   return codes;
 }
 
-const notebookCellCodes = atom((get) => {
-  const notebook = get(notebookAtom);
-  const codes = Objects.fromEntries(
-    notebook.cellIds.inOrderIds.map((id) => {
-      const handle = notebook.cellHandles[id];
-      return [
-        id,
-        handle?.current?.editorView
-          ? getEditorCodeAsPython(handle.current.editorView)
-          : "",
-      ];
-    }),
-  );
-  return codes;
-});
 const inOrderCellIdsAtom = atom((get) => {
   const notebook = get(notebookAtom);
   return notebook.cellIds.inOrderIds;

frontend/src/core/codemirror/lsp/__tests__/notebook-lsp.test.ts

@@ -8,7 +8,10 @@ import { beforeEach, describe, expect, it, type Mocked, vi } from "vitest";
 import * as LSP from "vscode-languageserver-protocol";
 import { cellId } from "@/__tests__/branded";
 import type { CellId } from "@/core/cells/ids";
+import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
+import { createCell } from "@/core/cells/types";
 import { store } from "@/core/state/jotai";
+import { MultiColumn } from "@/utils/id-tree";
 import { topologicalCodesAtom } from "../../copilot/getCodes";
 import { lspWorkspaceAtom } from "@/core/saving/file-state";
 import { languageAdapterState } from "../../language/extension";
@@ -62,6 +65,28 @@ function isPublishDiagnosticsNotification(
   );
 }
 
+describe("merged document with unmounted editors", () => {
+  it("should include the code of cells whose editors are not mounted", () => {
+    store.set(notebookAtom, {
+      ...initialNotebookState(),
+      cellIds: MultiColumn.from([[Cells.cell1, Cells.cell2]]),
+      cellData: {
+        [Cells.cell1]: createCell({ id: Cells.cell1, code: "import math" }),
+        [Cells.cell2]: createCell({
+          id: Cells.cell2,
+          code: "print(math.sqrt(4))",
+        }),
+      },
+    });
+
+    const { cellIds, codes } = store.get(topologicalCodesAtom);
+    const lens = createNotebookLens(cellIds, codes);
+
+    expect(lens.mergedText).toContain("import math");
+    expect(lens.mergedText).toContain("print(math.sqrt(4))");
+  });
+});
+
 describe("createNotebookLens", () => {
   it("should produce correct lens for same inputs", () => {
     // Use unique content for this test