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

package/src/components/slides/slide-cell-view.tsx

@@ -26,7 +26,7 @@ import { connectionAtom } from "@/core/network/connection";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { ReadonlyCode } from "../editor/code/readonly-python-code";
-import { languageAdapterFromCode } from "@/core/codemirror/language/extension";
+import { getReadonlyCodeDisplay } from "@/core/cells/readonly-code-display";
 
 type RuntimeCell = CellRuntimeState & CellData;
 
@@ -188,10 +188,7 @@ export const SlideCellReadOnlyView = ({ cell }: { cell: RuntimeCell }) => {
   const [userConfig] = useUserConfig();
   const cellOutputPosition = userConfig.display.cell_output;
 
-  const language = useMemo(() => {
-    const adapter = languageAdapterFromCode(cell.code.trim());
-    return adapter.type === "sql" ? "sql" : "python";
-  }, [cell.code]);
+  const display = useMemo(() => getReadonlyCodeDisplay(cell.code), [cell.code]);
 
   const output = (
     <CellOutputSlide
@@ -203,7 +200,11 @@ export const SlideCellReadOnlyView = ({ cell }: { cell: RuntimeCell }) => {
 
   const editor = (
     <div className="marimo-cell">
-      <ReadonlyCode code={cell.code} language={language} showHideCode={false} />
+      <ReadonlyCode
+        code={display.code}
+        language={display.language}
+        showHideCode={false}
+      />
     </div>
   );
 

package/src/core/ai/config.ts

@@ -60,7 +60,7 @@ export const useModelChange = () => {
       },
     };
 
-    saveConfig(newConfig);
+    await saveConfig(newConfig);
   };
 
   const saveModeChange = async (newMode: CopilotMode) => {
@@ -71,7 +71,7 @@ export const useModelChange = () => {
       },
     };
 
-    saveConfig(newConfig);
+    await saveConfig(newConfig);
   };
 
   return { saveModelChange, saveModeChange };

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

@@ -176,6 +176,33 @@ class FileContextProvider extends AIContextProvider<FileContextItem> {
   }
 }
 
+// Second mock provider with the same context type as MockContextProvider
+class SecondaryMockProvider extends AIContextProvider<MockContextItem> {
+  readonly title = "Secondary Mock Items";
+  readonly mentionPrefix = "@";
+  readonly contextType = "mock";
+
+  getItems(): MockContextItem[] {
+    return [
+      {
+        type: "mock",
+        uri: this.asURI("secondary-item"),
+        name: "Secondary Item",
+        description: "From the second mock provider",
+        data: { value: "secondary" },
+      },
+    ];
+  }
+
+  formatContext(item: MockContextItem): string {
+    return `Secondary: ${item.name}`;
+  }
+
+  formatCompletion(item: MockContextItem): Completion {
+    return this.createBasicCompletion(item);
+  }
+}
+
 describe("AIContextProvider", () => {
   let provider: MockContextProvider;
 
@@ -411,6 +438,95 @@ describe("AIContextRegistry", () => {
     });
   });
 
+  describe("resolveItems", () => {
+    let mockGetItems: ReturnType<typeof vi.spyOn>;
+    let fileGetItems: ReturnType<typeof vi.spyOn>;
+
+    beforeEach(() => {
+      registry.register(mockProvider);
+      registry.register(fileProvider);
+      mockGetItems = vi.spyOn(mockProvider, "getItems");
+      fileGetItems = vi.spyOn(fileProvider, "getItems");
+    });
+
+    it("should resolve valid IDs from the matching providers only", () => {
+      const resolved = registry.resolveItems([
+        "mock://item1",
+        "file://config.py",
+      ] as ContextLocatorId[]);
+
+      expect(resolved).toHaveLength(2);
+      expect(resolved.map((item) => item.uri)).toEqual([
+        "mock://item1",
+        "file://config.py",
+      ]);
+      expect(mockGetItems).toHaveBeenCalledTimes(1);
+      expect(fileGetItems).toHaveBeenCalledTimes(1);
+    });
+
+    it("should not call unrelated providers", () => {
+      registry.resolveItems(["mock://item1"] as ContextLocatorId[]);
+
+      expect(mockGetItems).toHaveBeenCalledTimes(1);
+      expect(fileGetItems).not.toHaveBeenCalled();
+    });
+
+    it("should ignore unknown or invalid URIs", () => {
+      const resolved = registry.resolveItems([
+        "mock://nonexistent",
+        "unknown://item",
+        "malformed" as ContextLocatorId,
+      ] as ContextLocatorId[]);
+
+      expect(resolved).toEqual([]);
+      expect(mockGetItems).toHaveBeenCalledTimes(1);
+      expect(fileGetItems).not.toHaveBeenCalled();
+    });
+
+    it("should return empty array for empty input", () => {
+      expect(registry.resolveItems([])).toEqual([]);
+      expect(mockGetItems).not.toHaveBeenCalled();
+      expect(fileGetItems).not.toHaveBeenCalled();
+    });
+
+    it("should preserve requested order when providers are interleaved", () => {
+      const resolved = registry.resolveItems([
+        "mock://item1",
+        "file://config.py",
+        "mock://item2",
+        "file://utils/helpers.py",
+      ] as ContextLocatorId[]);
+
+      expect(resolved.map((item) => item.uri)).toEqual([
+        "mock://item1",
+        "file://config.py",
+        "mock://item2",
+        "file://utils/helpers.py",
+      ]);
+      // Each provider is still queried only once, regardless of interleaving.
+      expect(mockGetItems).toHaveBeenCalledTimes(1);
+      expect(fileGetItems).toHaveBeenCalledTimes(1);
+    });
+
+    it("should resolve items from all providers that share a context type", () => {
+      const secondaryProvider = new SecondaryMockProvider();
+      const secondaryGetItems = vi.spyOn(secondaryProvider, "getItems");
+      registry.register(secondaryProvider);
+
+      const resolved = registry.resolveItems([
+        "mock://item1",
+        "mock://secondary-item",
+      ] as ContextLocatorId[]);
+
+      expect(resolved.map((item) => item.uri)).toEqual([
+        "mock://item1",
+        "mock://secondary-item",
+      ]);
+      expect(mockGetItems).toHaveBeenCalledTimes(1);
+      expect(secondaryGetItems).toHaveBeenCalledTimes(1);
+    });
+  });
+
   describe("getContextInfo", () => {
     beforeEach(() => {
       registry.register(mockProvider);
@@ -531,6 +647,19 @@ describe("AIContextRegistry", () => {
       expect(formatted).toContain("Mock: Item 1 (value1)");
       expect(formatted.split("\n\n")).toHaveLength(1);
     });
+
+    it("should format items from the provider that owns them", () => {
+      registry.register(new SecondaryMockProvider());
+
+      const formatted = registry.formatContextForAI([
+        "mock://item1",
+        "mock://secondary-item",
+      ] as ContextLocatorId[]);
+
+      expect(formatted).toBe(
+        "Mock: Item 1 (value1)\n\nSecondary: Secondary Item",
+      );
+    });
   });
 
   describe("edge cases and error handling", () => {
@@ -621,7 +750,7 @@ describe("AIContextRegistry", () => {
       expect(provider).toBe(errorProvider);
 
       const items = registry.getAllItems();
-      expect(items).toHaveLength(1);
+      expect(items).toHaveLength(3);
       expect(items[0].type).toBe("error");
       expect(items[0].name).toBe("Errors");
     });
@@ -654,10 +783,11 @@ describe("AIContextRegistry", () => {
       expect(provider).toBeDefined();
 
       const items = provider!.getItems();
-      expect(items).toHaveLength(1);
+      expect(items).toHaveLength(3);
 
       const completion = provider!.formatCompletion(items[0]);
       expect(completion.label).toBe("@Errors");
+      expect(completion.apply).toBe("@error://all");
       expect(completion.type).toBe("error");
     });
 
@@ -690,7 +820,7 @@ describe("AIContextRegistry", () => {
       const errorItems = items.filter((item) => item.type === "error");
       const mockItems = items.filter((item) => item.type === "mock");
 
-      expect(errorItems).toHaveLength(1);
+      expect(errorItems).toHaveLength(3);
       expect(mockItems).toHaveLength(3);
     });
 

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

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* oxlint-disable typescript/no-explicit-any */
 
+import { createStore } from "jotai";
 import { beforeEach, describe, expect, it } from "vitest";
 import { cellId, variableName } from "@/__tests__/branded";
 import type {
@@ -634,7 +635,8 @@ describe("DatasourceContextProvider", () => {
 
 describe("getDatasourceContext", () => {
   it("should return null if no cell ID is found", () => {
-    const context = getDatasourceContext(cellId("1"));
+    const store = createStore();
+    const context = getDatasourceContext(cellId("1"), store);
     expect(context).toBeNull();
   });
 });

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