@marimo-team/frontend 0.23.10 → 0.23.11-dev10

package/src/components/datasources/utils.ts

@@ -4,15 +4,54 @@ import { BigQueryDialect } from "@marimo-team/codemirror-sql/dialects";
 import { isKnownDialect } from "@/core/codemirror/language/languages/sql/utils";
 import type { SQLTableContext } from "@/core/datasets/data-source-connections";
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
-import type { DataTable, DataType } from "@/core/kernel/messages";
+import type {
+  Database,
+  DatabaseSchema,
+  DataTable,
+  DataType,
+} from "@/core/kernel/messages";
 import { logNever } from "@/utils/assertNever";
 import type { ColumnHeaderStatsKey } from "../data-table/types";
 
+/**
+ * Stable id for a table node in the datasources tree.
+ *
+ * schemaPath already includes the leaf schema for nested namespaces, so use it
+ * when present and fall back to the flat schema name otherwise (avoids
+ * duplicating the leaf, e.g. `top.nested.nested.table`).
+ */
+export function tableUniqueId(
+  sqlTableContext: SQLTableContext | undefined,
+  tableName: string,
+): string {
+  if (!sqlTableContext) {
+    return tableName;
+  }
+  const segments = (
+    sqlTableContext.schemaPath?.length
+      ? sqlTableContext.schemaPath
+      : [sqlTableContext.schema]
+  ).filter(Boolean);
+  return [sqlTableContext.database, ...segments, tableName].join(".");
+}
+
 // Some databases have no schemas, so we don't show it (eg. Clickhouse)
 export function isSchemaless(schemaName: string) {
   return schemaName === "";
 }
 
+// Lazy discovery: the `*_resolved` flags default to `true` and are only `false`
+// when enumeration was deferred. Helper functions to centralize logic
+export function areSchemasResolved(database: Database): boolean {
+  return database.schemas_resolved !== false;
+}
+export function areTablesResolved(schema: DatabaseSchema): boolean {
+  return schema.tables_resolved !== false;
+}
+export function areChildSchemasResolved(schema: DatabaseSchema): boolean {
+  return schema.child_schemas_resolved !== false;
+}
+
 interface SqlCodeFormatter {
   /**
    * Format the table path based on dialect-specific rules

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -27,7 +27,7 @@ import { Label } from "@/components/ui/label";
 import { Switch } from "@/components/ui/switch";
 import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
-import { stagedAICellsAtom } from "@/core/ai/staged-cells";
+import { stripWrappingBackticks } from "@/core/ai/strip-wrapping-backticks";
 import { type AiCompletionCell, includeOtherCellsAtom } from "@/core/ai/state";
 import type { CellId } from "@/core/cells/ids";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
@@ -45,6 +45,7 @@ import {
   RejectCompletionButton,
 } from "./completion-handlers";
 import { addContextCompletion, getAICompletionBody } from "./completion-utils";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
 
 const Original = CodeMirrorMerge.Original;
 const Modified = CodeMirrorMerge.Modified;
@@ -123,7 +124,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
     api: runtimeManager.getAiURL("completion").toString(),
     headers: runtimeManager.headers(),
     initialInput: initialPrompt,
-    streamProtocol: "text",
     // Throttle the messages and data updates to 100ms
     experimental_throttle: 100,
     body: {
@@ -143,13 +143,14 @@ export const AiCompletionEditor: React.FC<Props> = ({
       });
     },
     onFinish: (_prompt, completion) => {
-      // Remove trailing new lines
-      setCompletion(completion.trimEnd());
+      setCompletion(stripWrappingBackticks(completion).trimEnd());
     },
   });
 
   const inputRef = React.useRef<ReactCodeMirrorRef>(null);
-  const completion = untrimmedCompletion.trimEnd();
+  const completion = stripWrappingBackticks(untrimmedCompletion, {
+    streaming: isLoading,
+  }).trimEnd();
 
   const initialSubmit = useCallback(() => {
     if (triggerImmediately && !isLoading && initialPrompt) {

package/src/components/editor/connections/components.tsx

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { zodResolver } from "@hookform/resolvers/zod";
+import { useAtomValue } from "jotai";
 import React from "react";
 import { type DefaultValues, type FieldValues, useForm } from "react-hook-form";
 import type { z } from "zod";
@@ -16,8 +17,10 @@ import {
   SelectTrigger,
   SelectValue,
 } from "@/components/ui/select";
+import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
 import { useLastFocusedCellId } from "@/core/cells/focus";
+import { autoInstantiateAtom } from "@/core/config/config";
 import { ENV_RENDERER, SecretsProvider } from "./form-renderers";
 
 const RENDERERS: FormRenderer[] = [ENV_RENDERER];
@@ -106,8 +109,18 @@ export const ConnectionFormFooter = <L extends string>({
 export function useInsertCode() {
   const { createNewCell } = useCellActions();
   const lastFocusedCellId = useLastFocusedCellId();
+  const autoInstantiate = useAtomValue(autoInstantiateAtom);
 
   return (code: string) => {
+    // Ensure `mo` is importable when the generated code references it
+    if (/\bmo\./.test(code)) {
+      maybeAddMarimoImport({
+        autoInstantiate,
+        createNewCell,
+        fromCellId: lastFocusedCellId,
+      });
+    }
+
     createNewCell({
       code,
       before: false,

package/src/components/editor/connections/storage/__tests__/__snapshots__/as-code.test.ts.snap

@@ -85,14 +85,14 @@ store = GCSStore("my-bucket")"
 exports[`generateStorageCode > Google Drive > with default auth (no credentials) 1`] = `
 "from gdrive_fsspec import GoogleDriveFileSystem
 
-fs = GoogleDriveFileSystem(use_listings_cache=False)"
+fs = GoogleDriveFileSystem(use_listings_cache=False, skip_instance_cache=True)"
 `;
 
 exports[`generateStorageCode > Google Drive > with embedded auth (no credentials) 1`] = `
 "from gdrive_fsspec import GoogleDriveFileSystem
 
-fs = GoogleDriveFileSystem(use_listings_cache=False, auth_kwargs={"use_local_webserver": False})
-print("Google Drive connected! Important: Run this cell again to clear the console")"
+fs = GoogleDriveFileSystem(use_listings_cache=False, skip_instance_cache=True, auth_kwargs={"use_local_webserver": False})
+mo.output.clear_console()"
 `;
 
 exports[`generateStorageCode > Google Drive > with service account credentials 1`] = `
@@ -100,7 +100,7 @@ exports[`generateStorageCode > Google Drive > with service account credentials 1
 import json
 
 _creds = json.loads("""{"type": "service_account", "client_email": "test@test.iam.gserviceaccount.com"}""")
-fs = GoogleDriveFileSystem(creds=_creds, token="service_account", use_listings_cache=False)"
+fs = GoogleDriveFileSystem(creds=_creds, token="service_account", use_listings_cache=False, skip_instance_cache=True)"
 `;
 
 exports[`generateStorageCode > S3 > basic connection with all fields 1`] = `

package/src/components/editor/connections/storage/as-code.ts

@@ -168,6 +168,10 @@ function generateGDriveCode(
   connection: Extract<StorageConnection, { type: "gdrive" }>,
   options: { secrets: SecretContainer; isEmbedded?: boolean },
 ): { imports: Set<string>; code: string } {
+  /**
+   * Skip instance cache True so you can create multiple connections which don't reference the same creds.
+   * Use listings cache False so we don't get stale reads.
+   */
   const { secrets, isEmbedded = false } = options;
   const imports = new Set(["from gdrive_fsspec import GoogleDriveFileSystem"]);
 
@@ -179,18 +183,21 @@ function generateGDriveCode(
     );
     const code = dedent(`
       _creds = json.loads("""${connection.credentials_json?.startsWith("ENV:") ? `{${creds}}` : connection.credentials_json}""")
-      fs = GoogleDriveFileSystem(creds=_creds, token="service_account", use_listings_cache=False)
+      fs = GoogleDriveFileSystem(creds=_creds, token="service_account", use_listings_cache=False, skip_instance_cache=True)
     `);
     return { imports, code };
   }
 
+  // In the iframe (embedded) flow we authenticate via the console-based OOB
+  // flow, which prints an auth URL and reads the code from stdin. Clear the
+  // console afterwards so the (single-use) auth code doesn't linger.
   const code = isEmbedded
     ? dedent(`
-        fs = GoogleDriveFileSystem(use_listings_cache=False, auth_kwargs={"use_local_webserver": False})
-        print("Google Drive connected! Important: Run this cell again to clear the console")
+        fs = GoogleDriveFileSystem(use_listings_cache=False, skip_instance_cache=True, auth_kwargs={"use_local_webserver": False})
+        mo.output.clear_console()
       `)
     : dedent(`
-        fs = GoogleDriveFileSystem(use_listings_cache=False)
+        fs = GoogleDriveFileSystem(use_listings_cache=False, skip_instance_cache=True)
       `);
   return { imports, code };
 }

package/src/core/ai/__tests__/strip-wrapping-backticks.test.ts

@@ -0,0 +1,133 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { stripWrappingBackticks } from "../strip-wrapping-backticks";
+
+// Cases aligned with `without_wrapping_backticks` on a complete string (single chunk).
+const CASES: [string[], string][] = [
+  [["Hello", " world", "!"], "Hello world!"],
+  [["```", "print('hello')", "print('world')"], "print('hello')print('world')"],
+  [
+    ["print('hello')", "print('world')", "```"],
+    "print('hello')print('world')```",
+  ],
+  [["```", "print('hello')", "```"], "print('hello')"],
+  [["Hello", " ``` ", "world"], "Hello ``` world"],
+  [["``", "`print('hello')", "``", "`"], "print('hello')"],
+  [["``", "`", "\n", "print('hello')", "\n", "``", "`"], "print('hello')"],
+  [
+    ["```\n", "print('hello')", "print('world')", "\n```"],
+    "print('hello')print('world')",
+  ],
+  [
+    ["```\nprint('hello')\n", "print('world')\n```"],
+    "print('hello')\nprint('world')",
+  ],
+  [
+    ["```\n", "def test():\n    ", "return True\n```"],
+    "def test():\n    return True",
+  ],
+  [[], ""],
+  [["```idk```"], "idk"],
+  [["Hello world"], "Hello world"],
+  [
+    ["```python\n", "def hello():\n    ", "print('world')\n```"],
+    "def hello():\n    print('world')",
+  ],
+  [
+    ["```python", "\ndef hello():\n    ", "print('world')\n```"],
+    "def hello():\n    print('world')",
+  ],
+  [
+    ["```sql", "SELECT * FROM table", " WHERE id = 1", "```"],
+    "SELECT * FROM table WHERE id = 1",
+  ],
+  [
+    ["```sql\n", "SELECT * FROM table\n", "WHERE id = 1\n```"],
+    "SELECT * FROM table\nWHERE id = 1",
+  ],
+  [["```", "print('hello')", "```  "], "print('hello')  "],
+  [["```python\n", "print('hello')", "\n``` "], "print('hello') "],
+  [["```", "code", "```\t\n"], "code\t\n"],
+  [["  ```", "code", "```"], "code"],
+  [[" ```python\n", "code", "```"], "code"],
+  [["\t```", "code", "```"], "code"],
+  [["\n", "\n", "```\n", "code", "\n```\n"], "code\n"],
+  [["\n", "\n", "```python\n", "code", "\n```\n"], "code\n"],
+  [["\n``", "`python\n", "code", "\n```\n"], "code\n"],
+  [["\n`", "`", "`python\n", "code", "\n```\n"], "code\n"],
+  [["```python ", "code", "```"], " code"],
+  [["```python\t", "code", "```"], "\tcode"],
+  [["```\n", "```"], ""],
+  [["```python\n", "```"], ""],
+  [["```\n", "code"], "code"],
+  [["```python\n", "code"], "code"],
+  [["code", "\n```"], "code\n```"],
+  [
+    ["```\n", "x = 1\n", "```\n", "```\n", "y = 2\n", "```"],
+    "x = 1\n```\n```\ny = 2",
+  ],
+  [["```python\n", "s = 'use `backticks`'\n", "```"], "s = 'use `backticks`'"],
+  [["``", "`\n", "code\n", "``", "`"], "code"],
+  [["```", "python\n", "code\n", "```"], "code"],
+  [["```", "code", "```"], "code"],
+  [["prefix ", "```\n", "code\n", "```"], "prefix ```\ncode\n```"],
+  [["```\n", "code\n", "```", " suffix"], "code\n``` suffix"],
+  [["```\n\n", "code\n\n", "```"], "\ncode\n"],
+  [["```\n", "  code  \n", "```"], "  code  "],
+  [["```\n", "\tcode\t\n", "```"], "\tcode\t"],
+  [
+    ["```\n", "x\n", "```\n", "```python\n", "y\n", "```"],
+    "x\n```\n```python\ny",
+  ],
+  [["```javascript\n", "console.log()\n", "```"], "javascript\nconsole.log()"],
+  [["```markdown\n", "# Title\n", "```"], "# Title"],
+];
+
+describe("stripWrappingBackticks", () => {
+  it.each(CASES)("strips fences for %j", (chunks, expected) => {
+    expect(stripWrappingBackticks(chunks.join(""))).toBe(expected);
+  });
+});
+
+// In streaming mode, the opening fence is stripped as soon as it is
+// unambiguous, but the closing fence is left untouched (it may not have
+// arrived yet, or a trailing "```" may be content).
+const STREAMING_CASES: [string, string][] = [
+  // No fence: passthrough.
+  ["import pandas as pd", "import pandas as pd"],
+  // Plain opening fence stripped once the first line is terminated.
+  ["```\n", ""],
+  ["```\ncode", "code"],
+  ["```\ncode\nmore", "code\nmore"],
+  // Opening fence stripped, closing fence intentionally kept while streaming.
+  ["```\ncode\n```", "code\n```"],
+  ["```python\ncode\n```", "code\n```"],
+  // Language fences stripped as soon as the full identifier is present.
+  ["```python", ""],
+  ["```python\n", ""],
+  ["```python\ncode", "code"],
+  ["```sql\nSELECT 1", "SELECT 1"],
+  ["```markdown\n# Title", "# Title"],
+  // Leading whitespace before the fence.
+  ["  ```python\ncode", "code"],
+  // Partial language identifiers are left intact until they resolve.
+  ["```", "```"],
+  ["```p", "```p"],
+  ["```py", "```py"],
+  ["```pyth", "```pyth"],
+  ["```s", "```s"],
+  ["```mark", "```mark"],
+  // First-line content that is not a known language prefix is stripped as a
+  // plain fence even before a newline arrives.
+  ["```x", "x"],
+  ["```x = 1", "x = 1"],
+  // Unsupported language identifiers are kept as content (matches final mode).
+  ["```json\ncode", "json\ncode"],
+];
+
+describe("stripWrappingBackticks (streaming)", () => {
+  it.each(STREAMING_CASES)("strips opening fence for %j", (text, expected) => {
+    expect(stripWrappingBackticks(text, { streaming: true })).toBe(expected);
+  });
+});

package/src/core/ai/stream-completion-text.ts

@@ -0,0 +1,48 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { consumeStream, parseJsonEventStream, uiMessageChunkSchema } from "ai";
+import { stripWrappingBackticks } from "./strip-wrapping-backticks";
+
+/**
+ * Read an AI SDK UI message stream response and return the assistant text.
+ */
+export async function streamCompletionText(
+  response: Response,
+): Promise<string> {
+  if (!response.ok) {
+    throw new Error(await response.text());
+  }
+
+  if (!response.body) {
+    throw new Error("Failed to get response body");
+  }
+
+  let result = "";
+
+  await consumeStream({
+    stream: parseJsonEventStream({
+      stream: response.body,
+      schema: uiMessageChunkSchema,
+    }).pipeThrough(
+      new TransformStream({
+        transform(part) {
+          if (!part.success) {
+            throw part.error;
+          }
+
+          const streamPart = part.value;
+          if (streamPart.type === "text-delta") {
+            result += streamPart.delta;
+          } else if (streamPart.type === "error") {
+            throw new Error(streamPart.errorText);
+          }
+        },
+      }),
+    ),
+    onError: (error) => {
+      throw error;
+    },
+  });
+
+  return stripWrappingBackticks(result);
+}

package/src/core/ai/strip-wrapping-backticks.ts

@@ -0,0 +1,88 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+const LANGS = ["python", "sql", "markdown"] as const;
+
+interface StripOptions {
+  /**
+   * When true, the text is treated as a partial stream:
+   * - the opening fence is only stripped once we can be sure of its language
+   *   identifier (e.g. "```py" is left untouched since it may still become
+   *   "```python"), and
+   * - the closing fence is left in place, since it may not have arrived yet (or
+   *   a trailing "```" may be part of the content).
+   */
+  streaming?: boolean;
+}
+
+/**
+ * Removes wrapping markdown code fences from a completion string.
+ */
+export function stripWrappingBackticks(
+  text: string,
+  opts: StripOptions = {},
+): string {
+  const { streaming = false } = opts;
+  const leadingWhitespace = text.match(/^\s*/)?.[0] ?? "";
+  const rest = text.slice(leadingWhitespace.length);
+
+  let body: string | null = null;
+  let strippedOpening = false;
+
+  for (const lang of LANGS) {
+    if (rest.startsWith(`\`\`\`${lang}`)) {
+      strippedOpening = true;
+      body = rest.slice(3 + lang.length);
+      if (body.startsWith("\n")) {
+        body = body.slice(1);
+      }
+      break;
+    }
+  }
+
+  if (!strippedOpening && rest.startsWith("```")) {
+    const afterFence = rest.slice(3);
+    if (streaming && isPartialLanguageFence(afterFence)) {
+      return text;
+    }
+    strippedOpening = true;
+    body = afterFence;
+    if (body.startsWith("\n")) {
+      body = body.slice(1);
+    }
+  }
+
+  if (!strippedOpening || body === null) {
+    return text;
+  }
+
+  // While streaming, the closing fence may not have arrived yet, so leave the
+  // body as-is once the opening fence is removed.
+  if (streaming) {
+    return body;
+  }
+
+  const strippedEnd = body.trimEnd();
+  const trailingSpace = body.slice(strippedEnd.length);
+
+  if (strippedEnd.endsWith("\n```")) {
+    return strippedEnd.slice(0, -4) + trailingSpace;
+  }
+  if (strippedEnd.endsWith("```")) {
+    return strippedEnd.slice(0, -3) + trailingSpace;
+  }
+
+  return body;
+}
+
+/**
+ * Returns true if the text after an opening "```" has no terminating newline
+ * yet and could still become a known language identifier.
+ */
+function isPartialLanguageFence(afterFence: string): boolean {
+  // Once the first line is terminated, we already know it is not a known
+  // language fence (those are handled above), so treat it as a plain fence.
+  if (afterFence.includes("\n")) {
+    return false;
+  }
+  return LANGS.some((lang) => lang.startsWith(afterFence));
+}

package/src/core/cells/__tests__/cells.test.ts

@@ -2920,6 +2920,39 @@ describe("cell reducer", () => {
       expect(state.untouchedNewCells.has(newCellId)).toBe(true);
     });
 
+    it("adds interactively-created hidden cell with boilerplate code to untouchedNewCells", () => {
+      // Markdown cells are created with hideCode and non-empty default code
+      // (e.g. `mo.md(r"""\n""")`). They are user-initiated, so their editor
+      // should be shown until first blur.
+      actions.createNewCell({
+        cellId: "__end__",
+        before: false,
+        code: 'mo.md(r"""\n""")',
+        hideCode: true,
+        autoFocus: true,
+      });
+
+      const newCellId =
+        state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1];
+      expect(state.untouchedNewCells.has(newCellId)).toBe(true);
+    });
+
+    it("does not add programmatically-created hidden cell with code to untouchedNewCells", () => {
+      // Cells created by the kernel (e.g. via code_mode) carry code and
+      // autoFocus=false; their hide_code must take effect immediately.
+      actions.createNewCell({
+        cellId: "__end__",
+        before: false,
+        code: "x = 1",
+        hideCode: true,
+        autoFocus: false,
+      });
+
+      const newCellId =
+        state.cellIds.inOrderIds[state.cellIds.inOrderIds.length - 1];
+      expect(state.untouchedNewCells.has(newCellId)).toBe(false);
+    });
+
     it("does not add cell to untouchedNewCells when hideCode is false", () => {
       actions.createNewCell({
         cellId: "__end__",

package/src/core/cells/cells.ts

@@ -276,7 +276,7 @@ const {
       },
       scrollKey: autoFocus ? newCellId : null,
       untouchedNewCells:
-        hideCode && !code
+        hideCode && autoFocus
           ? new Set([...state.untouchedNewCells, newCellId])
           : state.untouchedNewCells,
     };

package/src/core/codemirror/ai/request.ts

@@ -2,6 +2,7 @@
 
 import { waitForConnectionOpen } from "@/core/network/connection";
 import type { AiCompletionRequest } from "@/core/network/types";
+import { streamCompletionText } from "@/core/ai/stream-completion-text";
 import { getRuntimeManager } from "@/core/runtime/config";
 import type { LanguageAdapterType } from "../language/types";
 
@@ -47,20 +48,7 @@ ${opts.codeAfter}
 
   const firstLineIndent = opts.selection.match(/^\s*/)?.[0] || "";
 
-  const reader = response.body?.getReader();
-  if (!reader) {
-    throw new Error("Failed to get response reader");
-  }
-
-  let result = "";
-  // oxlint-disable-next-line no-constant-condition
-  while (true) {
-    const { done, value } = await reader.read();
-    if (done) {
-      break;
-    }
-    result += new TextDecoder().decode(value);
-  }
+  let result = await streamCompletionText(response);
 
   // Add back the indent if it was stripped, which can happen with
   // LLM responses