@marimo-team/islands 0.23.14-dev5 → 0.23.14-dev50

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__/cell.test.ts

@@ -218,3 +218,42 @@ describe("transitionCell serialization", () => {
     );
   });
 });
+
+describe("transitionCell stdin resolution", () => {
+  function stdinOutput(): OutputMessage {
+    return {
+      channel: "stdin",
+      mimetype: "text/plain",
+      data: "(Pdb) ",
+      timestamp: Date.now() as Seconds,
+    } as OutputMessage;
+  }
+
+  it("resolves a dangling pdb prompt when the cell goes idle", () => {
+    // A paused-pdb cell shows an unanswered stdin prompt (the debugger box);
+    // once it finishes (e.g. after quitting with `q`), the box is removed.
+    const cell = createCellRuntimeState({
+      status: "running",
+      debuggerActive: true,
+      consoleOutputs: [stdinOutput()],
+    });
+    const next = transitionCell(cell, {
+      cell_id: "cell-1",
+      status: "idle",
+    } as CellMessage);
+    expect(next.debuggerActive).toBe(false);
+    expect(next.consoleOutputs[0].response).toBe("");
+  });
+
+  it("leaves an unanswered prompt interactive while still running", () => {
+    const cell = createCellRuntimeState({
+      status: "running",
+      consoleOutputs: [stdinOutput()],
+    });
+    const next = transitionCell(cell, {
+      cell_id: "cell-1",
+      status: null,
+    } as CellMessage);
+    expect(next.consoleOutputs[0].response).toBeUndefined();
+  });
+});

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/cell.ts

@@ -118,9 +118,11 @@ export function transitionCell(
   // Coalesce console outputs, which are streamed during execution.
   let consoleOutputs = cell.consoleOutputs;
 
-  // If interrupted on the incoming message,
-  // remove the debugger and resolve all stdin for previous console outputs
-  if (didInterruptFromThisMessage) {
+  // When the cell is interrupted or finishes (idle), there is no live
+  // stdin/pdb session, so resolve any dangling stdin prompt. This removes the
+  // debugger box — e.g. after quitting pdb with `q`, which stops the cell
+  // (idle) rather than interrupting it.
+  if (didInterruptFromThisMessage || message.status === "idle") {
     nextCell.debuggerActive = false;
     consoleOutputs = consoleOutputs.map((output) => {
       if (output.channel === "stdin") {

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

@@ -0,0 +1,185 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type { CellId } from "@/core/cells/ids";
+import { createMockObservable } from "@/core/state/__mocks__/mocks";
+import type { Observable } from "@/core/state/observable";
+import {
+  breakpointGutter,
+  debuggerLineHighlighter,
+} from "../debugger-decorations";
+
+vi.mock("../debugger-state", () => ({
+  toggleBreakpoint: vi.fn(),
+}));
+
+import { toggleBreakpoint } from "../debugger-state";
+
+const CODE = `def my_function():
+    x = 1
+    y = 2
+    return x + y
+
+result = my_function()`;
+
+function createLineHighlighterEditor(
+  lineObservable: Observable<number | null>,
+) {
+  const state = EditorState.create({
+    doc: CODE,
+    extensions: [python(), debuggerLineHighlighter(lineObservable)],
+  });
+  return new EditorView({ state, parent: document.body });
+}
+
+function createBreakpointGutterEditor(
+  cid: CellId,
+  breakpointsObservable: Observable<ReadonlySet<number>>,
+) {
+  const state = EditorState.create({
+    doc: CODE,
+    extensions: [python(), breakpointGutter(cid, breakpointsObservable)],
+  });
+  return new EditorView({ state, parent: document.body });
+}
+
+/** Flush the microtask queue and apply the resulting state changes to `view`. */
+async function flush(view: EditorView): Promise<void> {
+  await Promise.resolve();
+  view.dispatch({});
+}
+
+/**
+ * Count rendered breakpoint markers, excluding the always-present hidden
+ * spacer element (`initialSpacer`) that isn't an actual breakpoint.
+ */
+function visibleMarkerCount(view: EditorView): number {
+  const gutterElement = view.dom.querySelector(".cm-breakpoint-gutter");
+  if (!gutterElement) {
+    return 0;
+  }
+  return gutterElement.querySelectorAll(
+    '.cm-gutterElement:not([style*="visibility: hidden"]) .cm-breakpoint-marker',
+  ).length;
+}
+
+describe("debugger-decorations", () => {
+  let view: EditorView | null = null;
+
+  afterEach(() => {
+    if (view) {
+      view.destroy();
+      view = null;
+    }
+    document.body.innerHTML = "";
+    vi.clearAllMocks();
+  });
+
+  describe("debuggerLineHighlighter", () => {
+    it("renders no highlight when the observable starts null", () => {
+      const lineObservable = createMockObservable<number | null>(null);
+      view = createLineHighlighterEditor(lineObservable);
+
+      expect(view.dom.querySelector(".cm-debugger-current-line")).toBeNull();
+    });
+
+    it("highlights the line reported by the observable", () => {
+      const lineObservable = createMockObservable<number | null>(null);
+      view = createLineHighlighterEditor(lineObservable);
+
+      lineObservable.set(2);
+      view.dispatch({});
+
+      expect(
+        view.dom.querySelector(".cm-debugger-current-line"),
+      ).not.toBeNull();
+    });
+
+    it("clears the highlight when the line is set back to null", () => {
+      const lineObservable = createMockObservable<number | null>(2);
+      view = createLineHighlighterEditor(lineObservable);
+      view.dispatch({});
+      expect(
+        view.dom.querySelector(".cm-debugger-current-line"),
+      ).not.toBeNull();
+
+      lineObservable.set(null);
+      view.dispatch({});
+
+      expect(view.dom.querySelector(".cm-debugger-current-line")).toBeNull();
+    });
+
+    it("does not throw for an out-of-range line", () => {
+      const lineObservable = createMockObservable<number | null>(null);
+      view = createLineHighlighterEditor(lineObservable);
+
+      expect(() => {
+        lineObservable.set(9999);
+        view?.dispatch({});
+      }).not.toThrow();
+      expect(view.dom.querySelector(".cm-debugger-current-line")).toBeNull();
+    });
+  });
+
+  describe("breakpointGutter", () => {
+    it("renders no markers when there are no breakpoints", () => {
+      const cid = cellId("cell1");
+      const breakpointsObservable = createMockObservable<ReadonlySet<number>>(
+        new Set(),
+      );
+      view = createBreakpointGutterEditor(cid, breakpointsObservable);
+
+      expect(visibleMarkerCount(view)).toBe(0);
+    });
+
+    it("syncs pre-existing breakpoints on mount", async () => {
+      const cid = cellId("cell1");
+      const breakpointsObservable = createMockObservable<ReadonlySet<number>>(
+        new Set([2]),
+      );
+      view = createBreakpointGutterEditor(cid, breakpointsObservable);
+      await flush(view); // initial sync is deferred to a microtask
+
+      expect(visibleMarkerCount(view)).toBe(1);
+    });
+
+    it("adds a marker when the observable reports a new breakpoint", () => {
+      const cid = cellId("cell1");
+      const breakpointsObservable = createMockObservable<ReadonlySet<number>>(
+        new Set(),
+      );
+      view = createBreakpointGutterEditor(cid, breakpointsObservable);
+
+      breakpointsObservable.set(new Set([1, 3]));
+      view.dispatch({});
+
+      expect(visibleMarkerCount(view)).toBe(2);
+    });
+
+    it("toggles a breakpoint when an existing marker is clicked", async () => {
+      const cid = cellId("cell1");
+      // A single breakpoint on line 1 renders exactly one non-spacer
+      // gutter element, so the click lands unambiguously on that line.
+      const breakpointsObservable = createMockObservable<ReadonlySet<number>>(
+        new Set([1]),
+      );
+      view = createBreakpointGutterEditor(cid, breakpointsObservable);
+      await flush(view);
+
+      const gutterElement = view.dom.querySelector(".cm-breakpoint-gutter");
+      const lineElement = gutterElement?.querySelector(
+        '.cm-gutterElement:not([style*="visibility: hidden"])',
+      );
+      expect(lineElement).not.toBeNull();
+      lineElement?.dispatchEvent(
+        new MouseEvent("mousedown", { bubbles: true }),
+      );
+
+      expect(toggleBreakpoint).toHaveBeenCalledWith(cid, 1);
+    });
+  });
+});