@marimo-team/islands 0.25.3-dev322 → 0.25.3-dev348

package/dist/{useDateFormatter-2k4mHRBq.js → useDateFormatter-B4ygYuUU.js}

@@ -1,7 +1,7 @@
 import { a as e } from "./rolldown-runtime-BPCzCO3I.js";
 import { t } from "./react-Bgs_wQh3.js";
-import { h as n } from "./tooltip-DBv8QUvS.js";
-import { C as r } from "./select-Cv6lR39D.js";
+import { h as n } from "./tooltip-BVbJPkcl.js";
+import { C as r } from "./select-suaAxFuo.js";
 var i = n("chart-pie", [["path", {
 	d: "M21 12c.552 0 1.005-.449.95-.998a10 10 0 0 0-8.953-8.951c-.55-.055-.998.398-.998.95v8a1 1 0 0 0 1 1z",
 	key: "pzmjnu"

package/dist/{useDeepCompareMemoize-Cb2GRQIS.js → useDeepCompareMemoize-Bcu0DkCS.js}

@@ -1,6 +1,6 @@
 import { a as e } from "./rolldown-runtime-BPCzCO3I.js";
 import { t } from "./react-Bgs_wQh3.js";
-import { Mt as n } from "./use-restore-focus-BxJh5Zfa.js";
+import { Pt as n } from "./use-restore-focus-tPmARo1Q.js";
 //#region src/hooks/useDeepCompareMemoize.ts
 var r = /* @__PURE__ */ e(t(), 1);
 function i(e) {

package/dist/{useLifecycle-B2fV3c8N.js → useLifecycle-MTdtnR3-.js}

@@ -2,7 +2,7 @@ import { a as e } from "./rolldown-runtime-BPCzCO3I.js";
 import { h as t, r as n } from "./button-qQhINE7z.js";
 import { t as r } from "./react-Bgs_wQh3.js";
 import { t as i } from "./compiler-runtime-C5KHtL5D.js";
-import { h as a } from "./tooltip-DBv8QUvS.js";
+import { h as a } from "./tooltip-BVbJPkcl.js";
 import { t as o } from "./jsx-runtime-D8btoEHk.js";
 var s = a("calendar", [
 	["path", {

package/dist/{vega-component-do3stf0K.js → vega-component-BhBbP4S7.js}

@@ -2,22 +2,22 @@ import { a as e } from "./rolldown-runtime-BPCzCO3I.js";
 import { g as t, h as n, m as r, s as i } from "./button-qQhINE7z.js";
 import { t as a } from "./react-Bgs_wQh3.js";
 import { t as o } from "./compiler-runtime-C5KHtL5D.js";
-import { t as s } from "./tooltip-DBv8QUvS.js";
-import { l as c, t as l } from "./config-B171m9zn.js";
+import { t as s } from "./tooltip-BVbJPkcl.js";
+import { l as c, t as l } from "./config-CvBnR945.js";
 import { t as u } from "./jsx-runtime-D8btoEHk.js";
-import { u as d } from "./use-restore-focus-BxJh5Zfa.js";
-import { n as f } from "./error-banner-3_4Lc93Y.js";
+import { d } from "./use-restore-focus-tPmARo1Q.js";
+import { n as f } from "./error-banner-DUB_qUeL.js";
 import { t as p } from "./debounce-DrhhziJv.js";
 import { t as ee } from "./objects-DCmMvZtn.js";
 import { t as te } from "./useEvent-BoN7U38T.js";
-import { a as ne, i as re, n as ie, t as m } from "./formats-aUcw6MrS.js";
+import { a as ne, i as re, n as ie, t as m } from "./formats-hF6fETr6.js";
 import { a as h, t as g } from "./arrays-DoQ-AILx.js";
 import { t as ae } from "./useAsyncData-P5-Oovhq.js";
 import { t as _ } from "./semaphore-8KhQouXt.js";
 import { W as oe, n as se } from "./vega-loader.browser-Co9gYE8Z.js";
 import { t as ce } from "./react-vega-CGgQEf0O.js";
 import { a as le, n as v } from "./loader-Bin_dutt.js";
-import { t as ue } from "./useDeepCompareMemoize-Cb2GRQIS.js";
+import { t as ue } from "./useDeepCompareMemoize-Bcu0DkCS.js";
 import { t as de } from "./use-vega-container-remeasure-BrUpWQfF.js";
 //#region src/plugins/impl/vega/fix-relative-url.ts
 var fe = /* @__PURE__ */ e(a(), 1), pe = o();

package/dist/{vennDiagram-4TSXK5OY-BLq7sNjb.js → vennDiagram-4TSXK5OY-BrHqWB1U.js}

@@ -3,7 +3,7 @@ import { t } from "./src-D5uZaZib.js";
 import { n } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { $ as r, H as i, K as a, U as o, a as s, b as c, c as l, et as u, f as d, nt as f, tt as p, v as m, w as h, y as g } from "./chunk-DU6HZSFF-4SDIhUTY.js";
 import { t as _ } from "./rough.esm-CtV9AnV1.js";
-import { p as v } from "./mermaid-C7lQ_2LJ.js";
+import { p as v } from "./mermaid--2K0cmt9.js";
 //#region ../node_modules/.pnpm/khroma@2.1.0/node_modules/khroma/dist/methods/transparentize.js
 var y = (e, t) => p(e, "a", -t), b = 1e-10;
 function x(e, t) {

package/dist/{wardleyDiagram-VM6X3IG4-oyJEWWtb.js → wardleyDiagram-VM6X3IG4-Cms1K_3v.js}

@@ -2,7 +2,7 @@ import { i as e } from "./chunk-75Z2AOVW-76L6My3P.js";
 import { n as t } from "./src-D5uZaZib.js";
 import { n } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { D as r, H as i, K as a, U as o, a as s, b as c, c as l, v as u, w as d, x as f, y as p } from "./chunk-DU6HZSFF-4SDIhUTY.js";
-import { p as m } from "./mermaid-C7lQ_2LJ.js";
+import { p as m } from "./mermaid--2K0cmt9.js";
 import { n as h } from "./mermaid-parser.core-DtOnl7m4.js";
 import { t as g } from "./chunk-JWPE2WC7-DJvv1t8U.js";
 //#region ../node_modules/.pnpm/mermaid@11.17.2/node_modules/mermaid/dist/chunks/mermaid.core/wardleyDiagram-VM6X3IG4.mjs

package/dist/{xychartDiagram-S5SC5T6Z-DMjWgQPm.js → xychartDiagram-S5SC5T6Z-_kUiODzS.js}

@@ -9,7 +9,7 @@ import { n as s } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { D as c, H as l, K as u, U as d, a as f, b as p, c as m, f as h, v as g, w as _, y as v, z as y } from "./chunk-DU6HZSFF-4SDIhUTY.js";
 import "./chunk-PWAF6VOD-CRu_RGP2.js";
 import { t as b } from "./chunk-GMAD6QVW-BJGF8QZL.js";
-import { p as x } from "./mermaid-C7lQ_2LJ.js";
+import { p as x } from "./mermaid--2K0cmt9.js";
 //#region ../node_modules/.pnpm/d3-scale@4.0.2/node_modules/d3-scale/src/band.js
 function S() {
 	var e = t().unknown(void 0), r = e.domain, i = e.range, o = 0, s = 1, c, l, u = !1, d = 0, f = 0, p = .5;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.25.3-dev322",
+  "version": "0.25.3-dev348",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/__mocks__/requests.ts

@@ -19,9 +19,7 @@ export const MockRequestClient = {
       sendRestart: vi.fn().mockResolvedValue({}),
       saveLayout: vi.fn().mockResolvedValue({}),
       sendCopy: vi.fn().mockResolvedValue({}),
-      sendStdin: vi.fn().mockResolvedValue({}),
       sendFormat: vi.fn().mockResolvedValue({ codes: {} }),
-      sendInterrupt: vi.fn().mockResolvedValue({}),
       sendShutdown: vi.fn().mockResolvedValue({}),
       sendRunScratchpad: vi.fn().mockResolvedValue({}),
       sendInstantiate: vi.fn().mockResolvedValue({}),

package/src/components/chat/__tests__/notebook-context.test.ts

@@ -0,0 +1,199 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { type UIMessage, validateUIMessages } from "ai";
+import { describe, expect, it } from "vitest";
+import { withNotebookContext } from "../notebook-context";
+
+const question: UIMessage["parts"] = [{ type: "text", text: "Help me" }];
+
+function user(parts: UIMessage["parts"]): UIMessage {
+  return { id: "user", role: "user", parts };
+}
+
+function context(data: unknown) {
+  return { type: "data-marimo-notebook", data };
+}
+
+const initial = {
+  sessionId: "session-1",
+  cells: { a: "x = 1", b: "y = x + 1" },
+};
+
+function initialHistory() {
+  return [user(withNotebookContext(question, [], initial))];
+}
+
+describe("notebook context", () => {
+  it("captures a complete initial snapshot including empty cells", () => {
+    expect(
+      withNotebookContext(question, [], {
+        sessionId: "session-1",
+        cells: { a: "" },
+      }),
+    ).toEqual([
+      context({
+        sessionId: "session-1",
+        cellIds: ["a"],
+        cells: [{ id: "a", code: "" }],
+      }),
+      ...question,
+    ]);
+  });
+
+  it("sends only a snapshot marker after restoring unchanged history", async () => {
+    const savedHistory = JSON.stringify(initialHistory());
+    const restored = await validateUIMessages({
+      messages: JSON.parse(savedHistory),
+    });
+    expect(withNotebookContext(question, restored, initial)).toEqual([
+      context({ sessionId: "session-1", cells: [] }),
+      ...question,
+    ]);
+  });
+
+  it("sends only changed or added cells without mutating the history", () => {
+    const history = initialHistory();
+    const saved = structuredClone(history);
+    expect(
+      withNotebookContext(question, history, {
+        sessionId: "session-1",
+        cells: { a: "x = 2", b: "y = x + 1", c: "" },
+      }),
+    ).toEqual([
+      context({
+        sessionId: "session-1",
+        cellIds: ["a", "b", "c"],
+        cells: [
+          { id: "a", code: "x = 2" },
+          { id: "c", code: "" },
+        ],
+      }),
+      ...question,
+    ]);
+    expect(history).toEqual(saved);
+  });
+
+  it("tracks deletions, reordering, and re-added cells", () => {
+    const history = initialHistory();
+    const deleted = withNotebookContext(question, history, {
+      sessionId: "session-1",
+      cells: { b: "y = x + 1" },
+    });
+    expect(deleted).toEqual([
+      context({ sessionId: "session-1", cellIds: ["b"], cells: [] }),
+      ...question,
+    ]);
+    history.push(user(deleted));
+    const readded = withNotebookContext(question, history, initial);
+    expect(readded).toEqual([
+      context({
+        sessionId: "session-1",
+        cellIds: ["a", "b"],
+        cells: [{ id: "a", code: "x = 1" }],
+      }),
+      ...question,
+    ]);
+    history.push(user(readded));
+    expect(
+      withNotebookContext(question, history, {
+        sessionId: "session-1",
+        cells: { b: "y = x + 1", a: "x = 1" },
+      }),
+    ).toEqual([
+      context({ sessionId: "session-1", cellIds: ["b", "a"], cells: [] }),
+      ...question,
+    ]);
+  });
+
+  it("represents an empty notebook explicitly", () => {
+    expect(
+      withNotebookContext(question, initialHistory(), {
+        sessionId: "session-1",
+        cells: {},
+      }),
+    ).toEqual([
+      context({ sessionId: "session-1", cellIds: [], cells: [] }),
+      ...question,
+    ]);
+  });
+
+  it("uses only the retained history when editing a message", () => {
+    const history = initialHistory();
+    const changed = { sessionId: "session-1", cells: { a: "x = 2" } };
+    const next = withNotebookContext(question, history, changed);
+    expect(
+      withNotebookContext(question, [...history, user(next)], initial),
+    ).toEqual([
+      context({
+        sessionId: "session-1",
+        cellIds: ["a", "b"],
+        cells: [
+          { id: "a", code: "x = 1" },
+          { id: "b", code: "y = x + 1" },
+        ],
+      }),
+      ...question,
+    ]);
+    expect(withNotebookContext(question, history, initial)).toEqual([
+      context({ sessionId: "session-1", cells: [] }),
+      ...question,
+    ]);
+  });
+
+  it("updates session identity without repeating code", () => {
+    expect(
+      withNotebookContext(question, initialHistory(), {
+        ...initial,
+        sessionId: "session-2",
+      }),
+    ).toEqual([context({ sessionId: "session-2", cells: [] }), ...question]);
+  });
+
+  it("captures a full snapshot for chats created before notebook context existed", () => {
+    const legacyHistory = [user(question)];
+    expect(withNotebookContext(question, legacyHistory, initial)).toEqual(
+      initialHistory()[0].parts,
+    );
+  });
+
+  it("compares against user snapshots independently of tool results", () => {
+    const history = initialHistory();
+    history.push({
+      id: "tool-result",
+      role: "assistant",
+      parts: [
+        {
+          type: "tool-edit_cell",
+          toolCallId: "edit-a",
+          state: "output-available",
+          input: { id: "a", code: "x = 2" },
+          output: { code: "x = 2" },
+        },
+      ],
+    });
+    expect(
+      withNotebookContext(question, history, {
+        sessionId: "session-1",
+        cells: { a: "x = 2", b: "y = x + 1" },
+      }),
+    ).toEqual([
+      context({
+        sessionId: "session-1",
+        cells: [{ id: "a", code: "x = 2" }],
+      }),
+      ...question,
+    ]);
+  });
+
+  it("can enter and leave full access without sending its notebook code", () => {
+    const full = withNotebookContext(question, [], { sessionId: "session-1" });
+    expect(full).toEqual([context({ sessionId: "session-1" }), ...question]);
+    const read = withNotebookContext(question, [user(full)], initial);
+    expect(read).toEqual(initialHistory()[0].parts);
+    expect(
+      withNotebookContext(question, [user(full), user(read)], {
+        sessionId: "session-1",
+      }),
+    ).toEqual([context({ sessionId: "session-1" }), ...question]);
+  });
+});

package/src/components/chat/chat-panel.tsx

@@ -93,7 +93,9 @@ import {
   useFileState,
   useMessageQueue,
 } from "./chat-utils";
-import { getCodes } from "@/core/codemirror/copilot/getCodes";
+import { notebookCellCodes } from "@/core/codemirror/copilot/getCodes";
+import { getSessionId } from "@/core/kernel/session";
+import { withNotebookContext } from "./notebook-context";
 import { focusInputAndMoveToEnd } from "@/core/codemirror/utils";
 import ScrollToBottomButton from "./acp/scroll-to-bottom-button";
 import { useChatControllerId } from "./use-chat-controller-id";
@@ -501,7 +503,6 @@ const ChatPanelBody = () => {
 
           const completionBody = {
             uiMessages: options.messages,
-            includeOtherCode: getCodes(""),
             options: store.get(chatOptionsAtom),
           };
 
@@ -588,10 +589,21 @@ const ChatPanelBody = () => {
     },
   });
 
+  const contextualize = useEvent(
+    (parts: ChatMessagePart[], history: UIMessage[]) =>
+      withNotebookContext(parts, history, {
+        sessionId: getSessionId(),
+        cells:
+          store.get(aiAtom)?.mode === "code_mode"
+            ? undefined
+            : store.get(notebookCellCodes),
+      }),
+  );
+
   const sendUserMessage = useEvent((parts: ChatMessagePart[]) => {
     clearAbortState();
     turn.start();
-    sendMessage({ role: "user", parts });
+    sendMessage({ role: "user", parts: contextualize(parts, messages) });
   });
 
   const tryFlushQueuedMessages = useEvent(
@@ -753,12 +765,15 @@ const ChatPanelBody = () => {
       turn.start();
       sendMessage({
         role: "user",
-        parts: [
-          { type: "text" as const, text: initialMessage },
-          ...(contextPart ? [contextPart] : []),
-          ...(fileParts ?? []),
-          ...attachments,
-        ],
+        parts: contextualize(
+          [
+            { type: "text" as const, text: initialMessage },
+            ...(contextPart ? [contextPart] : []),
+            ...(fileParts ?? []),
+            ...attachments,
+          ],
+          [],
+        ),
       });
       clearFiles();
       setInput("");
@@ -812,12 +827,15 @@ const ChatPanelBody = () => {
       sendMessage({
         messageId: messageId, // replace the message
         role: "user",
-        parts: [
-          { type: "text", text: newValue },
-          ...(contextPart ? [contextPart] : []),
-          ...userFileParts,
-          ...attachments,
-        ],
+        parts: contextualize(
+          [
+            { type: "text", text: newValue },
+            ...(contextPart ? [contextPart] : []),
+            ...userFileParts,
+            ...attachments,
+          ],
+          messages.slice(0, index),
+        ),
       });
     },
   );

package/src/components/chat/notebook-context.ts

@@ -0,0 +1,80 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { DataUIPart, UIMessage } from "ai";
+import { z } from "zod";
+
+const notebookContextSchema = z.object({
+  sessionId: z.string(),
+  cells: z.array(z.object({ id: z.string(), code: z.string() })).optional(),
+  cellIds: z.array(z.string()).optional(),
+});
+
+type NotebookContext = z.infer<typeof notebookContextSchema>;
+type NotebookContextPart = DataUIPart<{
+  "marimo-notebook": NotebookContext;
+}>;
+
+export const NOTEBOOK_CONTEXT_PART_TYPE =
+  "data-marimo-notebook" as const satisfies NotebookContextPart["type"];
+
+interface NotebookSnapshot {
+  sessionId: string;
+  cells?: Readonly<Record<string, string>>;
+}
+
+export function withNotebookContext(
+  parts: UIMessage["parts"],
+  history: UIMessage[],
+  snapshot: NotebookSnapshot,
+): UIMessage["parts"] {
+  let cellIds: string[] | undefined;
+  const cells = new Map<string, string>();
+  for (const message of history) {
+    if (message.role !== "user") {
+      continue;
+    }
+    for (const part of message.parts) {
+      if (part.type !== NOTEBOOK_CONTEXT_PART_TYPE) {
+        continue;
+      }
+      const parsed = notebookContextSchema.safeParse(part.data);
+      if (!parsed.success) {
+        continue;
+      }
+      if (parsed.data.cellIds) {
+        cellIds = parsed.data.cellIds;
+        const retained = new Set(cellIds);
+        for (const id of cells.keys()) {
+          if (!retained.has(id)) {
+            cells.delete(id);
+          }
+        }
+      }
+      for (const cell of parsed.data.cells ?? []) {
+        cells.set(cell.id, cell.code);
+      }
+    }
+  }
+
+  const update: NotebookContext = { sessionId: snapshot.sessionId };
+  if (snapshot.cells) {
+    const nextIds = Object.keys(snapshot.cells);
+    if (
+      !cellIds ||
+      nextIds.length !== cellIds.length ||
+      nextIds.some((id, index) => id !== cellIds[index])
+    ) {
+      update.cellIds = nextIds;
+    }
+    const changed = Object.entries(snapshot.cells)
+      .filter(([id, code]) => cells.get(id) !== code)
+      .map(([id, code]) => ({ id, code }));
+    update.cells = changed;
+  }
+
+  const context: NotebookContextPart = {
+    type: NOTEBOOK_CONTEXT_PART_TYPE,
+    data: update,
+  };
+  return [context, ...parts];
+}

package/src/components/editor/cell/StopButton.tsx

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { SquareIcon } from "lucide-react";
 import type { JSX } from "react";
-import { useRequestClient } from "@/core/network/requests";
+import { interruptRuntime } from "@/core/executions";
+import { store } from "@/core/state/jotai";
 import type { RuntimeState } from "@/core/network/types";
 import { isAppInteractionDisabled } from "@/core/session/connection-utils";
 import type { ConnectionState } from "@/core/session/types";
@@ -15,7 +16,6 @@ export const StopButton = (props: {
   connectionState: ConnectionState;
 }): JSX.Element => {
   const { connectionState, status } = props;
-  const { sendInterrupt } = useRequestClient();
 
   const running = status === "running";
 
@@ -26,7 +26,7 @@ export const StopButton = (props: {
     <ToolbarItem
       tooltip={renderShortcut("global.interrupt")}
       disabled={isAppInteractionDisabled(connectionState) || !showInterrupt}
-      onClick={showInterrupt ? sendInterrupt : Functions.NOOP}
+      onClick={showInterrupt ? () => interruptRuntime(store) : Functions.NOOP}
       variant={showInterrupt ? "stale" : "disabled"}
       data-testid="run-button"
     >

package/src/components/editor/cell/cell-context-menu.tsx

@@ -20,6 +20,7 @@ import {
 import { menuItemVariants } from "@/components/ui/menu-items";
 import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
+import { useRestoreFocus } from "@/components/ui/use-restore-focus";
 import { useCellData, useCellRuntime } from "@/core/cells/cells";
 import { CellOutputId } from "@/core/cells/ids";
 import { isOutputEmpty } from "@/core/cells/outputs";
@@ -71,6 +72,8 @@ export const CellActionsContextMenu = ({
   const [imageRightClicked, setImageRightClicked] =
     React.useState<HTMLImageElement>();
   const suppressCloseAutoFocus = React.useRef(false);
+  const pendingAction = React.useRef<(() => void) | null>(null);
+  const restoreFocus = useRestoreFocus();
 
   const DEFAULT_CONTEXT_MENU_ITEMS: ActionButton[] = [
     {
@@ -194,7 +197,13 @@ export const CellActionsContextMenu = ({
     .filter((group) => group.length > 0);
 
   return (
-    <ContextMenu>
+    <ContextMenu
+      onOpenChange={(open) => {
+        if (open) {
+          restoreFocus.onOpenAutoFocus();
+        }
+      }}
+    >
       <ContextMenuTrigger
         disabled={disabled}
         onContextMenu={(evt) => {
@@ -212,6 +221,13 @@ export const CellActionsContextMenu = ({
         className="w-[300px]"
         scrollable={true}
         onCloseAutoFocus={(evt) => {
+          const action = pendingAction.current;
+          pendingAction.current = null;
+          if (action) {
+            restoreFocus.onCloseAutoFocus(evt);
+            evt.preventDefault();
+            action();
+          }
           if (suppressCloseAutoFocus.current) {
             evt.preventDefault();
             suppressCloseAutoFocus.current = false;
@@ -270,7 +286,7 @@ export const CellActionsContextMenu = ({
                           if (action.disableClick || action.disabled) {
                             return;
                           }
-                          action.handle(evt);
+                          pendingAction.current = () => action.handle(evt);
                         }}
                         variant={action.variant}
                       >

package/src/components/editor/cell/code/__tests__/language-toggle.test.tsx

@@ -7,7 +7,7 @@ import { createRef } from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import { switchLanguage } from "@/core/codemirror/language/extension";
-import { LanguageToggle } from "../language-toggle";
+import { LanguageToggle, LanguageToggles } from "../language-toggle";
 
 vi.mock("@/core/codemirror/language/extension", () => ({
   switchLanguage: vi.fn(),
@@ -51,4 +51,41 @@ describe("LanguageToggle", () => {
       view.destroy();
     }
   });
+
+  it("uses the correct callback for each language toggle", () => {
+    const view = new EditorView({ state: EditorState.create() });
+    const editorViewRef = createRef<EditorView>();
+    editorViewRef.current = view;
+
+    const onAfterToggleSQL = vi.fn();
+    const onAfterToggleMarkdown = vi.fn();
+
+    try {
+      render(
+        <TooltipProvider>
+          <LanguageToggles
+            editorViewRef={editorViewRef}
+            code=""
+            currentLanguageAdapter="python"
+            onAfterToggleSQL={onAfterToggleSQL}
+            onAfterToggleMarkdown={onAfterToggleMarkdown}
+          />
+        </TooltipProvider>,
+      );
+
+      const toggleButtons = screen.getAllByTestId("language-toggle-button");
+
+      fireEvent.click(toggleButtons[0]);
+
+      expect(onAfterToggleSQL).toHaveBeenCalledOnce();
+      expect(onAfterToggleMarkdown).not.toHaveBeenCalled();
+
+      fireEvent.click(toggleButtons[1]);
+
+      expect(onAfterToggleMarkdown).toHaveBeenCalledOnce();
+      expect(onAfterToggleSQL).toHaveBeenCalledOnce();
+    } finally {
+      view.destroy();
+    }
+  });
 });

package/src/components/editor/cell/code/cell-editor.tsx

@@ -511,7 +511,8 @@ const CellEditorInternal = ({
               code={code}
               editorViewRef={editorViewRef}
               currentLanguageAdapter={languageAdapter}
-              onAfterToggle={afterToggleMarkdown}
+              onAfterToggleSQL={afterToggleSQL}
+              onAfterToggleMarkdown={afterToggleMarkdown}
             />
           </div>
         )}

package/src/components/editor/cell/code/language-toggle.tsx

@@ -16,7 +16,8 @@ interface LanguageTogglesProps {
   editorViewRef: React.RefObject<EditorView | null>;
   code: string;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
-  onAfterToggle: () => void;
+  onAfterToggleSQL: () => void;
+  onAfterToggleMarkdown: () => void;
   /**
    * Classes for the wrapper element. Defaults to the absolutely-positioned,
    * hover-revealed placement used inside the notebook cell editor.
@@ -28,7 +29,8 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
   editorViewRef,
   code,
   currentLanguageAdapter,
-  onAfterToggle,
+  onAfterToggleSQL,
+  onAfterToggleMarkdown,
   className = "absolute right-3 top-2 z-20 flex hover-action gap-1",
 }) => {
   const canUseMarkdown = useMemo(
@@ -55,7 +57,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         }
         toType="sql"
         displayName="SQL"
-        onAfterToggle={onAfterToggle}
+        onAfterToggle={onAfterToggleSQL}
       />
       <LanguageToggle
         editorViewRef={editorViewRef}
@@ -72,7 +74,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         }
         toType="markdown"
         displayName="Markdown"
-        onAfterToggle={onAfterToggle}
+        onAfterToggle={onAfterToggleMarkdown}
       />
       <LanguageToggle
         editorViewRef={editorViewRef}