1 files changed, 7 insertions, 8 deletions
diff --git a/packages/plugin/src/tool.ts b/packages/plugin/src/tool.ts
index 7c1d3d7c5..2998a1e72 100644
--- a/packages/plugin/src/tool.ts
+++ b/packages/plugin/src/tool.ts
@@ -7,14 +7,13 @@ export type ToolContext = {
   abort: AbortSignal
 }
 
-export function tool<Args extends z.ZodRawShape>(
-  input: (zod: typeof z) => {
-    description: string
-    args: Args
-    execute: (args: z.infer<z.ZodObject<Args>>, ctx: ToolContext) => Promise<string>
-  },
-) {
-  return input(z)
+export function tool<Args extends z.ZodRawShape>(input: {
+  description: string
+  args: Args
+  execute(args: z.infer<z.ZodObject<Args>>, context: ToolContext): Promise<string>
+}) {
+  return input
 }
+tool.schema = z
 
 export type ToolDefinition = ReturnType<typeof tool>