@marimo-team/islands 0.16.4 → 0.17.0

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,10 +18,11 @@ 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";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 
 type NamedDatasource = Omit<
   DataSourceConnection,
@@ -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[];
 
@@ -53,24 +60,35 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
   }
 
   getItems(): DatasourceContextItem[] {
-    return [...this.connectionsMap.values()].map((connection) => {
-      let description = "Database schema.";
-      const data: DatasourceContextItem["data"] = {
-        connection: connection,
-      };
-      if (INTERNAL_SQL_ENGINES.has(connection.name)) {
-        data.tables = this.dataframes;
-        description = "Database schema and the dataframes that can be queried";
-      }
-
-      return {
-        uri: this.asURI(connection.name),
-        name: connection.name,
-        description: description,
-        type: this.contextType,
-        data: data,
-      };
-    });
+    return [...this.connectionsMap.values()]
+      .map((connection): DatasourceContextItem | null => {
+        let description = "Database schema.";
+        const data: DatasourceContextItem["data"] = {
+          connection: connection,
+        };
+
+        if (INTERNAL_SQL_ENGINES.has(connection.name)) {
+          data.tables = this.dataframes;
+          description =
+            "Database schema and the dataframes that can be queried";
+        }
+
+        // Hide empty datasources
+        const hasNoTables =
+          connection.databases.length === 0 && (data.tables?.length ?? 0) === 0;
+        if (hasNoTables) {
+          return null;
+        }
+
+        return {
+          uri: this.asURI(connection.name),
+          name: connection.name,
+          description: description,
+          type: this.contextType,
+          data: data,
+        };
+      })
+      .filter(Boolean);
   }
 
   formatContext(item: DatasourceContextItem): string {
@@ -115,7 +133,7 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
       detail: dbDisplayName(dataConnection.dialect),
       boost: Boosts.MEDIUM,
       type: this.contextType,
-      section: "Data Sources",
+      section: Sections.DATA_SOURCES,
       info: () => {
         const infoContainer = document.createElement("div");
         infoContainer.classList.add("mo-cm-tooltip", "docs-documentation");
@@ -129,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/context/providers/error.ts

@@ -9,7 +9,7 @@ import { logNever } from "@/utils/assertNever";
 import { PluralWord } from "@/utils/pluralize";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Sections } from "./common";
 
 export interface ErrorContextItem extends AIContextItem {
   type: "error";
@@ -111,10 +111,9 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
         label: "@Errors",
         displayLabel: "Errors",
         detail: `${item.data.errors.length} ${errorsTxt.pluralize(item.data.errors.length)}`,
-        boost: Boosts.ERROR,
         type: "error",
         apply: "@Errors",
-        section: "Errors",
+        section: Sections.ERROR,
         info: () => {
           const infoContainer = document.createElement("div");
           infoContainer.classList.add(
@@ -150,7 +149,7 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
     return {
       label: "Error",
       displayLabel: "Error",
-      boost: Boosts.ERROR,
+      section: Sections.ERROR,
     };
   }
 

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

@@ -15,7 +15,7 @@ import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 export interface FileContextItem extends AIContextItem {
   type: "file";
   data: {
@@ -190,7 +190,7 @@ export class FileContextProvider extends AIContextProvider<FileContextItem> {
     return {
       ...this.createBasicCompletion(item),
       type: "file",
-      section: "File",
+      section: Sections.FILE,
       boost: data.isDirectory ? Boosts.MEDIUM : Boosts.LOW,
       detail: data.path,
       displayLabel: `${icon} ${name}`,

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

@@ -9,7 +9,7 @@ import type { DataTable } from "@/core/kernel/messages";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Boosts, Sections } from "./common";
 
 export interface TableContextItem extends AIContextItem {
   type: "data";
@@ -38,21 +38,46 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
 
   formatContext(item: TableContextItem): string {
     const { data } = item;
-    const { columns, source, num_rows, num_columns, name } = data;
+    const { columns, source, num_rows, num_columns, name, variable_name } =
+      data;
+
+    // Build shape information
     const shape = [
-      num_rows == null ? undefined : `${num_rows} rows`,
-      num_columns == null ? undefined : `${num_columns} columns`,
+      num_rows != null ? `${num_rows} rows` : undefined,
+      num_columns != null ? `${num_columns} columns` : undefined,
     ]
       .filter(Boolean)
       .join(", ");
 
     let details = "";
+
+    // Add shape information
     if (shape) {
-      details += `\nShape: ${shape}`;
+      details += `Shape: ${shape}\n`;
+    }
+
+    // Add variable name if available
+    if (variable_name) {
+      details += `Variable: ${variable_name}\n`;
     }
 
+    // Add column information with sample values
     if (columns && columns.length > 0) {
-      details += `\nColumns:\n${columns.map((col) => `  - ${col.name}: ${col.type}`).join("\n")}`;
+      details += "Columns:\n";
+      for (const col of columns) {
+        let columnInfo = `  ${col.name} (${col.type})`;
+
+        // Add sample values if available
+        if (col.sample_values && col.sample_values.length > 0) {
+          const samples = col.sample_values
+            .slice(0, 3) // Limit to first 3 samples
+            .map((val) => (val === null ? "null" : String(val)))
+            .join(", ");
+          columnInfo += ` - samples: [${samples}]`;
+        }
+
+        details += `${columnInfo}\n`;
+      }
     }
 
     return contextToXml({
@@ -61,7 +86,7 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
         name: name,
         source: source ?? "unknown",
       },
-      details: details,
+      details: details.trim(),
     });
   }
 
@@ -78,7 +103,10 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
           : Boosts.REMOTE_TABLE,
       type: getTableType(table),
       apply: `@${tableName}`,
-      section: getTableType(table) === "dataframe" ? "Dataframe" : "Table",
+      section: {
+        name: getTableType(table) === "dataframe" ? "Dataframe" : "Table",
+        rank: Sections.TABLE.rank,
+      },
       info: () => this.createTableInfoElement(tableName, table),
     };
   }

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

@@ -6,7 +6,7 @@ import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
 import type { Variable, Variables } from "@/core/variables/types";
 import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
-import { Boosts } from "./common";
+import { Sections } from "./common";
 
 export interface VariableContextItem extends AIContextItem {
   type: "variable";
@@ -58,9 +58,8 @@ export class VariableContextProvider extends AIContextProvider<VariableContextIt
       label: `@${variable.name}`,
       displayLabel: variable.name,
       detail: variable.dataType ?? "",
-      boost: Boosts.VARIABLE,
       type: this.contextType,
-      section: "Variable",
+      section: Sections.VARIABLE,
       info: () => {
         return createVariableInfoElement(variable);
       },

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

package/src/core/cells/__tests__/cells.test.ts

@@ -12,7 +12,7 @@ import {
   it,
   vi,
 } from "vitest";
-import type { CellHandle } from "@/components/editor/Cell";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import { CellId } from "@/core/cells/ids";
 import { foldAllBulk, unfoldAllBulk } from "@/core/codemirror/editing/commands";
 import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
@@ -2094,9 +2094,9 @@ describe("cell reducer", () => {
     `);
   });
 
-  it("can create and update a setup cell", () => {
+  it("can create and noop-update a setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);
@@ -2106,17 +2106,17 @@ describe("cell reducer", () => {
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
 
     // Update the setup cell
-    actions.upsertSetupCell({ code: "# Updated setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Updated setup code" });
 
     // Check that the same setup cell was updated, not duplicated
-    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Updated setup code");
+    expect(state.cellData[SETUP_CELL_ID].code).toBe("# Setup code");
     expect(state.cellData[SETUP_CELL_ID].edited).toBe(true);
     expect(state.cellIds.inOrderIds).toContain(SETUP_CELL_ID);
   });
 
   it("can delete and undelete the setup cell", () => {
     // Create the setup cell
-    actions.upsertSetupCell({ code: "# Setup code" });
+    actions.addSetupCellIfDoesntExist({ code: "# Setup code" });
 
     // Check that setup cell was created
     expect(state.cellData[SETUP_CELL_ID].id).toBe(SETUP_CELL_ID);