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

package/src/core/codemirror/lsp/__tests__/transports.test.ts

@@ -0,0 +1,112 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
+import { store } from "@/core/state/jotai";
+import { createTransport } from "../transports";
+
+vi.mock("../../../network/connection", () => ({
+  waitForConnectionOpen: vi.fn(() => Promise.resolve()),
+}));
+
+vi.mock("../../../runtime/config", () => ({
+  getRuntimeManager: () => ({
+    getLSPURL: (server: string) => new URL(`ws://localhost:2718/lsp/${server}`),
+  }),
+}));
+
+/** Minimal stand-in for the browser WebSocket used by `WebSocketTransport`. */
+class FakeWebSocket {
+  static readonly instances: FakeWebSocket[] = [];
+  static readonly OPEN = 1;
+
+  readyState = 0;
+  readonly url: string;
+  private readonly listeners = new Map<string, Set<(event: unknown) => void>>();
+
+  constructor(url: string) {
+    this.url = url;
+    FakeWebSocket.instances.push(this);
+  }
+
+  addEventListener(type: string, listener: (event: unknown) => void) {
+    const listeners = this.listeners.get(type) ?? new Set();
+    listeners.add(listener);
+    this.listeners.set(type, listeners);
+  }
+
+  send() {
+    // no-op
+  }
+
+  close() {
+    this.readyState = 3;
+  }
+
+  open() {
+    this.readyState = FakeWebSocket.OPEN;
+    this.emit("open", {});
+  }
+
+  receive(message: unknown) {
+    this.emit("message", { data: JSON.stringify(message) });
+  }
+
+  private emit(type: string, event: unknown) {
+    for (const listener of this.listeners.get(type) ?? []) {
+      listener(event);
+    }
+  }
+}
+
+describe("createTransport", () => {
+  beforeEach(() => {
+    FakeWebSocket.instances.length = 0;
+    vi.stubGlobal("WebSocket", FakeWebSocket);
+    vi.spyOn(console, "log").mockImplementation(() => undefined);
+    store.set(notebookAtom, initialNotebookState());
+  });
+
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  it("forwards window/logMessage notifications to the logs panel", async () => {
+    const transport = createTransport("pylsp");
+    const connected = transport.connect();
+    await vi.waitFor(() => expect(FakeWebSocket.instances).toHaveLength(1));
+    const socket = FakeWebSocket.instances[0];
+    socket.open();
+    await connected;
+
+    socket.receive({
+      jsonrpc: "2.0",
+      method: "window/logMessage",
+      params: { type: 1, message: "pylsp failed to start" },
+    });
+    socket.receive({
+      jsonrpc: "2.0",
+      method: "window/logMessage",
+      params: { type: 3, message: "indexing workspace" },
+    });
+    // Unrelated notifications are left alone.
+    socket.receive({
+      jsonrpc: "2.0",
+      method: "textDocument/publishDiagnostics",
+      params: { uri: "file:///a", diagnostics: [] },
+    });
+
+    expect(store.get(notebookAtom).cellLogs).toMatchObject([
+      {
+        level: "stderr",
+        message: "pylsp failed to start",
+        source: { type: "lsp", name: "pylsp" },
+      },
+      {
+        level: "stdout",
+        message: "indexing workspace",
+        source: { type: "lsp", name: "pylsp" },
+      },
+    ]);
+  });
+});

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

@@ -2,7 +2,7 @@
 
 import type * as LSP from "vscode-languageserver-protocol";
 import { Objects } from "@/utils/objects";
-import type { ILanguageServerClient } from "./types";
+import type { ILanguageServerClient, RoutableLspMethod } from "./types";
 import { getLspDocumentUri } from "./utils";
 
 function removeFalseyValues<T extends object>(obj: T): T {
@@ -54,7 +54,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
         }
         return undefined;
       })
-      .filter((c) => c != null);
+      .filter((c): c is LSP.ClientCapabilities => c != null);
 
     return mergeDictsIgnoreFalsey<LSP.ClientCapabilities>(capabilities);
   }
@@ -72,7 +72,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   get capabilities(): LSP.ServerCapabilities | null {
     const capabilities = this.clients
       .map((client) => client.capabilities)
-      .filter((c) => c !== null);
+      .filter((c): c is LSP.ServerCapabilities => c !== null);
     return mergeDictsIgnoreFalsey<LSP.ServerCapabilities>(capabilities);
   }
 
@@ -92,16 +92,20 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
     });
   }
 
-  private firstWithCapability(
-    capability: keyof LSP.ServerCapabilities,
+  private firstWithMethod(
+    method: RoutableLspMethod,
   ): ILanguageServerClient | undefined {
-    return this.clients.find((client) => client.capabilities?.[capability]);
+    return this.clients.find((client) => client.hasCapability(method));
   }
 
-  private clientsWithCapability(
-    capability: keyof LSP.ServerCapabilities,
+  private clientsWithMethod(
+    method: RoutableLspMethod,
   ): ILanguageServerClient[] {
-    return this.clients.filter((client) => client.capabilities?.[capability]);
+    return this.clients.filter((client) => client.hasCapability(method));
+  }
+
+  hasCapability(method: string): boolean {
+    return this.clients.some((client) => client.hasCapability(method));
   }
 
   async initialize(): Promise<void> {
@@ -114,18 +118,16 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
 
   async textDocumentDidChange(
     params: LSP.DidChangeTextDocumentParams,
-  ): Promise<LSP.DidChangeTextDocumentParams> {
+  ): Promise<void> {
     await Promise.all(
       this.clients.map((client) => client.textDocumentDidChange(params)),
     );
-
-    return params;
   }
 
   async completionItemResolve(
     item: LSP.CompletionItem,
   ): Promise<LSP.CompletionItem> {
-    const client = this.firstWithCapability("completionProvider");
+    const client = this.firstWithMethod("textDocument/completion");
     if (client) {
       return client.completionItemResolve(item);
     }
@@ -135,17 +137,22 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   async textDocumentCodeAction(
     params: LSP.CodeActionParams,
   ): Promise<(LSP.Command | LSP.CodeAction)[] | null> {
-    const client = this.firstWithCapability("codeActionProvider");
+    const client = this.firstWithMethod("textDocument/codeAction");
     if (client) {
       return client.textDocumentCodeAction(params);
     }
     return null;
   }
 
+  async codeActionResolve(action: LSP.CodeAction): Promise<LSP.CodeAction> {
+    const client = this.firstWithMethod("textDocument/codeAction");
+    return client ? client.codeActionResolve(action) : action;
+  }
+
   async textDocumentRename(
     params: LSP.RenameParams,
   ): Promise<LSP.WorkspaceEdit | null> {
-    const client = this.firstWithCapability("renameProvider");
+    const client = this.firstWithMethod("textDocument/rename");
     if (client) {
       return client.textDocumentRename(params);
     }
@@ -155,7 +162,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   async textDocumentPrepareRename(
     params: LSP.PrepareRenameParams,
   ): Promise<LSP.PrepareRenameResult | null> {
-    const client = this.firstWithCapability("renameProvider");
+    const client = this.firstWithMethod("textDocument/prepareRename");
     if (client) {
       return client.textDocumentPrepareRename(params);
     }
@@ -165,7 +172,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   async textDocumentSignatureHelp(
     params: LSP.SignatureHelpParams,
   ): Promise<LSP.SignatureHelp | null> {
-    const client = this.firstWithCapability("signatureHelpProvider");
+    const client = this.firstWithMethod("textDocument/signatureHelp");
     if (client) {
       return client.textDocumentSignatureHelp(params);
     }
@@ -176,7 +183,7 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   async textDocumentCompletion(
     params: LSP.CompletionParams,
   ): Promise<LSP.CompletionList | LSP.CompletionItem[] | null> {
-    const clients = this.clientsWithCapability("completionProvider");
+    const clients = this.clientsWithMethod("textDocument/completion");
     const results = await Promise.allSettled(
       clients.map((client) => client.textDocumentCompletion(params)),
     );
@@ -187,24 +194,70 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
   async textDocumentDefinition(
     params: LSP.DefinitionParams,
   ): Promise<LSP.Definition | LSP.LocationLink[] | null> {
-    const client = this.firstWithCapability("definitionProvider");
+    const client = this.firstWithMethod("textDocument/definition");
     if (client) {
       return client.textDocumentDefinition(params);
     }
     return null;
   }
 
-  async textDocumentDidOpen(params: LSP.DidOpenTextDocumentParams) {
-    await Promise.all(
+  async textDocumentDidOpen(
+    params: LSP.DidOpenTextDocumentParams,
+  ): Promise<boolean> {
+    const results = await Promise.all(
       this.clients.map((client) => client.textDocumentDidOpen(params)),
     );
 
-    return params;
+    return results.some((result) => result !== false);
+  }
+
+  async textDocumentDidClose(
+    params: LSP.DidCloseTextDocumentParams,
+  ): Promise<void> {
+    await Promise.all(
+      this.clients.map((client) => client.textDocumentDidClose(params)),
+    );
+  }
+
+  async textDocumentWillSave(
+    params: LSP.WillSaveTextDocumentParams,
+  ): Promise<void> {
+    await Promise.all(
+      this.clients.map((client) => client.textDocumentWillSave(params)),
+    );
+  }
+
+  async textDocumentWillSaveWaitUntil(
+    params: LSP.WillSaveTextDocumentParams,
+  ): Promise<LSP.TextEdit[] | null> {
+    // This blocks the save, so query concurrently rather than paying the sum of
+    // the servers' latencies. The earliest client that returned edits still
+    // wins, so the result doesn't depend on who answers first.
+    const results = await Promise.allSettled(
+      this.clients.map((client) =>
+        client.textDocumentWillSaveWaitUntil(params),
+      ),
+    );
+
+    for (const result of results) {
+      if (result.status === "fulfilled" && result.value?.length) {
+        return result.value;
+      }
+    }
+    return null;
+  }
+
+  async textDocumentDidSave(
+    params: LSP.DidSaveTextDocumentParams,
+  ): Promise<void> {
+    await Promise.all(
+      this.clients.map((client) => client.textDocumentDidSave(params)),
+    );
   }
 
   async textDocumentHover(params: LSP.HoverParams): Promise<LSP.Hover> {
     for (const client of this.clients) {
-      if (!client.capabilities?.hoverProvider) {
+      if (!client.hasCapability("textDocument/hover")) {
         continue;
       }
       const result = await client.textDocumentHover(params);

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

@@ -0,0 +1,37 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { JSONRPCMessage } from "@marimo-team/codemirror-languageserver";
+import { addLogs } from "@/core/cells/cells";
+import { isRecord } from "@/utils/records";
+
+// LSP MessageType: 1 = Error, 2 = Warning, 3 = Info, 4 = Log, 5 = Debug
+const ERROR_MESSAGE_TYPES = new Set<unknown>([1, 2]);
+
+/**
+ * Forward a server's `window/logMessage` notifications to the Logs panel.
+ * Servers report failures they can't surface as diagnostics this way, such as
+ * a missing interpreter or a plugin that failed to load.
+ */
+export function handleLogMessage(
+  serverName: string,
+  message: JSONRPCMessage,
+): void {
+  if (!("method" in message) || message.method !== "window/logMessage") {
+    return;
+  }
+  const params: unknown = message.params;
+  if (!isRecord(params) || typeof params.message !== "string") {
+    return;
+  }
+
+  addLogs({
+    logs: [
+      {
+        timestamp: Date.now() / 1000,
+        level: ERROR_MESSAGE_TYPES.has(params.type) ? "stderr" : "stdout",
+        message: params.message,
+        source: { type: "lsp", name: serverName },
+      },
+    ],
+  });
+}

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

@@ -6,8 +6,10 @@ import type * as LSP from "vscode-languageserver-protocol";
 const INLINE_MATH_ROLE_REGEX = /(?<!`):math:`([^\n`]+)`/g;
 // Matches role variant with HTML code tag: :math:<code>...</code>
 const INLINE_MATH_ROLE_CODE_REGEX = /:math:\s*<code>(.+?)<\/code>/gs;
-// Matches display bracket delimiters: \[...\]
-const DISPLAY_BRACKET_REGEX = /\\\[(.+?)\\]/gs;
+// Matches display bracket delimiters: \[...\]. Must start a line or follow
+// whitespace, so escaped brackets in type annotations (`Optional\[Callable\]`)
+// are not read as math.
+const DISPLAY_BRACKET_REGEX = /(?<!\S)\\\[(.+?)\\]/gs;
 // Matches inline paren delimiters: \(...\)
 const INLINE_PAREN_REGEX = /\\\((.+?)\\\)/gs;
 // Matches display dollar delimiters: $$...$$
@@ -36,6 +38,8 @@ export function normalizeMarkdownMath(markdown: string): string {
     return markdown;
   }
 
+  // Segments are split on line boundaries, so the newline between two of them
+  // belongs to neither and has to be restored on join.
   return splitByFencedCodeBlocks(markdown)
     .map((segment) => {
       if (segment.isFenced) {
@@ -55,7 +59,7 @@ export function normalizeMarkdownMath(markdown: string): string {
         })
         .join("");
     })
-    .join("");
+    .join("\n");
 }
 
 /**
@@ -132,7 +136,7 @@ function containsMathSyntax(markdown: string): boolean {
     markdown.includes(".. math::") ||
     markdown.includes(":math:`") ||
     // Detects display bracket delimiters.
-    /\\\[[\S\s]+?\\]/.test(markdown) ||
+    /(?<!\S)\\\[[\S\s]+?\\]/.test(markdown) ||
     // Detects inline paren delimiters.
     /\\\([\S\s]+?\\\)/.test(markdown) ||
     // Detects display dollar delimiters.