feat: Create frontend tools framework and integrate to backend ai system

frontend/src/components/chat/chat-panel.tsx

@@ -44,6 +44,7 @@ import {
   type ChatId,
   chatStateAtom,
 } from "@/core/ai/state";
+import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
 import { aiAtom, aiEnabledAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
@@ -585,6 +586,7 @@ const ChatPanelBody = () => {
 
         return {
           body: {
+            tools: FRONTEND_TOOL_REGISTRY.getToolSchemas(),
             ...options,
             ...completionBody,
           },

frontend/src/components/chat/chat-utils.ts

@@ -2,6 +2,7 @@
 
 import type { components } from "@marimo-team/marimo-api";
 import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
 import type {
   InvokeAiToolRequest,
   InvokeAiToolResponse,
@@ -117,15 +118,29 @@ export async function handleToolCall({
   };
 }) {
   try {
-    const response = await invokeAiTool({
-      toolName: toolCall.toolName,
-      arguments: toolCall.input,
-    });
-    addToolResult({
-      tool: toolCall.toolName,
-      toolCallId: toolCall.toolCallId,
-      output: response.result || response.error,
-    });
+    if (FRONTEND_TOOL_REGISTRY.has(toolCall.toolName)) {
+      // Invoke the frontend tool
+      const response = await FRONTEND_TOOL_REGISTRY.invoke(
+        toolCall.toolName,
+        toolCall.input,
+      );
+      addToolResult({
+        tool: toolCall.toolName,
+        toolCallId: toolCall.toolCallId,
+        output: response,
+      });
+    } else {
+      // Invoke the backend/mcp tool
+      const response = await invokeAiTool({
+        toolName: toolCall.toolName,
+        arguments: toolCall.input,
+      });
+      addToolResult({
+        tool: toolCall.toolName,
+        toolCallId: toolCall.toolCallId,
+        output: response.result || response.error,
+      });
+    }
   } catch (error) {
     Logger.error("Tool call failed:", error);
     addToolResult({

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

@@ -0,0 +1,48 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { z } from "zod";
+import type { AnyZodObject } from "./registry";
+
+/**
+ * Minimal base class for frontend tools.
+ *
+ * Structural typing ensures instances are compatible with the Tool<TIn, TOut>
+ * interface used by the registry, without importing it here.
+ */
+export abstract class BaseTool<
+  TIn extends AnyZodObject,
+  TOut extends AnyZodObject,
+> {
+  public readonly name: string;
+  public readonly description: string;
+  public readonly schema: TIn;
+  public readonly outputSchema: TOut;
+  public readonly mode: ("manual" | "ask")[];
+
+  /**
+   * Handler exposed to the registry. Calls the subclass implementation.
+   */
+  public readonly handler: (
+    args: z.infer<TIn>,
+  ) => z.infer<TOut> | Promise<z.infer<TOut>>;
+
+  constructor(options: {
+    name: string;
+    description: string;
+    schema: TIn;
+    mode: ("manual" | "ask")[];
+    outputSchema: TOut;
+  }) {
+    this.name = options.name;
+    this.description = options.description;
+    this.schema = options.schema;
+    this.mode = options.mode;
+    this.outputSchema = options.outputSchema;
+    this.handler = (args) => Promise.resolve(this.handle(args));
+  }
+
+  /** Implement tool logic in subclasses */
+  protected abstract handle(
+    args: z.infer<TIn>,
+  ): z.infer<TOut> | Promise<z.infer<TOut>>;
+}

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

@@ -0,0 +1,132 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { type ZodObject, z } from "zod";
+import type { BaseTool } from "./base";
+import { testFrontendTool } from "./sample-tool";
+
+export type AnyZodObject = ZodObject<z.ZodRawShape>;
+
+interface StoredTool {
+  /** Generic type for to avoid type errors */
+  name: string;
+  description: string;
+  schema: AnyZodObject;
+  outputSchema: AnyZodObject;
+  mode: CopilotMode[];
+  handler: (args: unknown) => Promise<unknown>;
+}
+
+/** should be the same as marimo/_config/config.py > CopilotMode */
+export type CopilotMode = "manual" | "ask";
+
+export interface FrontendToolDefinition {
+  /** should be the same as marimo/_server/ai/tools/types.py > ToolDefinition */
+  name: string;
+  description: string;
+  parameters: Record<string, unknown>;
+  source: "frontend";
+  mode: CopilotMode[];
+}
+
+export class FrontendToolRegistry {
+  /** All registered tools */
+  private tools = new Map<string, StoredTool>();
+
+  registerAll<TIn extends AnyZodObject, TOut extends AnyZodObject>(
+    tools: BaseTool<TIn, TOut>[],
+  ) {
+    tools.forEach((tool) => {
+      this.register(tool);
+    });
+  }
+
+  private register<TIn extends AnyZodObject, TOut extends AnyZodObject>(
+    tool: BaseTool<TIn, TOut>,
+  ) {
+    // Make type generic to avoid type errors
+    // Let invoke() handle runtime type checking
+    const stored: StoredTool = {
+      name: tool.name,
+      description: tool.description,
+      schema: tool.schema,
+      outputSchema: tool.outputSchema,
+      mode: tool.mode,
+      handler: tool.handler as (args: unknown) => Promise<unknown>,
+    };
+    this.tools.set(tool.name, stored);
+  }
+
+  has(toolName: string) {
+    return this.tools.has(toolName);
+  }
+
+  private getTool(toolName: string): StoredTool {
+    const tool = this.tools.get(toolName);
+    if (!tool) {
+      throw new Error(`Tool ${toolName} not found`);
+    }
+    return tool;
+  }
+
+  async invoke<TName extends string>(
+    toolName: TName,
+    rawArgs: unknown,
+  ): Promise<unknown> {
+    const tool = this.getTool(toolName);
+    const handler = tool.handler;
+    const inputSchema = tool.schema;
+    const outputSchema = tool.outputSchema;
+
+    try {
+      // Parse input args
+      const inputResponse = await inputSchema.safeParseAsync(rawArgs);
+      if (inputResponse.error) {
+        const strError = z.prettifyError(inputResponse.error);
+        throw new Error(`Tool ${toolName} returned invalid input: ${strError}`);
+      }
+      const args = inputResponse.data;
+
+      // Call the handler
+      const rawOutput = await handler(args);
+
+      // Parse output
+      const response = await outputSchema.safeParseAsync(rawOutput);
+      if (response.error) {
+        const strError = z.prettifyError(response.error);
+        throw new Error(
+          `Tool ${toolName} returned invalid output: ${strError}`,
+        );
+      }
+      const output = response.data;
+      return output;
+    } catch (error) {
+      return {
+        status: "error",
+        code: "TOOL_ERROR",
+        message: error instanceof Error ? error.message : String(error),
+        suggestedFix: "Try again with valid arguments.",
+        meta: {
+          args: rawArgs,
+        },
+      };
+    }
+  }
+
+  getToolSchemas(): FrontendToolDefinition[] {
+    return [...this.tools.values()].map((tool) => ({
+      name: tool.name,
+      description: tool.description,
+      parameters: z.toJSONSchema(tool.schema),
+      source: "frontend",
+      mode: tool.mode,
+    }));
+  }
+}
+
+export const FRONTEND_TOOL_REGISTRY = new FrontendToolRegistry();
+
+/* Register all the frontend tools */
+FRONTEND_TOOL_REGISTRY.registerAll([
+  testFrontendTool,
+  // ADD MORE TOOLS HERE
+]);

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

@@ -0,0 +1,30 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import { BaseTool } from "./base";
+
+const schema = z.object({ name: z.string() });
+const outputSchema = z.object({ message: z.string() });
+
+/** A sample frontend tool that returns "hello world" */
+export class TestFrontendTool extends BaseTool<
+  typeof schema,
+  typeof outputSchema
+> {
+  constructor() {
+    super({
+      name: "test_frontend_tool",
+      description:
+        "A test frontend tool that returns hi with the name passed in",
+      schema,
+      outputSchema,
+      mode: ["ask"],
+    });
+  }
+
+  protected async handle({ name }: z.infer<typeof schema>) {
+    return { message: `Hello: ${name}` };
+  }
+}
+
+export const testFrontendTool = new TestFrontendTool();

frontend/tsconfig.json

@@ -23,7 +23,9 @@
     /* Style */
     "allowJs": true /* needed for storybook */,
     "paths": {
-      "@/*": ["./src/*"]
+      "@/*": ["./src/*"],
+      /* Pin zod to root (v4); v3 exists via @modelcontextprotocol/sdk */
+      "zod": ["./node_modules/zod"]
     },
     "forceConsistentCasingInFileNames": true,
     "useDefineForClassFields": true,

marimo/_cli/development/commands.py

@@ -49,6 +49,7 @@ def _generate_server_api_schema() -> dict[str, Any]:
     from marimo._messaging.mimetypes import KnownMimeType
     from marimo._plugins.core.web_component import JSONType
     from marimo._runtime.packages.package_manager import PackageDescription
+    from marimo._server.ai.tools.types import ToolDefinition
     from marimo._server.api.router import build_routes
     from marimo._utils.dataclass_to_openapi import (
         PythonTypeToOpenAPI,
@@ -136,6 +137,7 @@ def _generate_server_api_schema() -> dict[str, Any]:
         ops.MessageOperation,
         # ai
         ChatMessage,
+        ToolDefinition,
     ]
 
     # dataclass components used in requests/responses