@marimo-team/frontend 0.25.3-dev5 → 0.25.3-dev52

package/src/components/editor/actions/pair-with-agent-modal.tsx

@@ -17,18 +17,16 @@ import { Events } from "@/utils/events";
 import { Tooltip } from "@/components/ui/tooltip";
 import { asRemoteURL, useRuntimeManager } from "@/core/runtime/config";
 import { API } from "@/core/network/api";
-import { pairPreviewAtom } from "@/core/config/pair";
+import { pairConfigAtom } from "@/core/config/pair";
 import { stableSessionIdAtom } from "@/core/kernel/session";
 import {
   AGENT_LABELS,
   AGENT_TABS,
   type AgentTab,
   type ConnectionInfo,
-  getFileFromURL,
   getRawPrompt,
   getTerminalCommand,
   maskToken,
-  SKILL_INSTALL,
   TERMINAL_TABS,
 } from "./pair-with-agent-commands";
 
@@ -52,15 +50,64 @@ export const PairWithAgentModal: React.FC<{
 }> = ({ onClose }) => {
   const [activeTab, setActiveTab] = useState<AgentTab>("claude");
   const runtimeManager = useRuntimeManager();
-  const preview = useAtomValue(pairPreviewAtom);
+  const config = useAtomValue(pairConfigAtom);
+  const hostEnabled = config?.hostEnabled ?? false;
   const stableSessionId = useAtomValue(stableSessionIdAtom);
-  const commandStep = preview ? 1 : 2;
+  const [lookup, setLookup] = useState<
+    {
+      sessionId: string;
+    } & (
+      | { status: "ready"; notebookId: string | null }
+      | { status: "loading" | "error" }
+    )
+  >();
+  useEffect(() => {
+    if (!hostEnabled || !stableSessionId) {
+      setLookup(undefined);
+      return;
+    }
+    const sessionId = stableSessionId;
+    setLookup({ sessionId, status: "loading" });
+    let cancelled = false;
+    API.get<Record<string, { session_id: string; notebook_id: string | null }>>(
+      "/sessions",
+    )
+      .then((sessions) => {
+        if (!cancelled) {
+          const session = Object.values(sessions).find(
+            (session) => session.session_id === sessionId,
+          );
+          setLookup(
+            session
+              ? { sessionId, status: "ready", notebookId: session.notebook_id }
+              : { sessionId, status: "error" },
+          );
+        }
+      })
+      .catch(() => {
+        if (!cancelled) {
+          setLookup({ sessionId, status: "error" });
+        }
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [hostEnabled, stableSessionId]);
+  const currentLookup =
+    lookup?.sessionId === stableSessionId ? lookup : undefined;
+  const verified = !hostEnabled || currentLookup?.status === "ready";
+  const notebookId =
+    hostEnabled && currentLookup?.status === "ready"
+      ? currentLookup.notebookId
+      : undefined;
+  const targetReady = Boolean(stableSessionId) && verified;
   const authToken = useAuthToken();
   const hasToken = Boolean(authToken);
   const connection: ConnectionInfo = {
     url: runtimeManager.httpURL.toString(),
-    file: preview ? undefined : getFileFromURL(window.location.href),
-    session: preview ? (stableSessionId ?? undefined) : undefined,
+    notebook: notebookId ?? undefined,
+    session:
+      targetReady && !notebookId ? (stableSessionId ?? undefined) : undefined,
   };
 
   return (
@@ -83,6 +130,13 @@ export const PairWithAgentModal: React.FC<{
       </DialogHeader>
 
       <div className="flex flex-col gap-4 py-2">
+        {!targetReady && (
+          <output className="text-sm text-muted-foreground">
+            {currentLookup?.status === "error"
+              ? "Could not verify this notebook. Close this dialog and try again."
+              : "Waiting for the notebook connection…"}
+          </output>
+        )}
         <Tabs
           value={activeTab}
           onValueChange={(v) => setActiveTab(v as AgentTab)}
@@ -101,30 +155,27 @@ export const PairWithAgentModal: React.FC<{
               value={tab}
               className="mt-4 flex flex-col gap-4"
             >
-              {!preview && (
-                <Step
-                  index={1}
-                  title="Install the skill"
-                  hint="Run once per machine."
-                >
-                  <CommandBlock command={SKILL_INSTALL} />
-                </Step>
-              )}
-              <Step index={commandStep} title="Run in your terminal">
+              <Step
+                index={1}
+                title="Run in your terminal"
+                hint={
+                  hostEnabled && tab === "opencode"
+                    ? "Agent connections require OpenCode 2.x."
+                    : undefined
+                }
+              >
                 <CommandBlock
                   command={getTerminalCommand(
                     tab,
                     connection,
                     hasToken,
-                    preview,
+                    config,
                   )}
+                  disabled={!targetReady}
                 />
               </Step>
               {hasToken && authToken && (
-                <Step
-                  index={commandStep + 1}
-                  title="Paste when prompted for a token"
-                >
+                <Step index={2} title="Paste when prompted for a token">
                   <CommandBlock
                     command={authToken}
                     display={maskToken(authToken)}
@@ -136,17 +187,8 @@ export const PairWithAgentModal: React.FC<{
           ))}
 
           <TabsContent value="prompt" className="mt-4 flex flex-col gap-4">
-            {!preview && (
-              <Step
-                index={1}
-                title="Make sure the marimo-pair skill is available to your agent"
-                hint="Skip if your agent already has it."
-              >
-                <CommandBlock command={SKILL_INSTALL} />
-              </Step>
-            )}
             <Step
-              index={commandStep}
+              index={1}
               title="Copy this prompt into your agent"
               hint={
                 hasToken
@@ -155,13 +197,14 @@ export const PairWithAgentModal: React.FC<{
               }
             >
               <CommandBlock
-                command={getRawPrompt(connection, authToken, preview)}
+                command={getRawPrompt(connection, authToken, config)}
                 display={getRawPrompt(
                   connection,
                   authToken ? maskToken(authToken) : null,
-                  preview,
+                  config,
                 )}
                 multiline={true}
+                disabled={!targetReady}
               />
             </Step>
           </TabsContent>
@@ -199,11 +242,21 @@ const CommandBlock: React.FC<{
   display?: string;
   multiline?: boolean;
   copyLabel?: string;
-}> = ({ command, display, multiline = false, copyLabel = "Copy command" }) => {
+  disabled?: boolean;
+}> = ({
+  command,
+  display,
+  multiline = false,
+  copyLabel = "Copy command",
+  disabled = false,
+}) => {
   const [copied, setCopied] = useState(false);
 
   const copy = Events.stopPropagation(async (e) => {
     e.preventDefault();
+    if (disabled) {
+      return;
+    }
     await copyToClipboard(command);
     setCopied(true);
     setTimeout(() => setCopied(false), 2000);
@@ -218,6 +271,7 @@ const CommandBlock: React.FC<{
         <Tooltip content="Copied!" open={copied}>
           <Button
             aria-label="Copy prompt"
+            disabled={disabled}
             onClick={copy}
             size="xs"
             variant="ghost"
@@ -240,7 +294,13 @@ const CommandBlock: React.FC<{
         {display ?? command}
       </code>
       <Tooltip content="Copied!" open={copied}>
-        <Button aria-label={copyLabel} onClick={copy} size="xs" variant="ghost">
+        <Button
+          aria-label={copyLabel}
+          onClick={copy}
+          size="xs"
+          variant="ghost"
+          disabled={disabled}
+        >
           {copied ? (
             <CheckIcon size={14} strokeWidth={1.5} />
           ) : (

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -106,9 +106,9 @@ export const AddCellWithAI: React.FC<{
   const runtimeManager = useRuntimeManager();
 
   const inputRef = useRef<ReactCodeMirrorRef>(null);
-  const submissionController = useRef(
-    new StagedCellSubmissionController(),
-  ).current;
+  const [submissionController] = useState(
+    () => new StagedCellSubmissionController(),
+  );
 
   const fileInputRef = useRef<HTMLInputElement>(null);
   const { files, addFiles, removeFile } = useFileState();

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { DatabaseIcon, DiamondPlusIcon, PlusIcon } from "lucide-react";
-import { useState } from "react";
+import { useRef, useState } from "react";
 import { Button } from "@/components/editor/inputs/Inputs";
 import { MinimalHotkeys } from "@/components/shortcuts/renderShortcut";
 import {
@@ -36,7 +36,8 @@ export const CreateCellButton = ({
   const { createNewCell, addSetupCellIfDoesntExist } = useCellActions();
   const shortcut = `${oneClickShortcut}-Click`;
   const [open, setOpen] = useState(false);
-  const [justOpened, setJustOpened] = useState(false);
+  // Time the dropdown last opened, to ignore the click that opened it
+  const openedAt = useRef(0);
 
   const baseTooltipContent =
     getConnectionTooltip(connectionState) || tooltipContent;
@@ -83,9 +84,7 @@ export const CreateCellButton = ({
 
   const openDropdown = () => {
     setOpen(true);
-    setJustOpened(true);
-    // Allow interactions after a brief delay to prevent the dropdown items immediately being clicked
-    setTimeout(() => setJustOpened(false), 200);
+    openedAt.current = Date.now();
   };
 
   // We use onPointerDownCapture (not onPointerDown) to intercept events in
@@ -121,15 +120,14 @@ export const CreateCellButton = ({
   const handleOpenChange = (isOpen: boolean) => {
     setOpen(isOpen);
     if (isOpen) {
-      setJustOpened(true);
-      // Allow interactions after a brief delay
-      setTimeout(() => setJustOpened(false), 200);
+      openedAt.current = Date.now();
     }
   };
 
   const handleFirstItemClick = (e: React.MouseEvent) => {
-    // Hack to prevent the first item from being clicked when the dropdown is opened
-    if (justOpened) {
+    // Hack to prevent the first item from being clicked when the dropdown is
+    // opened: allow interactions only after a brief delay
+    if (Date.now() - openedAt.current < 200) {
       e.preventDefault();
       e.stopPropagation();
       return;

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

@@ -0,0 +1,54 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { fireEvent, render, screen } from "@testing-library/react";
+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";
+
+vi.mock("@/core/codemirror/language/extension", () => ({
+  switchLanguage: vi.fn(),
+}));
+
+describe("LanguageToggle", () => {
+  afterEach(() => {
+    vi.mocked(switchLanguage).mockClear();
+  });
+
+  it("uses the editor mounted after the button rendered", () => {
+    const editorViewRef = createRef<EditorView>();
+    const onAfterToggle = vi.fn();
+    render(
+      <TooltipProvider>
+        <LanguageToggle
+          editorViewRef={editorViewRef}
+          currentLanguageAdapter="python"
+          canSwitchToLanguage={true}
+          toType="markdown"
+          displayName="Markdown"
+          icon="Markdown"
+          onAfterToggle={onAfterToggle}
+        />
+      </TooltipProvider>,
+    );
+
+    fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+    expect(switchLanguage).not.toHaveBeenCalled();
+    expect(onAfterToggle).not.toHaveBeenCalled();
+
+    const view = new EditorView({ state: EditorState.create() });
+    editorViewRef.current = view;
+    try {
+      fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+      expect(switchLanguage).toHaveBeenCalledWith(view, {
+        language: "markdown",
+      });
+      expect(onAfterToggle).toHaveBeenCalledOnce();
+    } finally {
+      view.destroy();
+    }
+  });
+});

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

@@ -572,7 +572,7 @@ const CellEditorInternal = ({
           <div className="absolute top-1 right-5">
             <LanguageToggles
               code={code}
-              editorView={editorViewRef.current}
+              editorViewRef={editorViewRef}
               currentLanguageAdapter={languageAdapter}
               onAfterToggle={afterToggleMarkdown}
             />

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

@@ -13,7 +13,7 @@ import { Functions } from "@/utils/functions";
 import { MarkdownIcon, PythonIcon } from "./icons";
 
 interface LanguageTogglesProps {
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   code: string;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   onAfterToggle: () => void;
@@ -25,7 +25,7 @@ interface LanguageTogglesProps {
 }
 
 export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
-  editorView,
+  editorViewRef,
   code,
   currentLanguageAdapter,
   onAfterToggle,
@@ -43,7 +43,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
   return (
     <div className={className}>
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={canUseSQL && currentLanguageAdapter === "python"}
         icon={
@@ -58,7 +58,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={
           canUseMarkdown && currentLanguageAdapter === "python"
@@ -75,7 +75,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={true}
         icon={
@@ -95,7 +95,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
 
 interface Props {
   className?: string;
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   canSwitchToLanguage: boolean;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   toType: LanguageAdapter["type"];
@@ -105,7 +105,7 @@ interface Props {
 }
 
 export const LanguageToggle: React.FC<Props> = ({
-  editorView,
+  editorViewRef,
   currentLanguageAdapter,
   canSwitchToLanguage,
   icon,
@@ -114,6 +114,7 @@ export const LanguageToggle: React.FC<Props> = ({
   onAfterToggle,
 }) => {
   const handleClick = () => {
+    const editorView = editorViewRef.current;
     if (!editorView) {
       return;
     }