@marimo-team/islands 0.23.14-dev2 → 0.23.14-dev23

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/__tests__/setup.test.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { EditorState, type Extension } from "@codemirror/state";
-import { keymap } from "@codemirror/view";
+import { EditorView, keymap } from "@codemirror/view";
 import { describe, expect, test, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
@@ -163,6 +163,38 @@ test("auto_close_pairs: false removes closeBrackets keymaps", () => {
   expect(hasBracketPairHandler(withoutAutoClose)).toBe(false);
 });
 
+test("go to definition falls through to lower-priority LSP keymap", () => {
+  const lspGoToDefinition = vi.fn(() => true);
+  const goToDefinitionKey = getOpts().hotkeys.getHotkey(
+    "cell.goToDefinition",
+  ).key;
+
+  const view = new EditorView({
+    state: EditorState.create({
+      doc: "parser.add_argument('--foo')",
+      selection: { anchor: "parser.".length },
+      extensions: [
+        ...setup(),
+        keymap.of([{ key: goToDefinitionKey, run: lspGoToDefinition }]),
+      ],
+    }),
+    parent: document.body,
+  });
+
+  try {
+    const event = new KeyboardEvent("keydown", {
+      key: goToDefinitionKey,
+      bubbles: true,
+      cancelable: true,
+    });
+    view.contentDOM.dispatchEvent(event);
+
+    expect(lspGoToDefinition).toHaveBeenCalledOnce();
+  } finally {
+    view.destroy();
+  }
+});
+
 test("placeholder adds another extension", () => {
   const opts = getOpts();
   const withAI = new PythonLanguageAdapter()

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/codemirror/cells/extensions.ts

@@ -106,11 +106,8 @@ function cellKeymaps({
     keybindings.push(
       {
         key: hotkeys.getHotkey("cell.goToDefinition").key,
-        preventDefault: true,
-        stopPropagation: true,
         run: (ev) => {
-          goToDefinitionAtCursorPosition(ev);
-          return true;
+          return goToDefinitionAtCursorPosition(ev);
         },
       },
       {

package/src/core/codemirror/completion/__tests__/signature-hint.test.ts

@@ -0,0 +1,94 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { EditorState } from "@codemirror/state";
+import { EditorView, type Tooltip } from "@codemirror/view";
+import { describe, expect, it } from "vitest";
+import {
+  asSignatureHint,
+  setSignatureHintEffect,
+  signatureHintField,
+} from "../signature-hint";
+
+function fakeTooltip(pos: number): Tooltip {
+  return {
+    pos,
+    above: true,
+    create: () => ({ dom: document.createElement("div") }),
+  };
+}
+
+function stateWithHint(doc: string, pos: number): EditorState {
+  const state = EditorState.create({
+    doc,
+    extensions: [signatureHintField],
+  });
+  return state.update({ effects: setSignatureHintEffect.of(fakeTooltip(pos)) })
+    .state;
+}
+
+describe("signatureHintField", () => {
+  it("starts empty", () => {
+    const state = EditorState.create({ extensions: [signatureHintField] });
+    expect(state.field(signatureHintField)).toBeNull();
+  });
+
+  it("shows a tooltip when the effect is dispatched", () => {
+    const state = stateWithHint("plt.plot(", 9);
+    expect(state.field(signatureHintField)?.pos).toBe(9);
+  });
+
+  it("clears the tooltip when the effect dispatches null", () => {
+    let state = stateWithHint("plt.plot(", 9);
+    state = state.update({
+      effects: setSignatureHintEffect.of(null),
+    }).state;
+    expect(state.field(signatureHintField)).toBeNull();
+  });
+
+  it("dismisses the tooltip on a selection-only change", () => {
+    let state = stateWithHint("plt.plot(x)", 9);
+    state = state.update({ selection: { anchor: 0 } }).state;
+    expect(state.field(signatureHintField)).toBeNull();
+  });
+
+  it("keeps and re-anchors the tooltip across edits", () => {
+    let state = stateWithHint("plt.plot(", 9);
+    // Insert before the tooltip position; it should shift to stay anchored.
+    state = state.update({ changes: { from: 0, insert: "xy" } }).state;
+    expect(state.field(signatureHintField)?.pos).toBe(11);
+  });
+});
+
+describe("asSignatureHint", () => {
+  it("nests the content so descendant styling applies", () => {
+    const content = document.createElement("span");
+    content.classList.add("mo-cm-tooltip", "docs-documentation");
+    const base: Tooltip = {
+      pos: 0,
+      create: () => ({ dom: content, resize: false }),
+    };
+
+    const view = new EditorView({ state: EditorState.create({}) });
+    const { dom } = asSignatureHint(base).create(view);
+
+    // Outer wrapper carries the tooltip sizing class...
+    expect(dom.classList.contains("mo-cm-tooltip")).toBe(true);
+    // ...and the documentation content is a descendant (not the same node),
+    // so `.cm-tooltip .docs-documentation` padding/font rules match.
+    expect(dom).not.toBe(content);
+    expect(dom.querySelector(".docs-documentation")).toBe(content);
+
+    view.destroy();
+  });
+
+  it("preserves other tooltip fields", () => {
+    const base: Tooltip = {
+      pos: 5,
+      above: true,
+      create: () => ({ dom: document.createElement("span"), resize: false }),
+    };
+    const wrapped = asSignatureHint(base);
+    expect(wrapped.pos).toBe(5);
+    expect(wrapped.above).toBe(true);
+  });
+});

package/src/core/codemirror/completion/completer.ts

@@ -9,6 +9,7 @@ import { documentationAtom } from "@/core/documentation/state";
 import { store } from "@/core/state/jotai";
 import { Logger } from "../../../utils/Logger";
 import { AUTOCOMPLETER, Autocompleter } from "./Autocompleter";
+import { asSignatureHint, setSignatureHintEffect } from "./signature-hint";
 
 /**
  * Completion source for Python, using Jedi.
@@ -36,10 +37,12 @@ export const pythonCompletionSource: CompletionSource = async (
     cellId: cellId,
   });
   if (!result) {
+    context.view?.dispatch({ effects: setSignatureHintEffect.of(null) });
     return null;
   }
 
-  // If it is a tooltip, show it as a Tooltip instead of a completion
+  // If it is a tooltip (e.g. a signature after `(` or `,`), show it as a
+  // Tooltip instead of a completion.
   const tooltip = Autocompleter.asHoverTooltip({
     position: context.pos,
     message: result,
@@ -50,6 +53,14 @@ export const pythonCompletionSource: CompletionSource = async (
       documentation: tooltip.html ?? null,
     });
   }
+  // Surface the signature as a floating hint near the cursor (the LSP path has
+  // its own signature help), and clear any stale hint when we instead have a
+  // real completion list.
+  context.view?.dispatch({
+    effects: setSignatureHintEffect.of(
+      tooltip ? asSignatureHint(tooltip) : null,
+    ),
+  });
   if (tooltip) {
     return null;
   }

package/src/core/codemirror/completion/signature-hint.ts

@@ -0,0 +1,68 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { StateEffect, StateField } from "@codemirror/state";
+import { showTooltip, type Tooltip } from "@codemirror/view";
+
+/**
+ * Effect to set (or clear, with `null`) the floating signature hint.
+ */
+export const setSignatureHintEffect = StateEffect.define<Tooltip | null>();
+
+/**
+ * Wrap a tooltip so it renders like the completion popup's info box.
+ *
+ * CodeMirror adds `cm-tooltip` directly to the DOM node returned by `create`,
+ * so the documentation content ends up on the same element as `cm-tooltip`.
+ * Our styling for padding/font (`.cm-tooltip .docs-documentation`) is a
+ * descendant selector, so we nest the content one level deeper to make it
+ * apply — mirroring how CodeMirror nests completion info inside its own
+ * wrapper. The outer `mo-cm-tooltip` class picks up the shared tooltip sizing.
+ */
+export function asSignatureHint(tooltip: Tooltip): Tooltip {
+  return {
+    ...tooltip,
+    create: (view) => {
+      const { dom: content, ...rest } = tooltip.create(view);
+      const dom = document.createElement("div");
+      dom.classList.add("mo-cm-tooltip");
+      dom.append(content);
+      return { ...rest, dom };
+    },
+  };
+}
+
+/**
+ * Holds the floating "signature hint" shown after typing `(` or `,` inside a
+ * call on the non-LSP (Jedi) completion path.
+ *
+ * The LSP path has its own signature help; this fills the gap for users
+ * without a language server. The completion source (`pythonCompletionSource`)
+ * drives it: it dispatches `setSignatureHintEffect` with the tooltip when the
+ * backend returns a signature and with `null` otherwise. The hint is also
+ * cleared when the cursor moves via a selection-only change (e.g. clicking
+ * away or arrowing out of the call), and kept anchored across edits so it
+ * doesn't flicker while a fresh result is in flight.
+ */
+export const signatureHintField = StateField.define<Tooltip | null>({
+  create: () => null,
+  update(tooltip, tr) {
+    for (const effect of tr.effects) {
+      if (effect.is(setSignatureHintEffect)) {
+        return effect.value;
+      }
+    }
+    if (!tooltip) {
+      return null;
+    }
+    // Cursor moved without editing (click / arrow key): dismiss the hint.
+    if (tr.selection && !tr.docChanged) {
+      return null;
+    }
+    // Keep the hint anchored across edits; the completion source refreshes or
+    // clears it as new results arrive.
+    if (tr.docChanged) {
+      return { ...tooltip, pos: tr.changes.mapPos(tooltip.pos) };
+    }
+    return tooltip;
+  },
+  provide: (field) => showTooltip.from(field),
+});

package/src/core/codemirror/language/languages/python.ts

@@ -31,6 +31,7 @@ import { Logger } from "@/utils/Logger";
 import { once } from "@/utils/once";
 import { cellActionsState } from "../../cells/state";
 import { pythonCompletionSource } from "../../completion/completer";
+import { signatureHintField } from "../../completion/signature-hint";
 import type { PlaceholderType } from "../../config/types";
 import { FederatedLanguageServerClient } from "../../lsp/federated-lsp";
 import { createLspMarkdownRenderer } from "../../lsp/markdown-renderer";
@@ -376,10 +377,15 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
         ];
       }
 
-      return autocompletion({
-        ...autocompleteOptions,
-        override: [pythonCompletionSource],
-      });
+      return [
+        autocompletion({
+          ...autocompleteOptions,
+          override: [pythonCompletionSource],
+        }),
+        // The Jedi path has no built-in signature help; show a floating hint
+        // fed by `pythonCompletionSource` (the LSP path handles this itself).
+        signatureHintField,
+      ];
     };
 
     return [

package/src/core/codemirror/utils.ts

@@ -2,6 +2,7 @@
 import type { EditorState, Transaction } from "@codemirror/state";
 import type { EditorView, ViewUpdate } from "@codemirror/view";
 import { getCM } from "@replit/codemirror-vim";
+import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 
 export function isAtStartOfEditor(ev: { state: EditorState }) {
   const main = ev.state.selection.main;
@@ -37,6 +38,20 @@ export function moveToEndOfEditor(ev: EditorView | undefined) {
   });
 }
 
+/** We delay the focus and move to end of the editor until React has rendered the prefilled value. */
+export function focusInputAndMoveToEnd(
+  ref: React.RefObject<ReactCodeMirrorRef | null>,
+) {
+  requestAnimationFrame(() => {
+    const view = ref.current?.view;
+    if (!view) {
+      return;
+    }
+    view.focus();
+    moveToEndOfEditor(view);
+  });
+}
+
 export function isInVimMode(ev: EditorView): boolean {
   return getCM(ev)?.state.vim != null;
 }