zodvex 0.2.5 → 0.3.2

package/src/custom.ts

@@ -178,13 +178,15 @@ export function customFnBuilder<
 
   return function customBuilder(fn: any): any {
     const { args, handler = fn, returns: maybeObject, ...extra } = fn
+    const skipConvexValidation = fn.skipConvexValidation ?? false
 
     const returns =
       maybeObject && !(maybeObject instanceof z.ZodType) ? z.object(maybeObject) : maybeObject
+    // Only generate Convex return validator when not skipping Convex validation
     const returnValidator =
-      returns && !fn.skipConvexValidation ? { returns: zodToConvex(returns) } : undefined
+      returns && !skipConvexValidation ? { returns: zodToConvex(returns) } : undefined
 
-    if (args && !fn.skipConvexValidation) {
+    if (args) {
       let argsValidator = args
       let argsSchema: z.ZodObject<any>
 
@@ -202,9 +204,13 @@ export function customFnBuilder<
         argsSchema = z.object(argsValidator)
       }
 
-      const convexValidator = zodToConvexFields(argsValidator)
+      // Only generate Convex args validator when not skipping Convex validation
+      const convexArgs = skipConvexValidation
+        ? inputArgs
+        : { ...zodToConvexFields(argsValidator), ...inputArgs }
+
       return builder({
-        args: { ...convexValidator, ...inputArgs },
+        args: convexArgs,
         ...returnValidator,
         handler: async (ctx: Ctx, allArgs: any) => {
           const added: any = await customInput(
@@ -215,6 +221,7 @@ export function customFnBuilder<
           const argKeys = Object.keys(argsValidator)
           const rawArgs = pick(allArgs, argKeys)
           const decoded = fromConvexJS(rawArgs, argsSchema)
+          // Always run Zod validation, regardless of skipConvexValidation
           const parsed = argsSchema.safeParse(decoded)
           if (!parsed.success) {
             handleZodValidationError(parsed.error, 'args')
@@ -224,7 +231,8 @@ export function customFnBuilder<
           const addedArgs = (added?.args as Record<string, unknown>) ?? {}
           const finalArgs = { ...baseArgs, ...addedArgs }
           const ret = await handler(finalCtx, finalArgs)
-          if (returns && !fn.skipConvexValidation) {
+          // Always run Zod return validation when returns schema is provided
+          if (returns) {
             let validated: any
             try {
               validated = (returns as z.ZodTypeAny).parse(ret)
@@ -261,7 +269,8 @@ export function customFnBuilder<
         const addedArgs = (added?.args as Record<string, unknown>) ?? {}
         const finalArgs = { ...baseArgs, ...addedArgs }
         const ret = await handler(finalCtx, finalArgs)
-        if (returns && !fn.skipConvexValidation) {
+        // Always run Zod return validation when returns schema is provided
+        if (returns) {
           let validated: any
           try {
             validated = (returns as z.ZodTypeAny).parse(ret)

package/src/ids.ts

@@ -16,23 +16,21 @@ export const registryHelpers = {
 /**
  * Create a Zod validator for a Convex Id
  *
- * Uses the string → transform → brand pattern for proper type narrowing with ctx.db.get()
- * This aligns with Zod v4 best practices and matches convex-helpers implementation
+ * Compatible with AI SDK and other tools that don't support transforms.
+ * Uses type-level branding instead of runtime transforms for GenericId<T> compatibility.
+ *
+ * @param tableName - The Convex table name for this ID
+ * @returns A Zod string validator typed as GenericId<TableName>
  */
 export function zid<TableName extends string>(
   tableName: TableName
 ): z.ZodType<GenericId<TableName>> & { _tableName: TableName } {
-  // Use the string → transform → brand pattern (aligned with Zod v4 best practices)
+  // Create base string validator with refinement (no transform or brand)
   const baseSchema = z
     .string()
     .refine(val => typeof val === 'string' && val.length > 0, {
       message: `Invalid ID for table "${tableName}"`
     })
-    .transform(val => {
-      // Cast to GenericId while keeping the string value
-      return val as string & GenericId<TableName>
-    })
-    .brand(`ConvexId_${tableName}`) // Use native Zod v4 .brand() method
     // Add a human-readable marker for client-side introspection utilities
     // used in apps/native (e.g., to detect relationship fields in dynamic forms).
     .describe(`convexId:${tableName}`)
@@ -47,6 +45,8 @@ export function zid<TableName extends string>(
   const branded = baseSchema as any
   branded._tableName = tableName
 
+  // Type assertion provides GenericId<TableName> typing without runtime transform
+  // This maintains type safety while being compatible with AI SDK and similar tools
   return branded as z.ZodType<GenericId<TableName>> & { _tableName: TableName }
 }
 

package/src/registry.ts

@@ -56,3 +56,174 @@ export function isDateSchema(schema: any): boolean {
 
   return false
 }
+
+// ============================================================================
+// JSON Schema Override Support
+// ============================================================================
+// Zod's toJSONSchema doesn't support transforms, brands, and other "unrepresentable"
+// types by default. This module provides overrides for zodvex-managed types
+// so they can be used with AI SDKs and other JSON Schema-based tools.
+
+/**
+ * Checks if a schema is a zid (Convex ID) schema by looking at its description.
+ * zid schemas are marked with "convexId:{tableName}" in their description.
+ */
+export function isZidSchema(schema: z.ZodTypeAny): boolean {
+  const description = schema.description
+  return typeof description === 'string' && description.startsWith('convexId:')
+}
+
+/**
+ * Extracts the table name from a zid schema's description.
+ * Returns undefined if not a zid schema.
+ */
+export function getZidTableName(schema: z.ZodTypeAny): string | undefined {
+  const description = schema.description
+  if (typeof description === 'string' && description.startsWith('convexId:')) {
+    return description.slice('convexId:'.length)
+  }
+  return undefined
+}
+
+/**
+ * Context object passed to the JSON Schema override function.
+ * Uses 'any' types for compatibility with Zod's internal types.
+ */
+export interface JSONSchemaOverrideContext {
+  zodSchema: any // Zod's internal $ZodTypes
+  jsonSchema: any // Zod's JSONSchema.BaseSchema
+}
+
+/**
+ * Override function for z.toJSONSchema that handles zodvex-managed types.
+ *
+ * Handles:
+ * - zid schemas: Converts to { type: "string" } with convexId format
+ * - z.date(): Converts to { type: "string", format: "date-time" }
+ *
+ * @example
+ * ```ts
+ * import { z } from 'zod'
+ * import { zid, zodvexJSONSchemaOverride } from 'zodvex'
+ *
+ * const schema = z.object({
+ *   userId: zid('users'),
+ *   name: z.string()
+ * })
+ *
+ * const jsonSchema = z.toJSONSchema(schema, {
+ *   unrepresentable: 'any',
+ *   override: zodvexJSONSchemaOverride
+ * })
+ * // => { type: "object", properties: { userId: { type: "string" }, name: { type: "string" } } }
+ * ```
+ */
+export function zodvexJSONSchemaOverride(ctx: JSONSchemaOverrideContext): void {
+  const { zodSchema, jsonSchema } = ctx
+
+  // Handle zid schemas (transforms with convexId description)
+  if (isZidSchema(zodSchema)) {
+    const tableName = getZidTableName(zodSchema)
+    // Set our properties - don't clear existing ones set by user overrides
+    // When unrepresentable: 'any', Zod already gives us {} so no clearing needed
+    jsonSchema.type = 'string'
+    if (tableName) {
+      jsonSchema.format = `convex-id:${tableName}`
+    }
+    // Preserve the description from .describe() - this is what the LLM sees
+    if (zodSchema.description) {
+      jsonSchema.description = zodSchema.description
+    }
+    return
+  }
+
+  // Handle z.date() - convert to ISO 8601 string format
+  // Zod v4 passes real schema instances here (ZodDate has `type === 'date'`).
+  if (zodSchema instanceof z.ZodDate || (zodSchema as any).type === 'date') {
+    jsonSchema.type = 'string'
+    jsonSchema.format = 'date-time'
+    return
+  }
+}
+
+/**
+ * Composes multiple JSON Schema override functions into one.
+ * Overrides run in order - first override runs first.
+ *
+ * @example
+ * ```ts
+ * import { composeOverrides, zodvexJSONSchemaOverride } from 'zodvex'
+ *
+ * const myOverride = (ctx) => {
+ *   if (ctx.zodSchema.description?.startsWith('myType:')) {
+ *     ctx.jsonSchema.type = 'string'
+ *     ctx.jsonSchema.format = 'my-format'
+ *   }
+ * }
+ *
+ * // User's override runs first, then zodvex's
+ * z.toJSONSchema(schema, {
+ *   unrepresentable: 'any',
+ *   override: composeOverrides(myOverride, zodvexJSONSchemaOverride)
+ * })
+ * ```
+ */
+export function composeOverrides(
+  ...overrides: Array<((ctx: JSONSchemaOverrideContext) => void) | undefined>
+): (ctx: JSONSchemaOverrideContext) => void {
+  return (ctx: JSONSchemaOverrideContext) => {
+    for (const override of overrides) {
+      override?.(ctx)
+    }
+  }
+}
+
+/**
+ * Options for toJSONSchema, matching Zod's interface.
+ */
+export interface ToJSONSchemaOptions {
+  target?: 'draft-4' | 'draft-7' | 'draft-2020-12' | 'openapi-3.0'
+  unrepresentable?: 'throw' | 'any'
+  cycles?: 'ref' | 'throw'
+  reused?: 'ref' | 'inline'
+  io?: 'input' | 'output'
+  override?: (ctx: JSONSchemaOverrideContext) => void
+}
+
+/**
+ * Converts a Zod schema to JSON Schema with zodvex-aware overrides.
+ *
+ * This is a convenience wrapper around z.toJSONSchema that automatically
+ * handles zodvex-managed types like zid and dates.
+ *
+ * @example
+ * ```ts
+ * import { zid, toJSONSchema } from 'zodvex'
+ *
+ * const schema = z.object({
+ *   userId: zid('users'),
+ *   createdAt: z.date(),
+ *   name: z.string()
+ * })
+ *
+ * const jsonSchema = toJSONSchema(schema)
+ * // Works with AI SDK's generateObject, etc.
+ * ```
+ */
+export function toJSONSchema<T extends z.ZodTypeAny>(
+  schema: T,
+  options?: ToJSONSchemaOptions
+): Record<string, any> {
+  const userOverride = options?.override
+
+  return z.toJSONSchema(schema, {
+    ...options,
+    // Default to 'any' so transforms don't throw
+    unrepresentable: options?.unrepresentable ?? 'any',
+    // Chain our override with user's override
+    override: ctx => {
+      zodvexJSONSchemaOverride(ctx)
+      userOverride?.(ctx)
+    }
+  })
+}

package/src/tables.ts

@@ -1,8 +1,98 @@
+import { defineTable } from 'convex/server'
 import type { GenericId } from 'convex/values'
 import { Table } from 'convex-helpers/server'
 import { z } from 'zod'
 import { zid } from './ids'
-import { type ConvexValidatorFromZodFieldsAuto, zodToConvexFields } from './mapping'
+import { type ConvexValidatorFromZodFieldsAuto, zodToConvex, zodToConvexFields } from './mapping'
+
+/**
+ * Helper type for Convex system fields added to documents
+ */
+type SystemFields<TableName extends string> = {
+  _id: ReturnType<typeof zid<TableName>>
+  _creationTime: z.ZodNumber
+}
+
+/**
+ * Maps over union options, extending each ZodObject variant with system fields.
+ * Non-object variants are preserved as-is.
+ */
+type MapSystemFields<TableName extends string, Options extends readonly z.ZodTypeAny[]> = {
+  [K in keyof Options]: Options[K] extends z.ZodObject<infer Shape extends z.ZodRawShape>
+    ? z.ZodObject<Shape & SystemFields<TableName>>
+    : Options[K]
+}
+
+/**
+ * Adds Convex system fields (_id, _creationTime) to a Zod schema.
+ *
+ * For object schemas: extends with system fields
+ * For union schemas: adds system fields to each variant
+ *
+ * @param tableName - The Convex table name
+ * @param schema - The Zod schema (object or union)
+ * @returns Schema with system fields added
+ */
+// Overload 1: ZodObject - extends with system fields
+export function addSystemFields<TableName extends string, Shape extends z.ZodRawShape>(
+  tableName: TableName,
+  schema: z.ZodObject<Shape>
+): z.ZodObject<Shape & SystemFields<TableName>>
+
+// Overload 2: ZodUnion - maps system fields to each variant
+export function addSystemFields<TableName extends string, Options extends readonly z.ZodTypeAny[]>(
+  tableName: TableName,
+  schema: z.ZodUnion<Options>
+): z.ZodUnion<MapSystemFields<TableName, Options>>
+
+// Overload 3: ZodDiscriminatedUnion - maps system fields preserving discriminator
+// Note: Zod v4 signature is ZodDiscriminatedUnion<Options, Discriminator>
+export function addSystemFields<
+  TableName extends string,
+  Options extends readonly z.ZodObject<z.ZodRawShape>[],
+  Discriminator extends string
+>(
+  tableName: TableName,
+  schema: z.ZodDiscriminatedUnion<Options, Discriminator>
+): z.ZodDiscriminatedUnion<MapSystemFields<TableName, Options>, Discriminator>
+
+// Overload 4: Fallback for other ZodTypes - returns as-is
+export function addSystemFields<TableName extends string, S extends z.ZodTypeAny>(
+  tableName: TableName,
+  schema: S
+): S
+
+// Implementation
+export function addSystemFields<TableName extends string>(
+  tableName: TableName,
+  schema: z.ZodTypeAny
+): z.ZodTypeAny {
+  // Handle union schemas - add system fields to each variant
+  if (schema instanceof z.ZodUnion || schema instanceof z.ZodDiscriminatedUnion) {
+    const options = (schema as z.ZodUnion<any>).options.map((variant: z.ZodTypeAny) => {
+      if (variant instanceof z.ZodObject) {
+        return variant.extend({
+          _id: zid(tableName),
+          _creationTime: z.number()
+        })
+      }
+      // Non-object variants are returned as-is (shouldn't happen in practice)
+      return variant
+    })
+    return z.union(options as any)
+  }
+
+  // Handle object schemas
+  if (schema instanceof z.ZodObject) {
+    return schema.extend({
+      _id: zid(tableName),
+      _creationTime: z.number()
+    })
+  }
+
+  // Fallback: return schema as-is
+  return schema
+}
 
 // Helper to create a Zod schema for a Convex document
 export function zodDoc<
@@ -35,21 +125,43 @@ export function zodDocOrNull<
 }
 
 /**
- * Defines a Convex table using a raw Zod shape (an object mapping field names to Zod types).
+ * Helper to detect if input is an object shape (plain object with Zod validators)
+ */
+function isObjectShape(input: any): input is Record<string, z.ZodTypeAny> {
+  // Check if it's a plain object (not a Zod instance)
+  if (!input || typeof input !== 'object') return false
+
+  // If it's a Zod instance, it's not an object shape
+  if (input instanceof z.ZodType) return false
+
+  // Check if all values are Zod types
+  for (const key in input) {
+    if (!(input[key] instanceof z.ZodType)) {
+      return false
+    }
+  }
+
+  return true
+}
+
+/**
+ * Defines a Convex table using either:
+ * - A raw Zod shape (an object mapping field names to Zod types)
+ * - A Zod union schema (for polymorphic tables)
  *
- * This function intentionally accepts a raw shape instead of a ZodObject instance.
+ * For object shapes, this function intentionally accepts a raw shape instead of a ZodObject instance.
  * Accepting raw shapes allows TypeScript to infer field types more accurately and efficiently,
  * leading to better type inference and performance throughout the codebase.
- * This architectural decision is important for projects that rely heavily on type safety and
- * developer experience, as it avoids the type erasure that can occur when using ZodObject directly.
+ *
+ * For union schemas, this enables polymorphic tables with discriminated unions.
  *
  * Returns the Table definition along with Zod schemas for documents and arrays.
  *
  * @param name - The table name
- * @param shape - A raw object mapping field names to Zod validators
- * @returns A Table with attached shape, zDoc schema, and docArray helper
+ * @param schemaOrShape - Either a raw object shape or a Zod union schema
+ * @returns A Table with attached helpers (shape, schema, zDoc, docArray, withSystemFields)
  *
- * @example
+ * @example Object shape
  * ```ts
  * const Users = zodTable('users', {
  *   name: z.string(),
@@ -66,36 +178,132 @@ export function zodDocOrNull<
  *   { returns: Users.docArray }
  * )
  * ```
+ *
+ * @example Union schema (polymorphic table)
+ * ```ts
+ * const shapeSchema = z.union([
+ *   z.object({ kind: z.literal('circle'), r: z.number() }),
+ *   z.object({ kind: z.literal('rectangle'), width: z.number() })
+ * ])
+ *
+ * const Shapes = zodTable('shapes', shapeSchema)
+ *
+ * // Use in schema
+ * export default defineSchema({ shapes: Shapes.table })
+ *
+ * // Use for return types with system fields
+ * export const getShapes = zQuery(query, {},
+ *   async (ctx) => ctx.db.query('shapes').collect(),
+ *   { returns: Shapes.docArray }
+ * )
+ * ```
  */
+// Helper type to compute the result of addSystemFields for use in zodTable return type
+type AddSystemFieldsResult<
+  TableName extends string,
+  Schema extends z.ZodTypeAny
+> = Schema extends z.ZodObject<infer Shape extends z.ZodRawShape>
+  ? z.ZodObject<Shape & SystemFields<TableName>>
+  : Schema extends z.ZodUnion<infer Options extends readonly z.ZodTypeAny[]>
+    ? z.ZodUnion<MapSystemFields<TableName, Options>>
+    : Schema extends z.ZodDiscriminatedUnion<
+          infer Options extends readonly z.ZodObject<z.ZodRawShape>[],
+          infer Disc extends string
+        >
+      ? z.ZodDiscriminatedUnion<MapSystemFields<TableName, Options>, Disc>
+      : Schema
+
+// Overload 1: Object shape (most common case)
 export function zodTable<TableName extends string, Shape extends Record<string, z.ZodTypeAny>>(
   name: TableName,
   shape: Shape
-) {
-  // Convert fields with proper types
-  const convexFields = zodToConvexFields(shape) as ConvexValidatorFromZodFieldsAuto<Shape>
-
-  // Create the Table from convex-helpers with explicit type
-  const table = Table<ConvexValidatorFromZodFieldsAuto<Shape>, TableName>(name, convexFields)
-
-  // Create zDoc schema with system fields
-  const zDoc = zodDoc(name, z.object(shape))
-
-  // Create docArray helper for return types
-  const docArray = z.array(zDoc)
-
-  // Attach everything for comprehensive usage
-  return Object.assign(table, {
-    shape,
-    zDoc,
-    docArray
-  }) as typeof table & {
-    shape: Shape
-    zDoc: z.ZodObject<
+): ReturnType<typeof Table<ConvexValidatorFromZodFieldsAuto<Shape>, TableName>> & {
+  shape: Shape
+  zDoc: z.ZodObject<
+    Shape & {
+      _id: ReturnType<typeof zid<TableName>>
+      _creationTime: z.ZodNumber
+    }
+  >
+  docArray: z.ZodArray<
+    z.ZodObject<
       Shape & {
         _id: ReturnType<typeof zid<TableName>>
         _creationTime: z.ZodNumber
       }
     >
-    docArray: z.ZodArray<typeof zDoc>
+  >
+}
+
+// Overload 2: Union/schema types
+export function zodTable<TableName extends string, Schema extends z.ZodTypeAny>(
+  name: TableName,
+  schema: Schema
+): {
+  table: ReturnType<typeof defineTable>
+  tableName: TableName
+  validator: ReturnType<typeof zodToConvex<Schema>>
+  schema: Schema
+  docArray: z.ZodArray<AddSystemFieldsResult<TableName, Schema>>
+  withSystemFields: () => AddSystemFieldsResult<TableName, Schema>
+}
+
+export function zodTable<
+  TableName extends string,
+  SchemaOrShape extends z.ZodTypeAny | Record<string, z.ZodTypeAny>
+>(name: TableName, schemaOrShape: SchemaOrShape): any {
+  // Detect if it's an object shape or a schema
+  if (isObjectShape(schemaOrShape)) {
+    // Original object shape logic
+    const shape = schemaOrShape as Record<string, z.ZodTypeAny>
+
+    // Convert fields with proper types
+    const convexFields = zodToConvexFields(shape) as ConvexValidatorFromZodFieldsAuto<typeof shape>
+
+    // Create the Table from convex-helpers with explicit type
+    const table = Table<ConvexValidatorFromZodFieldsAuto<typeof shape>, TableName>(
+      name,
+      convexFields
+    )
+
+    // Create zDoc schema with system fields
+    const zDoc = zodDoc(name, z.object(shape))
+
+    // Create docArray helper for return types
+    const docArray = z.array(zDoc)
+
+    // Attach everything for comprehensive usage
+    return Object.assign(table, {
+      shape,
+      zDoc,
+      docArray
+    })
+  } else {
+    // Union or other schema type logic
+    const schema = schemaOrShape as z.ZodTypeAny
+
+    // Convert schema to Convex validator
+    const convexValidator = zodToConvex(schema)
+
+    // For unions, use defineTable directly (not Table helper which expects object fields)
+    // Note: TypeScript types don't reflect it, but Convex supports union validators in tables
+    const table = defineTable(convexValidator as any)
+
+    // Create document schema with system fields
+    const withFields = addSystemFields(name, schema)
+
+    // Create docArray helper
+    const docArray = z.array(withFields)
+
+    // Attach helpers for union tables
+    // Return structure similar to Table() but without fields-based helpers
+    return {
+      table,
+      tableName: name,
+      validator: convexValidator,
+      schema,
+      docArray,
+      withSystemFields: () => addSystemFields(name, schema)
+    }
   }
 }

package/src/types.ts

@@ -15,31 +15,19 @@ export type InferArgs<A> = A extends z.ZodObject<infer S>
       ? z.infer<A>
       : Record<string, never>
 
-// Return type inference with immediate bailout for unions/custom to avoid depth
-export type InferReturns<R> = R extends z.ZodUnion<any>
-  ? any
-  : // Bail immediately for unions
-    R extends z.ZodCustom<any>
+// Return type inference - uses z.output for Zod schemas
+// Previously had bailouts for unions/custom to avoid TypeScript depth errors,
+// but research (Issue #20) showed convex-helpers handles these without issues.
+// Removing bailouts fixes Issue #19 (Promise<any> return types).
+export type InferReturns<R> = R extends z.ZodType<any, any, any>
+  ? z.output<R>
+  : R extends undefined
     ? any
-    : // Bail immediately for custom
-      R extends z.ZodType<any, any, any>
-      ? z.output<R>
-      : // Use z.output for other schemas
-        R extends undefined
-        ? any
-        : R
+    : R
 
 // For handler authoring: what the handler returns before wrapper validation/encoding
-export type InferHandlerReturns<R> = R extends z.ZodUnion<any>
-  ? any
-  : // Bail immediately for unions
-    R extends z.ZodCustom<any>
-    ? any
-    : // Bail immediately for custom
-      R extends z.ZodType<any, any, any>
-      ? z.input<R>
-      : // Use z.input for other schemas
-        any
+// Uses z.input since this is what the handler produces before encoding
+export type InferHandlerReturns<R> = R extends z.ZodType<any, any, any> ? z.input<R> : any
 
 /**
  * Extract the visibility type from a Convex builder function

package/src/wrappers.ts

@@ -39,8 +39,8 @@ function containsCustom(schema: z.ZodTypeAny, maxDepth = 50, currentDepth = 0):
 
   let result = false
 
-  // Use _def.typeName instead of instanceof since ZodCustom is not exported in Zod v4
-  if ((schema as any)._def?.typeName === 'ZodCustom') {
+  // Zod v4 exports ZodCustom and instances expose `schema.type === "custom"`.
+  if (schema instanceof z.ZodCustom) {
     result = true
   } else if (schema instanceof z.ZodUnion) {
     result = (schema.options as z.ZodTypeAny[]).some(opt =>