@marimo-team/islands 0.19.7-dev27 → 0.19.7-dev30

package/dist/main.js

@@ -32291,8 +32291,13 @@ ${c.sqlString}
       transformRange: (e2, r2) => shiftRange(e2, f(r2)),
       reverseRange: (e2, r2) => shiftRange(e2, -f(r2)),
       isInRange: (e2, d2) => {
-        let f2 = r[d2].split("\n").length, _2 = c.get(d2) || 0, v = e2.start.line - _2, y = e2.end.line - _2;
-        return v >= 0 && y < f2;
+        let f2 = r[d2].split("\n").length, _2 = c.get(d2) || 0, v = e2.start.line - _2, y = e2.end.line - _2, S = v >= 0 && v < f2, w = y >= 0 && y < f2;
+        return S && !w && Logger.warn("[lsp] Cross-cell diagnostic detected: range starts in cell but ends outside", {
+          cellId: d2,
+          range: e2,
+          cellLines: f2,
+          offset: _2
+        }), S && w;
       },
       getEditsForNewText: (e2) => {
         let d2 = e2.split("\n"), f2 = _.split("\n");
@@ -32393,6 +32398,12 @@ ${c.sqlString}
         });
       }), this.snapshotter = new Snapshotter(this.getNotebookCode.bind(this));
     }
+    static pruneSeenCellUris(r) {
+      for (let c of _a.SEEN_CELL_DOCUMENT_URIS) {
+        let d = CellDocumentUri.parse(c);
+        r.has(d) || _a.SEEN_CELL_DOCUMENT_URIS.delete(c);
+      }
+    }
     onNotification(e) {
       return this.client.onNotification(e);
     }
@@ -32466,18 +32477,7 @@ ${c.sqlString}
       } : r;
     }
     async sync() {
-      let { lens: e, version: r, didChange: c } = this.snapshotter.snapshot();
-      return c ? this.client.textDocumentDidChange({
-        textDocument: {
-          uri: this.documentUri,
-          version: r
-        },
-        contentChanges: [
-          {
-            text: e.mergedText
-          }
-        ]
-      }) : {
+      let { lens: e, version: r, didChange: c } = this.snapshotter.snapshot(), d = {
         textDocument: {
           uri: this.documentUri,
           version: r
@@ -32488,15 +32488,22 @@ ${c.sqlString}
           }
         ]
       };
+      return c && await this.client.textDocumentDidChange(d), {
+        params: d,
+        lens: e
+      };
     }
     async textDocumentDidChange(e) {
-      return Logger.debug("[lsp] textDocumentDidChange", e), e.contentChanges.length === 1 ? this.sync() : (Logger.warn("[lsp] Unhandled textDocumentDidChange", e), this.client.textDocumentDidChange(e));
+      if (Logger.debug("[lsp] textDocumentDidChange", e), e.contentChanges.length === 1) {
+        let { params: e2 } = await this.sync();
+        return e2;
+      }
+      return Logger.warn("[lsp] Unhandled textDocumentDidChange", e), this.client.textDocumentDidChange(e);
     }
     async textDocumentDefinition(e) {
       let r = e.textDocument.uri;
       if (!CellDocumentUri.is(r)) return Logger.warn("Invalid cell document URI", r), null;
-      await this.sync();
-      let { lens: c } = this.snapshotter.getLatestSnapshot(), d = CellDocumentUri.parse(r);
+      let { lens: c } = await this.sync(), d = CellDocumentUri.parse(r);
       return this.client.textDocumentDefinition({
         ...e,
         textDocument: {
@@ -32508,8 +32515,7 @@ ${c.sqlString}
     async textDocumentSignatureHelp(e) {
       let r = e.textDocument.uri;
       if (!CellDocumentUri.is(r)) return Logger.warn("Invalid cell document URI", r), null;
-      await this.sync();
-      let { lens: c } = this.snapshotter.getLatestSnapshot(), d = CellDocumentUri.parse(r);
+      let { lens: c } = await this.sync(), d = CellDocumentUri.parse(r);
       return await this.client.textDocumentSignatureHelp({
         ...e,
         textDocument: {
@@ -32525,9 +32531,7 @@ ${c.sqlString}
       var _a2;
       let r = e.textDocument.uri;
       if (!CellDocumentUri.is(r)) return Logger.warn("Invalid cell document URI", r), null;
-      let c = CellDocumentUri.parse(r);
-      await this.sync();
-      let { lens: d } = this.snapshotter.getLatestSnapshot(), f = d.transformPosition(e.position, c), _ = await this.client.textDocumentRename({
+      let c = CellDocumentUri.parse(r), { lens: d } = await this.sync(), f = d.transformPosition(e.position, c), _ = await this.client.textDocumentRename({
         ...e,
         textDocument: {
           uri: this.documentUri
@@ -32547,20 +32551,21 @@ ${c.sqlString}
         let r2 = O.state.doc.line(e.position.line + 1).from + e.position.character, { startToken: c2, endToken: d2 } = getPositionAtWordBounds(O.state.doc, r2), f2 = O.state.doc.sliceString(c2, d2);
         M = isPrivateVariable(f2), M && Logger.debug("[lsp] Private variable rename detected, limiting to current cell", f2);
       }
+      let I = [], z = 0;
       for (let [e2, r2] of Objects.entries(E)) {
         if (M && e2 !== c) continue;
         let d2 = w.get(e2);
         if (d2 == null) {
-          Logger.warn("No new code for cell", e2);
+          Logger.warn("[lsp] No new code for cell during rename", e2), I.push(e2);
           continue;
         }
         if (!r2) {
-          Logger.warn("No view for plugin", e2);
+          Logger.warn("[lsp] No editor view for cell during rename", e2), I.push(e2);
           continue;
         }
-        getEditorCodeAsPython(r2) !== d2 && updateEditorCodeFromPython(r2, d2);
+        getEditorCodeAsPython(r2) !== d2 && (updateEditorCodeFromPython(r2, d2), z++);
       }
-      return {
+      return I.length > 0 && Logger.error(`[lsp] Rename partially failed: could not update ${I.length} cell(s)`, I), Logger.debug(`[lsp] Rename completed: updated ${z} cell(s)`), {
         ..._,
         documentChanges: [
           {
@@ -32582,8 +32587,7 @@ ${c.sqlString}
     async textDocumentPrepareRename(e) {
       let r = e.textDocument.uri;
       if (!CellDocumentUri.is(r)) return Logger.warn("Invalid cell document URI", r), null;
-      await this.sync();
-      let { lens: c } = this.snapshotter.getLatestSnapshot(), d = CellDocumentUri.parse(r);
+      let { lens: c } = await this.sync(), d = CellDocumentUri.parse(r);
       return this.client.textDocumentPrepareRename({
         ...e,
         textDocument: {
@@ -32593,8 +32597,8 @@ ${c.sqlString}
       });
     }
     async textDocumentHover(e) {
-      Logger.debug("[lsp] textDocumentHover", e), await this.sync();
-      let { lens: r } = this.snapshotter.getLatestSnapshot(), c = CellDocumentUri.parse(e.textDocument.uri), d = await this.client.textDocumentHover({
+      Logger.debug("[lsp] textDocumentHover", e);
+      let { lens: r } = await this.sync(), c = CellDocumentUri.parse(e.textDocument.uri), d = await this.client.textDocumentHover({
         ...e,
         textDocument: {
           uri: this.documentUri
@@ -32631,11 +32635,13 @@ ${c.sqlString}
             y.get(r2).push(c3);
             break;
           }
-          let S = new Set(_a.SEEN_CELL_DOCUMENT_URIS);
+          let S = new Set(_.cellIds);
+          _a.pruneSeenCellUris(S);
+          let w = new Set(_a.SEEN_CELL_DOCUMENT_URIS);
           for (let [e, d2] of y.entries()) {
             Logger.debug("[lsp] diagnostics for cell", e, d2);
             let f2 = CellDocumentUri.of(e);
-            S.delete(f2), r({
+            w.delete(f2), r({
               ...c2,
               params: {
                 ...c2.params,
@@ -32645,9 +32651,9 @@ ${c.sqlString}
               }
             });
           }
-          if (S.size > 0) {
-            Logger.debug("[lsp] clearing diagnostics", S);
-            for (let e of S) r({
+          if (w.size > 0) {
+            Logger.debug("[lsp] clearing diagnostics", w);
+            for (let e of w) r({
               method: "textDocument/publishDiagnostics",
               params: {
                 uri: e,
@@ -61387,7 +61393,7 @@ ${r}
   const defaultHtmlToImageOptions = {
     filter: (e) => {
       try {
-        return "classList" in e ? !e.classList.contains("mpl-toolbar") : true;
+        return !("classList" in e && (e.classList.contains("mpl-toolbar") || e.classList.contains("no-print")));
       } catch (e2) {
         return Logger.error("Error filtering node:", e2), true;
       }
@@ -61396,7 +61402,7 @@ ${r}
       Logger.error("Error loading image:", e);
     }
   };
-  function toPng(e, r) {
+  async function toPng(e, r) {
     return toPng$1(e, {
       ...defaultHtmlToImageOptions,
       ...r
@@ -62019,7 +62025,7 @@ Defaulting to \`null\`.`;
     e === "too_many" ? "Unknown" : prettifyRowCount(e, c),
     `${prettyNumber(r, c)} ${new PluralWord("column").pluralize(r)}`
   ].join(", "), TableActions = ({ enableSearch: e, onSearchQueryChange: r, isSearchEnabled: c, setIsSearchEnabled: d, pagination: f, totalColumns: _, selection: v, onRowSelectionChange: y, table: S, downloadAs: E, getRowIds: O, toggleDisplayHeader: M, showChartBuilder: I, showColumnExplorer: z, showRowExplorer: G, showPageSizeSelector: q, togglePanel: IY, isPanelOpen: LY, tableLoading: RY }) => (0, import_jsx_runtime.jsxs)("div", {
-    className: "flex items-center shrink-0 pt-1",
+    className: "flex items-center shrink-0 pt-1 no-print",
     children: [
       r && e && (0, import_jsx_runtime.jsx)(Tooltip, {
         content: "Search",
@@ -73076,7 +73082,7 @@ Image URL: ${r.imageUrl}`)), contextToXml({
       return Logger.warn("Failed to get version from mount config"), null;
     }
   }
-  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.19.7-dev27"), showCodeInRunModeAtom = atom(true);
+  const marimoVersionAtom = atom(getVersionFromMountConfig() || "0.19.7-dev30"), showCodeInRunModeAtom = atom(true);
   atom(null);
   var import_compiler_runtime$88 = require_compiler_runtime();
   function useKeydownOnElement(e, r) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.19.7-dev27",
+  "version": "0.19.7-dev30",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/components/data-table/TableActions.tsx

@@ -107,7 +107,7 @@ export const TableActions = <TData,>({
   };
 
   return (
-    <div className="flex items-center shrink-0 pt-1">
+    <div className="flex items-center shrink-0 pt-1 no-print">
       {onSearchQueryChange && enableSearch && (
         <Tooltip content="Search">
           <Button

package/src/components/editor/header/__tests__/filename-form.test.tsx

@@ -0,0 +1,124 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { fireEvent, render, waitFor } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { FilenameForm } from "../filename-form";
+
+// Mock the hooks
+vi.mock("@/core/saving/filename", () => ({
+  useUpdateFilename: vi.fn(),
+}));
+
+vi.mock("@/core/saving/save-component", () => ({
+  useSaveNotebook: vi.fn(),
+}));
+
+// Mock FilenameInput to make testing easier
+vi.mock("@/components/editor/header/filename-input", () => ({
+  FilenameInput: ({
+    onNameChange,
+    initialValue,
+  }: {
+    onNameChange: (value: string) => void;
+    initialValue?: string | null;
+  }) => (
+    <input
+      data-testid="filename-input"
+      data-initial-value={initialValue || ""}
+      onChange={(e) => onNameChange(e.target.value)}
+    />
+  ),
+}));
+
+const mockUseUpdateFilename = vi.mocked(
+  await import("@/core/saving/filename"),
+).useUpdateFilename;
+
+const mockUseSaveNotebook = vi.mocked(
+  await import("@/core/saving/save-component"),
+).useSaveNotebook;
+
+describe("FilenameForm", () => {
+  const mockUpdateFilename = vi.fn();
+  const mockSaveNotebook = vi.fn();
+  const mockSaveOrNameNotebook = vi.fn();
+  const mockSaveIfNotebookIsPersistent = vi.fn();
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockUseUpdateFilename.mockReturnValue(mockUpdateFilename);
+    mockUseSaveNotebook.mockReturnValue({
+      saveNotebook: mockSaveNotebook,
+      saveOrNameNotebook: mockSaveOrNameNotebook,
+      saveIfNotebookIsPersistent: mockSaveIfNotebookIsPersistent,
+    });
+  });
+
+  it("should call saveNotebook when creating a new file from an unnamed notebook", async () => {
+    // Setup: updateFilename resolves with the new filename
+    mockUpdateFilename.mockResolvedValue("/path/to/new-file.py");
+
+    const { getByTestId } = render(<FilenameForm filename={null} />);
+
+    const input = getByTestId("filename-input");
+
+    // Simulate name change (user entering a filename)
+    fireEvent.change(input, { target: { value: "new-file.py" } });
+
+    // Wait for the promise to resolve
+    await waitFor(() => {
+      expect(mockUpdateFilename).toHaveBeenCalledWith("new-file.py");
+    });
+
+    await waitFor(() => {
+      expect(mockSaveNotebook).toHaveBeenCalledWith(
+        "/path/to/new-file.py",
+        true,
+      );
+    });
+  });
+
+  it("should not call saveNotebook when renaming an existing file", async () => {
+    // Setup: updateFilename resolves with the new filename
+    mockUpdateFilename.mockResolvedValue("/path/to/renamed-file.py");
+
+    const { getByTestId } = render(
+      <FilenameForm filename="/path/to/existing-file.py" />,
+    );
+
+    const input = getByTestId("filename-input");
+
+    // Simulate name change
+    fireEvent.change(input, { target: { value: "renamed-file.py" } });
+
+    // Wait for the promise to resolve
+    await waitFor(() => {
+      expect(mockUpdateFilename).toHaveBeenCalledWith("renamed-file.py");
+    });
+
+    // saveNotebook should NOT be called because the file was already named
+    expect(mockSaveNotebook).not.toHaveBeenCalled();
+  });
+
+  it("should not call saveNotebook when updateFilename returns null", async () => {
+    // Setup: updateFilename resolves with null (e.g., user cancelled or error)
+    mockUpdateFilename.mockResolvedValue(null);
+
+    const { getByTestId } = render(<FilenameForm filename={null} />);
+
+    const input = getByTestId("filename-input");
+
+    // Simulate name change
+    fireEvent.change(input, { target: { value: "new-file.py" } });
+
+    // Wait for the promise to resolve
+    await waitFor(() => {
+      expect(mockUpdateFilename).toHaveBeenCalled();
+    });
+
+    // saveNotebook should NOT be called because updateFilename returned null
+    expect(mockSaveNotebook).not.toHaveBeenCalled();
+  });
+});

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

@@ -25,6 +25,30 @@ const Cells = {
 };
 
 describe("createNotebookLens", () => {
+  it("should produce correct lens for same inputs", () => {
+    // Use unique content for this test
+    const cellIds: CellId[] = [Cells.cell1, Cells.cell2];
+    const codes: Record<CellId, string> = {
+      [Cells.cell1]: "unique_memo_test_line1",
+      [Cells.cell2]: "unique_memo_test_line2",
+    };
+
+    const lens1 = createNotebookLens(cellIds, codes);
+    const lens2 = createNotebookLens(cellIds, codes);
+
+    // Same inputs should produce equivalent lens (same merged text)
+    expect(lens1.mergedText).toBe(lens2.mergedText);
+    expect(lens1.cellIds).toEqual(lens2.cellIds);
+
+    // Different content should produce a different lens
+    const differentCodes: Record<CellId, string> = {
+      [Cells.cell1]: "unique_memo_different",
+      [Cells.cell2]: "unique_memo_content",
+    };
+    const lens3 = createNotebookLens(cellIds, differentCodes);
+    expect(lens3.mergedText).not.toBe(lens1.mergedText);
+  });
+
   it("should calculate correct line offsets", () => {
     const cellIds: CellId[] = [Cells.cell1, Cells.cell2, Cells.cell3];
     const codes: Record<CellId, string> = {
@@ -111,6 +135,7 @@ describe("createNotebookLens", () => {
       ),
     ).toBe(true);
 
+    // This triggers the cross-cell diagnostic warning (start in range, end outside)
     expect(
       lens.isInRange(
         {
@@ -122,6 +147,24 @@ describe("createNotebookLens", () => {
     ).toBe(false);
   });
 
+  it("should detect cross-cell diagnostics where start is in range but end is outside", () => {
+    const cellIds: CellId[] = [Cells.cell1, Cells.cell2];
+    const codes: Record<CellId, string> = {
+      [Cells.cell1]: "line1\nline2",
+      [Cells.cell2]: "line3",
+    };
+    const lens = createNotebookLens(cellIds, codes);
+
+    // Range that starts in cell1 (line 1) but ends in cell2 (line 2 = first line of cell2)
+    // This should return false and log a warning
+    const crossCellRange = {
+      start: { line: 1, character: 0 }, // line 1 is in cell1
+      end: { line: 2, character: 0 }, // line 2 is in cell2
+    };
+
+    expect(lens.isInRange(crossCellRange, Cells.cell1)).toBe(false);
+  });
+
   it("should join all code into merged text", () => {
     const cellIds: CellId[] = [Cells.cell1, Cells.cell2, Cells.cell3];
     const codes: Record<CellId, string> = {
@@ -1120,6 +1163,83 @@ describe("NotebookLanguageServerClient", () => {
     });
   });
 
+  describe("sync returns lens to prevent race conditions", () => {
+    it("should return the lens used for synchronization", async () => {
+      mockClient.textDocumentDidChange = vi
+        .fn()
+        .mockImplementation((params) => params);
+
+      const result = await notebookClient.sync();
+
+      // sync() should return both params and the lens
+      expect(result).toHaveProperty("params");
+      expect(result).toHaveProperty("lens");
+      expect(result.lens.cellIds).toEqual([
+        Cells.cell1,
+        Cells.cell2,
+        Cells.cell3,
+      ]);
+      expect(result.lens.mergedText).toContain("# this is a comment");
+    });
+
+    it("should return consistent lens even when called multiple times rapidly", async () => {
+      mockClient.textDocumentDidChange = vi
+        .fn()
+        .mockImplementation((params) => params);
+
+      // Call sync multiple times rapidly
+      const [result1, result2, result3] = await Promise.all([
+        notebookClient.sync(),
+        notebookClient.sync(),
+        notebookClient.sync(),
+      ]);
+
+      // All should return the same lens content (memoized)
+      expect(result1.lens.mergedText).toBe(result2.lens.mergedText);
+      expect(result2.lens.mergedText).toBe(result3.lens.mergedText);
+    });
+  });
+
+  describe("SEEN_CELL_DOCUMENT_URIS memory management", () => {
+    it("should track opened cells in SEEN_CELL_DOCUMENT_URIS", async () => {
+      // Clear any existing state
+      const seenUris = (NotebookLanguageServerClient as any)
+        .SEEN_CELL_DOCUMENT_URIS;
+      seenUris.clear();
+
+      // Open some cells to add them to SEEN_CELL_DOCUMENT_URIS
+      await notebookClient.textDocumentDidOpen({
+        textDocument: {
+          uri: CellDocumentUri.of(Cells.cell1),
+          languageId: "python",
+          version: 1,
+          text: "code1",
+        },
+      });
+
+      await notebookClient.textDocumentDidOpen({
+        textDocument: {
+          uri: CellDocumentUri.of(Cells.cell2),
+          languageId: "python",
+          version: 1,
+          text: "code2",
+        },
+      });
+
+      // Verify cells were added
+      expect(seenUris.has(CellDocumentUri.of(Cells.cell1))).toBe(true);
+      expect(seenUris.has(CellDocumentUri.of(Cells.cell2))).toBe(true);
+      expect(seenUris.size).toBe(2);
+    });
+
+    it("should have pruneSeenCellUris static method", () => {
+      // Verify the method exists and is callable
+      expect(
+        typeof (NotebookLanguageServerClient as any).pruneSeenCellUris,
+      ).toBe("function");
+    });
+  });
+
   describe("initialization and configuration", () => {
     it("should send configuration after initialization", async () => {
       const configNotifications: any[] = [];

package/src/core/codemirror/lsp/lens.ts

@@ -68,12 +68,28 @@ export function createNotebookLens(
     reverseRange: (range: LSP.Range, cellId: CellId) =>
       shiftRange(range, -getCurrentLineOffset(cellId)),
 
+    /**
+     * Check if a range falls entirely within the given cell.
+     * Returns false for ranges that span multiple cells (cross-cell diagnostics).
+     */
     isInRange: (range: LSP.Range, cellId: CellId) => {
       const cellLines = codes[cellId].split("\n").length;
       const offset = cellLineOffsets.get(cellId) || 0;
       const startLine = range.start.line - offset;
       const endLine = range.end.line - offset;
-      return startLine >= 0 && endLine < cellLines;
+
+      const startInRange = startLine >= 0 && startLine < cellLines;
+      const endInRange = endLine >= 0 && endLine < cellLines;
+
+      // Warn about cross-cell diagnostics that start in this cell but end outside
+      if (startInRange && !endInRange) {
+        Logger.warn(
+          "[lsp] Cross-cell diagnostic detected: range starts in cell but ends outside",
+          { cellId, range, cellLines, offset },
+        );
+      }
+
+      return startInRange && endInRange;
     },
 
     getEditsForNewText: (newText: string) => {

package/src/core/codemirror/lsp/notebook-lsp.ts

@@ -113,8 +113,28 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     EditorView | null | undefined
   >;
   private readonly initialSettings: Record<string, unknown>;
+  /**
+   * Tracks which cell document URIs have been opened with the LSP server.
+   * Used to clear diagnostics for cells that no longer have any.
+   *
+   * This set is pruned when diagnostics are processed to only include
+   * cells that exist in the current notebook snapshot.
+   */
   private static readonly SEEN_CELL_DOCUMENT_URIS = new Set<CellDocumentUri>();
 
+  /**
+   * Remove cell URIs that are no longer in the notebook.
+   * Called during diagnostic processing to prevent memory leaks.
+   */
+  private static pruneSeenCellUris(currentCellIds: Set<CellId>): void {
+    for (const uri of NotebookLanguageServerClient.SEEN_CELL_DOCUMENT_URIS) {
+      const cellId = CellDocumentUri.parse(uri);
+      if (!currentCellIds.has(cellId)) {
+        NotebookLanguageServerClient.SEEN_CELL_DOCUMENT_URIS.delete(uri);
+      }
+    }
+  }
+
   /**
    * Cache of completion items to avoid jitter while typing in the same completion item
    */
@@ -274,26 +294,30 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     };
   }
 
-  async sync(): Promise<LSP.DidChangeTextDocumentParams> {
+  /**
+   * Synchronize the document with the LSP server and return the lens used.
+   * This ensures the caller uses the same lens that was sent to the server,
+   * avoiding race conditions if cells change between sync and subsequent operations.
+   */
+  async sync(): Promise<{
+    params: LSP.DidChangeTextDocumentParams;
+    lens: NotebookLens;
+  }> {
     const { lens, version, didChange } = this.snapshotter.snapshot();
-    if (!didChange) {
-      return {
-        textDocument: {
-          uri: this.documentUri,
-          version: version,
-        },
-        contentChanges: [{ text: lens.mergedText }],
-      };
-    }
-
-    // Update changes for merged doc, etc.
-    return this.client.textDocumentDidChange({
+    const params: LSP.DidChangeTextDocumentParams = {
       textDocument: {
         uri: this.documentUri,
         version: version,
       },
       contentChanges: [{ text: lens.mergedText }],
-    });
+    };
+
+    if (didChange) {
+      // Update changes for merged doc
+      await this.client.textDocumentDidChange(params);
+    }
+
+    return { params, lens };
   }
 
   public async textDocumentDidChange(params: LSP.DidChangeTextDocumentParams) {
@@ -301,7 +325,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     // We know how to only handle single content changes
     // But that is all we expect to receive
     if (params.contentChanges.length === 1) {
-      return this.sync();
+      const { params: syncParams } = await this.sync();
+      return syncParams;
     }
 
     Logger.warn("[lsp] Unhandled textDocumentDidChange", params);
@@ -319,9 +344,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       return null;
     }
 
-    // This LSP method has no version, so lets sync and then get the latest snapshot
-    await this.sync();
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    // Sync and use the same lens that was sent to the server
+    const { lens } = await this.sync();
 
     // Find the lens for this cell
     const cellId = CellDocumentUri.parse(cellDocumentUri);
@@ -345,9 +369,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       return null;
     }
 
-    // This LSP method has no version, so lets sync and then get the latest snapshot
-    await this.sync();
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    // Sync and use the same lens that was sent to the server
+    const { lens } = await this.sync();
 
     const cellId = CellDocumentUri.parse(cellDocumentUri);
 
@@ -367,6 +390,23 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     return response;
   }
 
+  /**
+   * Code actions are currently disabled because mapping code action edits
+   * back to individual cells is complex and error-prone.
+   *
+   * The LSP server returns edits in merged document coordinates, but we need
+   * to apply them to individual cell editors. Unlike simple position transforms
+   * (hover, completion), code actions can include arbitrary text edits that
+   * may span multiple cells or change line counts, making the reverse mapping
+   * unreliable.
+   *
+   * To enable this, we would need to:
+   * 1. Transform edit ranges from merged doc coordinates to cell coordinates
+   * 2. Handle edits that span cell boundaries (split or reject them)
+   * 3. Apply edits atomically across multiple cell editors
+   *
+   * See textDocumentRename for a similar workaround that manually applies edits.
+   */
   textDocumentCodeAction(
     params: LSP.CodeActionParams,
   ): Promise<(LSP.Command | LSP.CodeAction)[] | null> {
@@ -378,12 +418,23 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   }
 
   /**
-   * HACK
-   * This whole function is a hack to work around the fact that we don't have
-   * a great way to map the edits back to the appropriate LSP view plugin.
+   * Rename implementation with manual edit application.
+   *
+   * This is a workaround because the standard LSP rename flow doesn't work well
+   * with our notebook architecture. The LSP server returns a WorkspaceEdit with
+   * edits in merged document coordinates, but CodeMirror's LSP plugin expects
+   * edits for individual cell documents.
+   *
+   * Instead of trying to transform the WorkspaceEdit (which is complex because
+   * edits can span cells, change line counts, etc.), we:
+   * 1. Request the rename from the LSP server
+   * 2. Extract the new merged text from the response
+   * 3. Split it back into per-cell text using the lens
+   * 4. Manually update each cell's editor
    *
-   * Instead, we parse out the new text from the response and then update the
-   * code in the plugins manually, instead of using the LSP.
+   * This approach is simpler and more reliable, though it bypasses the normal
+   * LSP edit application flow. The trade-off is that we lose undo grouping
+   * across cells.
    */
   async textDocumentRename(
     params: LSP.RenameParams,
@@ -397,9 +448,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
 
     const cellId = CellDocumentUri.parse(cellDocumentUri);
 
-    // This LSP method has no version, so lets sync and then get the latest snapshot
-    await this.sync();
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    // Sync and use the same lens that was sent to the server
+    const { lens } = await this.sync();
 
     const transformedPosition = lens.transformPosition(params.position, cellId);
 
@@ -468,6 +518,9 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       }
     }
 
+    const failedCells: CellId[] = [];
+    let updatedCount = 0;
+
     for (const [currentCellId, ev] of Objects.entries(editors)) {
       // For private variable renames, only update the originating cell
       if (isPrivateRename && currentCellId !== cellId) {
@@ -476,12 +529,17 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
 
       const newCode = editsToNewCode.get(currentCellId);
       if (newCode == null) {
-        Logger.warn("No new code for cell", currentCellId);
+        Logger.warn("[lsp] No new code for cell during rename", currentCellId);
+        failedCells.push(currentCellId);
         continue;
       }
 
       if (!ev) {
-        Logger.warn("No view for plugin", currentCellId);
+        Logger.warn(
+          "[lsp] No editor view for cell during rename",
+          currentCellId,
+        );
+        failedCells.push(currentCellId);
         continue;
       }
 
@@ -489,9 +547,19 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       const code = getEditorCodeAsPython(ev);
       if (code !== newCode) {
         updateEditorCodeFromPython(ev, newCode);
+        updatedCount++;
       }
     }
 
+    if (failedCells.length > 0) {
+      Logger.error(
+        `[lsp] Rename partially failed: could not update ${failedCells.length} cell(s)`,
+        failedCells,
+      );
+    }
+
+    Logger.debug(`[lsp] Rename completed: updated ${updatedCount} cell(s)`);
+
     return {
       ...response,
       documentChanges: [
@@ -531,9 +599,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
       return null;
     }
 
-    // This LSP method has no version, so lets sync and then get the latest snapshot
-    await this.sync();
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    // Sync and use the same lens that was sent to the server
+    const { lens } = await this.sync();
 
     const cellId = CellDocumentUri.parse(cellDocumentUri);
 
@@ -549,9 +616,8 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   public async textDocumentHover(params: LSP.HoverParams) {
     Logger.debug("[lsp] textDocumentHover", params);
 
-    // This LSP method has no version, so lets sync and then get the latest snapshot
-    await this.sync();
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    // Sync and use the same lens that was sent to the server
+    const { lens } = await this.sync();
 
     const cellId = CellDocumentUri.parse(params.textDocument.uri);
 
@@ -655,6 +721,10 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
           }
         }
 
+        // Prune any cell URIs for cells that no longer exist
+        const currentCellIds = new Set(lens.cellIds);
+        NotebookLanguageServerClient.pruneSeenCellUris(currentCellIds);
+
         const cellsToClear = new Set(
           NotebookLanguageServerClient.SEEN_CELL_DOCUMENT_URIS,
         );

package/src/utils/download.ts

@@ -94,6 +94,8 @@ function prepareCellElementForScreenshot(
   };
 }
 
+const THRESHOLD_TIME_MS = 500;
+
 /**
  * Capture a cell output as a PNG data URL.
  *
@@ -119,7 +121,17 @@ export async function getImageDataUrlForCell(
   const cleanup = prepareCellElementForScreenshot(element, enablePrintMode);
 
   try {
-    return await toPng(element);
+    const startTime = Date.now();
+    const dataUrl = await toPng(element);
+    const timeTaken = Date.now() - startTime;
+    if (timeTaken > THRESHOLD_TIME_MS) {
+      Logger.debug(
+        "toPng operation for element",
+        element,
+        `took ${timeTaken} ms (exceeds threshold)`,
+      );
+    }
+    return dataUrl;
   } finally {
     cleanup();
   }

package/src/utils/html-to-image.ts

@@ -13,7 +13,13 @@ export const defaultHtmlToImageOptions: HtmlToImageOptions = {
     try {
       if ("classList" in node) {
         // Filter out matplotlib toolbars
-        return !node.classList.contains("mpl-toolbar");
+        if (node.classList.contains("mpl-toolbar")) {
+          return false;
+        }
+
+        if (node.classList.contains("no-print")) {
+          return false;
+        }
       }
       return true;
     } catch (error) {
@@ -30,7 +36,7 @@ export const defaultHtmlToImageOptions: HtmlToImageOptions = {
  * Convert an HTML element to a PNG data URL.
  * This is a wrapper around html-to-image's toPng with default options applied.
  */
-export function toPng(
+export async function toPng(
   element: HTMLElement,
   options?: HtmlToImageOptions,
 ): Promise<string> {