@marimo-team/frontend 0.23.9-dev3 → 0.23.9-dev30

package/src/components/editor/output/MarimoTracebackOutput.tsx

@@ -5,17 +5,13 @@ import { useAtomValue } from "jotai";
 import {
   BugPlayIcon,
   ChevronDown,
+  ChevronRight,
   CopyIcon,
   ExternalLinkIcon,
   MessageCircleIcon,
   SearchIcon,
 } from "lucide-react";
 import { type JSX, useState } from "react";
-import {
-  Accordion,
-  AccordionContent,
-  AccordionItem,
-} from "@/components/ui/accordion";
 import { Button } from "@/components/ui/button";
 import {
   DropdownMenu,
@@ -36,13 +32,17 @@ import { isWasm } from "@/core/wasm/utils";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import { sanitizeHtml } from "@/plugins/core/sanitize-html";
 import { copyToClipboard } from "@/utils/copy";
+import {
+  containsMangledLocal,
+  splitMangledLocals,
+} from "@/utils/local-variables";
 import {
   elementContainsMarimoCellFile,
   extractAllTracebackInfo,
   getTracebackInfo,
 } from "@/utils/traceback";
-import { cn } from "../../../utils/cn";
 import { AIFixButton } from "../errors/auto-fix";
+import { MangledSegments } from "../errors/mangled-local-chip";
 import { CellLinkTraceback } from "../links/cell-link";
 import type { OnRefactorWithAI } from "../Output";
 
@@ -52,8 +52,6 @@ interface Props {
   onRefactorWithAI?: OnRefactorWithAI;
 }
 
-const KEY = "item";
-
 /**
  * List of errors due to violations of Marimo semantics.
  */
@@ -64,9 +62,12 @@ export const MarimoTracebackOutput = ({
 }: Props): JSX.Element => {
   const htmlTraceback = renderHTML({
     html: traceback,
-    additionalReplacements: [replaceTracebackFilenames, replaceTracebackPrefix],
+    additionalReplacements: [
+      replaceTracebackFilenames,
+      replaceTracebackPrefix,
+      replaceMangledLocal,
+    ],
   });
-  const [expanded, setExpanded] = useState(true);
 
   const lastTracebackLine = lastLine(traceback);
   const aiEnabled = useAtomValue(aiEnabledAtom);
@@ -86,6 +87,8 @@ export const MarimoTracebackOutput = ({
 
   const showSearch = !isStaticNotebook();
 
+  const [isOpen, setIsOpen] = useState(true);
+
   const handleRefactorWithAI = (triggerImmediately: boolean) => {
     onRefactorWithAI?.({
       prompt: `My code gives the following error:\n\n${lastTracebackLine}`,
@@ -93,32 +96,29 @@ export const MarimoTracebackOutput = ({
     });
   };
 
-  const [error, errorMessage] = lastTracebackLine.split(":", 2);
-
   return (
     <div className="flex flex-col gap-2 min-w-full w-fit">
-      <Accordion type="single" collapsible={true} value={expanded ? KEY : ""}>
-        <AccordionItem value={KEY} className="border-none">
-          <div
-            className="flex gap-2 h-10 px-2 hover:bg-muted rounded-sm select-none items-center cursor-pointer transition-all"
-            onClick={() => setExpanded((prev) => !prev)}
-          >
-            <ChevronDown
-              className={cn(
-                "h-4 w-4 text-muted-foreground transition-transform duration-200 shrink-0",
-                expanded ? "rotate-180" : "rotate-0",
-              )}
-            />
-            <div className="text-sm inline font-mono">
-              <span className="text-destructive">{error || "Error"}:</span>{" "}
-              {errorMessage}
-            </div>
-          </div>
-          <AccordionContent className="text-muted-foreground px-4 pt-2 text-xs overflow-auto">
-            {htmlTraceback}
-          </AccordionContent>
-        </AccordionItem>
-      </Accordion>
+      <button
+        type="button"
+        onClick={() => setIsOpen(!isOpen)}
+        aria-expanded={isOpen}
+        aria-label={isOpen ? "Collapse traceback" : "Expand traceback"}
+        className="self-start flex items-center gap-1 pt-2 text-muted-foreground/70 hover:text-muted-foreground transition-colors"
+      >
+        {isOpen ? (
+          <ChevronDown className="h-3 w-3" />
+        ) : (
+          <ChevronRight className="h-3 w-3" />
+        )}
+        <span className="text-[0.6875rem] uppercase tracking-wider">
+          Traceback
+        </span>
+      </button>
+      {isOpen && (
+        <div className="text-muted-foreground pr-4 text-xs overflow-auto">
+          {htmlTraceback}
+        </div>
+      )}
       <div className="flex gap-2">
         {showAIFix && (
           <AIFixButton
@@ -249,6 +249,23 @@ export const replaceTracebackFilenames = (domNode: DOMNode) => {
   }
 };
 
+/**
+ * Replace any cell-local mangled name (`_cell_<id>_<name>`) inside a text
+ * node with a {@link MangledLocalChip}. The mangled name appears in both
+ * the final `NameError:` line and inside compiled-cell source lines because
+ * the compiler rewrites underscore-prefixed references at AST-visit time.
+ */
+export const replaceMangledLocal = (domNode: DOMNode) => {
+  if (!(domNode instanceof Text) || !domNode.nodeValue) {
+    return;
+  }
+  if (!containsMangledLocal(domNode.nodeValue)) {
+    return;
+  }
+  const segments = splitMangledLocals(domNode.nodeValue);
+  return <MangledSegments segments={segments} />;
+};
+
 export const replaceTracebackPrefix = (domNode: DOMNode) => {
   if (
     domNode instanceof Text &&

package/src/components/ui/number-field.tsx

@@ -15,10 +15,17 @@ import { maxFractionalDigits } from "@/utils/numbers";
 export interface NumberFieldProps extends AriaNumberFieldProps {
   placeholder?: string;
   variant?: "default" | "xs";
+  /**
+   * Fires on every keystroke with the raw text currently in the input.
+   * React-aria's `onChange` only fires on commit (blur/Enter/stepper); use this
+   * when callers need to peek at the in-progress value, e.g. to enable an Apply
+   * button while the user is still typing.
+   */
+  onInputText?: (text: string) => void;
 }
 
 export const NumberField = React.forwardRef<HTMLInputElement, NumberFieldProps>(
-  ({ placeholder, variant = "default", ...props }, ref) => {
+  ({ placeholder, variant = "default", onInputText, ...props }, ref) => {
     const { locale } = useLocale();
     return (
       <AriaNumberField
@@ -48,6 +55,11 @@ export const NumberField = React.forwardRef<HTMLInputElement, NumberFieldProps>(
                 e.stopPropagation();
               }
             }}
+            onInput={
+              onInputText
+                ? (e) => onInputText(e.currentTarget.value)
+                : undefined
+            }
             className={cn(
               "flex-1",
               "w-full",

package/src/core/ai/__tests__/model-registry.test.ts

@@ -1,46 +1,52 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
-// Mock the models.json import
 vi.mock("@marimo-team/llm-info/models.json", () => {
-  const models: AiModel[] = [
-    {
-      name: "GPT-4",
-      model: "gpt-4",
-      description: "OpenAI GPT-4 model",
-      providers: ["openai"],
-      roles: ["chat", "edit"],
-      thinking: false,
-    },
-    {
-      name: "Claude 3",
-      model: "claude-3-sonnet",
-      description: "Anthropic Claude 3 Sonnet",
-      providers: ["anthropic"],
-      roles: ["chat", "edit"],
-      thinking: false,
-    },
-    {
-      name: "Gemini Pro",
-      model: "gemini-pro",
-      description: "Google Gemini Pro model",
-      providers: ["google"],
-      roles: ["chat", "edit"],
-      thinking: false,
-    },
-    {
-      name: "Multi Provider Model",
-      model: "multi-model",
-      description: "Model available on multiple providers",
-      providers: ["openai", "anthropic"],
-      roles: ["chat", "edit"],
-      thinking: false,
-    },
-  ];
-
-  return {
-    models: models,
+  const make = (
+    overrides: Partial<AiModel> & Pick<AiModel, "name" | "model">,
+  ): AiModel => ({
+    description: "",
+    roles: ["chat", "edit"],
+    capabilities: [],
+    input_types: [],
+    output_types: [],
+    release_date: "1970-01-01",
+    ...overrides,
+  });
+
+  const multiModel = make({
+    name: "Multi Provider Model",
+    model: "multi-model",
+    description: "Model available on multiple providers",
+  });
+
+  const models: Record<string, AiModel[]> = {
+    openai: [
+      make({
+        name: "GPT-4",
+        model: "gpt-4",
+        description: "OpenAI GPT-4 model",
+      }),
+      multiModel,
+    ],
+    anthropic: [
+      make({
+        name: "Claude 3",
+        model: "claude-3-sonnet",
+        description: "Anthropic Claude 3 Sonnet",
+      }),
+      multiModel,
+    ],
+    google: [
+      make({
+        name: "Gemini Pro",
+        model: "gemini-pro",
+        description: "Google Gemini Pro model",
+      }),
+    ],
   };
+
+  return { models };
 });
 
 import type { AiModel } from "@marimo-team/llm-info";
@@ -107,14 +113,15 @@ describe("AiModelRegistry", () => {
       });
 
       const ids = [...registry.getModelsMap().keys()];
-      // Include custom and all default ones.
+      // Include custom and all default ones; iteration follows provider
+      // sections in the source data (openai → anthropic → google).
       expect(ids).toEqual([
         "openai/custom-gpt",
         "openai/gpt-4",
-        "anthropic/claude-3-sonnet",
-        "google/gemini-pro",
         "openai/multi-model",
+        "anthropic/claude-3-sonnet",
         "anthropic/multi-model",
+        "google/gemini-pro",
       ]);
     });
   });
@@ -125,9 +132,9 @@ describe("AiModelRegistry", () => {
       const openaiModels = registry.getModelsByProvider("openai");
 
       expect(openaiModels).toHaveLength(2); // gpt-4 and multi-model
-      expect(
-        openaiModels.every((model) => model.providers.includes("openai")),
-      ).toBe(true);
+      expect(openaiModels.every((model) => model.provider === "openai")).toBe(
+        true,
+      );
     });
 
     it("should return empty array for provider with no models", () => {
@@ -147,9 +154,9 @@ describe("AiModelRegistry", () => {
       expect(customModel?.name).toBe("custom-gpt");
       expect(customModel?.model).toBe("custom-gpt");
       expect(customModel?.description).toBe("Custom model");
-      expect(customModel?.providers).toEqual(["openai"]);
+      expect(customModel?.provider).toBe("openai");
       expect(customModel?.roles).toEqual([]);
-      expect(customModel?.thinking).toBe(false);
+      expect(customModel?.capabilities).toEqual([]);
     });
 
     it("should filter models based on displayed models", () => {
@@ -309,7 +316,10 @@ describe("AiModelRegistry", () => {
 
       expect(multiModelInOpenai).toBeDefined();
       expect(multiModelInAnthropic).toBeDefined();
-      expect(multiModelInOpenai).toEqual(multiModelInAnthropic);
+      // Same model id, but each entry belongs to its own provider.
+      expect(multiModelInOpenai?.provider).toBe("openai");
+      expect(multiModelInAnthropic?.provider).toBe("anthropic");
+      expect(multiModelInOpenai?.name).toBe(multiModelInAnthropic?.name);
     });
 
     it("should handle displayed models filter with non-existent models", () => {
@@ -339,20 +349,20 @@ describe("AiModelRegistry", () => {
           expect(model).toHaveProperty("name");
           expect(model).toHaveProperty("model");
           expect(model).toHaveProperty("description");
-          expect(model).toHaveProperty("providers");
+          expect(model).toHaveProperty("provider");
           expect(model).toHaveProperty("roles");
-          expect(model).toHaveProperty("thinking");
+          expect(model).toHaveProperty("capabilities");
           expect(model).toHaveProperty("custom");
 
           expect(typeof model.name).toBe("string");
           expect(typeof model.model).toBe("string");
           expect(typeof model.description).toBe("string");
-          expect(Array.isArray(model.providers)).toBe(true);
+          expect(typeof model.provider).toBe("string");
           expect(Array.isArray(model.roles)).toBe(true);
-          expect(typeof model.thinking).toBe("boolean");
+          expect(Array.isArray(model.capabilities)).toBe(true);
           expect(typeof model.custom).toBe("boolean");
 
-          expect(model.providers).toContain(provider);
+          expect(model.provider).toBe(provider);
         }
       }
     });
@@ -367,31 +377,33 @@ describe("AiModelRegistry", () => {
 
       expect(customModel).toMatchInlineSnapshot(`
         {
+          "capabilities": [],
           "custom": true,
           "description": "Custom model",
+          "input_types": [],
           "model": "custom-gpt",
           "name": "custom-gpt",
-          "providers": [
-            "openai",
-          ],
+          "output_types": [],
+          "provider": "openai",
+          "release_date": "1970-01-01",
           "roles": [],
-          "thinking": false,
         }
       `);
       expect(defaultModel).toMatchInlineSnapshot(`
         {
+          "capabilities": [],
           "custom": false,
           "description": "OpenAI GPT-4 model",
+          "input_types": [],
           "model": "gpt-4",
           "name": "GPT-4",
-          "providers": [
-            "openai",
-          ],
+          "output_types": [],
+          "provider": "openai",
+          "release_date": "1970-01-01",
           "roles": [
             "chat",
             "edit",
           ],
-          "thinking": false,
         }
       `);
     });

package/src/core/ai/model-registry.ts

@@ -1,12 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type {
-  AiModel as AiModelType,
-  AiProvider,
-  Role,
-} from "@marimo-team/llm-info";
-import { models } from "@marimo-team/llm-info/models.json";
-import { providers } from "@marimo-team/llm-info/providers.json";
+import type { AiModel as AiModelType, AiProvider } from "@marimo-team/llm-info";
+import { models as modelsJson } from "@marimo-team/llm-info/models.json";
+import { providers as providersJson } from "@marimo-team/llm-info/providers.json";
 import { Logger } from "@/utils/Logger";
 import { MultiMap } from "@/utils/multi-map";
 import { once } from "@/utils/once";
@@ -14,13 +10,19 @@ import type { ProviderId } from "./ids/ids";
 import { AiModelId, type QualifiedModelId, type ShortModelId } from "./ids/ids";
 
 export interface AiModel extends AiModelType {
-  roles: Role[];
   model: ShortModelId;
-  providers: ProviderId[];
+  /** The provider this entry belongs to. */
+  provider: ProviderId;
   /** Whether this is a custom model. */
   custom: boolean;
 }
 
+// JSON shape matches the `AiModel` schema (Zod-validated at codegen time).
+const models = modelsJson as unknown as Partial<
+  Record<ProviderId, AiModelType[]>
+>;
+const providers = providersJson as unknown as readonly AiProvider[];
+
 interface KnownModelMaps {
   /** Map of qualified model ID to model info */
   modelMap: ReadonlyMap<QualifiedModelId, AiModel>;
@@ -32,24 +34,25 @@ export const getKnownModelMaps = once((): KnownModelMaps => {
   const modelMap = new Map<QualifiedModelId, AiModel>();
   const defaultModelByProvider = new Map<ProviderId, QualifiedModelId>();
 
-  for (const model of models) {
-    const modelId = model.model as ShortModelId;
-    const modelInfo: AiModel = {
-      ...model,
-      model: model.model as ShortModelId,
-      roles: model.roles.map((role) => role as Role),
-      providers: model.providers as ProviderId[],
-      custom: false,
-    };
-
-    const supportsChatOrEdit =
-      modelInfo.roles.includes("chat") || modelInfo.roles.includes("edit");
+  for (const [providerKey, providerModels] of Object.entries(models)) {
+    if (!providerModels) {
+      continue;
+    }
+    const provider = providerKey as ProviderId;
+    for (const raw of providerModels) {
+      const modelId = raw.model as ShortModelId;
+      const modelInfo: AiModel = {
+        ...raw,
+        model: modelId,
+        provider,
+        custom: false,
+      };
 
-    for (const provider of modelInfo.providers) {
       const qualifiedModelId: QualifiedModelId = `${provider}/${modelId}`;
       modelMap.set(qualifiedModelId, modelInfo);
 
-      // Track first model per provider that supports chat or edit
+      const supportsChatOrEdit =
+        modelInfo.roles.includes("chat") || modelInfo.roles.includes("edit");
       if (supportsChatOrEdit && !defaultModelByProvider.has(provider)) {
         defaultModelByProvider.set(provider, qualifiedModelId);
       }
@@ -67,9 +70,8 @@ const getProviderMap = once(
     const providerMap = new Map<ProviderId, AiProvider>();
     const providerToOrderIdx = new Map<ProviderId, number>();
     providers.forEach((provider, idx) => {
-      const providerId = provider.id as ProviderId;
-      providerMap.set(providerId, provider);
-      providerToOrderIdx.set(providerId, idx);
+      providerMap.set(provider.id, provider);
+      providerToOrderIdx.set(provider.id, idx);
     });
     return { providerMap, providerToOrderIdx };
   },
@@ -158,9 +160,12 @@ export class AiModelRegistry {
         name: modelId.shortModelId,
         model: modelId.shortModelId,
         description: "Custom model",
-        providers: [modelId.providerId],
+        provider: modelId.providerId,
         roles: [],
-        thinking: false,
+        capabilities: [],
+        input_types: [],
+        output_types: [],
+        release_date: "1970-01-01",
         custom: true,
       };
       customModelsMap.set(model, modelInfo);

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

@@ -0,0 +1,48 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
+import { notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { createCellRuntimeState } from "@/core/cells/types";
+import { store } from "@/core/state/jotai";
+import { notebookScrollToRunning } from "../actions";
+
+vi.mock("@/components/editor/links/cell-link", () => ({
+  scrollAndHighlightCell: vi.fn(),
+}));
+
+describe("notebookScrollToRunning", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    store.set(notebookAtom, MockNotebook.notebookState({ cellData: {} }));
+  });
+
+  it("scrolls to the first running cell in notebook order", () => {
+    const runtimeOnlyCellId = "runtime-only" as CellId;
+    const idleCellId = "idle-cell" as CellId;
+    const runningCellId = "running-cell" as CellId;
+
+    const notebook = MockNotebook.notebookState({
+      cellData: {
+        [idleCellId]: {},
+        [runningCellId]: {},
+      },
+      cellRuntime: {
+        [idleCellId]: { status: "idle" },
+        [runningCellId]: { status: "running" },
+      },
+    });
+    notebook.cellRuntime = {
+      [runtimeOnlyCellId]: createCellRuntimeState({ status: "running" }),
+      ...notebook.cellRuntime,
+    };
+    store.set(notebookAtom, notebook);
+
+    notebookScrollToRunning();
+
+    expect(scrollAndHighlightCell).toHaveBeenCalledOnce();
+    expect(scrollAndHighlightCell).toHaveBeenCalledWith(runningCellId, "focus");
+  });
+});

package/src/core/cells/actions.ts

@@ -1,7 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
-import { Objects } from "@/utils/objects";
 import { store } from "../state/jotai";
 import { notebookAtom } from "./cells";
 
@@ -10,12 +9,12 @@ import { notebookAtom } from "./cells";
  */
 export function notebookScrollToRunning() {
   // find cell that is currently in "running" state
-  const { cellRuntime } = store.get(notebookAtom);
-  const cell = Objects.entries(cellRuntime).find(
-    ([cellid, runtimestate]) => runtimestate.status === "running",
+  const { cellIds, cellRuntime } = store.get(notebookAtom);
+  const cellId = cellIds.inOrderIds.find(
+    (id) => cellRuntime[id]?.status === "running",
   );
-  if (!cell) {
+  if (!cellId) {
     return;
   }
-  scrollAndHighlightCell(cell[0], "focus");
+  scrollAndHighlightCell(cellId, "focus");
 }

package/src/core/codemirror/__tests__/setup.test.ts

@@ -134,6 +134,35 @@ describe("snapshot all duplicate keymaps", () => {
   });
 });
 
+test("auto_close_pairs: false removes closeBrackets keymaps", () => {
+  const withAutoClose = EditorState.create({
+    extensions: setup(),
+  });
+  const withoutAutoClose = EditorState.create({
+    extensions: setup({
+      completionConfig: {
+        ...getOpts().completionConfig,
+        auto_close_pairs: false,
+      },
+    }),
+  });
+
+  const keysWith = withAutoClose.facet(keymap).flat();
+  const keysWithout = withoutAutoClose.facet(keymap).flat();
+
+  // closeBracketsKeymap contributes Backspace and Enter handlers
+  expect(keysWith.length).toBeGreaterThan(keysWithout.length);
+
+  const hasBracketPairHandler = (state: EditorState) =>
+    state
+      .facet(keymap)
+      .flat()
+      .some((k) => k.run?.name === "deleteBracketPair");
+
+  expect(hasBracketPairHandler(withAutoClose)).toBe(true);
+  expect(hasBracketPairHandler(withoutAutoClose)).toBe(false);
+});
+
 test("placeholder adds another extension", () => {
   const opts = getOpts();
   const withAI = new PythonLanguageAdapter()

package/src/core/codemirror/cells/traceback-decorations.ts

@@ -147,7 +147,7 @@ export function errorLineHighlighter(
         backgroundColor: "color-mix(in srgb, var(--red-4) 40%, transparent)",
       },
       "&.cm-focused .cm-error-line.cm-activeLine": {
-        backgroundColor: "color-mix(in srgb, var(--red-6) 40%, transparent)",
+        backgroundColor: "color-mix(in srgb, var(--red-5) 40%, transparent)",
       },
     }),
   ];

package/src/core/codemirror/cm.ts

@@ -182,6 +182,7 @@ export const basicBundle = (opts: CodeMirrorSetupOpts): Extension[] => {
     diagnosticsConfig,
   } = opts;
   const placeholderType = getPlaceholderType(opts);
+  const autoClosePairs = completionConfig.auto_close_pairs !== false;
 
   return [
     ///// View
@@ -208,10 +209,10 @@ export const basicBundle = (opts: CodeMirrorSetupOpts): Extension[] => {
     copilotBundle(completionConfig),
     foldGutter(),
     stringsAutoCloseBraces(),
-    closeBrackets(),
+    autoClosePairs ? closeBrackets() : [],
     completionKeymap(acceptCompletionOnEnter),
     // to avoid clash with charDeleteBackward keymap
-    Prec.high(keymap.of(closeBracketsKeymap)),
+    autoClosePairs ? Prec.high(keymap.of(closeBracketsKeymap)) : [],
     bracketMatching(),
     indentOnInput(),
     indentUnit.of("    "),

package/src/core/codemirror/format.ts

@@ -205,6 +205,7 @@ async function getSqlFormatterDialect(
     case "mongodb":
     case "timescaledb":
     case "datafusion":
+    case "dremio":
       return sql;
     case "databricks":
       return spark;