@marimo-team/frontend 0.23.16-dev35 → 0.23.16-dev36

package/src/core/codemirror/language/languages/python.ts

@@ -38,7 +38,11 @@ import { createLspMarkdownRenderer } from "../../lsp/markdown-renderer";
 import { NotebookLanguageServerClient } from "../../lsp/notebook-lsp";
 import { createTransport } from "../../lsp/transports";
 import { CellDocumentUri, type ILanguageServerClient } from "../../lsp/types";
-import { getLspRootUri, getLspWorkspaceFolders } from "../../lsp/utils";
+import {
+  getLspRootUri,
+  getLspWorkspaceFolders,
+  isKnownCellDocumentUri,
+} from "../../lsp/utils";
 import {
   clickablePlaceholderExtension,
   smartPlaceholderExtension,
@@ -338,6 +342,21 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
           clients.length === 1
             ? (clients[0] as NotebookLanguageServerClient)
             : new FederatedLanguageServerClient(clients);
+        const openExternalLocation = (result: { uri: string }) => {
+          Logger.debug("openExternalLocation", result);
+          // Neither the merged document nor a cell is a file on disk. The
+          // plugin handles same-cell locations itself; revealing a location in
+          // another cell isn't supported yet.
+          if (
+            client.documentUri === result.uri ||
+            isKnownCellDocumentUri(result.uri)
+          ) {
+            return;
+          }
+          getRequestClient().openFile({
+            path: result.uri.replace("file://", ""),
+          });
+        };
 
         return [
           languageServerWithClient({
@@ -346,6 +365,7 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
             allowHTMLContent: true,
             markdownRenderer: createLspMarkdownRenderer(),
             useSnippetOnCompletion: true,
+            clientSideFiltering: true,
             hoverConfig: hoverOptions,
             completionConfig: autocompleteOptions,
             // Default to false
@@ -362,16 +382,8 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
               rename: hotkeys.getHotkey("cell.renameSymbol").key,
             },
             completionMatchBefore,
-            onGoToDefinition: (result) => {
-              Logger.debug("onGoToDefinition", result);
-              if (client.documentUri === result.uri) {
-                // Local definition
-                return;
-              }
-              getRequestClient().openFile({
-                path: result.uri.replace("file://", ""),
-              });
-            },
+            onGoToDefinition: openExternalLocation,
+            onShowLocation: openExternalLocation,
           }),
           documentUri.of(CellDocumentUri.of(cellId)),
         ];

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

@@ -0,0 +1,161 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it, type Mocked, vi } from "vitest";
+import type * as LSP from "vscode-languageserver-protocol";
+import { FederatedLanguageServerClient } from "../federated-lsp";
+import type { ILanguageServerClient } from "../types";
+
+function createClient(
+  supportedMethods: string[] = [],
+): Mocked<ILanguageServerClient> {
+  const client: Mocked<ILanguageServerClient> = {
+    ready: true,
+    capabilities: {},
+    initializePromise: Promise.resolve(),
+    clientCapabilities: {},
+    initialize: vi.fn(),
+    close: vi.fn(),
+    hasCapability: vi.fn((method) => supportedMethods.includes(method)),
+    textDocumentDidOpen: vi.fn().mockResolvedValue(false),
+    textDocumentDidChange: vi.fn(),
+    textDocumentDidClose: vi.fn(),
+    textDocumentWillSave: vi.fn(),
+    textDocumentWillSaveWaitUntil: vi.fn().mockResolvedValue(null),
+    textDocumentDidSave: vi.fn(),
+    textDocumentHover: vi.fn().mockResolvedValue({ contents: [] }),
+    textDocumentCompletion: vi.fn().mockResolvedValue(null),
+    completionItemResolve: vi.fn(),
+    textDocumentDefinition: vi.fn().mockResolvedValue(null),
+    textDocumentCodeAction: vi.fn().mockResolvedValue(null),
+    codeActionResolve: vi.fn(),
+    textDocumentRename: vi.fn().mockResolvedValue(null),
+    textDocumentPrepareRename: vi.fn().mockResolvedValue(null),
+    textDocumentSignatureHelp: vi.fn().mockResolvedValue(null),
+    onNotification: vi.fn().mockReturnValue(() => true),
+  };
+  return client;
+}
+
+describe("FederatedLanguageServerClient", () => {
+  it("routes requests using dynamic method capabilities", async () => {
+    const staticOnlyClient = createClient();
+    const dynamicClient = createClient(["textDocument/definition"]);
+    const definition: LSP.Location = {
+      uri: "file:///definition.py",
+      range: {
+        start: { line: 1, character: 2 },
+        end: { line: 1, character: 5 },
+      },
+    };
+    dynamicClient.textDocumentDefinition.mockResolvedValue(definition);
+    const client = new FederatedLanguageServerClient([
+      staticOnlyClient,
+      dynamicClient,
+    ]);
+    const params: LSP.DefinitionParams = {
+      textDocument: { uri: "file:///notebook.py" },
+      position: { line: 0, character: 0 },
+    };
+
+    await expect(client.textDocumentDefinition(params)).resolves.toEqual(
+      definition,
+    );
+    expect(client.hasCapability("textDocument/definition")).toBe(true);
+    expect(staticOnlyClient.textDocumentDefinition).not.toHaveBeenCalled();
+    expect(dynamicClient.textDocumentDefinition).toHaveBeenCalledWith(params);
+  });
+
+  it("prefers the first client's willSaveWaitUntil edits regardless of latency", async () => {
+    const slowEdits: LSP.TextEdit[] = [
+      {
+        range: {
+          start: { line: 0, character: 0 },
+          end: { line: 0, character: 1 },
+        },
+        newText: "first",
+      },
+    ];
+    const first = createClient();
+    const second = createClient();
+    // Answers last, but still wins
+    first.textDocumentWillSaveWaitUntil.mockImplementation(
+      () => new Promise((resolve) => setTimeout(() => resolve(slowEdits), 5)),
+    );
+    second.textDocumentWillSaveWaitUntil.mockResolvedValue([
+      {
+        range: {
+          start: { line: 1, character: 0 },
+          end: { line: 1, character: 1 },
+        },
+        newText: "second",
+      },
+    ]);
+    const client = new FederatedLanguageServerClient([first, second]);
+
+    await expect(
+      client.textDocumentWillSaveWaitUntil({
+        textDocument: { uri: "file:///notebook.py" },
+        reason: 1,
+      }),
+    ).resolves.toEqual(slowEdits);
+    expect(second.textDocumentWillSaveWaitUntil).toHaveBeenCalled();
+  });
+
+  it("skips clients that reject willSaveWaitUntil", async () => {
+    const failing = createClient();
+    const working = createClient();
+    const edits: LSP.TextEdit[] = [
+      {
+        range: {
+          start: { line: 0, character: 0 },
+          end: { line: 0, character: 1 },
+        },
+        newText: "ok",
+      },
+    ];
+    failing.textDocumentWillSaveWaitUntil.mockRejectedValue(
+      new Error("server exploded"),
+    );
+    working.textDocumentWillSaveWaitUntil.mockResolvedValue(edits);
+    const client = new FederatedLanguageServerClient([failing, working]);
+
+    await expect(
+      client.textDocumentWillSaveWaitUntil({
+        textDocument: { uri: "file:///notebook.py" },
+        reason: 1,
+      }),
+    ).resolves.toEqual(edits);
+  });
+
+  it("fans out document lifecycle and save notifications", async () => {
+    const first = createClient();
+    const second = createClient();
+    second.textDocumentDidOpen.mockResolvedValue(true);
+    const client = new FederatedLanguageServerClient([first, second]);
+    const openParams: LSP.DidOpenTextDocumentParams = {
+      textDocument: {
+        uri: "file:///notebook.py",
+        languageId: "python",
+        version: 1,
+        text: "value = 1",
+      },
+    };
+    const closeParams: LSP.DidCloseTextDocumentParams = {
+      textDocument: { uri: "file:///notebook.py" },
+    };
+    const saveParams: LSP.DidSaveTextDocumentParams = {
+      textDocument: { uri: "file:///notebook.py" },
+      text: "value = 1",
+    };
+
+    await expect(client.textDocumentDidOpen(openParams)).resolves.toBe(true);
+    await client.textDocumentDidClose(closeParams);
+    await client.textDocumentDidSave(saveParams);
+
+    for (const child of [first, second]) {
+      expect(child.textDocumentDidOpen).toHaveBeenCalledWith(openParams);
+      expect(child.textDocumentDidClose).toHaveBeenCalledWith(closeParams);
+      expect(child.textDocumentDidSave).toHaveBeenCalledWith(saveParams);
+    }
+  });
+});

package/src/core/codemirror/lsp/__tests__/log-messages.test.ts

@@ -0,0 +1,53 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { JSONRPCMessage } from "@marimo-team/codemirror-languageserver";
+import { beforeEach, describe, expect, it } from "vitest";
+import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
+import { store } from "@/core/state/jotai";
+import { handleLogMessage } from "../log-messages";
+
+function logMessage(type: number, message: string): JSONRPCMessage {
+  return {
+    jsonrpc: "2.0",
+    method: "window/logMessage",
+    params: { type, message },
+  } as JSONRPCMessage;
+}
+
+describe("handleLogMessage", () => {
+  beforeEach(() => {
+    store.set(notebookAtom, initialNotebookState());
+  });
+
+  it("adds a log for each message, erroring on Error and Warning", () => {
+    handleLogMessage("pylsp", logMessage(1, "failed to start"));
+    handleLogMessage("pylsp", logMessage(2, "plugin not installed"));
+    handleLogMessage("pylsp", logMessage(3, "indexing"));
+
+    expect(store.get(notebookAtom).cellLogs).toMatchObject([
+      { level: "stderr", message: "failed to start" },
+      { level: "stderr", message: "plugin not installed" },
+      {
+        level: "stdout",
+        message: "indexing",
+        source: { type: "lsp", name: "pylsp" },
+      },
+    ]);
+  });
+
+  it("ignores messages that aren't logs", () => {
+    handleLogMessage("pylsp", { jsonrpc: "2.0", id: 1 } as JSONRPCMessage);
+    handleLogMessage("pylsp", {
+      jsonrpc: "2.0",
+      method: "textDocument/publishDiagnostics",
+      params: { uri: "file:///a", diagnostics: [] },
+    } as JSONRPCMessage);
+    handleLogMessage("pylsp", {
+      jsonrpc: "2.0",
+      method: "window/logMessage",
+      params: { type: 1 },
+    } as JSONRPCMessage);
+
+    expect(store.get(notebookAtom).cellLogs).toEqual([]);
+  });
+});

package/src/core/codemirror/lsp/__tests__/normalize-markdown-math.test.ts

@@ -90,6 +90,31 @@ Use \`:math:\` as text and :math:\`z\` as math.
     expect(normalized).not.toContain(":math:`z`");
     expect(normalized).toContain("||(z||)");
   });
+
+  it("leaves markdown-escaped brackets in type annotations alone", () => {
+    const markdown =
+      'on\\_change (Optional\\[Callable\\]) and editable\\_columns (Union\\[list\\[str\\], Literal\\["all"\\]\\])';
+
+    expect(normalizeMarkdownMath(markdown)).toBe(markdown);
+  });
+
+  it("keeps the newlines around fenced code blocks", () => {
+    const markdown = `
+\`\`\`python
+class data_editor()
+\`\`\`
+---
+Inline \\[z^2\\] math forces normalization.
+
+Done.
+`.trim();
+
+    // A closing fence glued to the next line stops closing the block, and the
+    // rest of the document is swallowed into it.
+    expect(normalizeMarkdownMath(markdown)).toContain(
+      "class data_editor()\n```\n---\n",
+    );
+  });
 });
 
 describe("normalizeLspDocumentation", () => {