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

package/dist/index.html

@@ -66,7 +66,7 @@
     <marimo-server-token data-token="{{ server_token }}" hidden></marimo-server-token>
     <!-- /TODO -->
     <title>{{ title }}</title>
-    <script type="module" crossorigin src="./assets/index-CNmQbGXF.js"></script>
+    <script type="module" crossorigin src="./assets/index-BcxGcwkJ.js"></script>
     <link rel="modulepreload" crossorigin href="./assets/preload-helper-BW0IMuFq.js">
     <link rel="modulepreload" crossorigin href="./assets/chunk-LvLJmgfZ.js">
     <link rel="modulepreload" crossorigin href="./assets/react-Bj1aDYRI.js">
@@ -133,7 +133,7 @@
     <link rel="modulepreload" crossorigin href="./assets/debounce-OiXznToA.js">
     <link rel="modulepreload" crossorigin href="./assets/database-zap-DIlSqcdp.js">
     <link rel="modulepreload" crossorigin href="./assets/main-DmhaVp6Y.js">
-    <link rel="modulepreload" crossorigin href="./assets/cells-DENsmH9t.js">
+    <link rel="modulepreload" crossorigin href="./assets/cells-RJTD-euN.js">
     <link rel="modulepreload" crossorigin href="./assets/kbd-UVzWBHfi.js">
     <link rel="modulepreload" crossorigin href="./assets/useInstallPackage-C1Dd0aIZ.js">
     <link rel="modulepreload" crossorigin href="./assets/alert-dialog-CC0gG4Av.js">
@@ -146,40 +146,40 @@
     <link rel="modulepreload" crossorigin href="./assets/usePress-BSFUeqnS.js">
     <link rel="modulepreload" crossorigin href="./assets/input-DB2LiTiQ.js">
     <link rel="modulepreload" crossorigin href="./assets/ImperativeModal-B62Q6S9l.js">
-    <link rel="modulepreload" crossorigin href="./assets/cell-link-exzcyUhT.js">
+    <link rel="modulepreload" crossorigin href="./assets/cell-link-CBRggkzf.js">
     <link rel="modulepreload" crossorigin href="./assets/config-DMNo9K-z.js">
     <link rel="modulepreload" crossorigin href="./assets/alert-1y_DnRWn.js">
     <link rel="modulepreload" crossorigin href="./assets/chevron-right-C7RH0Nc7.js">
     <link rel="modulepreload" crossorigin href="./assets/dropdown-menu-DfGLyCbP.js">
     <link rel="modulepreload" crossorigin href="./assets/links-CCkJpphF.js">
-    <link rel="modulepreload" crossorigin href="./assets/useRunCells-CTpvxsqp.js">
+    <link rel="modulepreload" crossorigin href="./assets/useRunCells-D5aAUcAa.js">
     <link rel="modulepreload" crossorigin href="./assets/copy-D8MtpXgN.js">
     <link rel="modulepreload" crossorigin href="./assets/copy-IBNGDypD.js">
     <link rel="modulepreload" crossorigin href="./assets/copy-icon-EHMT-qlr.js">
-    <link rel="modulepreload" crossorigin href="./assets/RenderHTML-v15QgtRg.js">
-    <link rel="modulepreload" crossorigin href="./assets/datasource-BCOYpQ69.js">
-    <link rel="modulepreload" crossorigin href="./assets/state-LM4pQmYL.js">
+    <link rel="modulepreload" crossorigin href="./assets/RenderHTML-B9mw7TD1.js">
+    <link rel="modulepreload" crossorigin href="./assets/datasource-RP7m6Xrh.js">
+    <link rel="modulepreload" crossorigin href="./assets/state-DHqaTLr2.js">
     <link rel="modulepreload" crossorigin href="./assets/hat-glasses-DlK0E4SX.js">
     <link rel="modulepreload" crossorigin href="./assets/package-BTIV9-Dk.js">
     <link rel="modulepreload" crossorigin href="./assets/sparkles-COaEc8Qn.js">
-    <link rel="modulepreload" crossorigin href="./assets/MarimoErrorOutput-D07aAGkQ.js">
+    <link rel="modulepreload" crossorigin href="./assets/MarimoErrorOutput-BokuSnpU.js">
     <link rel="modulepreload" crossorigin href="./assets/ErrorBoundary-Bo313Z9X.js">
     <link rel="modulepreload" crossorigin href="./assets/eye-off-DI0EQXw0.js">
     <link rel="modulepreload" crossorigin href="./assets/refresh-cw-C9FftltI.js">
-    <link rel="modulepreload" crossorigin href="./assets/tree-actions-Bnmh7YGo.js">
+    <link rel="modulepreload" crossorigin href="./assets/tree-actions-DLySKXUq.js">
     <link rel="modulepreload" crossorigin href="./assets/spinner-BbI9qk88.js">
-    <link rel="modulepreload" crossorigin href="./assets/html-to-image-B6eCuh-2.js">
-    <link rel="modulepreload" crossorigin href="./assets/focus-CZUEQpxX.js">
+    <link rel="modulepreload" crossorigin href="./assets/html-to-image-B8ja7wwq.js">
+    <link rel="modulepreload" crossorigin href="./assets/focus-BViZMdHR.js">
     <link rel="modulepreload" crossorigin href="./assets/useAsyncData-DAQcLiyE.js">
     <link rel="modulepreload" crossorigin href="./assets/LazyAnyLanguageCodeMirror-DpraoNlO.js">
     <link rel="modulepreload" crossorigin href="./assets/micromark-factory-space-BUQpMdx2.js">
     <link rel="modulepreload" crossorigin href="./assets/chunk-5FQGJX7Z-OScactr4.js">
-    <link rel="modulepreload" crossorigin href="./assets/markdown-renderer-CcLl4-oB.js">
+    <link rel="modulepreload" crossorigin href="./assets/markdown-renderer-jghC8PFL.js">
     <link rel="modulepreload" crossorigin href="./assets/command-DRHbYHUP.js">
     <link rel="modulepreload" crossorigin href="./assets/field-DlNY-dDH.js">
     <link rel="modulepreload" crossorigin href="./assets/popover-DMNO6Poc.js">
     <link rel="modulepreload" crossorigin href="./assets/errors-9jMmWHd9.js">
-    <link rel="modulepreload" crossorigin href="./assets/download-4NoBestd.js">
+    <link rel="modulepreload" crossorigin href="./assets/download-CJ6tdZ-1.js">
     <link rel="modulepreload" crossorigin href="./assets/table-CnW850Wt.js">
     <link rel="modulepreload" crossorigin href="./assets/useIframeCapabilities-B0dNZiKK.js">
     <link rel="modulepreload" crossorigin href="./assets/error-banner-BSQVN46J.js">
@@ -199,21 +199,21 @@
     <link rel="modulepreload" crossorigin href="./assets/plus-ZXRHY1up.js">
     <link rel="modulepreload" crossorigin href="./assets/trash-2-CqTUwVuU.js">
     <link rel="modulepreload" crossorigin href="./assets/react-resizable-panels.browser.esm-Ce2ksurd.js">
-    <link rel="modulepreload" crossorigin href="./assets/JsonOutput-C3c9GGaV.js">
+    <link rel="modulepreload" crossorigin href="./assets/JsonOutput-CYIo9GSP.js">
     <link rel="modulepreload" crossorigin href="./assets/chart-no-axes-column-DVEV2j1H.js">
     <link rel="modulepreload" crossorigin href="./assets/square-function-Dx0pKHko.js">
     <link rel="modulepreload" crossorigin href="./assets/spec-CDgWI0fT.js">
-    <link rel="modulepreload" crossorigin href="./assets/column-preview-BPgYXdDd.js">
+    <link rel="modulepreload" crossorigin href="./assets/column-preview-dW7xtQ7S.js">
     <link rel="modulepreload" crossorigin href="./assets/icons-avFV3lBF.js">
-    <link rel="modulepreload" crossorigin href="./assets/radio-group-B8OldelF.js">
-    <link rel="modulepreload" crossorigin href="./assets/floating-outline--OlHG_yX.js">
+    <link rel="modulepreload" crossorigin href="./assets/radio-group-DFkF8CeQ.js">
+    <link rel="modulepreload" crossorigin href="./assets/floating-outline-BOkj3eC_.js">
     <link rel="modulepreload" crossorigin href="./assets/objectWithoutPropertiesLoose-BkUn3I0-.js">
     <link rel="modulepreload" crossorigin href="./assets/esm-X3Tr2pur.js">
-    <link rel="modulepreload" crossorigin href="./assets/readonly-python-code-CoFR55Yt.js">
+    <link rel="modulepreload" crossorigin href="./assets/readonly-python-code-Cla-uJpn.js">
     <link rel="modulepreload" crossorigin href="./assets/file-headphone-CUN2DC78.js">
     <link rel="modulepreload" crossorigin href="./assets/file-BvoBtS_V.js">
     <link rel="modulepreload" crossorigin href="./assets/image-BzUcK6ho.js">
-    <link rel="modulepreload" crossorigin href="./assets/file-icons-BFLcd4J7.js">
+    <link rel="modulepreload" crossorigin href="./assets/file-icons-CGJ6uWXm.js">
     <link rel="modulepreload" crossorigin href="./assets/switch-DvDNxcZP.js">
     <link rel="modulepreload" crossorigin href="./assets/events-XR_LpvHc.js">
     <link rel="modulepreload" crossorigin href="./assets/globals-Bx8Km3LH.js">
@@ -222,11 +222,11 @@
     <link rel="modulepreload" crossorigin href="./assets/memoize-Tp7rARFe.js">
     <link rel="modulepreload" crossorigin href="./assets/get-C-qh_et5.js">
     <link rel="modulepreload" crossorigin href="./assets/_baseSet-CsW0tABf.js">
-    <link rel="modulepreload" crossorigin href="./assets/state-CQn0dE8X.js">
+    <link rel="modulepreload" crossorigin href="./assets/state-CXzKtmF5.js">
     <link rel="modulepreload" crossorigin href="./assets/label-CFvCMYuj.js">
-    <link rel="modulepreload" crossorigin href="./assets/textarea-Tswrj-fB.js">
+    <link rel="modulepreload" crossorigin href="./assets/textarea-B5Bp08T3.js">
     <link rel="modulepreload" crossorigin href="./assets/refresh-ccw-Bee7JLOl.js">
-    <link rel="modulepreload" crossorigin href="./assets/form-CJ3pB1ow.js">
+    <link rel="modulepreload" crossorigin href="./assets/form-BbqpGicG.js">
     <link rel="modulepreload" crossorigin href="./assets/renderShortcut-LpVarhnQ.js">
     <link rel="modulepreload" crossorigin href="./assets/useBoolean-BMzj74q1.js">
     <link rel="modulepreload" crossorigin href="./assets/useDeepCompareMemoize-CWcgQCbT.js">

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/frontend",
-  "version": "0.23.14-dev19",
+  "version": "0.23.14-dev21",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",
@@ -56,7 +56,7 @@
     "@lezer/python": "^1.1.18",
     "@marimo-team/codemirror-ai": "^0.3.7",
     "@marimo-team/codemirror-languageserver": "^1.16.12",
-    "@marimo-team/codemirror-mcp": "^0.1.5",
+    "@marimo-team/codemirror-mcp": "^0.1.7",
     "@marimo-team/codemirror-sql": "^0.2.8",
     "@marimo-team/llm-info": "workspace:*",
     "@marimo-team/marimo-api": "workspace:*",

package/src/components/editor/errors/__tests__/auto-fix.test.ts

@@ -1,119 +1,23 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { describe, expect, it } from "vitest";
 import { cellId } from "@/__tests__/branded";
-import type { MarimoError } from "@/core/kernel/messages";
+import { buildFixInChatPrompt } from "../auto-fix";
 
-const getDatasourceContext = vi.fn<(id: unknown) => string | null>(() => null);
-
-vi.mock("@/core/ai/context/providers/datasource", () => ({
-  getDatasourceContext: (id: unknown) => getDatasourceContext(id),
-}));
-
-const { buildFixPrompt, buildFixPromptFromText } = await import("../auto-fix");
-
-describe("buildFixPromptFromText", () => {
-  beforeEach(() => {
-    getDatasourceContext.mockReset();
-    getDatasourceContext.mockReturnValue(null);
-  });
-
-  it("includes the cell id in the header when provided", () => {
-    const prompt = buildFixPromptFromText("boom", cellId("cell-1"));
-    expect(prompt).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom",
-    );
-  });
-
-  it("uses a generic header when no cell id is provided", () => {
-    const prompt = buildFixPromptFromText("boom");
-    expect(prompt).toBe(
-      "My code gives the following error. Please fix it:\n\nboom",
+describe("buildFixInChatPrompt", () => {
+  it("references the per-cell error context URI", () => {
+    expect(buildFixInChatPrompt(cellId("cell-1"))).toBe(
+      "@error://cell-1\n\nPlease fix this error.",
     );
   });
 
-  it("appends datasource context when explicitly requested for the cell", () => {
-    getDatasourceContext.mockReturnValue("@datasource://my_db");
-    const prompt = buildFixPromptFromText("boom", cellId("cell-1"), {
-      includeDatasourceContext: true,
-    });
-    expect(prompt).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom\n\nDatabase schema: @datasource://my_db",
+  it("falls back to inline error text when no cell id", () => {
+    expect(buildFixInChatPrompt(undefined, "ValueError: boom")).toBe(
+      "My code gives the following error. Please fix it:\n\nValueError: boom",
     );
   });
 
-  it("does not append datasource context by default", () => {
-    getDatasourceContext.mockReturnValue("@datasource://my_db");
-    const prompt = buildFixPromptFromText("boom", cellId("cell-1"));
-    expect(prompt).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom",
-    );
-    expect(getDatasourceContext).not.toHaveBeenCalled();
-  });
-
-  it("does not look up datasource context without a cell id", () => {
-    buildFixPromptFromText("boom", undefined, {
-      includeDatasourceContext: true,
-    });
-    expect(getDatasourceContext).not.toHaveBeenCalled();
-  });
-});
-
-describe("buildFixPrompt", () => {
-  beforeEach(() => {
-    getDatasourceContext.mockReset();
-    getDatasourceContext.mockReturnValue(null);
-  });
-
-  it("uses the error message when present", () => {
-    const errors: MarimoError[] = [
-      { type: "sql-error", msg: "syntax error", sql_statement: "SELECT" },
-    ];
-    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nsyntax error",
-    );
-  });
-
-  it("joins multiple errors with newlines", () => {
-    const errors: MarimoError[] = [
-      { type: "syntax", msg: "bad syntax" },
-      {
-        type: "exception",
-        exception_type: "ValueError",
-        msg: "bad value",
-        raising_cell: null,
-      },
-    ];
-    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nbad syntax\nbad value",
-    );
-  });
-
-  it("falls back to the error type when there is no message", () => {
-    const errors: MarimoError[] = [
-      { type: "multiple-defs", name: "foo", cells: [cellId("foo")] },
-    ];
-    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nmultiple-defs",
-    );
-  });
-
-  it("appends datasource context for SQL errors", () => {
-    getDatasourceContext.mockReturnValue("@datasource://my_db");
-    const errors: MarimoError[] = [
-      { type: "sql-error", msg: "syntax error", sql_statement: "SELECT" },
-    ];
-    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nsyntax error\n\nDatabase schema: @datasource://my_db",
-    );
-  });
-
-  it("does not append datasource context for non-SQL errors", () => {
-    getDatasourceContext.mockReturnValue("@datasource://my_db");
-    const errors: MarimoError[] = [{ type: "syntax", msg: "bad syntax" }];
-    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
-      "My cell (id: cell-1) produced the following error. Please fix it:\n\nbad syntax",
-    );
-    expect(getDatasourceContext).not.toHaveBeenCalled();
+  it("uses a generic prompt when no cell id and no error text", () => {
+    expect(buildFixInChatPrompt(undefined)).toBe("Please fix this error.");
   });
 });

package/src/components/editor/errors/auto-fix.tsx

@@ -21,44 +21,24 @@ import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { aiFeaturesEnabledAtom } from "@/core/config/config";
-import { getDatasourceContext } from "@/core/ai/context/providers/datasource";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { cn } from "@/utils/cn";
 import { useOpenAiAssistant } from "../chrome/wrapper/useOpenAiAssistant";
 import { type FixMode, useFixMode } from "./fix-mode";
 
-export function buildFixPromptFromText(
-  errorText: string,
-  cellId?: CellId,
-  {
-    includeDatasourceContext = false,
-  }: { includeDatasourceContext?: boolean } = {},
+export function buildFixInChatPrompt(
+  cellId: CellId | undefined,
+  fallbackErrorText?: string,
 ): string {
-  const header =
-    cellId != null
-      ? `My cell (id: ${cellId}) produced the following error. Please fix it:`
-      : "My code gives the following error. Please fix it:";
-  let prompt = `${header}\n\n${errorText}`;
-  if (cellId != null && includeDatasourceContext) {
-    const datasourceContext = getDatasourceContext(cellId);
-    if (datasourceContext) {
-      prompt += `\n\nDatabase schema: ${datasourceContext}`;
-    }
+  if (cellId != null) {
+    return `@error://${cellId}\n\nPlease fix this error.`;
   }
-  return prompt;
-}
-
-export function buildFixPrompt(errors: MarimoError[], cellId: CellId): string {
-  const errorText = errors
-    .map((error) => ("msg" in error && error.msg ? error.msg : error.type))
-    .join("\n");
-  const includeDatasourceContext = errors.some(
-    (error) => error.type === "sql-error",
-  );
-  return buildFixPromptFromText(errorText, cellId, {
-    includeDatasourceContext,
-  });
+  const errorText = fallbackErrorText?.trim();
+  if (errorText) {
+    return `My code gives the following error. Please fix it:\n\n${errorText}`;
+  }
+  return "Please fix this error.";
 }
 
 export const AutoFixButton = ({
@@ -112,7 +92,7 @@ export const AutoFixButton = ({
 
   const openAISidebar = () => {
     openAiAssistant({
-      prompt: buildFixPrompt(errors, cellId),
+      prompt: buildFixInChatPrompt(cellId),
     });
   };
 

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

@@ -42,7 +42,7 @@ import {
   getTracebackInfo,
 } from "@/utils/traceback";
 import { useOpenAiAssistant } from "../chrome/wrapper/useOpenAiAssistant";
-import { AIFixButton, buildFixPromptFromText } from "../errors/auto-fix";
+import { AIFixButton, buildFixInChatPrompt } from "../errors/auto-fix";
 import { MangledSegments } from "../errors/mangled-local-chip";
 import { CellLinkTraceback } from "../links/cell-link";
 import type { OnRefactorWithAI } from "../Output";
@@ -101,7 +101,7 @@ export const MarimoTracebackOutput = ({
 
   const openAISidebar = () => {
     openAiAssistant({
-      prompt: buildFixPromptFromText(lastTracebackLine, cellId),
+      prompt: buildFixInChatPrompt(cellId, lastTracebackLine),
     });
   };
 

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