@marimo-team/frontend 0.25.2 → 0.25.3-dev100

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

@@ -11,10 +11,10 @@ import {
 import { Provider } from "jotai";
 import type React from "react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { PAIR_PREVIEW } from "@/__tests__/fixtures/pair-preview";
+import { PAIR_CONFIG } from "@/__tests__/fixtures/pair-config";
 import { Dialog } from "@/components/ui/dialog";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import { pairPreviewAtom } from "@/core/config/pair";
+import { pairConfigAtom } from "@/core/config/pair";
 import {
   type StableSessionId,
   stableSessionIdAtom,
@@ -26,6 +26,7 @@ import {
 } from "@/core/runtime/config";
 import { store } from "@/core/state/jotai";
 import { PairWithAgentModal } from "../pair-with-agent-modal";
+import { Deferred } from "@/utils/Deferred";
 import { copyToClipboard } from "@/utils/copy";
 
 vi.mock("@/utils/copy", () => ({
@@ -43,11 +44,24 @@ function wrapper({ children }: { children: React.ReactNode }) {
 }
 
 const STABLE_ID = "sess-stable-1" as StableSessionId;
+const NOTEBOOK_ID = "nb-one";
+function responseFor(url: string | URL | Request, token: string | null = null) {
+  return Promise.resolve(
+    new Response(
+      JSON.stringify(
+        String(url).endsWith("/api/sessions")
+          ? { one: { session_id: STABLE_ID, notebook_id: NOTEBOOK_ID } }
+          : { token },
+      ),
+      { headers: { "Content-Type": "application/json" } },
+    ),
+  );
+}
 
 describe("PairWithAgentModal", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    store.set(pairPreviewAtom, PAIR_PREVIEW);
+    store.set(pairConfigAtom, PAIR_CONFIG);
     store.set(stableSessionIdAtom, STABLE_ID);
     store.set(filenameAtom, "/project/my notebook.py");
     store.set(runtimeConfigAtom, {
@@ -57,19 +71,19 @@ describe("PairWithAgentModal", () => {
     window.history.replaceState({}, "", "/?file=notebook.py");
     vi.stubGlobal(
       "fetch",
-      vi.fn().mockResolvedValue(new Response(JSON.stringify({ token: null }))),
+      vi.fn().mockImplementation((url) => responseFor(url)),
     );
   });
 
   afterEach(() => {
-    store.set(pairPreviewAtom, undefined);
+    store.set(pairConfigAtom, undefined);
     store.set(stableSessionIdAtom, null);
     store.set(filenameAtom, null);
     vi.unstubAllGlobals();
   });
 
   it.each(["Claude", "Codex", "OpenCode", "Prompt"])(
-    "omits skill installation and targets the stable session in %s preview",
+    "targets the verified host Notebook in %s",
     async (tab) => {
       await act(async () => {
         render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
@@ -81,10 +95,13 @@ describe("PairWithAgentModal", () => {
 
       const panel = screen.getByRole("tabpanel");
       expect(within(panel).queryAllByText(/skill/i)).toEqual([]);
-      expect(panel).toHaveTextContent(STABLE_ID);
+      expect(panel).toHaveTextContent(NOTEBOOK_ID);
       expect(panel).not.toHaveTextContent("--file");
       expect(panel).not.toHaveTextContent("File:");
       expect(panel).not.toHaveTextContent("notebook.py");
+      expect(
+        panel.textContent?.includes("Agent connections require OpenCode 2.x"),
+      ).toBe(tab === "OpenCode");
       expect(panel).toHaveTextContent(
         tab === "Prompt" ? "1. Copy this prompt" : "1. Run in your terminal",
       );
@@ -95,13 +112,15 @@ describe("PairWithAgentModal", () => {
       );
       await waitFor(() => expect(copyToClipboard).toHaveBeenCalledOnce());
       const copied = vi.mocked(copyToClipboard).mock.calls[0][0];
-      expect(copied).toContain(STABLE_ID);
+      expect(copied).toContain(NOTEBOOK_ID);
       expect(copied).not.toContain("notebook.py");
+      expect(copied).not.toContain(STABLE_ID);
+      expect(panel).not.toHaveTextContent(STABLE_ID);
     },
   );
 
   it.each(["Claude", "Prompt"])(
-    "omits the session in %s preview before kernel-ready arrives",
+    "waits for kernel-ready before copying the %s target",
     async (tab) => {
       store.set(stableSessionIdAtom, null);
       await act(async () => {
@@ -113,17 +132,21 @@ describe("PairWithAgentModal", () => {
       });
 
       const panel = screen.getByRole("tabpanel");
-      expect(panel).not.toHaveTextContent("--session");
-      expect(panel).not.toHaveTextContent("Session:");
+      expect(panel).not.toHaveTextContent("--notebook");
+      expect(panel).not.toHaveTextContent("Notebook:");
       expect(panel).toHaveTextContent("http://localhost:8000");
+      expect(
+        within(panel).getByRole("button", {
+          name: tab === "Prompt" ? "Copy prompt" : "Copy command",
+        }),
+      ).toBeDisabled();
     },
   );
 
   it.each(["Claude", "Codex", "OpenCode", "Prompt"])(
-    "preserves the %s skill flow without preview configuration",
+    "uses the editor CLI in %s when host is disabled",
     async (tab) => {
-      store.set(pairPreviewAtom, undefined);
-      window.history.replaceState({}, "", "/");
+      store.set(pairConfigAtom, { ...PAIR_CONFIG, hostEnabled: false });
       await act(async () => {
         render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
       });
@@ -131,27 +154,33 @@ describe("PairWithAgentModal", () => {
         button: 0,
         ctrlKey: false,
       });
-
       const panel = screen.getByRole("tabpanel");
-      expect(
-        within(panel).getByText("npx skills add marimo-team/marimo-pair"),
-      ).toBeVisible();
-      expect(panel).not.toHaveTextContent(STABLE_ID);
-      expect(panel).not.toHaveTextContent("--file");
+      expect(panel).toHaveTextContent(STABLE_ID);
+      expect(panel).not.toHaveTextContent(NOTEBOOK_ID);
+      expect(panel).not.toHaveTextContent("skill");
+      expect(panel).not.toHaveTextContent("Agent connections require");
       expect(panel).not.toHaveTextContent("MARIMO_PAIR_NEXT");
-      expect(panel).not.toHaveTextContent("/project/my notebook.py");
-      expect(panel).toHaveTextContent(
-        tab === "Prompt" ? "2. Copy this prompt" : "2. Run in your terminal",
+      fireEvent.click(
+        within(panel).getByRole("button", {
+          name: tab === "Prompt" ? "Copy prompt" : "Copy command",
+        }),
       );
+      await waitFor(() => expect(copyToClipboard).toHaveBeenCalledOnce());
+      expect(vi.mocked(copyToClipboard).mock.calls[0][0]).toContain(STABLE_ID);
+      expect(
+        vi
+          .mocked(fetch)
+          .mock.calls.some(([url]) => String(url).endsWith("/api/sessions")),
+      ).toBe(false);
     },
   );
 
   it.each([true, false])(
-    "masks the raw token but copies it intact (preview=%s)",
-    async (preview) => {
-      store.set(pairPreviewAtom, preview ? PAIR_PREVIEW : undefined);
-      vi.mocked(fetch).mockResolvedValue(
-        new Response(JSON.stringify({ token: "secret-token" })),
+    "masks the raw token but copies it intact (host mode=%s)",
+    async (hostEnabled) => {
+      store.set(pairConfigAtom, { ...PAIR_CONFIG, hostEnabled });
+      vi.mocked(fetch).mockImplementation((url) =>
+        responseFor(url, "secret-token"),
       );
       await act(async () => {
         render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
@@ -169,7 +198,7 @@ describe("PairWithAgentModal", () => {
       );
       await waitFor(() => expect(copyToClipboard).toHaveBeenCalledOnce());
       expect(vi.mocked(copyToClipboard).mock.calls[0][0]).toContain(
-        preview ? "export MARIMO_TOKEN=secret-token" : "--token secret-token",
+        "export MARIMO_TOKEN=secret-token",
       );
       expect(vi.mocked(copyToClipboard).mock.calls[0][0]).not.toContain(
         "********oken",
@@ -177,9 +206,156 @@ describe("PairWithAgentModal", () => {
     },
   );
 
-  it("copies a preview terminal command without the token and lets the user copy the token separately", async () => {
-    vi.mocked(fetch).mockResolvedValue(
-      new Response(JSON.stringify({ token: "secret-token" })),
+  it("uses the CLI when an older server omits Pair configuration", async () => {
+    store.set(pairConfigAtom, undefined);
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    const panel = screen.getByRole("tabpanel");
+    expect(panel).toHaveTextContent(`--session ${STABLE_ID}`);
+    expect(panel).not.toHaveTextContent("skill");
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeEnabled();
+  });
+
+  it("ignores a previous Session's late Notebook lookup", async () => {
+    const previous = new Deferred<Response>();
+    const nextId = "sess-next" as StableSessionId;
+    let lookups = 0;
+    vi.mocked(fetch).mockImplementation((url) => {
+      if (!String(url).endsWith("/api/sessions")) {
+        return responseFor(url);
+      }
+      if (lookups++ === 0) {
+        return previous.promise;
+      }
+      return Promise.resolve(
+        new Response(
+          JSON.stringify({
+            next: { session_id: nextId, notebook_id: "nb-next" },
+          }),
+          { headers: { "Content-Type": "application/json" } },
+        ),
+      );
+    });
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    await act(async () => store.set(stableSessionIdAtom, nextId));
+    expect(screen.getByRole("tabpanel")).toHaveTextContent("nb-next");
+    await act(async () =>
+      previous.resolve(
+        new Response(
+          JSON.stringify({
+            previous: { session_id: STABLE_ID, notebook_id: NOTEBOOK_ID },
+          }),
+          { headers: { "Content-Type": "application/json" } },
+        ),
+      ),
+    );
+    expect(screen.getByRole("tabpanel")).toHaveTextContent("nb-next");
+    expect(screen.getByRole("tabpanel")).not.toHaveTextContent(NOTEBOOK_ID);
+  });
+
+  it("disables target copying while a new Session's Notebook is unverified", async () => {
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(NOTEBOOK_ID);
+    const pending = new Deferred<Response>();
+    vi.mocked(fetch).mockImplementation((url) =>
+      String(url).endsWith("/api/sessions")
+        ? pending.promise
+        : responseFor(url),
+    );
+    const nextId = "sess-stable-2" as StableSessionId;
+    act(() => store.set(stableSessionIdAtom, nextId));
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeDisabled();
+    expect(screen.getByRole("tabpanel")).not.toHaveTextContent(NOTEBOOK_ID);
+    fireEvent.click(screen.getByRole("button", { name: "Copy command" }));
+    expect(copyToClipboard).not.toHaveBeenCalled();
+    await act(async () =>
+      pending.resolve(
+        new Response(
+          JSON.stringify({
+            next: { session_id: nextId, notebook_id: "nb-two" },
+          }),
+          { headers: { "Content-Type": "application/json" } },
+        ),
+      ),
+    );
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeEnabled();
+    fireEvent.click(screen.getByRole("button", { name: "Copy command" }));
+    await waitFor(() => expect(copyToClipboard).toHaveBeenCalledOnce());
+    expect(vi.mocked(copyToClipboard).mock.calls[0][0]).toContain("nb-two");
+    expect(vi.mocked(copyToClipboard).mock.calls[0][0]).not.toContain(nextId);
+  });
+
+  it("uses the Session after verifying that a notebook is unsaved", async () => {
+    vi.mocked(fetch).mockImplementation((url) =>
+      String(url).endsWith("/api/sessions")
+        ? Promise.resolve(
+            new Response(
+              JSON.stringify({
+                one: { session_id: STABLE_ID, notebook_id: null },
+              }),
+              { headers: { "Content-Type": "application/json" } },
+            ),
+          )
+        : responseFor(url),
+    );
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    expect(screen.queryByRole("status")).not.toBeInTheDocument();
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeEnabled();
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(
+      `--session ${STABLE_ID}`,
+    );
+    expect(screen.getByRole("tabpanel")).not.toHaveTextContent("--notebook");
+  });
+
+  it("reverifies the target when host mode is enabled again for the same Session", async () => {
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(NOTEBOOK_ID);
+    act(() =>
+      store.set(pairConfigAtom, { ...PAIR_CONFIG, hostEnabled: false }),
+    );
+
+    const pending = new Deferred<Response>();
+    vi.mocked(fetch).mockImplementation((url) =>
+      String(url).endsWith("/api/sessions")
+        ? pending.promise
+        : responseFor(url),
+    );
+    act(() => store.set(pairConfigAtom, PAIR_CONFIG));
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeDisabled();
+    expect(screen.getByRole("tabpanel")).not.toHaveTextContent(NOTEBOOK_ID);
+
+    await act(async () => pending.resolve(await responseFor("/api/sessions")));
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeEnabled();
+    expect(screen.getByRole("tabpanel")).toHaveTextContent(NOTEBOOK_ID);
+  });
+
+  it("keeps copying disabled when Notebook verification fails", async () => {
+    vi.mocked(fetch).mockImplementation((url) =>
+      String(url).endsWith("/api/sessions")
+        ? Promise.reject(new Error("offline"))
+        : responseFor(url),
+    );
+    await act(async () => {
+      render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
+    });
+    expect(screen.getByRole("status")).toHaveTextContent(
+      "Could not verify this notebook",
+    );
+    expect(screen.getByRole("button", { name: "Copy command" })).toBeDisabled();
+  });
+
+  it("copies a host mode terminal command without the token and lets the user copy the token separately", async () => {
+    vi.mocked(fetch).mockImplementation((url) =>
+      responseFor(url, "secret-token"),
     );
     await act(async () => {
       render(<PairWithAgentModal onClose={vi.fn()} />, { wrapper });
@@ -193,9 +369,9 @@ describe("PairWithAgentModal", () => {
     );
     await waitFor(() =>
       expect(copyToClipboard).toHaveBeenCalledWith(
-        String.raw`claude "$(MARIMO_PAIR_NEXT=1 uvx marimo@latest pair prompt \
+        String.raw`claude "$(uvx marimo@latest pair prompt \
   --url http://localhost:8000/ \
-  --session ${STABLE_ID} \
+  --notebook ${NOTEBOOK_ID} \
   --with-token)"`,
       ),
     );

package/src/components/editor/actions/pair-with-agent-commands.ts

@@ -1,9 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { assertNever } from "@/utils/assertNever";
-import { KnownQueryParams } from "@/core/constants";
 import { shellQuote } from "@/utils/shell";
-import type { PairPreviewConfig } from "@/core/config/pair";
+import type { PairConfig } from "@/core/config/pair";
 
 export type AgentTab = "claude" | "codex" | "opencode" | "prompt";
 
@@ -18,30 +17,18 @@ export const AGENT_LABELS: Record<AgentTab, string> = {
   prompt: "Prompt",
 };
 
-export const SKILL_INSTALL = "npx skills add marimo-team/marimo-pair";
-
 /** How to invoke marimo: from the local checkout in dev, else via uvx. */
 export function getMarimoCommand(): string {
   return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest";
 }
 
-/** Return the server file key from a page URL, preserving its decoded value. */
-export function getFileFromURL(href: string): string | undefined {
-  const file = new URL(href).searchParams.get(KnownQueryParams.filePath);
-  return file === null || file === "" ? undefined : file;
-}
-
-function getFileFlag(file: string | undefined): string {
-  return file ? ` --file ${shellQuote(file)}` : "";
-}
-
 /** Identifies the specific running notebook to pair on. */
 export interface ConnectionInfo {
   url: string;
-  /** The server's file key, when known. */
-  file?: string;
-  /** The stable session ID, when known. */
+  /** Stable Session ID for ordinary editor execution. */
   session?: string;
+  /** The host notebook ID, when known. */
+  notebook?: string;
 }
 
 /**
@@ -50,21 +37,23 @@ export interface ConnectionInfo {
  */
 export function getTerminalCommand(
   agent: Exclude<AgentTab, "prompt">,
-  { url, session }: ConnectionInfo,
+  { url, notebook, session }: ConnectionInfo,
   withToken: boolean,
-  preview?: PairPreviewConfig,
+  config?: PairConfig,
 ): string {
-  const command = preview?.command ?? getMarimoCommand();
-  const prefix = preview ? "MARIMO_PAIR_NEXT=1 " : "";
+  const command = config?.command ?? getMarimoCommand();
   const base = [
-    `${prefix}${command} pair prompt`,
+    `${command} pair prompt`,
     `--url ${shellQuote(url)}`,
-    preview && session ? `--session ${shellQuote(session)}` : "",
+    notebook
+      ? `--notebook ${shellQuote(notebook)}`
+      : session
+        ? `--session ${shellQuote(session)}`
+        : "",
     withToken ? "--with-token" : "",
-    preview ? "" : `--${agent}`,
   ]
     .filter(Boolean)
-    .join(preview ? " \\\n  " : " ");
+    .join(" \\\n  ");
   switch (agent) {
     case "claude":
       return `claude "$(${base})"`;
@@ -83,37 +72,39 @@ export function getTerminalCommand(
  * an agent behaves the same as the terminal commands.
  */
 export function getRawPrompt(
-  { url, file, session }: ConnectionInfo,
+  { url, notebook, session }: ConnectionInfo,
   token: string | null,
-  preview?: PairPreviewConfig,
+  config?: PairConfig,
 ): string {
-  if (preview) {
-    const { command, templates } = preview;
-    return formatPrompt(templates.prompt, {
+  const command = config?.command ?? getMarimoCommand();
+  const templates = config?.templates;
+  const authentication = token
+    ? formatPrompt(
+        templates?.token ??
+          "\n\nFor authenticated notebook discovery, set `export MARIMO_TOKEN={token}` in the shell that runs marimo.",
+        { token: shellQuote(token) },
+      )
+    : "";
+  return formatPrompt(
+    templates?.prompt ??
+      "Pair with me on this running marimo notebook.\n\nURL: {url}\n{notebook}{session}\nRun `{command} pair --help` first.\nUse `{command}` for all marimo commands.{authentication}",
+    {
       command,
       url,
-      session: session ? formatPrompt(templates.session, { session }) : "",
-      authentication: token
-        ? formatPrompt(templates.token, { token: shellQuote(token) })
+      notebook: notebook
+        ? formatPrompt(templates?.notebook ?? "Notebook: {notebook}\n", {
+            notebook,
+          })
         : "",
-    });
-  }
-
-  const fileFlag = getFileFlag(file);
-  const fileHint = file ? ` (file ${file})` : "";
-  const executeCmd = `execute-code.sh --url ${shellQuote(url)}${fileFlag}`;
-  const tokenHint = token
-    ? `\n\nUse this auth token when calling \`execute-code.sh\`: \`${executeCmd} --token ${shellQuote(token)}\`.`
-    : "";
-  return [
-    "Use the /marimo-pair skill to pair-program on a running marimo notebook.",
-    "",
-    `Connect to the notebook at: ${url}${fileHint}`,
-    "",
-    `Use \`${executeCmd}\` from the marimo-pair skill to execute code in the notebook.${tokenHint}`,
-    "",
-    "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair.",
-  ].join("\n");
+      session:
+        session && !notebook
+          ? formatPrompt(templates?.session ?? "Session: {session}\n", {
+              session,
+            })
+          : "",
+      authentication,
+    },
+  );
 }
 
 function formatPrompt(