@marimo-team/frontend 0.23.10 → 0.23.11-dev2

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/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/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

package/src/core/datasets/__tests__/data-source.test.ts

@@ -2,9 +2,12 @@
 import { beforeEach, describe, expect, it } from "vitest";
 import { variableName } from "@/__tests__/branded";
 import type { DatabaseSchema, DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
 import type { VariableName } from "@/core/variables/types";
 import {
+  allTablesAtom,
   type DataSourceConnection,
+  dataSourceConnectionsAtom,
   type DataSourceState,
   exportedForTesting,
   type SQLTableContext,
@@ -575,3 +578,226 @@ describe("add table", () => {
     expect(db1?.schemas.length).toBe(1);
   });
 });
+
+describe("nested namespaces", () => {
+  // Iceberg-style: database "top" with a schemaless schema and a nested
+  // namespace "nested" that has a deferred child namespace "deep".
+  const nestedConnections: DataSourceConnection[] = [
+    {
+      name: "ice" as ConnectionName,
+      source: "iceberg",
+      display_name: "Iceberg",
+      dialect: "iceberg",
+      databases: [
+        {
+          name: "top",
+          dialect: "iceberg",
+          schemas_resolved: true,
+          schemas: [
+            { name: "", tables: [], tables_resolved: true },
+            {
+              name: "nested",
+              tables: [],
+              tables_resolved: false,
+              child_schemas: [],
+              child_schemas_resolved: false,
+            },
+          ],
+        },
+      ],
+    },
+  ];
+
+  let baseState: DataSourceState;
+
+  beforeEach(() => {
+    baseState = addConnection(nestedConnections, initialState());
+  });
+
+  const findSchema = (
+    state: DataSourceState,
+    path: string[],
+  ): DatabaseSchema | undefined => {
+    const conn = state.connectionsMap.get("ice" as ConnectionName);
+    const db = conn?.databases.find((d) => d.name === "top");
+    let schemas = db?.schemas ?? [];
+    let found: DatabaseSchema | undefined;
+    for (const segment of path) {
+      found = schemas.find((s) => s.name === segment);
+      schemas = found?.child_schemas ?? [];
+    }
+    return found;
+  };
+
+  it("sets child namespaces at a nested path", () => {
+    const children: DatabaseSchema[] = [
+      { name: "deep", tables: [], tables_resolved: false },
+    ];
+    const newState = reducer(baseState, {
+      type: "addSchemaList",
+      payload: {
+        schemas: children,
+        sqlSchemaContext: {
+          engine: "ice",
+          database: "top",
+          schemaPath: ["nested"],
+        },
+      },
+    });
+
+    const nested = findSchema(newState, ["nested"]);
+    expect(nested?.child_schemas_resolved).toBe(true);
+    expect(nested?.child_schemas?.map((s) => s.name)).toEqual(["deep"]);
+    // The schemaless sibling is untouched.
+    expect(findSchema(newState, [""])?.name).toBe("");
+  });
+
+  it("sets tables at a nested path", () => {
+    const tables: DataTable[] = [
+      {
+        name: "table4",
+        columns: [],
+        num_columns: 0,
+        num_rows: 0,
+        variable_name: null,
+        source: "iceberg",
+        source_type: "catalog",
+        type: "table",
+      },
+    ];
+    const newState = reducer(baseState, {
+      type: "addTableList",
+      payload: {
+        tables,
+        sqlTableContext: {
+          engine: "ice",
+          database: "top",
+          schema: "nested",
+          dialect: "iceberg",
+          schemaPath: ["nested"],
+        },
+      },
+    });
+
+    const nested = findSchema(newState, ["nested"]);
+    expect(nested?.tables_resolved).toBe(true);
+    expect(nested?.tables.map((t) => t.name)).toEqual(["table4"]);
+  });
+
+  it("replaces a single table at a nested path", () => {
+    const makeTable = (numRows: number): DataTable => ({
+      name: "table4",
+      columns: [],
+      num_columns: 0,
+      num_rows: numRows,
+      variable_name: null,
+      source: "iceberg",
+      source_type: "catalog",
+      type: "table",
+    });
+    const context = {
+      engine: "ice",
+      database: "top",
+      schema: "nested",
+      dialect: "iceberg",
+      schemaPath: ["nested"],
+    };
+    let state = reducer(baseState, {
+      type: "addTableList",
+      payload: { tables: [makeTable(1)], sqlTableContext: context },
+    });
+    state = reducer(state, {
+      type: "addTable",
+      payload: { table: makeTable(42), sqlTableContext: context },
+    });
+
+    const nested = findSchema(state, ["nested"]);
+    expect(nested?.tables).toHaveLength(1);
+    expect(nested?.tables[0].num_rows).toBe(42);
+  });
+
+  it("does not change anything for a missing nested path", () => {
+    const newState = reducer(baseState, {
+      type: "addSchemaList",
+      payload: {
+        schemas: [{ name: "deep", tables: [] }],
+        sqlSchemaContext: {
+          engine: "ice",
+          database: "top",
+          schemaPath: ["does_not_exist"],
+        },
+      },
+    });
+    // The result is unchanged: nested namespace stays unresolved and the
+    // database keeps its two schemas (schemaless + nested).
+    const newDb = newState.connectionsMap
+      .get("ice" as ConnectionName)
+      ?.databases.find((d) => d.name === "top");
+    expect(findSchema(newState, ["nested"])?.child_schemas_resolved).toBe(
+      false,
+    );
+    expect(newDb?.schemas.length).toBe(2);
+  });
+});
+
+describe("allTablesAtom with nested namespaces", () => {
+  it("enumerates tables from nested namespaces", () => {
+    const table = (name: string): DataTable => ({
+      name,
+      columns: [],
+      num_columns: 0,
+      num_rows: 0,
+      variable_name: null,
+      source: "iceberg",
+      source_type: "catalog",
+      type: "table",
+    });
+
+    const state = addConnection(
+      [
+        {
+          name: "ice" as ConnectionName,
+          source: "iceberg",
+          display_name: "Iceberg",
+          dialect: "iceberg",
+          databases: [
+            {
+              name: "top",
+              dialect: "iceberg",
+              schemas_resolved: true,
+              schemas: [
+                {
+                  name: "",
+                  tables: [table("toptable")],
+                  tables_resolved: true,
+                },
+                {
+                  name: "nested",
+                  tables: [table("nestedtable")],
+                  tables_resolved: true,
+                  child_schemas_resolved: true,
+                  child_schemas: [
+                    {
+                      name: "deep",
+                      tables: [table("deeptable")],
+                      tables_resolved: true,
+                      child_schemas: [],
+                      child_schemas_resolved: true,
+                    },
+                  ],
+                },
+              ],
+            },
+          ],
+        },
+      ],
+      initialState(),
+    );
+
+    store.set(dataSourceConnectionsAtom, state);
+    const names = [...store.get(allTablesAtom).values()].map((t) => t.name);
+    expect(names).toContain("toptable");
+    expect(names).toContain("nestedtable");
+    expect(names).toContain("deeptable");
+  });
+});