@marimo-team/islands 0.17.2 → 0.17.4

package/src/core/ai/tools/run-cells-tool.ts

@@ -21,11 +21,12 @@ import {
 } from "./base";
 import type { CopilotMode } from "./registry";
 
+const POST_EXECUTION_DELAY = 200;
+const WAIT_FOR_CELLS_TIMEOUT = 30_000;
+
 interface CellOutput {
   consoleOutput?: string;
-  // consoleAttachments?: FileUIPart[];
   cellOutput?: string;
-  // cellAttachments?: FileUIPart[];
 }
 
 // Must use Record instead of Map because Map serializes to JSON as {}
@@ -67,6 +68,12 @@ export class RunStaleCellsTool
   }) satisfies z.ZodType<RunStaleCellsOutput>;
   readonly mode: CopilotMode[] = ["agent"];
 
+  private readonly postExecutionDelay: number;
+
+  constructor(opts?: { postExecutionDelay?: number }) {
+    this.postExecutionDelay = opts?.postExecutionDelay ?? POST_EXECUTION_DELAY;
+  }
+
   handler = async (
     _args: EmptyToolInput,
     toolContext: ToolNotebookContext,
@@ -91,7 +98,12 @@ export class RunStaleCellsTool
     });
 
     // Wait for all cells to finish executing
-    const allCellsFinished = await this.waitForCellsToFinish(store, staleCells);
+    const allCellsFinished = await this.waitForCellsToFinish(
+      store,
+      staleCells,
+      WAIT_FOR_CELLS_TIMEOUT,
+      this.postExecutionDelay,
+    );
     if (!allCellsFinished) {
       return {
         status: "success",
@@ -105,6 +117,7 @@ export class RunStaleCellsTool
 
     const cellsToOutput = new Map<CellId, CellOutput | null>();
     let resultMessage = "";
+    let outputHasErrors = false;
 
     for (const cellId of staleCells) {
       const cellContextData = getCellContextData(cellId, updatedNotebook, {
@@ -112,13 +125,13 @@ export class RunStaleCellsTool
       });
 
       let cellOutputString: string | undefined;
-      // let cellAttachments: FileUIPart[] | undefined;
       let consoleOutputString: string | undefined;
-      // let consoleAttachments: FileUIPart[] | undefined;
 
       const cellOutput = cellContextData.cellOutput;
       const consoleOutputs = cellContextData.consoleOutputs;
-      if (!cellOutput && !consoleOutputs) {
+      const hasConsoleOutput = consoleOutputs && consoleOutputs.length > 0;
+
+      if (!cellOutput && !hasConsoleOutput) {
         // Set null to show no output
         cellsToOutput.set(cellId, null);
         continue;
@@ -126,31 +139,26 @@ export class RunStaleCellsTool
 
       if (cellOutput) {
         cellOutputString = this.formatOutputString(cellOutput);
-        // cellAttachments = await getAttachmentsForOutputs(
-        //   [cellOutput],
-        //   cellId,
-        //   cellContextData.cellName,
-        // );
+        if (this.outputHasErrors(cellOutput)) {
+          outputHasErrors = true;
+        }
       }
 
-      if (consoleOutputs) {
-        // consoleAttachments = await getAttachmentsForOutputs(
-        //   consoleOutputs,
-        //   cellId,
-        //   cellContextData.cellName,
-        // );
+      if (hasConsoleOutput) {
         consoleOutputString = consoleOutputs
           .map((output) => this.formatOutputString(output))
           .join("\n");
         resultMessage +=
           "Console output represents the stdout or stderr of the cell (eg. print statements).";
+
+        if (consoleOutputs.some((output) => this.outputHasErrors(output))) {
+          outputHasErrors = true;
+        }
       }
 
       cellsToOutput.set(cellId, {
         cellOutput: cellOutputString,
         consoleOutput: consoleOutputString,
-        // cellAttachments: cellAttachments,
-        // consoleAttachments: consoleAttachments,
       });
     }
 
@@ -161,16 +169,32 @@ export class RunStaleCellsTool
       };
     }
 
+    const nextSteps = [
+      "Review the output of the cells. The CellId is the key of the result object.",
+      outputHasErrors
+        ? "There are errors in the cells. Please fix them by using the edit notebook tool and the given CellIds."
+        : "You may edit the notebook further with the given CellIds.",
+    ];
+
     return {
       status: "success",
       cellsToOutput: Object.fromEntries(cellsToOutput),
       message: resultMessage === "" ? undefined : resultMessage,
-      next_steps: [
-        "Review the output of the cells, if you need to make any changes, you can run this tool again, the CellId is the key of the result object if you want to edit the cell.",
-      ],
+      next_steps: nextSteps,
     };
   };
 
+  private outputHasErrors(cellOutput: BaseOutput): boolean {
+    const { output } = cellOutput;
+    if (
+      output.mimetype === "application/vnd.marimo+error" ||
+      output.mimetype === "application/vnd.marimo+traceback"
+    ) {
+      return true;
+    }
+    return false;
+  }
+
   private formatOutputString(cellOutput: BaseOutput): string {
     let outputString = "";
     const { outputType, processedContent, imageUrl, output } = cellOutput;
@@ -199,7 +223,8 @@ export class RunStaleCellsTool
   private async waitForCellsToFinish(
     store: JotaiStore,
     cellIds: CellId[],
-    timeout = 30_000,
+    timeout: number,
+    postExecutionDelay: number,
   ): Promise<boolean> {
     const checkAllFinished = (
       notebook: ReturnType<typeof notebookAtom.read>,
@@ -212,9 +237,18 @@ export class RunStaleCellsTool
       });
     };
 
-    // If already finished, return immediately
-    if (checkAllFinished(store.get(notebookAtom))) {
+    // Add a small delay after cells finish to allow console outputs to arrive
+    // Console outputs are streamed and might still be in-flight
+    const delayForConsoleOutputs = async () => {
+      if (postExecutionDelay > 0) {
+        await new Promise((resolve) => setTimeout(resolve, postExecutionDelay));
+      }
       return true;
+    };
+
+    // Return immediately if all cells are finished
+    if (checkAllFinished(store.get(notebookAtom))) {
+      return await delayForConsoleOutputs();
     }
 
     // Wait for notebook state changes with timeout
@@ -225,7 +259,7 @@ export class RunStaleCellsTool
           setTimeout(() => reject(new Error("timeout")), timeout),
         ),
       ]);
-      return true;
+      return await delayForConsoleOutputs();
     } catch {
       return false;
     }

package/src/core/cells/ids.ts

@@ -94,7 +94,7 @@ export const HTMLCellId = {
  */
 export function findCellId(element: HTMLElement): CellId | null {
   let cellId: CellId | null = null;
-  const cellContainer = HTMLCellId.findElement(element);
+  const cellContainer = HTMLCellId.findElementThroughShadowDOMs(element);
   if (cellContainer) {
     cellId = HTMLCellId.parse(cellContainer.id);
   }

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

@@ -163,23 +163,19 @@ describe("switchLanguage", () => {
     // Switch back to python
     switchLanguage(mockEditor, { language: "python", keepCodeAsIs: false });
     expect(mockEditor.state.doc.toString()).toMatchInlineSnapshot(`
-      "mo.md(
-          r"""
+      "mo.md(r"""
       print('Hello')
       print('Goodbye')
-      """
-      )"
+      """)"
     `);
 
     // Switch to sql
     switchLanguage(mockEditor, { language: "sql", keepCodeAsIs: false });
     expect(mockEditor.state.doc.toString()).toMatchInlineSnapshot(`
-      "mo.md(
-          r"""
+      "mo.md(r"""
       print('Hello')
       print('Goodbye')
-      """
-      )"
+      """)"
     `);
 
     // Switch back to python
@@ -187,12 +183,10 @@ describe("switchLanguage", () => {
     expect(mockEditor.state.doc.toString()).toMatchInlineSnapshot(`
       "_df = mo.sql(
           f"""
-          mo.md(
-              r\\"""
+          mo.md(r\\"""
           print('Hello')
           print('Goodbye')
-          \\"""
-          )
+          \\""")
           """
       )"
     `);

package/src/core/codemirror/language/__tests__/markdown.test.ts

@@ -185,11 +185,10 @@ describe("MarkdownLanguageAdapter", () => {
     });
 
     it("should preserve indentation in f-strings", () => {
-      const pythonCode =
-        'mo.md(\n    f"""\n```python\n{some_variable}\n```\n"""\n)';
+      const pythonCode = 'mo.md(f"""\n```python\n{some_variable}\n```\n""")';
       const [innerCode, offset, metadata] = adapter.transformIn(pythonCode);
       expect(innerCode).toBe("```python\n{some_variable}\n```");
-      expect(offset).toBe(15);
+      expect(offset).toBe(10);
       expect(metadata.quotePrefix).toBe("f");
 
       // Transform out
@@ -198,7 +197,7 @@ describe("MarkdownLanguageAdapter", () => {
         metadata,
       );
       expect(outerCode).toMatch(pythonCode);
-      expect(outerOffset).toBe(17);
+      expect(outerOffset).toBe(12);
     });
 
     it("should handle f-strings", () => {
@@ -252,41 +251,35 @@ describe("MarkdownLanguageAdapter", () => {
       const code = '"Hello" world';
       const [wrappedCode, offset] = adapter.transformOut(code, metadata);
       expect(wrappedCode).toMatchInlineSnapshot(`
-        "mo.md(
-            """
+        "mo.md("""
         "Hello" world
-        """
-        )"
+        """)"
       `);
-      expect(offset).toBe(16);
+      expect(offset).toBe(11);
     });
 
     it("ends with quote", () => {
       const code = 'Hello "world"';
       const [wrappedCode, offset] = adapter.transformOut(code, metadata);
       expect(wrappedCode).toMatchInlineSnapshot(`
-        "mo.md(
-            """
+        "mo.md("""
         Hello "world"
-        """
-        )"
+        """)"
       `);
-      expect(offset).toBe(16);
+      expect(offset).toBe(11);
     });
 
     it("should wrap Markdown code with triple double-quoted string format", () => {
       const code = "# Markdown Title\n\nSome content here.";
       const [wrappedCode, offset] = adapter.transformOut(code, metadata);
       expect(wrappedCode).toMatchInlineSnapshot(`
-        "mo.md(
-            """
+        "mo.md("""
         # Markdown Title
 
         Some content here.
-        """
-        )"
+        """)"
       `);
-      expect(offset).toBe(16);
+      expect(offset).toBe(11);
     });
 
     it("should escape triple quotes in the Markdown code", () => {
@@ -312,14 +305,12 @@ describe("MarkdownLanguageAdapter", () => {
       metadata.quotePrefix = "f";
       const [wrappedCode, offset] = adapter.transformOut(code, metadata);
       expect(wrappedCode).toMatchInlineSnapshot(`
-        "mo.md(
-            f"""
+        "mo.md(f"""
         # Title
         {some_variable}
-        """
-        )"
+        """)"
       `);
-      expect(offset).toBe(17);
+      expect(offset).toBe(12);
     });
 
     it("should handle rf-strings in transformOut", () => {
@@ -327,14 +318,12 @@ describe("MarkdownLanguageAdapter", () => {
       metadata.quotePrefix = "rf";
       const [wrappedCode, offset] = adapter.transformOut(code, metadata);
       expect(wrappedCode).toMatchInlineSnapshot(`
-        "mo.md(
-            rf"""
+        "mo.md(rf"""
         # Title
         {some_variable}
-        """
-        )"
+        """)"
       `);
-      expect(offset).toBe(18);
+      expect(offset).toBe(13);
     });
   });
 
@@ -369,13 +358,11 @@ describe("MarkdownLanguageAdapter", () => {
 
     it("should return true for complex nested markdown", () => {
       const pythonCode = String.raw`
-      mo.md(
-        rf"""
-        \`\`\`python
-        {pathlib.Path(__file__).read_text(encoding="utf-8")}
-        \`\`\`
-        """
-      )
+      mo.md(rf"""
+      \`\`\`python
+      {pathlib.Path(__file__).read_text(encoding="utf-8")}
+      \`\`\`
+      """)
       `;
       expect(adapter.isSupported(pythonCode)).toBe(true);
     });

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

@@ -44,8 +44,15 @@ import {
 import type { LanguageAdapter } from "../types";
 
 const pylspClient = once((lspConfig: LSPConfig) => {
+  // Create a mutable reference for the resync callback
+  let resyncCallback: (() => Promise<void>) | undefined;
+
+  const transport = createTransport("pylsp", async () => {
+    await resyncCallback?.();
+  });
+
   const lspClientOpts = {
-    transport: createTransport("pylsp"),
+    transport,
     rootUri: getLSPDocumentRootUri(),
     workspaceFolders: [],
   };
@@ -126,47 +133,74 @@ const pylspClient = once((lspConfig: LSPConfig) => {
 
   // We wrap the client in a NotebookLanguageServerClient to add some
   // additional functionality to handle multiple cells
-  return new NotebookLanguageServerClient(
+  const notebookClient = new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
     }),
     settings,
   );
+
+  // Set the resync callback now that the client exists
+  resyncCallback = () => notebookClient.resyncAllDocuments();
+
+  return notebookClient;
 });
 
 const tyLspClient = once((_: LSPConfig) => {
+  let resyncCallback: (() => Promise<void>) | undefined;
+
+  const transport = createTransport("ty", async () => {
+    await resyncCallback?.();
+  });
+
   const lspClientOpts = {
-    transport: createTransport("ty"),
+    transport,
     rootUri: getLSPDocumentRootUri(),
     workspaceFolders: [],
   };
 
   // We wrap the client in a NotebookLanguageServerClient to add some
   // additional functionality to handle multiple cells
-  return new NotebookLanguageServerClient(
+  const notebookClient = new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
       getWorkspaceConfiguration: (_) => [{ disableLanguageServices: true }],
     }),
     {},
   );
+
+  // Set the resync callback now that the client exists
+  resyncCallback = () => notebookClient.resyncAllDocuments();
+
+  return notebookClient;
 });
 
 const pyrightClient = once((_: LSPConfig) => {
+  let resyncCallback: (() => Promise<void>) | undefined;
+
+  const transport = createTransport("basedpyright", async () => {
+    await resyncCallback?.();
+  });
+
   const lspClientOpts = {
-    transport: createTransport("basedpyright"),
+    transport,
     rootUri: getLSPDocumentRootUri(),
     workspaceFolders: [],
   };
 
   // We wrap the client in a NotebookLanguageServerClient to add some
   // additional functionality to handle multiple cells
-  return new NotebookLanguageServerClient(
+  const notebookClient = new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
     }),
     {},
   );
+
+  // Set the resync callback now that the client exists
+  resyncCallback = () => notebookClient.resyncAllDocuments();
+
+  return notebookClient;
 });
 
 /**

package/src/core/codemirror/language/languages/sql/completion-store.ts

@@ -134,7 +134,7 @@ class SQLCompletionStore {
     if (!connection) {
       return ModifiedStandardSQL;
     }
-    return guessDialect(connection) ?? ModifiedStandardSQL;
+    return guessDialect(connection);
   }
 
   getCompletionSource(connectionName: ConnectionName): SQLConfig | null {
@@ -152,7 +152,7 @@ class SQLCompletionStore {
     const schema = this.cache.getOrCreate(connection);
 
     return {
-      dialect: guessDialect(connection) ?? ModifiedStandardSQL,
+      dialect: guessDialect(connection),
       schema: schema.shouldAddLocalTables
         ? { ...schema.schema, ...getTablesMap() }
         : schema.schema,

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

@@ -38,6 +38,7 @@ import type { HotkeyProvider } from "@/core/hotkeys/hotkeys";
 import type { ValidateSQLResult } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { resolvedThemeAtom } from "@/theme/useTheme";
+import { logNever } from "@/utils/assertNever";
 import { Logger } from "@/utils/Logger";
 import { variableCompletionSource } from "../../embedded/embedded-python";
 import { languageMetadataField } from "../../metadata";
@@ -52,6 +53,7 @@ import {
 } from "./completion-sources";
 import { SCHEMA_CACHE } from "./completion-store";
 import { getSQLMode, type SQLMode } from "./sql-mode";
+import { isKnownDialect } from "./utils";
 
 const DEFAULT_DIALECT = DuckDBDialect;
 const DEFAULT_PARSER_DIALECT = "DuckDB";
@@ -353,6 +355,11 @@ function connectionNameToParserDialect(
 ): ParserDialects | null {
   const dialect =
     SCHEMA_CACHE.getInternalDialect(connectionName)?.toLowerCase();
+
+  if (!dialect || !isKnownDialect(dialect)) {
+    return null;
+  }
+
   switch (dialect) {
     case "postgresql":
     case "postgres":
@@ -385,8 +392,15 @@ function connectionNameToParserDialect(
     case "flink":
       return "FlinkSQL";
     case "mongodb":
+    case "noql":
       return "Noql";
+    case "oracle":
+    case "oracledb":
+    case "timescaledb":
+      Logger.debug("Unsupported dialect", { dialect });
+      return null;
     default:
+      logNever(dialect);
       return null;
   }
 }

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

@@ -17,14 +17,52 @@ import {
   DuckDBDialect,
 } from "@marimo-team/codemirror-sql/dialects";
 import type { DataSourceConnection } from "@/core/kernel/messages";
+import { logNever } from "@/utils/assertNever";
+import { Logger } from "@/utils/Logger";
+
+const KNOWN_DIALECTS_ARRAY = [
+  "postgresql",
+  "postgres",
+  "db2",
+  "mysql",
+  "sqlite",
+  "mssql",
+  "sqlserver",
+  "duckdb",
+  "mariadb",
+  "cassandra",
+  "noql",
+  "athena",
+  "bigquery",
+  "hive",
+  "redshift",
+  "snowflake",
+  "flink",
+  "mongodb",
+  "oracle",
+  "oracledb",
+  "timescaledb",
+] as const;
+const KNOWN_DIALECTS: ReadonlySet<string> = new Set(KNOWN_DIALECTS_ARRAY);
+type KnownDialect = (typeof KNOWN_DIALECTS_ARRAY)[number];
+
+export function isKnownDialect(dialect: string): dialect is KnownDialect {
+  return KNOWN_DIALECTS.has(dialect);
+}
 
 /**
  * Guess the CodeMirror SQL dialect from the backend connection dialect.
+ * If unknown, return the standard SQL dialect.
  */
 export function guessDialect(
   connection: Pick<DataSourceConnection, "dialect">,
-): SQLDialect | undefined {
-  switch (connection.dialect) {
+): SQLDialect {
+  const dialect = connection.dialect;
+  if (!isKnownDialect(dialect)) {
+    return ModifiedStandardSQL;
+  }
+
+  switch (dialect) {
     case "postgresql":
     case "postgres":
       return PostgreSQL;
@@ -46,8 +84,21 @@ export function guessDialect(
       return PLSQL;
     case "bigquery":
       return BigQueryDialect;
+    case "timescaledb":
+      return PostgreSQL; // TimescaleDB is a PostgreSQL dialect
+    case "athena":
+    case "db2":
+    case "hive":
+    case "redshift":
+    case "snowflake":
+    case "flink":
+    case "mongodb":
+    case "noql":
+      Logger.debug("Unsupported dialect", { dialect });
+      return ModifiedStandardSQL;
     default:
-      return undefined;
+      logNever(dialect);
+      return ModifiedStandardSQL;
   }
 }
 

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

@@ -99,9 +99,17 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     CellId,
     EditorView | null | undefined
   >;
-
+  private readonly initialSettings: Record<string, unknown>;
   private static readonly SEEN_CELL_DOCUMENT_URIS = new Set<CellDocumentUri>();
 
+  /**
+   * Cache of completion items to avoid jitter while typing in the same completion item
+   */
+  private readonly completionItemCache = new LRUCache<
+    string,
+    Promise<LSP.CompletionItem>
+  >(10);
+
   constructor(
     client: ILanguageServerClient,
     initialSettings: Record<string, unknown>,
@@ -112,7 +120,7 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   ) {
     this.documentUri = getLSPDocument();
     this.getNotebookEditors = getNotebookEditors;
-
+    this.initialSettings = initialSettings;
     this.client = client;
     this.patchProcessNotification();
 
@@ -177,6 +185,37 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     this.client.close();
   }
 
+  /**
+   * Re-synchronize all open documents with the LSP server.
+   * This is called after a WebSocket reconnection to restore document state.
+   */
+  public async resyncAllDocuments(): Promise<void> {
+    invariant(
+      isClientWithNotify(this.client),
+      "notify is not a method on the client",
+    );
+    await this.client.initialize();
+    this.client.notify("workspace/didChangeConfiguration", {
+      settings: this.initialSettings,
+    });
+
+    // Get the current document state
+    const { lens, version } = this.snapshotter.snapshot();
+
+    // Re-open the merged document with the LSP server
+    // This sends a textDocument/didOpen for the entire notebook
+    await this.client.textDocumentDidOpen({
+      textDocument: {
+        languageId: "python", // Default to Python for marimo notebooks
+        text: lens.mergedText,
+        uri: this.documentUri,
+        version: version,
+      },
+    });
+
+    Logger.log("[lsp] Document re-synchronization complete");
+  }
+
   private getNotebookCode() {
     return store.get(topologicalCodesAtom);
   }
@@ -428,10 +467,19 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     };
   }
 
-  completionItemResolve(
+  async completionItemResolve(
     params: LSP.CompletionItem,
   ): Promise<LSP.CompletionItem> {
-    return this.client.completionItemResolve(params);
+    // Used cached result to avoid jitter while typing in the same completion item
+    const key = JSON.stringify(params);
+    const cached = this.completionItemCache.get(key);
+    if (cached) {
+      return cached;
+    }
+
+    const resolved = this.client.completionItemResolve(params);
+    this.completionItemCache.set(key, resolved);
+    return resolved;
   }
 
   async textDocumentPrepareRename(