@marimo-team/frontend 0.25.1-dev26 → 0.25.1-dev28

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

@@ -0,0 +1,107 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { fireEvent, render, screen } from "@testing-library/react";
+import { Provider } from "jotai";
+import type React from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { ModalProvider } from "@/components/modal/ImperativeModal";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { store } from "@/core/state/jotai";
+import { PairWithAgentBanner, PairWithAgentButton } from "../pair-with-agent";
+
+const mocks = vi.hoisted(() => ({
+  isWasm: vi.fn(() => false),
+}));
+
+vi.mock("@/core/wasm/utils", () => ({
+  isWasm: mocks.isWasm,
+}));
+
+function wrapper({ children }: { children: React.ReactNode }) {
+  return (
+    <Provider store={store}>
+      <TooltipProvider>
+        <ModalProvider>{children}</ModalProvider>
+      </TooltipProvider>
+    </Provider>
+  );
+}
+
+function stubTokenFetch() {
+  vi.stubGlobal(
+    "fetch",
+    vi.fn().mockResolvedValue(new Response(JSON.stringify({ token: null }))),
+  );
+}
+
+describe("PairWithAgent", () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+    mocks.isWasm.mockReset();
+    mocks.isWasm.mockReturnValue(false);
+  });
+
+  it("shows local agent options", () => {
+    render(<PairWithAgentBanner />, { wrapper });
+
+    expect(
+      screen.getByRole("group", { name: "Supported agents" }),
+    ).toBeVisible();
+    expect(screen.getByRole("figure", { name: "Claude Code" })).toBeVisible();
+    expect(screen.getByRole("figure", { name: "Codex" })).toBeVisible();
+    expect(screen.getByRole("figure", { name: "Cursor" })).toBeVisible();
+    expect(screen.getByRole("figure", { name: "Gemini CLI" })).toBeVisible();
+    expect(screen.getByRole("figure", { name: "OpenCode" })).toBeVisible();
+    expect(
+      screen.getByRole("figure", { name: "GitHub Copilot" }),
+    ).toBeVisible();
+    expect(
+      screen.getByRole("figure", { name: "Any local agent" }),
+    ).toBeVisible();
+    expect(
+      screen.getByRole("button", {
+        name: "Learn more.",
+      }),
+    ).toBeVisible();
+  });
+
+  it("opens the pairing dialog from the banner description", () => {
+    stubTokenFetch();
+    render(<PairWithAgentBanner />, { wrapper });
+
+    fireEvent.click(
+      screen.getByRole("button", {
+        name: "Learn more.",
+      }),
+    );
+
+    expect(
+      screen.getByRole("dialog", { name: "Pair with an agent" }),
+    ).toBeVisible();
+  });
+
+  it("opens the pairing dialog and preserves a passed click handler", () => {
+    const onClick = vi.fn();
+    stubTokenFetch();
+    render(<PairWithAgentButton onClick={onClick} />, { wrapper });
+
+    fireEvent.click(screen.getByRole("button", { name: "Connect your agent" }));
+
+    expect(onClick).toHaveBeenCalledOnce();
+    expect(
+      screen.getByRole("dialog", { name: "Pair with an agent" }),
+    ).toBeVisible();
+  });
+
+  it("hides pairing controls in WASM", () => {
+    mocks.isWasm.mockReturnValue(true);
+
+    const { container, rerender } = render(<PairWithAgentButton />, {
+      wrapper,
+    });
+    expect(container).toBeEmptyDOMElement();
+
+    rerender(<PairWithAgentBanner />);
+    expect(container).toBeEmptyDOMElement();
+  });
+});

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

@@ -0,0 +1,119 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { BotIcon } from "lucide-react";
+import React from "react";
+import { AiProviderIcon } from "@/components/ai/ai-provider-icon";
+import { useImperativeModal } from "@/components/modal/ImperativeModal";
+import { Button, type ButtonProps } from "@/components/ui/button";
+import { isWasm } from "@/core/wasm/utils";
+import { cn } from "@/utils/cn";
+import { PairWithAgentModal } from "./pair-with-agent-modal";
+
+const FEATURED_AGENTS = [
+  { id: "claude", label: "Claude Code" },
+  { id: "codex", label: "Codex" },
+  { id: "cursor", label: "Cursor" },
+  { id: "gemini", label: "Gemini CLI" },
+  { id: "opencode", label: "OpenCode" },
+  { id: "github", label: "GitHub Copilot" },
+  { id: "openai-compatible", label: "Any local agent" },
+] as const satisfies ReadonlyArray<{
+  id: React.ComponentProps<typeof AiProviderIcon>["provider"];
+  label: string;
+}>;
+
+export const usePairWithAgentModal = () => {
+  const { openModal, closeModal } = useImperativeModal();
+
+  return React.useCallback(
+    () => openModal(<PairWithAgentModal onClose={closeModal} />),
+    [openModal, closeModal],
+  );
+};
+
+export const PairWithAgentButton: React.FC<
+  ButtonProps & { label?: string }
+> = ({ className, label = "Connect your agent", onClick, ...props }) => {
+  const openPairWithAgentModal = usePairWithAgentModal();
+
+  if (isWasm()) {
+    return null;
+  }
+
+  return (
+    <Button
+      type="button"
+      variant="link"
+      size="xs"
+      className={cn("h-auto gap-1 px-0", className)}
+      onClick={(event) => {
+        onClick?.(event);
+        if (!event.defaultPrevented) {
+          openPairWithAgentModal();
+        }
+      }}
+      {...props}
+    >
+      {label}
+    </Button>
+  );
+};
+
+export const PairWithAgentBanner: React.FC<{
+  label?: string;
+  className?: string;
+}> = ({ label = "Connect your agent", className }) => {
+  const openPairWithAgentModal = usePairWithAgentModal();
+
+  if (isWasm()) {
+    return null;
+  }
+
+  return (
+    <div
+      className={cn(
+        "rounded-lg border border-(--blue-6) bg-(--blue-2)/50 p-3",
+        className,
+      )}
+    >
+      <div className="flex items-center gap-2 text-sm font-semibold">
+        {label}
+        <fieldset
+          className="m-0 flex min-w-0 -space-x-1.5 border-0 p-0 mb-0.5"
+          aria-label="Supported agents"
+        >
+          {FEATURED_AGENTS.map(({ id, label }) => (
+            <figure
+              key={id}
+              title={label}
+              aria-label={label}
+              className="flex size-6 items-center justify-center rounded-full border bg-background shadow-xs"
+            >
+              {id === "openai-compatible" ? (
+                <BotIcon className="size-3.5" aria-hidden={true} />
+              ) : (
+                <AiProviderIcon
+                  provider={id}
+                  className="size-3.5"
+                  aria-hidden={true}
+                />
+              )}
+            </figure>
+          ))}
+        </fieldset>
+      </div>
+      <div className="mt-1 text-xs leading-relaxed">
+        <p className="inline">
+          Use your AI coding agent to pair-program on this notebook.
+        </p>
+        <button
+          type="button"
+          onClick={openPairWithAgentModal}
+          className="text-link hover:underline ml-1"
+        >
+          Learn more.
+        </button>
+      </div>
+    </div>
+  );
+};

package/src/components/editor/actions/useNotebookActions.tsx

@@ -297,11 +297,12 @@ export function useNotebookActions({
 
     {
       icon: <SparklesIcon size={14} strokeWidth={1.5} />,
-      label: "Pair with an agent",
+      label: "Connect your agent",
       hidden: isWasm(),
       handle: async () => {
         openModal(<PairWithAgentModal onClose={closeModal} />);
       },
+      additionalKeywords: ["connect"],
     },
 
     {

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

@@ -29,6 +29,7 @@ import useEvent from "react-use-event-hook";
 import { DefaultChatTransport } from "ai";
 import { z } from "zod";
 import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
+import { PairWithAgentButton } from "@/components/editor/actions/pair-with-agent";
 import {
   AddContextButton,
   AttachFileButton,
@@ -301,6 +302,7 @@ export const AddCellWithAI: React.FC<{
           showAddCustomModelDocs={true}
         />
         {languageDropdown}
+        <PairWithAgentButton />
       </div>
       <div className="flex flex-row items-center">
         {files.length > 0 && (

package/src/components/editor/ai/staged-cell-submission.ts

@@ -14,13 +14,13 @@ interface SubmissionAttempt {
 export class StagedCellSubmissionController {
   private activeAttempt: SubmissionAttempt | null = null;
 
-  cancel() {
+  public cancel() {
     if (this.activeAttempt) {
       this.activeAttempt.cancelled = true;
     }
   }
 
-  async run<T>({ prepare, submit, onError }: SubmissionCallbacks<T>) {
+  public async run<T>({ prepare, submit, onError }: SubmissionCallbacks<T>) {
     if (this.activeAttempt !== null) {
       return;
     }

package/src/components/editor/ai/transport/chat-transport.tsx

@@ -15,7 +15,7 @@ export class StreamingChunkTransport<
 > extends DefaultChatTransport<UI_MESSAGE> {
   private onChunkReceived: (chunk: UIMessageChunk) => void;
 
-  constructor(
+  public constructor(
     options: HttpChatTransportInitOptions<UI_MESSAGE>,
     onChunkReceived: (chunk: UIMessageChunk) => void,
   ) {

package/src/components/editor/connections/database/as-code.ts

@@ -36,7 +36,7 @@ abstract class CodeGenerator<T extends DatabaseConnection["type"]> {
   protected orm: ConnectionLibrary;
   protected secrets: SecretContainer;
 
-  constructor(
+  public constructor(
     connection: Extract<DatabaseConnection, { type: T }>,
     orm: ConnectionLibrary,
     secrets: SecretContainer,
@@ -46,7 +46,7 @@ abstract class CodeGenerator<T extends DatabaseConnection["type"]> {
     this.secrets = secrets;
   }
 
-  get imports(): Set<string> {
+  public get imports(): Set<string> {
     const imports = new Set<string>(this.generateImports());
     switch (this.orm) {
       case "sqlalchemy":
@@ -67,7 +67,7 @@ abstract class CodeGenerator<T extends DatabaseConnection["type"]> {
 
   protected abstract generateImports(): string[];
 
-  abstract generateConnectionCode(): string;
+  public abstract generateConnectionCode(): string;
 }
 
 const makePrivate = (name: string) => `_${name}`;
@@ -75,14 +75,14 @@ const makePrivate = (name: string) => `_${name}`;
 class SecretContainer {
   private secrets: Record<string, string> = {};
 
-  get imports(): Set<string> {
+  public get imports(): Set<string> {
     if (Object.keys(this.secrets).length === 0) {
       return new Set<string>();
     }
     return new Set<string>(["import os"]);
   }
 
-  print(
+  public print(
     varName: string,
     secretKeyOrValue: string | number | boolean,
     defaultValue?: string | undefined,
@@ -119,7 +119,7 @@ class SecretContainer {
     return `"${secretKeyOrValue}"`;
   }
 
-  printInFString(
+  public printInFString(
     varName: string,
     secretKeyOrValue: string | number | undefined | boolean,
     defaultValue?: string | undefined,
@@ -171,7 +171,7 @@ class SecretContainer {
    * f"db://sample:access_token@sample.com"
    * ```
    */
-  printPassword(
+  public printPassword(
     password: string | undefined,
     passwordPlaceholder: string,
     inFString: boolean,
@@ -187,11 +187,11 @@ class SecretContainer {
       : printMethod(variableName || "password", passwordPlaceholder, password);
   }
 
-  getSecrets(): Record<string, string> {
+  public getSecrets(): Record<string, string> {
     return this.secrets;
   }
 
-  formatSecrets(): string {
+  public formatSecrets(): string {
     if (Object.keys(this.secrets).length === 0) {
       return "";
     }
@@ -203,11 +203,11 @@ class SecretContainer {
 }
 
 class PostgresGenerator extends CodeGenerator<"postgres"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const ssl = this.connection.ssl
       ? ", connect_args={'sslmode': 'require'}"
       : "";
@@ -235,11 +235,11 @@ class PostgresGenerator extends CodeGenerator<"postgres"> {
 }
 
 class MySQLGenerator extends CodeGenerator<"mysql"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const ssl = this.connection.ssl
       ? ", connect_args={'ssl': {'ssl-mode': 'preferred'}}"
       : "";
@@ -267,11 +267,11 @@ class MySQLGenerator extends CodeGenerator<"mysql"> {
 }
 
 class SQLiteGenerator extends CodeGenerator<"sqlite"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const database = this.connection.database
       ? this.secrets.printInFString("database", this.connection.database)
       : "";
@@ -289,11 +289,11 @@ class SQLiteGenerator extends CodeGenerator<"sqlite"> {
 }
 
 class SnowflakeGenerator extends CodeGenerator<"snowflake"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return ["from snowflake.sqlalchemy import URL"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const { authType, account, database, warehouse, schema, role } =
       this.connection;
     const baseParams: Record<string, string | undefined> = {
@@ -422,11 +422,11 @@ class BigQueryGenerator extends CodeGenerator<"bigquery"> {
     );
   }
 
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return this.credential.kind === "path" ? [] : ["import json"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const project = this.secrets.printInFString(
       "project",
       this.connection.project,
@@ -451,11 +451,11 @@ class BigQueryGenerator extends CodeGenerator<"bigquery"> {
 }
 
 class DuckDBGenerator extends CodeGenerator<"duckdb"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const database = this.secrets.printInFString(
       "database",
       this.connection.database || ":memory:",
@@ -469,11 +469,11 @@ class DuckDBGenerator extends CodeGenerator<"duckdb"> {
 }
 
 class MotherDuckGenerator extends CodeGenerator<"motherduck"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const database = this.secrets.printInFString(
       "database",
       this.connection.database,
@@ -497,11 +497,11 @@ class MotherDuckGenerator extends CodeGenerator<"motherduck"> {
 }
 
 class ClickHouseGenerator extends CodeGenerator<"clickhouse_connect"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return ["import clickhouse_connect"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const password = this.secrets.printPassword(
       this.connection.password,
       "CLICKHOUSE_PASSWORD",
@@ -530,11 +530,11 @@ ${formatUrlParams(params, (inner) => `        ${inner}`)},
 }
 
 class TimeplusGenerator extends CodeGenerator<"timeplus"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const password = this.secrets.printPassword(
       this.connection.password,
       "TIMEPLUS_PASSWORD",
@@ -555,11 +555,11 @@ class TimeplusGenerator extends CodeGenerator<"timeplus"> {
 }
 
 class ChDBGenerator extends CodeGenerator<"chdb"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return ["import chdb"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const database =
       this.secrets.print("database", this.connection.database) || '""';
 
@@ -570,13 +570,13 @@ class ChDBGenerator extends CodeGenerator<"chdb"> {
 }
 
 class TrinoGenerator extends CodeGenerator<"trino"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return this.connection.async_support
       ? ["import aiotrino"]
       : ["import trino.sqlalchemy"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const trinoExtension = this.connection.async_support ? "aiotrino" : "trino";
     const schema = this.connection.schema ? `/${this.connection.schema}` : "";
 
@@ -614,7 +614,7 @@ class PyIcebergGenerator extends CodeGenerator<"iceberg"> {
     );
   }
 
-  generateImports(): string[] {
+  public generateImports(): string[] {
     if (Object.keys(this.getCatalogOptions()).length === 0) {
       return ["from pyiceberg.catalog import load_catalog"];
     }
@@ -635,7 +635,7 @@ class PyIcebergGenerator extends CodeGenerator<"iceberg"> {
     }
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const name = `"${this.connection.name}"`;
     const options = this.getCatalogOptions();
     if (Object.keys(options).length === 0) {
@@ -705,12 +705,12 @@ class PyIcebergGenerator extends CodeGenerator<"iceberg"> {
 }
 
 class DataFusionGenerator extends CodeGenerator<"datafusion"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     // To trigger installation of ibis-datafusion
     return ["from datafusion import SessionContext"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     if (this.connection.sessionContext) {
       return dedent(`
         ctx = SessionContext()
@@ -727,11 +727,11 @@ class DataFusionGenerator extends CodeGenerator<"datafusion"> {
 }
 
 class PySparkGenerator extends CodeGenerator<"pyspark"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return ["from pyspark.sql import SparkSession"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     if (this.connection.host || this.connection.port) {
       const host = this.secrets.printInFString("host", this.connection.host);
       const port = this.secrets.printInFString("port", this.connection.port);
@@ -747,11 +747,11 @@ class PySparkGenerator extends CodeGenerator<"pyspark"> {
 }
 
 class RedshiftGenerator extends CodeGenerator<"redshift"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return ["import redshift_connector"];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const host = this.secrets.print("host", this.connection.host);
     const port = this.secrets.print("port", this.connection.port);
     const database = this.secrets.print("database", this.connection.database);
@@ -819,11 +819,11 @@ ${formatUrlParams(params, (inner) => `          ${inner}`)},
 }
 
 class DatabricksGenerator extends CodeGenerator<"databricks"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const useFString = this.orm !== "ibis";
 
     const accessToken = this.secrets.printPassword(
@@ -883,14 +883,14 @@ class DatabricksGenerator extends CodeGenerator<"databricks"> {
 }
 
 class SupabaseGenerator extends CodeGenerator<"supabase"> {
-  generateImports(): string[] {
+  public generateImports(): string[] {
     if (this.connection.disable_client_pooling) {
       return ["from sqlalchemy.pool import NullPool"];
     }
     return [];
   }
 
-  generateConnectionCode(): string {
+  public generateConnectionCode(): string {
     const password = this.secrets.printPassword(
       this.connection.password,
       "SUPABASE_PASSWORD",
@@ -921,7 +921,7 @@ class SupabaseGenerator extends CodeGenerator<"supabase"> {
 class CodeGeneratorFactory {
   public secrets = new SecretContainer();
 
-  createGenerator(
+  public createGenerator(
     connection: DatabaseConnection,
     orm: ConnectionLibrary,
   ): CodeGenerator<DatabaseConnection["type"]> {

package/src/components/editor/connections/storage/as-code.ts

@@ -26,14 +26,14 @@ export const StorageLibraryDisplayNames: Record<StorageLibrary, string> = {
 class SecretContainer {
   private secrets: Record<string, string> = {};
 
-  get imports(): Set<string> {
+  public get imports(): Set<string> {
     if (Object.keys(this.secrets).length === 0) {
       return new Set<string>();
     }
     return new Set<string>(["import os"]);
   }
 
-  print(varName: string, value: string | undefined): string {
+  public print(varName: string, value: string | undefined): string {
     if (value === undefined || value === "") {
       return "";
     }
@@ -46,7 +46,7 @@ class SecretContainer {
     return `"${escapePythonString(value)}"`;
   }
 
-  formatSecrets(): string {
+  public formatSecrets(): string {
     if (Object.keys(this.secrets).length === 0) {
       return "";
     }