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

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

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { useAtomValue } from "jotai";
 import { Button } from "@/components/ui/button";
 import {
@@ -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,67 @@ 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 lookupIdentity = useMemo(
+    () => ({ hostEnabled, stableSessionId }),
+    [hostEnabled, stableSessionId],
+  );
+  const [lookup, setLookup] = useState<
+    {
+      identity: typeof lookupIdentity;
+    } & ({ status: "ready"; notebookId: string | null } | { status: "error" })
+  >();
+  useEffect(() => {
+    const { hostEnabled, stableSessionId: sessionId } = lookupIdentity;
+    if (!hostEnabled || !sessionId) {
+      return;
+    }
+    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
+              ? {
+                  identity: lookupIdentity,
+                  status: "ready",
+                  notebookId: session.notebook_id,
+                }
+              : { identity: lookupIdentity, status: "error" },
+          );
+        }
+      })
+      .catch(() => {
+        if (!cancelled) {
+          setLookup({ identity: lookupIdentity, status: "error" });
+        }
+      });
+    return () => {
+      cancelled = true;
+    };
+  }, [lookupIdentity]);
+  const currentLookup =
+    lookup?.identity === lookupIdentity ? 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 +133,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 +158,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 +190,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 +200,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 +245,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 +274,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 +297,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/__tests__/useShouldShowInterrupt.test.ts

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useShouldShowInterrupt } from "../useShouldShowInterrupt";
+
+describe("useShouldShowInterrupt", () => {
+  beforeEach(() => vi.useFakeTimers());
+  afterEach(() => vi.useRealTimers());
+
+  it("waits for each run and hides immediately when it stops", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: false,
+    });
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(true);
+    rerender(false);
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(199));
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(1));
+    expect(result.current).toBe(true);
+  });
+
+  it("cancels the delay for a short run", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: true,
+    });
+    act(() => vi.advanceTimersByTime(100));
+    rerender(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(false);
+  });
+});

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;
     }

package/src/components/editor/cell/useShouldShowInterrupt.ts

@@ -1,5 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useEffect, useState } from "react";
+import { useDelayElapsed } from "@/hooks/useDelayElapsed";
 
 const DELAY_MS = 200;
 
@@ -10,17 +10,7 @@ const DELAY_MS = 200;
 export function useShouldShowInterrupt(running: boolean) {
   // Start a timer when the run starts.
   // After 200ms, show the interrupt button to avoid flickering.
-  const [hasRunLongEnough, setHasRunLongEnough] = useState(false);
-  useEffect(() => {
-    if (!running) {
-      return;
-    }
-    setHasRunLongEnough(false);
-    const timeout = setTimeout(() => {
-      setHasRunLongEnough(true);
-    }, DELAY_MS);
-    return () => clearTimeout(timeout);
-  }, [running]);
+  const hasRunLongEnough = useDelayElapsed(running ? DELAY_MS : 0);
 
   return running && hasRunLongEnough;
 }

package/src/components/editor/chrome/panels/__tests__/packages-panel.test.tsx

@@ -18,7 +18,10 @@ import type {
   DependencyTreeNode,
   DependencyTreeResponse,
 } from "@/core/network/types";
-import { withPackageInvalidation } from "@/core/packages/package-data";
+import {
+  invalidatePackageData,
+  withPackageInvalidation,
+} from "@/core/packages/package-data";
 import { useInstallPackages } from "@/core/packages/useInstallPackage";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
 import { store } from "@/core/state/jotai";
@@ -143,6 +146,32 @@ describe("PackagesPanel", () => {
     expect(screen.queryByText("cycle")).not.toBeInTheDocument();
   });
 
+  it("collapses expanded nodes when the dependency tree is refreshed", async () => {
+    const { client } = renderPanel(
+      { kind: "sandbox", backend: "pixi" },
+      populatedTree,
+    );
+    fireEvent.click(await screen.findByRole("treeitem", { name: /polars/ }));
+    expect(screen.getByText("already in tree")).toBeInTheDocument();
+
+    vi.mocked(client.getDependencyTree).mockResolvedValue({
+      tree: structuredClone(populatedTree),
+      context: { kind: "sandbox", backend: "pixi" },
+    });
+    act(() => invalidatePackageData());
+
+    await waitFor(() =>
+      expect(client.getDependencyTree).toHaveBeenCalledTimes(2),
+    );
+    await waitFor(() =>
+      expect(screen.queryByText("already in tree")).not.toBeInTheDocument(),
+    );
+    expect(screen.getByRole("treeitem", { name: /polars/ })).toHaveAttribute(
+      "aria-expanded",
+      "false",
+    );
+  });
+
   it("keeps the configured package manager changeable outside a sandbox", async () => {
     const { getPackageList } = renderPanel({
       kind: "package-manager",

package/src/components/editor/chrome/panels/packages-panel.tsx

@@ -31,6 +31,7 @@ import { usePackageAction } from "@/core/packages/usePackageAction";
 import { usePackageDependencies } from "@/core/packages/usePackageDependencies";
 import { useInstallPackages } from "@/core/packages/useInstallPackage";
 import { isWasm } from "@/core/wasm/utils";
+import { useStateForKey } from "@/hooks/useStateForKey";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { copyToClipboard } from "@/utils/copy";
@@ -440,15 +441,11 @@ const DependencyTree: React.FC<{
   error?: Error | null;
   sandboxBackend: "uv" | "pixi" | null;
 }> = ({ tree, error, sandboxBackend }) => {
-  const [expandedNodes, setExpandedNodes] = React.useState<Set<string>>(
-    new Set(),
+  const [expandedNodes, setExpandedNodes] = useStateForKey(
+    tree,
+    () => new Set<string>(),
   );
 
-  // Reset tree to collapsed state when tree data changes (including refetches)
-  React.useEffect(() => {
-    setExpandedNodes(new Set());
-  }, [tree]);
-
   if (error) {
     return <ErrorBanner error={error} />;
   }