@marimo-team/islands 0.25.2-dev14 → 0.25.2-dev140

package/src/core/codemirror/language/embedded/__tests__/embedded-python.test.ts

@@ -2,13 +2,17 @@
 
 import type { CompletionContext } from "@codemirror/autocomplete";
 import { python } from "@codemirror/lang-python";
+import { StreamLanguage } from "@codemirror/language";
+import { stexMath } from "@codemirror/legacy-modes/mode/stex";
 import { EditorState } from "@codemirror/state";
-import type { InlineContext } from "@lezer/markdown";
-import { beforeEach, describe, expect, it } from "vitest";
+import { classHighlighter, highlightTree } from "@lezer/highlight";
+import { parser as markdownParser } from "@lezer/markdown";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { cellId, variableName } from "@/__tests__/branded";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import { parsePython, variableCompletionSource } from "../embedded-python";
+import { parseLatex } from "../latex";
 
 const IS_ACTIVE = () => true;
 
@@ -23,25 +27,133 @@ describe("parsePython", () => {
     ]);
   });
 
-  it("should not parse double curly braces", () => {
+  const cases = [
+    "{math.exp(3)}",
+    '${math.exp(3) + len("$}")} $',
+    "${math.exp(3)} + {math.sin(3)}$",
+    "${math.exp(3)}$",
+    "$$ {math.exp(3)} $$",
+    String.raw`\({math.exp(3)}\)`,
+    String.raw`\[{math.exp(3)}\]`,
+    String.raw`\\({math.exp(3)}\\)`,
+    String.raw`\\[{math.exp(3)}\\]`,
+    String.raw`$\frac{{{math.exp(3)}}}{{2}}$`,
+    String.raw`$\{math.exp(3)}$`,
+    '${ {"}": math.exp(3)} }$',
+    "${math.exp(3):{width}.{precision}f}$",
+  ];
+
+  function highlightedProperties(doc: string, prefix: string) {
+    const parser = markdownParser.configure([
+      parseLatex(StreamLanguage.define(stexMath).parser),
+      parsePython(python().language.parser, () => prefix.includes("f")),
+    ]);
+    const properties: string[] = [];
+    highlightTree(parser.parse(doc), classHighlighter, (from, to, classes) => {
+      if (classes.split(" ").includes("tok-propertyName")) {
+        properties.push(doc.slice(from, to));
+      }
+    });
+    return properties;
+  }
+
+  it.each(["f", "rf", "fr"])(
+    "highlights interpolation with %s strings",
+    (prefix) => {
+      for (const doc of cases) {
+        expect(highlightedProperties(doc, prefix)).toEqual(
+          doc.includes("math.sin") ? ["exp", "sin"] : ["exp"],
+        );
+      }
+    },
+  );
+
+  it.each(["", "r"])("does not interpolate %s strings", (prefix) => {
+    for (const doc of cases) {
+      expect(highlightedProperties(doc, prefix)).toEqual([]);
+    }
+  });
+
+  it("parses replacement boundaries once for a paragraph with many interpolations", () => {
     const pythonParser = python().language.parser;
-    const config = parsePython(pythonParser, IS_ACTIVE);
+    const parse = vi.spyOn(pythonParser, "parse");
+    try {
+      const doc = "{math.exp(3)} ".repeat(1000);
+      const parser = markdownParser.configure([
+        parseLatex(StreamLanguage.define(stexMath).parser),
+        parsePython(pythonParser, IS_ACTIVE),
+      ]);
+      const properties: string[] = [];
+      highlightTree(
+        parser.parse(doc),
+        classHighlighter,
+        (from, to, classes) => {
+          if (classes.split(" ").includes("tok-propertyName")) {
+            properties.push(doc.slice(from, to));
+          }
+        },
+      );
+      expect(properties).toEqual(Array.from({ length: 1000 }, () => "exp"));
+      expect(parse).toHaveBeenCalledTimes(1);
+    } finally {
+      parse.mockRestore();
+    }
+  });
 
-    // Create a mock context
-    const mockContext = {
-      slice: (from: number, to: number) => {
-        if (from === 5 && to === 6) {
-          return "{";
-        }
-        return "";
-      },
-      addDelimiter: () => -1,
-    } as unknown as InlineContext;
+  it("preserves raw backslashes between replacements", () => {
+    expect(
+      highlightedProperties(String.raw`{math.exp(3)} \N{math.sin(3)}`, "rf"),
+    ).toEqual(["exp", "sin"]);
+  });
+
+  it("finds replacements after triple quotes in markdown text", () => {
+    expect(
+      highlightedProperties('{math.exp(3)} """ {math.sin(3)}', "f"),
+    ).toEqual(["exp", "sin"]);
+  });
 
-    // Test with double curly braces
-    const result = config.parseInline![0].parse(mockContext, 123, 6); // 123 is OPEN_BRACE
-    expect(result).toBe(-1);
+  it("keeps replacement positions local to each paragraph", () => {
+    expect(
+      highlightedProperties("${math.exp(3)}$\n\n${math.sin(3)}$", "rf"),
+    ).toEqual(["exp", "sin"]);
   });
+
+  it("preserves LaTeX highlighting around Python replacements", () => {
+    const doc = String.raw`$\frac{{{math.exp(3)}}}{{2}} + \alpha$`;
+    const parser = markdownParser.configure([
+      parseLatex(StreamLanguage.define(stexMath).parser),
+      parsePython(python().language.parser, IS_ACTIVE),
+    ]);
+    const highlighted: string[] = [];
+    highlightTree(parser.parse(doc), classHighlighter, (from, to) => {
+      highlighted.push(doc.slice(from, to));
+    });
+    expect(highlighted).toEqual(
+      expect.arrayContaining(["\\frac", "exp", "\\alpha"]),
+    );
+  });
+
+  it.each(["${math.exp(3)$", "${math.exp(3)", "${{math.exp(3)}$"])(
+    "does not treat incomplete or escaped braces as replacements: %s",
+    (doc) => {
+      expect(highlightedProperties(doc, "f")).toEqual([]);
+    },
+  );
+
+  it.each(["f", "rf", "fr"])(
+    "keeps doubled braces literal in %s strings",
+    (prefix) => {
+      for (const doc of [
+        "{{math.exp(3)}}",
+        "${{math.exp(3)}}$",
+        "$ {{{{math.exp(3)}}}} $",
+        String.raw`$\{{math.exp(3)}}$`,
+        String.raw`$\frac{{1}}{{2}}$`,
+      ]) {
+        expect(highlightedProperties(doc, prefix)).toEqual([]);
+      }
+    },
+  );
 });
 
 describe("variableCompletionSource", () => {
@@ -72,10 +184,10 @@ describe("variableCompletionSource", () => {
     const context: Partial<CompletionContext> = {
       pos: 8,
       explicit: false,
-      matchBefore: (expr: RegExp) => ({ from: 7, to: 8, text: "v" }),
+      matchBefore: (_expr: RegExp) => ({ from: 7, to: 8, text: "v" }),
       state,
       aborted: false,
-      tokenBefore: (types: readonly string[]) => null,
+      tokenBefore: (_types: readonly string[]) => null,
     };
 
     const result = variableCompletionSource(context as CompletionContext);
@@ -92,10 +204,10 @@ describe("variableCompletionSource", () => {
     const context: Partial<CompletionContext> = {
       pos: 7,
       explicit: false,
-      matchBefore: (expr: RegExp) => ({ from: 6, to: 7, text: "v" }),
+      matchBefore: (_expr: RegExp) => ({ from: 6, to: 7, text: "v" }),
       state,
       aborted: false,
-      tokenBefore: (types: readonly string[]) => null,
+      tokenBefore: (_types: readonly string[]) => null,
     };
 
     const result = variableCompletionSource(context as CompletionContext);
@@ -110,10 +222,10 @@ describe("variableCompletionSource", () => {
     const context: Partial<CompletionContext> = {
       pos: 13,
       explicit: false,
-      matchBefore: (expr: RegExp) => ({ from: 12, to: 13, text: "v" }),
+      matchBefore: (_expr: RegExp) => ({ from: 12, to: 13, text: "v" }),
       state,
       aborted: false,
-      tokenBefore: (types: readonly string[]) => null,
+      tokenBefore: (_types: readonly string[]) => null,
     };
 
     const result = variableCompletionSource(context as CompletionContext);
@@ -131,7 +243,7 @@ describe("variableCompletionSource", () => {
       matchBefore: () => null,
       state,
       aborted: false,
-      tokenBefore: (types: readonly string[]) => null,
+      tokenBefore: (_types: readonly string[]) => null,
     };
 
     const result = variableCompletionSource(context as CompletionContext);
@@ -146,10 +258,10 @@ describe("variableCompletionSource", () => {
     const context: Partial<CompletionContext> = {
       pos: 9,
       explicit: false,
-      matchBefore: (expr: RegExp) => ({ from: 8, to: 9, text: "v" }),
+      matchBefore: (_expr: RegExp) => ({ from: 8, to: 9, text: "v" }),
       state,
       aborted: false,
-      tokenBefore: (types: readonly string[]) => null,
+      tokenBefore: (_types: readonly string[]) => null,
     };
 
     const result = variableCompletionSource(context as CompletionContext);

package/src/core/codemirror/language/embedded/embedded-python.ts

@@ -16,11 +16,8 @@ import { getVariableCompletions } from "../../completion/variable-completions";
 const PYTHON = "Python";
 
 const OPEN_BRACE = 123; // {
-const CLOSE_BRACE = 125; // }
 const DELIMITER_LENGTH = 1;
 
-const MARK = { mark: `${PYTHON}Mark`, resolve: PYTHON };
-
 /**
  * Define a Python code block parser for Markdown.
  *
@@ -31,6 +28,11 @@ export function parsePython(
   pythonParser: Parser,
   isActivated: () => boolean,
 ): MarkdownConfig {
+  const replacements = new WeakMap<
+    InlineContext,
+    { ends: Map<number, number>; parsedTo: number }
+  >();
+
   const defineNodes = [
     { name: PYTHON, style: tags.emphasis },
     { name: `${PYTHON}Mark`, style: tags.processingInstruction },
@@ -41,26 +43,58 @@ export function parsePython(
     parseInline: [
       {
         name: PYTHON,
+        before: "Escape",
         parse(cx: InlineContext, next: number, pos: number): number {
           if (!isActivated()) {
             return -1;
           }
 
-          if (next !== OPEN_BRACE && next !== CLOSE_BRACE) {
+          // A backslash does not escape a brace in Python f-strings.
+          if (next === 92 && cx.char(pos + 1) === OPEN_BRACE) {
+            return pos + 1;
+          }
+          if (next !== OPEN_BRACE) {
             return -1;
           }
+          if (cx.char(pos + 1) === OPEN_BRACE) {
+            return pos + 2;
+          }
 
-          // Check for double curly braces
-          if (next === OPEN_BRACE && cx.slice(pos - 1, pos) === "{") {
+          let cached = replacements.get(cx);
+          if (!cached || pos >= cached.parsedTo) {
+            // Cache all replacements so each brace does not reparse the paragraph.
+            // Keep Markdown backslashes from consuming braces as Python escapes.
+            const prefix = 'rf"""';
+            const tree = pythonParser.parse(
+              `${prefix}${cx.slice(pos, cx.end)}\n"""`,
+            );
+            const formatString = tree.topNode
+              .getChild("ExpressionStatement")
+              ?.getChild("FormatString");
+            const offset = pos - prefix.length;
+            const ends = new Map<number, number>();
+            for (const replacement of formatString?.getChildren(
+              "FormatReplacement",
+            ) ?? []) {
+              if (replacement.lastChild?.name === "}") {
+                ends.set(offset + replacement.from, offset + replacement.to);
+              }
+            }
+            cached = {
+              ends,
+              parsedTo: formatString ? offset + formatString.to : cx.end,
+            };
+            replacements.set(cx, cached);
+          }
+          const end = cached.ends.get(pos);
+          if (end === undefined) {
             return -1;
           }
-
-          return cx.addDelimiter(
-            MARK,
-            pos,
-            pos + DELIMITER_LENGTH,
-            next === OPEN_BRACE,
-            next === CLOSE_BRACE,
+          return cx.addElement(
+            cx.elt(PYTHON, pos, end, [
+              cx.elt(`${PYTHON}Mark`, pos, pos + 1),
+              cx.elt(`${PYTHON}Mark`, end - 1, end),
+            ]),
           );
         },
       },

package/src/core/codemirror/language/embedded/latex.ts

@@ -29,7 +29,7 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-import { type Parser, parseMixed } from "@lezer/common";
+import { IterMode, type Parser, parseMixed } from "@lezer/common";
 import { tags } from "@lezer/highlight";
 import type {
   DelimiterType,
@@ -231,24 +231,31 @@ export function parseLatex(latexParser: Parser): MarkdownConfig {
       // Test if the node type is one of the math expression
       const delimiterLength = DELIMITER_LENGTH[node.type.name];
       if (delimiterLength) {
-        const from = node.from + delimiterLength;
-        const to = node.to - delimiterLength;
+        const from = node.node.firstChild?.to ?? node.from + delimiterLength;
+        const to = node.node.lastChild?.from ?? node.to - delimiterLength;
 
         // Ensure valid ranges
         if (from >= to || from < 0 || to < 0) {
           return null;
         }
 
-        return {
-          parser: latexParser,
-          // Remove delimiter from LaTeX parser otherwise it won't be highlighted
-          overlay: [
-            {
-              from,
-              to,
-            },
-          ],
-        };
+        // Leave Python replacements to their own parser and highlighter.
+        const overlay: { from: number; to: number }[] = [];
+        let start = from;
+        node.node.cursor(IterMode.IgnoreMounts).iterate((child) => {
+          if (child.name !== "Python") {
+            return;
+          }
+          if (start < child.from) {
+            overlay.push({ from: start, to: child.from });
+          }
+          start = child.to;
+          return false;
+        });
+        if (start < to) {
+          overlay.push({ from: start, to });
+        }
+        return overlay.length ? { parser: latexParser, overlay } : null;
       }
 
       return null;

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

@@ -359,11 +359,111 @@ describe("NotebookLanguageServerClient", () => {
       return undefined;
     });
 
-    notebookClient = new NotebookLanguageServerClient(mockClient, {}, () => ({
-      [Cells.cell1]: new EditorView({ doc: "# this is a comment" }),
-      [Cells.cell2]: new EditorView({ doc: "import math\nimport numpy" }),
-      [Cells.cell3]: new EditorView({ doc: "print(math.sqrt(4))" }),
-    }));
+    notebookClient = new NotebookLanguageServerClient(
+      mockClient,
+      {},
+      () => ({}),
+    );
+  });
+
+  describe("live editor snapshots", () => {
+    function createEditor(code: string) {
+      return new EditorView({
+        doc: code,
+        extensions: [
+          languageAdapterState.init(() => new PythonLanguageAdapter()),
+          languageMetadataField.init(() => ({})),
+        ],
+      });
+    }
+
+    it("sends edits and their removal without invalidating the notebook atom", async () => {
+      const cachedCodes = store.get(topologicalCodesAtom);
+      const editor = createEditor("print(math.sqrt(4))");
+      replaceNotebookEditors(notebookClient, () => ({ [Cells.cell3]: editor }));
+      const uri = CellDocumentUri.of(Cells.cell3);
+      try {
+        await notebookClient.textDocumentDidOpen({
+          textDocument: {
+            uri,
+            languageId: "python",
+            version: 1,
+            text: editor.state.doc.toString(),
+          },
+        });
+        for (const [index, code] of [
+          "math.missing()",
+          "print(math.sqrt(4))",
+        ].entries()) {
+          editor.dispatch({
+            changes: { from: 0, to: editor.state.doc.length, insert: code },
+          });
+          await notebookClient.textDocumentDidChange({
+            textDocument: { uri, version: index + 2 },
+            contentChanges: [{ text: code }],
+          });
+          expect(store.get(topologicalCodesAtom)).toEqual(cachedCodes);
+          expect(mockClient.textDocumentDidChange).toHaveBeenNthCalledWith(
+            index + 1,
+            {
+              textDocument: {
+                uri: notebookClient.documentUri,
+                version: index + 2,
+              },
+              contentChanges: [
+                {
+                  text: `# this is a comment\nimport math\nimport numpy\n${code}`,
+                },
+              ],
+            },
+          );
+        }
+      } finally {
+        editor.destroy();
+      }
+    });
+
+    it("syncs pending edits before completion and uses their updated line offsets", async () => {
+      const editor = createEditor("import math\nimport numpy");
+      replaceNotebookEditors(notebookClient, () => ({ [Cells.cell2]: editor }));
+      try {
+        await notebookClient.textDocumentDidOpen({
+          textDocument: {
+            uri: CellDocumentUri.of(Cells.cell2),
+            languageId: "python",
+            version: 1,
+            text: editor.state.doc.toString(),
+          },
+        });
+        editor.dispatch({
+          changes: { from: editor.state.doc.length, insert: "\nextra = 1" },
+        });
+        mockClient.textDocumentCompletion.mockImplementation(async () => {
+          expect(mockClient.textDocumentDidChange).toHaveBeenCalledWith({
+            textDocument: { uri: notebookClient.documentUri, version: 2 },
+            contentChanges: [
+              {
+                text: "# this is a comment\nimport math\nimport numpy\nextra = 1\nprint(math.sqrt(4))",
+              },
+            ],
+          });
+          return null;
+        });
+        const params: LSP.CompletionParams = {
+          textDocument: { uri: CellDocumentUri.of(Cells.cell3) },
+          position: { line: 0, character: 11 },
+        };
+        await notebookClient.textDocumentCompletion(params);
+        expect(mockClient.textDocumentCompletion).toHaveBeenCalledWith({
+          textDocument: { uri: notebookClient.documentUri },
+          position: { line: 4, character: 11 },
+        });
+        await notebookClient.textDocumentCompletion(params);
+        expect(mockClient.textDocumentDidChange).toHaveBeenCalledTimes(1);
+      } finally {
+        editor.destroy();
+      }
+    });
   });
 
   describe("document lifecycle", () => {
@@ -1893,6 +1993,7 @@ describe("NotebookLanguageServerClient", () => {
       const isolatedNotebookClient = new NotebookLanguageServerClient(
         isolatedClient,
         {},
+        () => ({}),
       );
       await isolatedNotebookClient.sync();
       isolatedNotebookClient.patchProcessNotification();

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

@@ -343,7 +343,21 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   }
 
   private getNotebookCode() {
-    return store.get(topologicalCodesAtom);
+    const { cellIds, codes } = store.get(topologicalCodesAtom);
+    const editors = this.getNotebookEditors();
+    // EditorView edits do not invalidate the cached notebook atom.
+    return {
+      cellIds,
+      codes: Objects.fromEntries(
+        cellIds.map((id) => {
+          const editor = editors[id];
+          return [
+            id,
+            editor ? getEditorCodeAsPython(editor) : (codes[id] ?? ""),
+          ];
+        }),
+      ),
+    };
   }
 
   private assertCellDocumentUri(
@@ -870,7 +884,7 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   }
 
   public async textDocumentCompletion(params: LSP.CompletionParams) {
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    const { lens } = await this.sync();
 
     // Check if URI is valid
     if (!CellDocumentUri.is(params.textDocument.uri)) {

package/src/core/codemirror/readonly/__tests__/extension.test.ts

@@ -4,12 +4,12 @@ import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { createStore } from "jotai";
 import { describe, expect, it } from "vitest";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import { connectionAtom } from "../../../network/connection";
 import { dynamicReadonly, isEditorReadonly } from "../extension";
 
 function makeStoreWithConnection(
-  state: typeof WebSocketState.CONNECTING | typeof WebSocketState.OPEN,
+  state: typeof ConnectionState.CONNECTING | typeof ConnectionState.OPEN,
 ) {
   const store = createStore();
   store.set(connectionAtom, { state });
@@ -18,7 +18,7 @@ function makeStoreWithConnection(
 
 describe("dynamicReadonly", () => {
   it("should be readonly when connection is not OPEN", () => {
-    const store = makeStoreWithConnection(WebSocketState.CONNECTING);
+    const store = makeStoreWithConnection(ConnectionState.CONNECTING);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -27,7 +27,7 @@ describe("dynamicReadonly", () => {
   });
 
   it("should be editable when connection is OPEN", async () => {
-    const store = makeStoreWithConnection(WebSocketState.OPEN);
+    const store = makeStoreWithConnection(ConnectionState.OPEN);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -37,7 +37,7 @@ describe("dynamicReadonly", () => {
   });
 
   it("should toggle readonly when connection state changes", () => {
-    const store = makeStoreWithConnection(WebSocketState.CONNECTING);
+    const store = makeStoreWithConnection(ConnectionState.CONNECTING);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -45,7 +45,7 @@ describe("dynamicReadonly", () => {
     const view = new EditorView({ state });
     expect(isEditorReadonly(state)).toBe(true);
     // Simulate connection opening
-    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
     // The extension uses a plugin to update the state, so we need to create a view to trigger it
     view.requestMeasure();
     // The state should now be editable

package/src/core/codemirror/readonly/extension.ts

@@ -7,7 +7,7 @@ import {
   isAppConnecting,
 } from "@/core/websocket/connection-utils";
 import { connectionAtom } from "../../network/connection";
-import { WebSocketState } from "../../websocket/types";
+import { ConnectionState } from "../../websocket/types";
 
 /**
  * State effect for updating readonly status based on connection
@@ -111,7 +111,7 @@ export function dynamicReadonly(store: ReturnType<typeof createStore>) {
     // Subscribe to connectionAtom and manage listeners
     const unsubscribe = store.sub(connectionAtom, () => {
       const connection = store.get(connectionAtom);
-      if (connection.state === WebSocketState.CONNECTING) {
+      if (connection.state === ConnectionState.CONNECTING) {
         attachListeners();
       } else {
         detachListeners();

package/src/core/config/pair.ts

@@ -7,7 +7,6 @@ export const pairPreviewSchema = z.object({
   command: z.string(),
   templates: z.object({
     prompt: z.string(),
-    file: z.string(),
     session: z.string(),
     token_file: z.string(),
     token: z.string(),