diff --git a/packages/core/src/promptdom.ts b/packages/core/src/promptdom.ts
index 0456148a8..71d9f5f63 100644
--- a/packages/core/src/promptdom.ts
+++ b/packages/core/src/promptdom.ts
@@ -362,13 +362,13 @@ export function createImageNode(
 // Function to create a schema node.
 export function createSchemaNode(
     name: string,
-    value: JSONSchema,
+    value: JSONSchema | ZodTypeLike,
     options?: DefSchemaOptions
 ): PromptSchemaNode {
     assert(!!name)
     assert(value !== undefined)
     // auto zod conversion
-    value = tryZodToJsonSchema(value) ?? value
+    value = tryZodToJsonSchema(value as ZodTypeLike) ?? (value as JSONSchema)
     return { type: "schema", name, value, options }
 }
 
diff --git a/packages/core/src/types/prompt_template.d.ts b/packages/core/src/types/prompt_template.d.ts
index 6ff01de74..f6481a577 100644
--- a/packages/core/src/types/prompt_template.d.ts
+++ b/packages/core/src/types/prompt_template.d.ts
@@ -2604,10 +2604,12 @@ interface McpServerConfig {
 
 type McpServersConfig = Record<string, Omit<McpServerConfig, "id" | "options">>
 
+type ZodTypeLike = { _def: any, safeParse: any, refine: any }
+
 interface ChatGenerationContext extends ChatTurnGenerationContext {
     defSchema(
         name: string,
-        schema: JSONSchema | object,
+        schema: JSONSchema | ZodTypeLike,
         options?: DefSchemaOptions
     ): string
     defImages(
diff --git a/packages/core/src/types/prompt_type.d.ts b/packages/core/src/types/prompt_type.d.ts
index 529770149..56c7042f3 100644
--- a/packages/core/src/types/prompt_type.d.ts
+++ b/packages/core/src/types/prompt_type.d.ts
@@ -239,7 +239,7 @@ declare function fetchText(
  */
 declare function defSchema(
     name: string,
-    schema: JSONSchema | object,
+    schema: JSONSchema | ZodTypeLike,
     options?: DefSchemaOptions
 ): string
 
diff --git a/packages/core/src/zod.ts b/packages/core/src/zod.ts
index 68f2a94da..f3fe5a600 100644
--- a/packages/core/src/zod.ts
+++ b/packages/core/src/zod.ts
@@ -6,17 +6,14 @@ import { zodToJsonSchema as _zodToJsonSchema } from "zod-to-json-schema"
  * @param options
  * @returns
  */
-export function tryZodToJsonSchema(z: object, options?: object): JSONSchema {
-    if (!z) return undefined
-    // instanceof not working, test for some existing methoid
-    if (!(z as any)._def) return undefined
-    try {
-        const schema = _zodToJsonSchema(z as any, {
-            target: "openAi",
-            ...(options || {}),
-        })
-        return structuredClone(schema) as JSONSchema
-    } catch (e) {
-        return undefined
-    }
+export function tryZodToJsonSchema(
+    z: ZodTypeLike,
+    options?: object
+): JSONSchema {
+    if (!z || !z._def || !z.refine || !z.safeParse) return undefined
+    const schema = _zodToJsonSchema(z as any, {
+        target: "openAi",
+        ...(options || {}),
+    })
+    return structuredClone(schema) as JSONSchema
 }