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

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/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,

package/src/core/errors/errors.ts

@@ -6,6 +6,7 @@ import type { AiCompletionCell } from "../ai/state";
 import type { CellId } from "../cells/ids";
 import type { MarimoError } from "../kernel/messages";
 import { wrapInFunction } from "./utils";
+import { store } from "../state/jotai";
 
 interface AIFix {
   setAiCompletionCell: (opts: AiCompletionCell) => void;
@@ -84,7 +85,7 @@ export function getAutoFixes(
         description: "Fix the SQL statement",
         fixType: "ai",
         onFix: async (ctx) => {
-          const datasourceContext = getDatasourceContext(ctx.cellId);
+          const datasourceContext = getDatasourceContext(ctx.cellId, store);
           let initialPrompt = `Fix the SQL statement: ${error.msg}.`;
           if (datasourceContext) {
             initialPrompt += `\nDatabase schema: ${datasourceContext}`;

package/src/core/mime.ts

@@ -1,9 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { OutputMessage } from "./kernel/messages";
 
+export function isMarimoErrorsMime(
+  mime: OutputMessage["mimetype"] | undefined,
+) {
+  return mime === "application/vnd.marimo+error";
+}
+
+export function isTracebackMime(mime: OutputMessage["mimetype"] | undefined) {
+  return mime === "application/vnd.marimo+traceback";
+}
+
 export function isErrorMime(mime: OutputMessage["mimetype"] | undefined) {
-  return (
-    mime === "application/vnd.marimo+error" ||
-    mime === "application/vnd.marimo+traceback"
-  );
+  return isMarimoErrorsMime(mime) || isTracebackMime(mime);
 }

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

@@ -1,3 +0,0 @@
-// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
-
-exports[`ErrorContextProvider > formatContext > should format context for basic errors > basic-error-context 1`] = `"<error name="Cell 1" description="Invalid syntax"></error>"`;