feat: Create frontend tools framework and integrate to backend ai system (#6609)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->

Create frontend tools framework to manage and send tools to backend via
ai body. Update openapi to include ToolDefinition in body. Remove any
frontend tool handing from ToolManager.

<img width="504" height="951" alt="Screenshot 2025-09-29 at 9 12 24 PM"
src="https://github.com/user-attachments/assets/6ddf33e6-0382-416e-8d2b-98c6b86a9bfa"
/>


## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->
- Create frontend tool registry to handle registering and invoking tools
from the frontend in registry.ts
- Add BaseTool to handle frontend tool shape validation and centralize
logic in base.ts
- Handle zod v4 collision in tsconfig.json
- Add sample tools as example of how to create a frontend tool (maybe
remove before merge?)
- Send back frontend tools via body in chat-panel.tsx
- Update invokeAiTool to handle frontend tools if in list in
chat-utils.ts
- Add ToolDefinition to ChatRequest body and OpenAPI schemas
- Inject tools through stream_completion if tools is avaliable in ai.py
and providers.py
- Remove all left over frontend tool logic and handling in
tool_manager.py

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [x] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [ ] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> Introduce a frontend tools framework and pass tool schemas through
chat requests to providers, while removing frontend-tool handling from
the backend and updating OpenAPI/types.
> 
> - **Frontend**:
> - **Tools Framework**: Add `AiTool` interface (`base.ts`),
`FrontendToolRegistry` with zod validation and memoized schemas
(`registry.ts`), and a sample tool (`sample-tool.ts`); add tests for
registry behavior.
> - **Chat Integration**: Send `FRONTEND_TOOL_REGISTRY.getToolSchemas()`
in chat request body and handle frontend tool calls client-side before
falling back to backend (`chat-panel.tsx`, `chat-utils.ts`).
>   - **Config**: Pin `zod` path in `tsconfig.json`.
> - **Backend**:
> - **Providers**: Extend `stream_completion` to accept
`additional_tools` and merge with configured tools for OpenAI,
Anthropic, Google, and Bedrock (`providers.py`).
> - **Endpoints**: Pass `body.tools` to providers in `POST
/api/ai/chat`; no tools for completion/inline completion (`ai.py`).
> - **ToolManager**: Drop frontend tool registration/lookup; warn on
frontend source requests; simplify init log; tighten error messages
(`tool_manager.py`).
> - **API/OpenAPI**:
> - Add `ToolDefinition` schema and include optional `tools` in
`ChatRequest`; regenerate TS types (`api.yaml`, `src/api.ts`).
> - Include `ToolDefinition` in dev OpenAPI generation (`commands.py`).
> - **Tests**:
> - Update provider tests to new `stream_completion` signature
(`test_providers.py`).
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
d7836e8. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Myles Scolnick <[email protected]>

Author: 3 people

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/__tests__/registry.test.ts

@@ -0,0 +1,53 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { FrontendToolRegistry } from "../registry";
+import { TestFrontendTool } from "../sample-tool";
+
+describe("FrontendToolRegistry", () => {
+  it("registers tools via constructor and supports has()", () => {
+    const registry = new FrontendToolRegistry([new TestFrontendTool()]);
+    expect(registry.has("test_frontend_tool")).toBe(true);
+    expect(registry.has("nonexistent_tool" as string)).toBe(false);
+  });
+
+  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" });
+  });
+
+  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");
+  });
+
+  it("returns tool schemas with expected shape and memoizes the result", () => {
+    const registry = new FrontendToolRegistry([new TestFrontendTool()]);
+
+    const schemas1 = registry.getToolSchemas();
+    expect(Array.isArray(schemas1)).toBe(true);
+    expect(schemas1.length).toBe(1);
+
+    const def = schemas1[0];
+    expect(def.name).toBe("test_frontend_tool");
+    expect(def.source).toBe("frontend");
+    expect(def.mode).toEqual(["ask"]);
+    expect(typeof def.description).toBe("string");
+
+    const params = def.parameters as Record<string, unknown>;
+    expect(params && typeof params === "object").toBe(true);
+    const properties = (params as { properties?: Record<string, unknown> })
+      .properties;
+    expect(properties && typeof properties === "object").toBe(true);
+    expect("name" in (properties ?? {})).toBe(true);
+
+    const schemas2 = registry.getToolSchemas();
+    expect(schemas2).toBe(schemas1);
+  });
+});

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

@@ -0,0 +1,19 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { z } from "zod";
+import type { CopilotMode } from "./registry";
+
+/**
+ * Contract for a frontend tool.
+ *
+ * Implementations can be plain objects or classes. The registry consumes this
+ * interface without caring about the underlying implementation.
+ */
+export interface AiTool<TIn, TOut> {
+  name: string;
+  description: string;
+  schema: z.ZodType<TIn>;
+  outputSchema: z.ZodType<TOut>;
+  mode: CopilotMode[];
+  handler: (args: TIn) => TOut | Promise<TOut>;
+}

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

@@ -0,0 +1,102 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+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 { TestFrontendTool } from "./sample-tool";
+
+export type AnyZodObject = ZodObject<z.ZodRawShape>;
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type StoredTool = AiTool<any, any>;
+
+/** should be the same as marimo/_config/config.py > CopilotMode */
+
+type ToolDefinition = components["schemas"]["ToolDefinition"];
+export type CopilotMode = ToolDefinition["mode"][number];
+
+export interface FrontendToolDefinition extends ToolDefinition {
+  source: "frontend";
+}
+
+export class FrontendToolRegistry {
+  /** All registered tools */
+  private tools = new Map<string, StoredTool>();
+
+  constructor(tools: StoredTool[] = []) {
+    this.tools = new Map(tools.map((tool) => [tool.name, tool]));
+  }
+
+  has(toolName: string) {
+    return this.tools.has(toolName);
+  }
+
+  private getToolOrThrow(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.getToolOrThrow(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,
+        },
+      };
+    }
+  }
+
+  @Memoize()
+  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([
+  ...(import.meta.env.DEV ? [new TestFrontendTool()] : []),
+  // ADD MORE TOOLS HERE
+]);

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

@@ -0,0 +1,36 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import type { AiTool } from "./base";
+import type { CopilotMode } from "./registry";
+
+const description = `
+Test frontend tool that returns a greeting message.
+
+Args:
+- name (string): The name to include in the greeting.
+
+Returns:
+- { message: string } — The greeting message, e.g., "Hello: Alice".
+`;
+
+interface Input {
+  name: string;
+}
+
+interface Output {
+  message: string;
+}
+
+/** A sample frontend tool that returns "hello world" */
+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 mode: CopilotMode[] = ["ask"];
+
+  async handler({ name }: Input) {
+    return { message: `Hello: ${name}` };
+  }
+}

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