@marimo-team/frontend 0.23.16-dev4 → 0.23.16-dev41

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/language/languages/sql/completion-store.ts

@@ -109,7 +109,11 @@ class SQLCompletionStore {
     return {
       shouldAddLocalTables: true,
       schema: builder.build(),
-      defaultSchema: default_schema ?? undefined,
+      defaultSchema:
+        default_schema ??
+        (defaultDb?.schemas.length === 1
+          ? defaultDb.schemas[0]?.name
+          : undefined),
     };
   }
 

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

@@ -12,12 +12,17 @@ import {
   type ViewUpdate,
 } from "@codemirror/view";
 import {
+  aliasColumnCompletionSource,
+  createCteCompletionSource,
   defaultSqlHoverTheme,
   NodeSqlParser,
   type NodeSqlParserResult,
+  QueryContextAnalyzer,
   type SupportedDialects as ParserDialects,
   type SqlParseError,
   sqlExtension,
+  SqlStructureAnalyzer,
+  unqualifiedColumnCompletionSource,
 } from "@marimo-team/codemirror-sql";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import { type SQLMetadata, SQLParser } from "@marimo-team/smart-cells";
@@ -142,6 +147,7 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
     _placeholderType: PlaceholderType,
     lspConfig: LSPConfig & { diagnostics: DiagnosticsConfig },
   ): Extension[] {
+    const analysis = createSQLAnalysis();
     const extensions = [
       // This can be updated with a dispatch effect
       sqlConfigCompartment.of(sql({ dialect: DEFAULT_DIALECT })),
@@ -167,6 +173,10 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
           variableCompletionSource,
           // Completions for dialect keywords
           customKeywordCompletionSource(),
+          // Completions based on the current query
+          analysis.cteCompletionSource,
+          analysis.aliasCompletionSource,
+          analysis.columnCompletionSource,
         ],
       }),
     ];
@@ -185,10 +195,23 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
 
       extensions.push(
         sqlExtension({
+          schema: getSchema,
           enableLinting: true,
           linterConfig: {
             delay: 250, // Delay before running validation
             parser: parser,
+            // CustomSqlParser performs backend DuckDB validation in validateSql.
+            perStatement: false,
+          },
+          enableSemanticLinting: true,
+          semanticLinterConfig: {
+            parser: analysis.parser,
+            structureAnalyzer: analysis.structureAnalyzer,
+            severity: {
+              unknownTable: "off",
+              unknownColumn: "warning",
+              ambiguousColumn: "warning",
+            },
           },
           enableGutterMarkers: true,
           gutterConfig: {
@@ -198,14 +221,21 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
             parser: parser,
           },
           hoverConfig: {
-            schema: getSchema, // Use the same schema as autocomplete
             hoverTime: 300, // 300ms hover delay
             enableKeywords: true, // Show keyword information
             enableTables: true, // Show table information
             enableColumns: true, // Show column information
-            parser: parser,
+            parser: analysis.parser,
+            contextAnalyzer: analysis.contextAnalyzer,
             theme: defaultSqlHoverTheme(theme),
           },
+          enableNavigation: true,
+          navigationConfig: {
+            contextAnalyzer: analysis.contextAnalyzer,
+            keymap: false,
+            parser: analysis.parser,
+            structureAnalyzer: analysis.structureAnalyzer,
+          },
         }),
         EditorView.updateListener.of((update) => {
           if (update.focusChanged) {
@@ -223,6 +253,123 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
   }
 }
 
+function createSQLAnalysis() {
+  const parser = new NodeSqlParser({
+    getParserOptions: (state: EditorState) => {
+      return {
+        database: getAnalysisDialect(state),
+      };
+    },
+  });
+  const contextAnalyzer = new DialectAwareQueryContextAnalyzer(parser);
+  const structureAnalyzer = new DialectAwareSqlStructureAnalyzer(parser);
+
+  return {
+    parser,
+    contextAnalyzer,
+    structureAnalyzer,
+    cteCompletionSource: createCteCompletionSource({
+      parser,
+      contextAnalyzer,
+      structureAnalyzer,
+    }),
+    aliasCompletionSource: aliasColumnCompletionSource({
+      schema: getSchema,
+      parser,
+      contextAnalyzer,
+    }),
+    columnCompletionSource: unqualifiedColumnCompletionSource({
+      schema: getSchema,
+      parser,
+      contextAnalyzer,
+    }),
+  };
+}
+
+function getAnalysisDialect(state: EditorState): ParserDialects {
+  return guessParserDialect(state) ?? DEFAULT_PARSER_DIALECT;
+}
+
+class DialectAwareQueryContextAnalyzer extends QueryContextAnalyzer {
+  private readonly analyzers = new Map<ParserDialects, QueryContextAnalyzer>();
+  private readonly analysisParser: NodeSqlParser;
+
+  constructor(analysisParser: NodeSqlParser) {
+    super(analysisParser);
+    this.analysisParser = analysisParser;
+  }
+
+  override getContext(
+    sql: string,
+    opts: { state: EditorState },
+  ): ReturnType<QueryContextAnalyzer["getContext"]> {
+    return this.getAnalyzer(opts.state).getContext(sql, opts);
+  }
+
+  override clearCache(): void {
+    for (const analyzer of this.analyzers.values()) {
+      analyzer.clearCache();
+    }
+  }
+
+  private getAnalyzer(state: EditorState): QueryContextAnalyzer {
+    const dialect = getAnalysisDialect(state);
+    let analyzer = this.analyzers.get(dialect);
+    if (!analyzer) {
+      analyzer = new QueryContextAnalyzer(this.analysisParser);
+      this.analyzers.set(dialect, analyzer);
+    }
+    return analyzer;
+  }
+}
+
+class DialectAwareSqlStructureAnalyzer extends SqlStructureAnalyzer {
+  private readonly analyzers = new Map<ParserDialects, SqlStructureAnalyzer>();
+  private readonly analysisParser: NodeSqlParser;
+
+  constructor(analysisParser: NodeSqlParser) {
+    super(analysisParser);
+    this.analysisParser = analysisParser;
+  }
+
+  override analyzeDocument(
+    state: EditorState,
+  ): ReturnType<SqlStructureAnalyzer["analyzeDocument"]> {
+    return this.getAnalyzer(state).analyzeDocument(state);
+  }
+
+  override getStatementAtPosition(
+    state: EditorState,
+    position: number,
+  ): ReturnType<SqlStructureAnalyzer["getStatementAtPosition"]> {
+    return this.getAnalyzer(state).getStatementAtPosition(state, position);
+  }
+
+  override getStatementsInRange(
+    state: EditorState,
+    from: number,
+    to: number,
+  ): ReturnType<SqlStructureAnalyzer["getStatementsInRange"]> {
+    return this.getAnalyzer(state).getStatementsInRange(state, from, to);
+  }
+
+  override clearCache(): void {
+    for (const analyzer of this.analyzers.values()) {
+      analyzer.clearCache();
+    }
+  }
+
+  private getAnalyzer(state: EditorState): SqlStructureAnalyzer {
+    const dialect = getAnalysisDialect(state);
+    let analyzer = this.analyzers.get(dialect);
+    if (!analyzer) {
+      analyzer = new SqlStructureAnalyzer(this.analysisParser);
+      this.analyzers.set(dialect, analyzer);
+    }
+    return analyzer;
+  }
+}
+
 class CustomSqlParser extends NodeSqlParser {
   private validationTimeout: number | null = null;
   private readonly VALIDATION_DELAY_MS = 300; // Wait 300ms after user stops typing
@@ -344,7 +491,61 @@ function getSchema(view: EditorView): SQLNamespace {
     return {};
   }
 
-  return config.schema;
+  const schema = config.schema;
+  if (!isNamespaceMap(schema)) {
+    return schema;
+  }
+
+  const defaultSchema = config.defaultSchema;
+  const defaultSchemaNamespace = defaultSchema
+    ? schema[defaultSchema]
+    : undefined;
+
+  // The completion schema contains both default-schema and fully qualified
+  // paths. Promote default-schema tables so unqualified references resolve to
+  // one exact match instead of appearing ambiguous.
+  if (
+    defaultSchemaNamespace &&
+    !Array.isArray(defaultSchemaNamespace) &&
+    typeof defaultSchemaNamespace === "object" &&
+    "children" in defaultSchemaNamespace &&
+    isNamespaceMap(defaultSchemaNamespace.children)
+  ) {
+    return {
+      ...schema,
+      ...Object.fromEntries(
+        Object.entries(defaultSchemaNamespace.children).filter(
+          ([name]) => !Object.hasOwn(schema, name),
+        ),
+      ),
+    };
+  }
+
+  return schema;
+}
+
+function isNamespaceMap(
+  namespace: SQLNamespace,
+): namespace is Record<string, SQLNamespace> {
+  if (Array.isArray(namespace)) {
+    return false;
+  }
+
+  if (!("self" in namespace) || !("children" in namespace)) {
+    return true;
+  }
+
+  // A self/children wrapper stores a Completion in `self`. A table or schema
+  // literally named "self" stores another namespace there.
+  const self = namespace.self;
+  const isSelfChildrenWrapper =
+    typeof self === "object" &&
+    self !== null &&
+    !Array.isArray(self) &&
+    "label" in self &&
+    typeof self.label === "string";
+
+  return !isSelfChildrenWrapper;
 }
 
 function guessParserDialect(state: EditorState): ParserDialects | null {
@@ -529,3 +730,9 @@ async function validateSQL(
   }
   return result;
 }
+
+export const exportedForTesting = {
+  createSQLAnalysis,
+  CustomSqlParser,
+  getSchema,
+};

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", () => {