@marimo-team/frontend 0.23.14-dev2 → 0.23.14-dev21

package/src/core/ai/context/providers/error.ts

@@ -1,26 +1,133 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { Completion } from "@codemirror/autocomplete";
-import { cellErrorsAtom } from "@/core/cells/cells";
+import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import type { MarimoError } from "@/core/kernel/messages";
+import { displayCellName } from "@/core/cells/names";
+import type { MarimoError, OutputMessage } from "@/core/kernel/messages";
+import { isErrorMime, isMarimoErrorsMime, isTracebackMime } from "@/core/mime";
 import type { JotaiStore } from "@/core/state/jotai";
 import { logNever } from "@/utils/assertNever";
+import { parseHtmlContent } from "@/utils/dom";
 import { PluralWord } from "@/utils/pluralize";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
 import { Sections } from "./common";
+import { formatDatasourceContextForCell } from "./datasource";
+
+export interface CellErrorEntry {
+  cellId: CellId;
+  cellName: string;
+  cellCode: string;
+  errorData: MarimoError[];
+  tracebackHtml?: string;
+}
+
+export type ErrorContextItemData =
+  | {
+      type: "all-errors";
+      errors: CellErrorEntry[];
+    }
+  | {
+      type: "cell-error";
+      error: CellErrorEntry;
+    };
 
 export interface ErrorContextItem extends AIContextItem {
   type: "error";
-  data: {
-    type: "all-errors";
-    errors: {
-      cellId: CellId;
-      cellName: string;
-      errorData: MarimoError[];
-    }[];
-  };
+  data: ErrorContextItemData;
+}
+
+function parseCellErrorOutput(
+  output: OutputMessage,
+): Pick<CellErrorEntry, "errorData" | "tracebackHtml"> | null {
+  if (isMarimoErrorsMime(output.mimetype)) {
+    if (!Array.isArray(output.data)) {
+      return null;
+    }
+    const errorData = output.data.filter(
+      (error) => !error.type.includes("ancestor"),
+    );
+    if (errorData.length === 0) {
+      return null;
+    }
+    return { errorData };
+  }
+
+  if (isTracebackMime(output.mimetype)) {
+    if (typeof output.data !== "string" || output.data.length === 0) {
+      return null;
+    }
+    return { errorData: [], tracebackHtml: output.data };
+  }
+
+  return null;
+}
+
+function errorDataHasTraceback(errorData: MarimoError[]): boolean {
+  return errorData.some((error) => "traceback" in error && error.traceback);
+}
+
+function getTracebackFromConsole(
+  consoleOutputs: OutputMessage[] | undefined,
+): string | undefined {
+  const tracebackOutput = consoleOutputs?.find((output) =>
+    isTracebackMime(output.mimetype),
+  );
+  if (
+    !tracebackOutput ||
+    typeof tracebackOutput.data !== "string" ||
+    tracebackOutput.data.length === 0
+  ) {
+    return undefined;
+  }
+  return tracebackOutput.data;
+}
+
+function resolveTracebackHtml(
+  parsed: Pick<CellErrorEntry, "errorData" | "tracebackHtml">,
+  consoleOutputs: OutputMessage[] | undefined,
+): string | undefined {
+  if (parsed.tracebackHtml) {
+    return parsed.tracebackHtml;
+  }
+  if (errorDataHasTraceback(parsed.errorData)) {
+    return undefined;
+  }
+  return getTracebackFromConsole(consoleOutputs);
+}
+
+export function getCellErrorEntries(store: JotaiStore): CellErrorEntry[] {
+  const { cellIds, cellRuntime, cellData } = store.get(notebookAtom);
+  const entries: CellErrorEntry[] = [];
+
+  for (const [cellIndex, cellId] of cellIds.inOrderIds.entries()) {
+    const cell = cellRuntime[cellId];
+    const output = cell.output;
+    if (!output || !isErrorMime(output.mimetype)) {
+      continue;
+    }
+
+    const parsed = parseCellErrorOutput(output);
+    if (!parsed) {
+      continue;
+    }
+
+    const cellName = displayCellName(cellData[cellId].name, cellIndex);
+    // Prefer the code from the last execution: runtime errors correspond to
+    // the previous run, while `code` may already contain unsaved edits.
+    const cellCode = cellData[cellId].lastCodeRun ?? cellData[cellId].code;
+
+    entries.push({
+      cellId,
+      cellName,
+      cellCode,
+      errorData: parsed.errorData,
+      tracebackHtml: resolveTracebackHtml(parsed, cell.consoleOutputs),
+    });
+  }
+
+  return entries;
 }
 
 function describeError(error: MarimoError): string {
@@ -67,6 +174,77 @@ function describeError(error: MarimoError): string {
   return "Unknown error";
 }
 
+function formatSingleError(error: MarimoError): string {
+  let text = describeError(error);
+  if ("traceback" in error && error.traceback) {
+    text += `\n\nTraceback:\n${parseHtmlContent(error.traceback)}`;
+  }
+  return text;
+}
+
+function formatCellErrorDetails(
+  entry: CellErrorEntry,
+  store: JotaiStore,
+): string {
+  const parts: string[] = [`Code:\n${entry.cellCode}`];
+
+  if (entry.tracebackHtml) {
+    parts.push(`Traceback:\n${parseHtmlContent(entry.tracebackHtml)}`);
+  }
+
+  if (entry.errorData.length > 0) {
+    parts.push(
+      entry.errorData.map((error) => formatSingleError(error)).join("\n\n"),
+    );
+  }
+
+  if (entry.errorData.some((error) => error.type === "sql-error")) {
+    const datasourceContext = formatDatasourceContextForCell(
+      entry.cellId,
+      store,
+    );
+    if (datasourceContext) {
+      parts.push(`Database schema:\n${datasourceContext}`);
+    }
+  }
+
+  return parts.join("\n\n");
+}
+
+function formatCellErrorXml(entry: CellErrorEntry, store: JotaiStore): string {
+  return contextToXml({
+    type: "error",
+    data: {
+      name: entry.cellName || `Cell ${entry.cellId}`,
+      cellId: entry.cellId,
+    },
+    details: formatCellErrorDetails(entry, store),
+  });
+}
+
+function summarizeCellError(entry: CellErrorEntry): string {
+  if (entry.errorData.length === 1) {
+    return describeError(entry.errorData[0]);
+  }
+  if (entry.errorData.length > 1) {
+    return `${entry.errorData.length} errors`;
+  }
+  if (entry.tracebackHtml) {
+    const text = parseHtmlContent(entry.tracebackHtml).trim();
+    const firstLine = text
+      .split("\n")
+      .find((line) => line.trim())
+      ?.trim();
+    return firstLine || "traceback";
+  }
+  return "error";
+}
+
+function errorContextName(entry: CellErrorEntry): string {
+  const cellName = entry.cellName || `Cell ${entry.cellId}`;
+  return `Error: ${cellName}`;
+}
+
 const errorsTxt = new PluralWord("error", "errors");
 export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
   readonly title = "Errors";
@@ -80,39 +258,50 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
   }
 
   getItems(): ErrorContextItem[] {
-    const errors = this.store.get(cellErrorsAtom);
+    const errors = getCellErrorEntries(this.store);
+
     if (errors.length === 0) {
       return [];
     }
 
-    return [
+    const items: ErrorContextItem[] = [
       {
         uri: this.asURI("all"),
         name: "Errors",
         type: this.contextType,
         data: {
           type: "all-errors",
-          errors: errors.map((error) => ({
-            cellId: error.cellId,
-            cellName: error.cellName,
-            errorData: error.output.data,
-          })),
+          errors,
         },
         description: "All errors in the notebook",
       },
     ];
 
-    // TODO: maybe handle single errors or grouped by types
+    for (const error of errors) {
+      items.push({
+        uri: this.asURI(error.cellId),
+        name: errorContextName(error),
+        type: this.contextType,
+        data: {
+          type: "cell-error",
+          error,
+        },
+        description: summarizeCellError(error),
+      });
+    }
+
+    return items;
   }
 
   formatCompletion(item: ErrorContextItem): Completion {
     if (item.data.type === "all-errors") {
+      const errorCount = item.data.errors.length;
       return {
         label: "@Errors",
         displayLabel: "Errors",
-        detail: `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`,
+        detail: `${errorCount} ${errorsTxt.pluralize(errorCount)}`,
         type: "error",
-        apply: "@Errors",
+        apply: "@error://all",
         section: Sections.ERROR,
         info: () => {
           const infoContainer = document.createElement("div");
@@ -136,7 +325,7 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
 
           const descriptionDiv = document.createElement("div");
           descriptionDiv.classList.add("text-sm", "text-muted-foreground");
-          descriptionDiv.textContent = `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`;
+          descriptionDiv.textContent = `${errorCount} ${errorsTxt.pluralize(errorCount)}`;
           headerDiv.append(descriptionDiv);
 
           infoContainer.append(headerDiv);
@@ -146,6 +335,17 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
       };
     }
 
+    if (item.data.type === "cell-error") {
+      const { error } = item.data;
+      return {
+        ...this.createBasicCompletion(item, {
+          detail: summarizeCellError(error),
+          type: "error",
+        }),
+        section: Sections.ERROR,
+      };
+    }
+
     return {
       label: "Error",
       displayLabel: "Error",
@@ -154,21 +354,11 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
   }
 
   formatContext(item: ErrorContextItem): string {
-    const { data } = item;
-
-    const xmls = data.errors.map((err) => {
-      return contextToXml({
-        type: this.contextType,
-        data: {
-          name: err.cellName || `Cell ${err.cellId}`,
-          description: err.errorData
-            .map((err) => describeError(err))
-            .join("\n"),
-        },
-      });
-    });
+    const entries =
+      item.data.type === "all-errors" ? item.data.errors : [item.data.error];
 
-    const xml = xmls.join("\n\n");
-    return xml;
+    return entries
+      .map((entry) => formatCellErrorXml(entry, this.store))
+      .join("\n\n");
   }
 }

package/src/core/ai/context/registry.ts

@@ -9,11 +9,19 @@ import { MultiMap } from "@/utils/multi-map";
 import type { TypedString } from "@/utils/typed";
 
 /**
- * Unique identifier for a context item in the format "type:id"
+ * Unique identifier for a context item in the format "type://id"
  * e.g., "variable://my_var", "data://users", "file://config.py"
  */
 export type ContextLocatorId = TypedString<"ContextLocatorId">;
 
+function parseContextType(uri: ContextLocatorId): string | undefined {
+  const separator = uri.indexOf("://");
+  if (separator === -1) {
+    return undefined;
+  }
+  return uri.slice(0, separator);
+}
+
 /**
  * Base interface for context items that can be mentioned in AI prompts
  */
@@ -142,40 +150,80 @@ export class AIContextRegistry<T extends AIContextItem> {
     );
   }
 
+  private findProviderForUri(
+    uri: ContextLocatorId,
+  ): AIContextProvider | undefined {
+    const type = parseContextType(uri);
+    if (!type) {
+      return undefined;
+    }
+    return [...this.providers].find(
+      (provider) =>
+        provider.contextType === type &&
+        provider.getItems().some((item) => item.uri === uri),
+    );
+  }
+
   /**
-   * Get context information for mentioned items
+   * Resolve only the requested context items, querying each matching provider
    */
-  getContextInfo(contextIds: ContextLocatorId[]): T[] {
-    const contextInfo: T[] = [];
-    const allItems = new Map<ContextLocatorId, T>(
-      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
-    );
+  resolveItems(contextIds: ContextLocatorId[]): T[] {
+    if (contextIds.length === 0) {
+      return [];
+    }
+
+    const idsByType = new MultiMap<string, ContextLocatorId>();
+    for (const contextId of contextIds) {
+      const type = parseContextType(contextId);
+      if (type) {
+        idsByType.add(type, contextId);
+      }
+    }
+
+    const itemsById = new Map<ContextLocatorId, T>();
+    for (const [type, ids] of idsByType.entries()) {
+      const providers = [...this.providers].filter(
+        (provider) => provider.contextType === type,
+      );
+      for (const provider of providers) {
+        const itemsByUri = new Map<ContextLocatorId, T>(
+          provider
+            .getItems()
+            .map((item) => [item.uri as ContextLocatorId, item as T]),
+        );
+        for (const id of ids) {
+          const item = itemsByUri.get(id);
+          if (item) {
+            itemsById.set(id, item);
+          }
+        }
+      }
+    }
 
+    // Preserve the order in which the ids were requested, so formatted context
+    // matches the order the user mentioned them in the prompt.
+    const results: T[] = [];
     for (const contextId of contextIds) {
-      const item = allItems.get(contextId);
+      const item = itemsById.get(contextId);
       if (item) {
-        contextInfo.push(item);
+        results.push(item);
       }
     }
+    return results;
+  }
 
-    return contextInfo;
+  /**
+   * Get context information for mentioned items
+   */
+  getContextInfo(contextIds: ContextLocatorId[]): T[] {
+    return this.resolveItems(contextIds);
   }
 
   /**
    * Format context for AI prompt inclusion
    */
   formatContextForAI(contextIds: ContextLocatorId[]): string {
-    const allItems = new Map<ContextLocatorId, T>(
-      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
-    );
-
-    const contextInfo: T[] = [];
-    for (const contextId of contextIds) {
-      const item = allItems.get(contextId);
-      if (item) {
-        contextInfo.push(item);
-      }
-    }
+    const contextInfo = this.resolveItems(contextIds);
 
     if (contextInfo.length === 0) {
       return "";
@@ -183,7 +231,7 @@ export class AIContextRegistry<T extends AIContextItem> {
 
     return contextInfo
       .map((item) => {
-        const provider = this.getProvider(item.type);
+        const provider = this.findProviderForUri(item.uri as ContextLocatorId);
         return provider?.formatContext(item) || "";
       })
       .join("\n\n");
@@ -195,36 +243,22 @@ export class AIContextRegistry<T extends AIContextItem> {
   async getAttachmentsForContext(
     contextIds: ContextLocatorId[],
   ): Promise<FileUIPart[]> {
-    const allItems = new Map<ContextLocatorId, T>(
-      this.getAllItems().map((item) => [item.uri as ContextLocatorId, item]),
-    );
-
-    const contextInfo: T[] = [];
-    for (const contextId of contextIds) {
-      const item = allItems.get(contextId);
-      if (item) {
-        contextInfo.push(item);
-      }
-    }
+    const contextInfo = this.resolveItems(contextIds);
 
     if (contextInfo.length === 0) {
       return [];
     }
 
-    // Group items by provider type to batch attachment requests
-    const itemsByProvider = new MultiMap<string, T>();
+    const itemsByProvider = new MultiMap<AIContextProvider, T>();
     for (const item of contextInfo) {
-      const providerType = item.type;
-      itemsByProvider.add(providerType, item);
+      const provider = this.findProviderForUri(item.uri as ContextLocatorId);
+      if (provider) {
+        itemsByProvider.add(provider, item);
+      }
     }
 
-    // Collect attachments from all providers
     const attachmentPromises = [...itemsByProvider.entries()].map(
-      async ([providerType, items]) => {
-        const provider = this.getProvider(providerType);
-        if (!provider) {
-          return [];
-        }
+      async ([provider, items]) => {
         try {
           return await provider.getAttachments(items);
         } catch (error) {

package/src/core/ai/state.ts

@@ -21,6 +21,17 @@ export interface AiCompletionCell {
 
 export const aiCompletionCellAtom = atom<AiCompletionCell | null>(null);
 
+/**
+ * A prompt queued to be delivered to the AI assistant panel.
+ * The active panel consumes this on mount/when ready, then clears it.
+ */
+export interface PendingAiPrompt {
+  prompt: string;
+  submit: boolean;
+}
+
+export const pendingAiPromptAtom = atom<PendingAiPrompt | null>(null);
+
 const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
 export const includeOtherCellsAtom = atomWithStorage<boolean>(
   INCLUDE_OTHER_CELLS_KEY,

package/src/core/cells/__tests__/readonly-code-display.test.ts

@@ -0,0 +1,46 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { getReadonlyCodeDisplay } from "../readonly-code-display";
+
+describe("getReadonlyCodeDisplay", () => {
+  it("unwraps SQL cells to their inner query and marks them as sql", () => {
+    const result = getReadonlyCodeDisplay(
+      'my_table = mo.sql("""SELECT 1 AS id""")',
+    );
+    expect(result.language).toBe("sql");
+    expect(result.code).toBe("SELECT 1 AS id");
+  });
+
+  it("unwraps f-string SQL cells", () => {
+    const result = getReadonlyCodeDisplay(
+      'my_table = mo.sql(f"""SELECT 1 AS id""")',
+    );
+    expect(result.language).toBe("sql");
+    expect(result.code).toBe("SELECT 1 AS id");
+  });
+
+  it("leaves plain Python cells untouched", () => {
+    const code = "x = 1 + 2\nprint(x)";
+    const result = getReadonlyCodeDisplay(code);
+    expect(result.language).toBe("python");
+    expect(result.code).toBe(code);
+  });
+
+  it("unwraps markdown cells to their inner content", () => {
+    const result = getReadonlyCodeDisplay('mo.md("""## Heading""")');
+    expect(result.language).toBe("markdown");
+    expect(result.code).toBe("## Heading");
+  });
+
+  it("treats empty or whitespace-only code as python", () => {
+    expect(getReadonlyCodeDisplay("")).toEqual({
+      code: "",
+      language: "python",
+    });
+    expect(getReadonlyCodeDisplay("   \n\t  ")).toEqual({
+      code: "   \n\t  ",
+      language: "python",
+    });
+  });
+});

package/src/core/cells/readonly-code-display.ts

@@ -0,0 +1,35 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { MarkdownParser, SQLParser } from "@marimo-team/smart-cells";
+import type { LanguageAdapterType } from "@/core/codemirror/language/types";
+
+export interface ReadonlyCodeDisplay {
+  code: string;
+  language: LanguageAdapterType;
+}
+
+const markdownParser = new MarkdownParser();
+const sqlParser = new SQLParser();
+
+/**
+ * Unwrap SQL and markdown cells so read-only views show inner content with the
+ * correct syntax highlighting instead of the raw Python wrapper.
+ */
+export function getReadonlyCodeDisplay(code: string): ReadonlyCodeDisplay {
+  const trimmed = code.trim();
+  if (!trimmed) {
+    return { code, language: "python" };
+  }
+
+  if (markdownParser.isSupported(trimmed)) {
+    return {
+      code: markdownParser.transformIn(trimmed).code,
+      language: "markdown",
+    };
+  }
+
+  if (sqlParser.isSupported(trimmed)) {
+    return { code: sqlParser.transformIn(trimmed).code, language: "sql" };
+  }
+  return { code, language: "python" };
+}

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

@@ -11,6 +11,7 @@ import {
   resourceDecorations,
   resourceInputFilter,
   resourcesField,
+  resourceSync,
   resourceTheme,
 } from "@marimo-team/codemirror-mcp";
 import {
@@ -38,25 +39,28 @@ const NONE_RESOURCE_FORMAT_COMPLETION = {
   },
 };
 
+function getRegistryResources(store: JotaiStore): Resource[] {
+  const registry = getAIContextRegistry(store);
+  const resources = registry.getAllItems();
+  if (resources.length === 0) {
+    return NONE_RESOURCE;
+  }
+  return resources;
+}
+
 export function resourceExtension(opts: {
   language: Language;
   store: JotaiStore;
   onAddFiles?: (files: File[]) => void;
 }): Extension[] {
   const { language, store, onAddFiles } = opts;
+  const getResources = () => getRegistryResources(store);
 
   return [
     language.data.of({
       // Resource completion for static resources (variables, tables, etc.)
       autocomplete: resourceCompletion(
-        async (): Promise<Resource[]> => {
-          const registry = getAIContextRegistry(store);
-          const resources = registry.getAllItems();
-          if (resources.length === 0) {
-            return NONE_RESOURCE;
-          }
-          return resources;
-        },
+        async (): Promise<Resource[]> => getResources(),
         (resource) => {
           if (resource.type === NONE_RESOURCE_TYPE) {
             return NONE_RESOURCE_FORMAT_COMPLETION;
@@ -92,9 +96,10 @@ export function resourceExtension(opts: {
       : []),
     resourceDecorations,
     resourceInputFilter,
+    // Resolve @-mention chips for programmatic prefills.
+    resourceSync(getResources, { logger: Logger }),
     resourcesField.init(() => {
-      const registry = getAIContextRegistry(store);
-      const resources = registry.getAllItems();
+      const resources = getRegistryResources(store);
       return new Map(resources.map((resource) => [resource.uri, resource]));
     }),
     resourceTheme,