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

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

@@ -57,7 +57,7 @@ describe("ErrorContextProvider", () => {
 
       const items = provider.getItems();
 
-      expect(items).toHaveLength(1);
+      expect(items).toHaveLength(3);
       expect(items[0]).toMatchObject({
         name: "Errors",
         type: "error",
@@ -78,6 +78,18 @@ describe("ErrorContextProvider", () => {
           ],
         },
       });
+      expect(items[1]).toMatchObject({
+        uri: `error://${cellId1}`,
+        name: "Error: Cell 1",
+        description: "Invalid syntax",
+        data: { type: "cell-error", error: { cellId: cellId1 } },
+      });
+      expect(items[2]).toMatchObject({
+        uri: `error://${cellId2}`,
+        name: "Error: Cell 2",
+        description: "Runtime error",
+        data: { type: "cell-error", error: { cellId: cellId2 } },
+      });
     });
 
     it("should handle cells without names", () => {
@@ -92,7 +104,12 @@ describe("ErrorContextProvider", () => {
       ]);
 
       const items = provider.getItems();
-      expect(items[0].data.errors[0].cellName).toBe("");
+      expect(items[1].data.type).toBe("cell-error");
+      if (items[1].data.type === "cell-error") {
+        expect(items[1].data.error.cellName).toBe("cell-0");
+      }
+      expect(items[1].name).toBe("Error: cell-0");
+      expect(items[1].description).toBe("Invalid syntax");
     });
   });
 
@@ -119,7 +136,7 @@ describe("ErrorContextProvider", () => {
 
       expect(completion).toMatchInlineSnapshot(`
         {
-          "apply": "@Errors",
+          "apply": "@error://all",
           "detail": "2 errors",
           "displayLabel": "Errors",
           "info": [Function],
@@ -170,6 +187,24 @@ describe("ErrorContextProvider", () => {
       expect(completion.detail).toBe("1 error");
     });
 
+    it("formats completion for a single cell error", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const completion = provider.formatCompletion(items[1]);
+      expect(completion.apply).toBe(`@error://${cellId}`);
+      expect(completion.displayLabel).toBe("Error: Cell 1");
+      expect(items[1].name).toBe("Error: Cell 1");
+    });
+
     it("should handle fallback for unknown error types", () => {
       const item = {
         uri: "error://unknown",
@@ -210,7 +245,8 @@ describe("ErrorContextProvider", () => {
 
       const items = provider.getItems();
       const context = provider.formatContext(items[0]);
-      expect(context).toMatchSnapshot("basic-error-context");
+      expect(context).toContain("Invalid syntax");
+      expect(context).toContain("Code:");
     });
 
     it("should format context for multiple error types", () => {
@@ -239,14 +275,12 @@ describe("ErrorContextProvider", () => {
       const items = provider.getItems();
       const context = provider.formatContext(items[0]);
 
-      // Check for expected content instead of exact snapshot match due to random cell IDs
-      expect(context).toMatchInlineSnapshot(`
-        "<error name="Cell 1" description="Invalid syntax
-        Runtime error"></error>
-
-        <error name="Cell cell-2" description="This cell is in a cycle
-        The variable 'variable_x' was defined by another cell"></error>"
-      `);
+      expect(context).toContain("Invalid syntax");
+      expect(context).toContain("Runtime error");
+      expect(context).toContain("This cell is in a cycle");
+      expect(context).toContain(
+        "The variable 'variable_x' was defined by another cell",
+      );
     });
 
     it("should handle cells without names", () => {
@@ -261,8 +295,158 @@ describe("ErrorContextProvider", () => {
       ]);
 
       const items = provider.getItems();
-      const context = provider.formatContext(items[0]);
-      expect(context).toContain(`Cell ${cellId}`);
+      const context = provider.formatContext(items[1]);
+      expect(context).toContain(`cellId="${cellId}"`);
+      expect(context).toContain("Invalid syntax");
+    });
+
+    it("includes exception tracebacks in formatted context", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [
+            {
+              type: "exception",
+              msg: "boom",
+              exception_type: "ValueError",
+              raising_cell: null,
+              traceback: "<pre>line 1\nline 2</pre>",
+            },
+          ],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[1]);
+      expect(context).toContain("boom");
+      expect(context).toContain("line 1");
+      expect(context).toContain("line 2");
+    });
+
+    it("includes traceback-only cell outputs", () => {
+      const cellId = CellIdClass.create();
+      const notebookState = MockNotebook.notebookState({
+        cellData: {
+          [cellId]: {
+            name: "Cell 1",
+            code: "raise ValueError('boom')",
+          },
+        },
+      });
+      notebookState.cellRuntime[cellId] = {
+        ...notebookState.cellRuntime[cellId],
+        output: {
+          channel: "marimo-error",
+          data: "<pre>ValueError: boom\n  line 1</pre>",
+          mimetype: "application/vnd.marimo+traceback",
+          timestamp: Date.now(),
+        },
+      };
+      store.set(notebookAtom, notebookState);
+
+      const items = provider.getItems();
+      expect(items).toHaveLength(2);
+      expect(items[1].description).toBe("ValueError: boom");
+      const context = provider.formatContext(items[1]);
+      expect(context).toContain("raise ValueError('boom')");
+      expect(context).toContain("ValueError: boom");
+      expect(context).toContain("line 1");
+    });
+
+    it("includes console tracebacks for marimo error outputs", () => {
+      const cellId = CellIdClass.create();
+      const notebookState = MockNotebook.notebookState({
+        cellData: {
+          [cellId]: {
+            name: "Cell 1",
+            code: "raise ValueError('boom')",
+          },
+        },
+      });
+      notebookState.cellRuntime[cellId] = {
+        ...notebookState.cellRuntime[cellId],
+        output: {
+          channel: "marimo-error",
+          data: [
+            {
+              type: "exception",
+              msg: "boom",
+              exception_type: "ValueError",
+              raising_cell: null,
+            },
+          ],
+          mimetype: "application/vnd.marimo+error",
+          timestamp: Date.now(),
+        },
+        consoleOutputs: [
+          {
+            channel: "stderr",
+            data: '<pre>ValueError: boom\n  File "notebook.py", line 1</pre>',
+            mimetype: "application/vnd.marimo+traceback",
+            timestamp: Date.now(),
+          },
+        ],
+      };
+      store.set(notebookAtom, notebookState);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[1]);
+      expect(context).toContain("boom");
+      expect(context).toContain("ValueError: boom");
+      expect(context).toContain('File "notebook.py", line 1');
+      expect(items[1].description).toBe("boom");
+    });
+
+    it("does not duplicate console traceback when error already has one", () => {
+      const cellId = CellIdClass.create();
+      const notebookState = MockNotebook.notebookState({
+        cellData: {
+          [cellId]: {
+            name: "Cell 1",
+            code: "raise ValueError('boom')",
+          },
+        },
+      });
+      notebookState.cellRuntime[cellId] = {
+        ...notebookState.cellRuntime[cellId],
+        output: {
+          channel: "marimo-error",
+          data: [
+            {
+              type: "exception",
+              msg: "boom",
+              exception_type: "ValueError",
+              raising_cell: null,
+              traceback: "<pre>embedded line 1</pre>",
+            },
+          ],
+          mimetype: "application/vnd.marimo+error",
+          timestamp: Date.now(),
+        },
+        consoleOutputs: [
+          {
+            channel: "stderr",
+            data: "<pre>console line 1</pre>",
+            mimetype: "application/vnd.marimo+traceback",
+            timestamp: Date.now(),
+          },
+        ],
+      };
+      store.set(notebookAtom, notebookState);
+
+      const items = provider.getItems();
+      if (items[1].data.type !== "cell-error") {
+        throw new Error("Expected cell-error item");
+      }
+      expect(items[1].data.error.tracebackHtml).toBeUndefined();
+
+      const context = provider.formatContext(items[1]);
+      expect(context).toContain("embedded line 1");
+      expect(context).not.toContain("console line 1");
+      expect(items[1].description).toBe("boom");
     });
   });
 
@@ -319,7 +503,8 @@ describe("ErrorContextProvider", () => {
         testStore.set(notebookAtom, notebookState);
 
         const items = testProvider.getItems();
-        expect(items).toHaveLength(1);
+        // One "all-errors" item plus one per-cell error item.
+        expect(items).toHaveLength(2);
         const context = testProvider.formatContext(items[0]);
         expect(context).toContain(expected);
       }

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

@@ -8,6 +8,7 @@ import type { CellId } from "@/core/cells/ids";
 import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { renderDatasourceInfo } from "@/core/codemirror/language/languages/sql/renderers";
 import {
+  allTablesAtom,
   type ConnectionsMap,
   type DatasetTablesMap,
   dataSourceConnectionsAtom,
@@ -18,7 +19,7 @@ import {
   INTERNAL_SQL_ENGINES,
 } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
-import { store } from "@/core/state/jotai";
+import type { JotaiStore } from "@/core/state/jotai";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
@@ -148,7 +149,10 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
   }
 }
 
-export function getDatasourceContext(cellId: CellId): string | null {
+export function getDatasourceContext(
+  cellId: CellId,
+  store: JotaiStore,
+): string | null {
   const cellData = store.get(cellDataAtom(cellId));
   const code = cellData?.code;
   if (!code || code.trim() === "") {
@@ -164,3 +168,24 @@ export function getDatasourceContext(cellId: CellId): string | null {
   }
   return null;
 }
+
+export function formatDatasourceContextForCell(
+  cellId: CellId,
+  store: JotaiStore,
+): string | null {
+  const mention = getDatasourceContext(cellId, store);
+  if (!mention) {
+    return null;
+  }
+
+  const connectionsMap = store.get(dataSourceConnectionsAtom).connectionsMap;
+  const tablesMap = store.get(allTablesAtom);
+  const provider = new DatasourceContextProvider(connectionsMap, tablesMap);
+  const uri = mention.slice(1);
+  const item = provider.getItems().find((candidate) => candidate.uri === uri);
+  if (!item) {
+    return null;
+  }
+
+  return provider.formatContext(item);
+}

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");
   }
 }