diff --git a/frontend/src/core/ai/tools/__tests__/registry.test.ts b/frontend/src/core/ai/tools/__tests__/registry.test.ts index 8604724235f..1b4e19f6d95 100644 --- a/frontend/src/core/ai/tools/__tests__/registry.test.ts +++ b/frontend/src/core/ai/tools/__tests__/registry.test.ts @@ -13,8 +13,22 @@ 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 result = await registry.invoke("test_frontend_tool", { + name: "Alice", + }); + + expect(result).toMatchObject({ + status: "success", + data: { + greeting: "Hello: Alice", + }, + next_steps: expect.arrayContaining([expect.any(String)]), + }); + + // Verify timestamp is present and valid + const output = 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 () => { diff --git a/frontend/src/core/ai/tools/base.ts b/frontend/src/core/ai/tools/base.ts index 0e2c4c366fe..a33600d92e9 100644 --- a/frontend/src/core/ai/tools/base.ts +++ b/frontend/src/core/ai/tools/base.ts @@ -1,8 +1,43 @@ /* 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"; + +/** + * 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; +} + +/** + * 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. * diff --git a/frontend/src/core/ai/tools/sample-tool.ts b/frontend/src/core/ai/tools/sample-tool.ts index 454ff3f7d2b..bf8ab826303 100644 --- a/frontend/src/core/ai/tools/sample-tool.ts +++ b/frontend/src/core/ai/tools/sample-tool.ts @@ -1,7 +1,7 @@ /* Copyright 2024 Marimo. All rights reserved. */ import { z } from "zod"; -import type { AiTool } from "./base"; +import { type AiTool, type ToolOutputBase, toolOutputBaseSchema } from "./base"; import type { CopilotMode } from "./registry"; const description = ` @@ -11,26 +11,46 @@ 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 { 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 { + 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", + ], + }; } }