@marimo-team/islands 0.16.5 → 0.17.0

package/src/components/editor/output/__tests__/json-output.test.ts

@@ -218,6 +218,49 @@ describe("getCopyValue", () => {
     `,
     );
   });
+
+  it("should handle bigint", () => {
+    const bigint = String(BigInt(2 ** 64));
+    const value = `text/plain+bigint:${bigint}`;
+    const result = getCopyValue(value);
+    expect(result).toMatchInlineSnapshot(`"18446744073709551616"`);
+
+    const nestedBigInt = {
+      key1: bigint, // this will be just a string
+      key2: `text/plain+bigint:${bigint}`, // this will convert to number
+      key3: true,
+    };
+    const nestedResult = getCopyValue(nestedBigInt);
+    expect(nestedResult).toMatchInlineSnapshot(
+      `
+      "{
+        "key1": "18446744073709551616",
+        "key2": 18446744073709551616,
+        "key3": True
+      }"
+      `,
+    );
+
+    const bigintRaw = BigInt(2 ** 64);
+    const bigintRawResult = getCopyValue(bigintRaw);
+    expect(bigintRawResult).toMatchInlineSnapshot(`"18446744073709551616"`);
+
+    const nestedBigIntRaw = {
+      key1: bigintRaw, // raw number
+      key2: `text/plain+bigint:${bigintRaw}`,
+      key3: true,
+    };
+    const nestedBigIntRawResult = getCopyValue(nestedBigIntRaw);
+    expect(nestedBigIntRawResult).toMatchInlineSnapshot(
+      `
+      "{
+        "key1": 18446744073709551616,
+        "key2": 18446744073709551616,
+        "key3": True
+      }"
+      `,
+    );
+  });
 });
 
 describe("determineMaxDisplayLength", () => {

package/src/components/editor/package-alert.tsx

@@ -463,6 +463,7 @@ const ExtrasSelector: React.FC<ExtrasSelectorProps> = ({
               <button
                 className="hover:bg-muted/50 rounded text-sm px-1 transition-colors border border-muted-foreground/30 hover:border-muted-foreground/60 min-w-0 flex-1 truncate text-left"
                 title={`Selected extras: ${selectedExtras.join(", ")}`}
+                type="button"
               >
                 {selectedExtras.join(",")}
               </button>
@@ -523,6 +524,7 @@ const ExtrasSelector: React.FC<ExtrasSelectorProps> = ({
                 !canSelectExtras && "opacity-50 cursor-not-allowed",
               )}
               title={canSelectExtras ? "Add extras" : "Loading extras..."}
+              type="button"
             >
               <PlusIcon className="w-3 h-3 shrink-0" />
             </button>
@@ -619,6 +621,7 @@ const StreamingLogsViewer: React.FC<StreamingLogsViewerProps> = ({
       <button
         onClick={() => setIsExpanded(!isExpanded)}
         className="flex items-center gap-2 text-sm font-medium text-muted-foreground hover:text-foreground transition-colors"
+        type="button"
       >
         {isExpanded ? (
           <ChevronDownIcon className="w-4 h-4" />

package/src/components/forms/__tests__/form-utils.test.ts

@@ -189,7 +189,8 @@ describe("getDefaults", () => {
     const schema = z.object({
       map: z.map(z.string(), z.number()).default(new Map([["a", 1]])),
     });
-    const result = getDefaults(schema) as { map: Map<string, number> };
+    const defaults = getDefaults(schema);
+    const result = defaults as { map: Map<string, number> };
     expect(result.map instanceof Map).toBe(true);
     expect([...result.map.entries()]).toEqual([["a", 1]]);
   });
@@ -198,7 +199,8 @@ describe("getDefaults", () => {
     const schema = z.object({
       set: z.set(z.string()).default(new Set(["a", "b"])),
     });
-    const result = getDefaults(schema) as { set: Set<string> };
+    const defaults = getDefaults(schema);
+    const result = defaults as { set: Set<string> };
     expect(result.set instanceof Set).toBe(true);
     expect([...result.set]).toEqual(["a", "b"]);
   });

package/src/components/icons/copy-icon.tsx

@@ -9,10 +9,11 @@ import { Tooltip } from "../ui/tooltip";
 import { toast } from "../ui/use-toast";
 
 interface Props {
-  value: string | (() => string);
+  value: string | ((event: React.MouseEvent) => string);
   className?: string;
-  tooltip?: string | false;
+  tooltip?: React.ReactNode | false;
   toastTitle?: string;
+  ariaLabel?: string;
 }
 
 export const CopyClipboardIcon: React.FC<Props> = ({
@@ -20,11 +21,12 @@ export const CopyClipboardIcon: React.FC<Props> = ({
   className,
   tooltip,
   toastTitle,
+  ariaLabel,
 }) => {
   const [isCopied, setIsCopied] = useState(false);
 
-  const handleCopy = Events.stopPropagation(async () => {
-    const valueToCopy = typeof value === "function" ? value() : value;
+  const handleCopy = Events.stopPropagation(async (event: React.MouseEvent) => {
+    const valueToCopy = typeof value === "function" ? value(event) : value;
     await copyToClipboard(valueToCopy).then(() => {
       setIsCopied(true);
       setTimeout(() => setIsCopied(false), 2000);
@@ -35,7 +37,11 @@ export const CopyClipboardIcon: React.FC<Props> = ({
   });
 
   const button = (
-    <button type="button" onClick={handleCopy} aria-label="Copy to clipboard">
+    <button
+      type="button"
+      onClick={handleCopy}
+      aria-label={ariaLabel ?? "Copy to clipboard"}
+    >
       {isCopied ? (
         <CheckIcon className={cn(className, "text-(--grass-11)")} />
       ) : (

package/src/components/shortcuts/renderShortcut.tsx

@@ -102,7 +102,9 @@ export const MinimalHotkeys: React.FC<{
   );
 };
 
-function prettyPrintHotkey(key: string): [label: string, symbol?: string] {
+export function prettyPrintHotkey(
+  key: string,
+): [label: string, symbol?: string] {
   const platform = isPlatformMac() ? "mac" : "default";
 
   const lowerKey = key.toLowerCase();

package/src/components/ui/button.tsx

@@ -70,6 +70,8 @@ const buttonVariants = cva(
         link: "underline-offset-4 hover:underline text-link",
         linkDestructive:
           "underline-offset-4 hover:underline text-destructive underline-destructive",
+        outlineDestructive:
+          "border border-destructive text-destructive hover:bg-destructive/10",
       },
       size: {
         default: "h-10 py-2 px-4",

package/src/components/ui/confirmation-button.tsx

@@ -0,0 +1,85 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import React from "react";
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogDestructiveAction,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+  AlertDialogTrigger,
+} from "./alert-dialog";
+
+interface ConfirmationButtonProps {
+  /**
+   * The button element to trigger the confirmation dialog
+   */
+  children: React.ReactElement;
+  /**
+   * Title of the confirmation dialog
+   */
+  title: string;
+  /**
+   * Description/message of the confirmation dialog
+   */
+  description: string;
+  /**
+   * Callback when the user confirms the action
+   */
+  onConfirm: () => void;
+  /**
+   * Text for the confirm button (default: "Continue")
+   */
+  confirmText?: string;
+  /**
+   * Text for the cancel button (default: "Cancel")
+   */
+  cancelText?: string;
+  /**
+   * Whether to use destructive styling for the confirm button
+   */
+  destructive?: boolean;
+}
+
+export const ConfirmationButton: React.FC<ConfirmationButtonProps> = ({
+  children,
+  title,
+  description,
+  onConfirm,
+  confirmText = "Continue",
+  cancelText = "Cancel",
+  destructive = false,
+}) => {
+  const [open, setOpen] = React.useState(false);
+
+  const handleConfirm = () => {
+    onConfirm();
+    setOpen(false);
+  };
+
+  const ActionComponent = destructive
+    ? AlertDialogDestructiveAction
+    : AlertDialogAction;
+
+  return (
+    <AlertDialog open={open} onOpenChange={setOpen}>
+      <AlertDialogTrigger asChild>{children}</AlertDialogTrigger>
+      <AlertDialogContent>
+        <AlertDialogHeader>
+          <AlertDialogTitle>{title}</AlertDialogTitle>
+          <AlertDialogDescription>{description}</AlertDialogDescription>
+        </AlertDialogHeader>
+        <AlertDialogFooter>
+          <AlertDialogCancel>{cancelText}</AlertDialogCancel>
+          <ActionComponent onClick={handleConfirm}>
+            {confirmText}
+          </ActionComponent>
+        </AlertDialogFooter>
+      </AlertDialogContent>
+    </AlertDialog>
+  );
+};

package/src/components/ui/switch.tsx

@@ -12,6 +12,7 @@ const switchVariants = cva(
       size: {
         default: "h-6 w-11 mb-1",
         sm: "h-4.5 w-8.5",
+        xs: "h-4 w-7",
       },
     },
     defaultVariants: {
@@ -28,6 +29,7 @@ const switchThumbVariants = cva(
         default:
           "h-5 w-5 data-[state=checked]:translate-x-5 data-[state=unchecked]:translate-x-0",
         sm: "h-3.5 w-3.5 data-[state=checked]:translate-x-4 data-[state=unchecked]:translate-x-0",
+        xs: "h-3 w-3 data-[state=checked]:translate-x-3 data-[state=unchecked]:translate-x-0",
       },
     },
     defaultVariants: {

package/src/core/ai/context/providers/__tests__/datasource.test.ts

@@ -2,6 +2,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import { beforeEach, describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
 import type {
   ConnectionsMap,
   DatasetTablesMap,
@@ -9,7 +10,7 @@ import type {
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
 import { Boosts, Sections } from "../common";
-import { DatasourceContextProvider } from "../datasource";
+import { DatasourceContextProvider, getDatasourceContext } from "../datasource";
 
 // Mock data for testing
 const createMockDataSourceConnection = (
@@ -630,3 +631,10 @@ describe("DatasourceContextProvider", () => {
     });
   });
 });
+
+describe("getDatasourceContext", () => {
+  it("should return null if no cell ID is found", () => {
+    const context = getDatasourceContext("1" as CellId);
+    expect(context).toBeNull();
+  });
+});

package/src/core/ai/context/providers/datasource.ts

@@ -3,10 +3,14 @@
 import type { Completion } from "@codemirror/autocomplete";
 import { createRoot } from "react-dom/client";
 import { dbDisplayName } from "@/components/databases/display";
+import { cellDataAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { renderDatasourceInfo } from "@/core/codemirror/language/languages/sql/renderers";
 import {
   type ConnectionsMap,
   type DatasetTablesMap,
+  dataSourceConnectionsAtom,
   getTableType,
 } from "@/core/datasets/data-source-connections";
 import {
@@ -14,6 +18,7 @@ import {
   INTERNAL_SQL_ENGINES,
 } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
@@ -37,10 +42,12 @@ export interface DatasourceContextItem extends AIContextItem {
   };
 }
 
+const CONTEXT_TYPE = "datasource";
+
 export class DatasourceContextProvider extends AIContextProvider<DatasourceContextItem> {
   readonly title = "Datasource";
   readonly mentionPrefix = "@";
-  readonly contextType = "datasource";
+  readonly contextType = CONTEXT_TYPE;
   private connectionsMap: ConnectionsMap;
   private dataframes: DataTable[];
 
@@ -140,3 +147,20 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
     };
   }
 }
+
+export function getDatasourceContext(cellId: CellId): string | null {
+  const cellData = store.get(cellDataAtom(cellId));
+  const code = cellData?.code;
+  if (!code || code.trim() === "") {
+    return null;
+  }
+
+  const [_sqlStatement, _, metadata] = LanguageAdapters.sql.transformIn(code);
+  const datasourceSchema = store
+    .get(dataSourceConnectionsAtom)
+    .connectionsMap.get(metadata.engine);
+  if (datasourceSchema) {
+    return `@${CONTEXT_TYPE}://${datasourceSchema.name}`;
+  }
+  return null;
+}

package/src/core/ai/state.ts

@@ -13,10 +13,13 @@ const KEY = "marimo:ai:chatState:v5";
 
 export type ChatId = TypedString<"ChatId">;
 
-export const aiCompletionCellAtom = atom<{
+export interface AiCompletionCell {
   cellId: CellId;
   initialPrompt?: string;
-} | null>(null);
+  triggerImmediately?: boolean;
+}
+
+export const aiCompletionCellAtom = atom<AiCompletionCell | null>(null);
 
 const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
 export const includeOtherCellsAtom = atomWithStorage<boolean>(
@@ -86,7 +89,7 @@ export const activeChatAtom = atom(
     }
     return state.chats.get(state.activeChatId);
   },
-  (get, set, chatId: ChatId | null) => {
+  (_get, set, chatId: ChatId | null) => {
     set(chatStateAtom, (prev) => ({
       ...prev,
       activeChatId: chatId,

package/src/core/ai/tools/__tests__/registry.test.ts

@@ -13,18 +13,43 @@ describe("FrontendToolRegistry", () => {
 
   it("invokes a tool with valid args and validates input/output", async () => {
     const registry = new FrontendToolRegistry([new TestFrontendTool()]);
-    const ok = await registry.invoke("test_frontend_tool", { name: "Alice" });
-    expect(ok).toEqual({ message: "Hello: Alice" });
+    const response = await registry.invoke("test_frontend_tool", {
+      name: "Alice",
+    });
+
+    // Check InvokeResult wrapper
+    expect(response.tool_name).toBe("test_frontend_tool");
+    expect(response.error).toBeNull();
+
+    // Check the actual tool output
+    expect(response.result).toMatchObject({
+      status: "success",
+      data: {
+        greeting: "Hello: Alice",
+      },
+      next_steps: expect.arrayContaining([expect.any(String)]),
+    });
+
+    // Verify timestamp is present and valid
+    const output = response.result as { data: { timestamp: string } };
+    expect(output.data.timestamp).toBeDefined();
+    expect(typeof output.data.timestamp).toBe("string");
   });
 
   it("returns a structured error on invalid args", async () => {
     const registry = new FrontendToolRegistry([new TestFrontendTool()]);
-    const err = (await registry.invoke("test_frontend_tool", {})) as Record<
-      string,
-      unknown
-    >;
-    expect(err.status).toBe("error");
-    expect(err.code).toBe("TOOL_ERROR");
+    const response = await registry.invoke("test_frontend_tool", {});
+
+    // Check InvokeResult wrapper
+    expect(response.tool_name).toBe("test_frontend_tool");
+    expect(response.result).toBeNull();
+    expect(response.error).toBeDefined();
+    expect(typeof response.error).toBe("string");
+
+    // Verify error message contains expected prefix
+    expect(response.error).toContain("Error invoking tool ToolExecutionError:");
+    expect(response.error).toContain('"code":"TOOL_ERROR"');
+    expect(response.error).toContain('"is_retryable":false');
   });
 
   it("returns tool schemas with expected shape and memoizes the result", () => {

package/src/core/ai/tools/base.ts

@@ -1,8 +1,88 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { z } from "zod";
+import { z } from "zod";
 import type { CopilotMode } from "./registry";
 
+/**
+ * Status value for tool responses, mirroring status value in marimo/_ai/_tools/types.py
+ */
+export type StatusValue = "success" | "error" | "warning";
+
+/**
+ * Structured error for tool execution failures.
+ * Mirrors the ToolExecutionError dataclass from marimo/_ai/_tools/utils/exceptions.py
+ *
+ * @example
+ * throw new ToolExecutionError(
+ *   "Failed to fetch data",
+ *   "FETCH_ERROR",
+ *   true,
+ *   "Check your network connection"
+ * );
+ */
+export class ToolExecutionError extends Error {
+  readonly code: string;
+  readonly isRetryable: boolean;
+  readonly suggestedFix?: string;
+  readonly meta?: Record<string, unknown>;
+
+  constructor(
+    message: string,
+    code = "TOOL_ERROR",
+    isRetryable = false,
+    suggestedFix?: string,
+    meta?: Record<string, unknown>,
+  ) {
+    super(message);
+    this.name = "ToolExecutionError";
+    this.code = code;
+    this.isRetryable = isRetryable;
+    this.suggestedFix = suggestedFix;
+    this.meta = meta;
+  }
+
+  toStructuredString(): string {
+    const stringError = JSON.stringify({
+      message: this.message,
+      code: this.code,
+      is_retryable: this.isRetryable,
+      suggested_fix: this.suggestedFix,
+      meta: this.meta ?? {},
+    });
+    return `Error invoking tool ${this.name}: ${stringError}`;
+  }
+}
+
+/**
+ * Base interface for tool output responses.
+ * Mirrors the SuccessResult dataclass from marimo/_ai/_tools/types.py
+ *
+ * Tool outputs should extend this interface to include standardized
+ * metadata like next_steps, messages, and status information.
+ */
+export interface ToolOutputBase {
+  status?: StatusValue;
+  auth_required?: boolean;
+  next_steps?: string[];
+  action_url?: string;
+  message?: string;
+  meta?: Record<string, unknown>;
+}
+
+/**
+ * Base Zod schema for tool outputs.
+ *
+ * Tool output schemas should extend this using .extend() to add their specific fields.
+ */
+export const toolOutputBaseSchema = z.object({
+  status: z.enum(["success", "error", "warning"]).optional(),
+  auth_required: z.boolean().optional(),
+  next_steps: z.array(z.string()).optional(),
+  action_url: z.string().optional(),
+  message: z.string().optional(),
+  meta: z.record(z.string(), z.unknown()).optional(),
+});
+
 /**
  * Contract for a frontend tool.
  *

package/src/core/ai/tools/registry.ts

@@ -3,7 +3,7 @@
 import type { components } from "@marimo-team/marimo-api";
 import { Memoize } from "typescript-memoize";
 import { type ZodObject, z } from "zod";
-import type { AiTool } from "./base";
+import { type AiTool, ToolExecutionError } from "./base";
 import { TestFrontendTool } from "./sample-tool";
 
 export type AnyZodObject = ZodObject<z.ZodRawShape>;
@@ -11,6 +11,12 @@ export type AnyZodObject = ZodObject<z.ZodRawShape>;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type StoredTool = AiTool<any, any>;
 
+interface InvokeResult<TName> {
+  tool_name: TName;
+  result: unknown;
+  error: string | null;
+}
+
 /** should be the same as marimo/_config/config.py > CopilotMode */
 
 type ToolDefinition = components["schemas"]["ToolDefinition"];
@@ -43,7 +49,7 @@ export class FrontendToolRegistry {
   async invoke<TName extends string>(
     toolName: TName,
     rawArgs: unknown,
-  ): Promise<unknown> {
+  ): Promise<InvokeResult<TName>> {
     const tool = this.getToolOrThrow(toolName);
     const handler = tool.handler;
     const inputSchema = tool.schema;
@@ -69,17 +75,37 @@ export class FrontendToolRegistry {
           `Tool ${toolName} returned invalid output: ${strError}`,
         );
       }
-      const output = response.data;
-      return output;
-    } catch (error) {
+      const result = response.data;
       return {
-        status: "error",
-        code: "TOOL_ERROR",
-        message: error instanceof Error ? error.message : String(error),
-        suggestedFix: "Try again with valid arguments.",
-        meta: {
+        tool_name: toolName,
+        result,
+        error: null,
+      };
+    } catch (error) {
+      // Handle structured errors
+      if (error instanceof ToolExecutionError) {
+        return {
+          tool_name: toolName,
+          result: null,
+          error: error.toStructuredString(),
+        };
+      }
+
+      // Handle unknown/generic errors
+      const genericError = new ToolExecutionError(
+        error instanceof Error ? error.message : String(error),
+        "TOOL_ERROR",
+        false,
+        "Check the error message and try again with valid arguments.",
+        {
           args: rawArgs,
+          errorType: error?.constructor?.name ?? typeof error,
         },
+      );
+      return {
+        tool_name: toolName,
+        result: null,
+        error: genericError.toStructuredString(),
       };
     }
   }

package/src/core/ai/tools/sample-tool.ts

@@ -1,7 +1,12 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { z } from "zod";
-import type { AiTool } from "./base";
+import {
+  type AiTool,
+  ToolExecutionError,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
 import type { CopilotMode } from "./registry";
 
 const description = `
@@ -11,26 +16,57 @@ Args:
 - name (string): The name to include in the greeting.
 
 Returns:
-- { message: string } — The greeting message, e.g., "Hello: Alice".
+- Output with data containing the greeting message.
 `;
 
 interface Input {
   name: string;
 }
 
-interface Output {
-  message: string;
+interface GreetingData {
+  greeting: string;
+  timestamp: string;
 }
 
-/** A sample frontend tool that returns "hello world" */
+interface Output extends ToolOutputBase {
+  data: GreetingData;
+}
+
+/** A sample frontend tool that demonstrates real tool output structure */
 export class TestFrontendTool implements AiTool<Input, Output> {
   readonly name = "test_frontend_tool";
   readonly description = description;
   readonly schema = z.object({ name: z.string() });
-  readonly outputSchema = z.object({ message: z.string() });
+  readonly outputSchema = toolOutputBaseSchema.extend({
+    data: z.object({
+      greeting: z.string(),
+      timestamp: z.string(),
+    }),
+  });
   readonly mode: CopilotMode[] = ["ask"];
 
-  async handler({ name }: Input) {
-    return { message: `Hello: ${name}` };
+  async handler({ name }: Input): Promise<Output> {
+    // Example: Validate input and throw ToolExecutionError on invalid data
+    if (!name.trim()) {
+      throw new ToolExecutionError(
+        "Name cannot be empty",
+        "INVALID_INPUT",
+        false,
+        "Please provide a non-empty name",
+        { field: "name", received: name },
+      );
+    }
+
+    return {
+      status: "success",
+      data: {
+        greeting: `Hello: ${name}`,
+        timestamp: new Date().toISOString(),
+      },
+      next_steps: [
+        "You can now proceed with your next task",
+        "Try calling another tool if needed",
+      ],
+    };
   }
 }

package/src/core/cache/requests.ts

@@ -0,0 +1,5 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { atom } from "jotai";
+import type { CacheInfoFetched } from "../kernel/messages";
+
+export const cacheInfoAtom = atom<CacheInfoFetched | null>(null);