Zod-like type inference, but for JSONSchema (#490)

* Zod-like type inference for our JsonSchema variant, including Cell support
* add type inference to recipe, lift, handler, .asCell, .asSchema and many more
* added `schema` helper so we don't have to write `as const as JSONSchema` + allow `example` in schema
* added `asStream` support

Author: seefeldb

typescript/packages/common-builder/src/index.ts

@@ -58,6 +58,7 @@ export {
   unsafe_parentRecipe,
   type UnsafeBinding,
 } from "./types.ts";
+export { type Schema, schema } from "./schema-to-ts.ts";
 
 // This should be a separate package, but for now it's easier to keep it here.
 export {

typescript/packages/common-builder/src/module.ts

@@ -17,6 +17,7 @@ import {
   traverseValue,
 } from "./utils.ts";
 import { getTopFrame } from "./recipe.ts";
+import { Schema, SchemaWithoutCell } from "./schema-to-ts.ts";
 import { z } from "zod";
 import { zodToJsonSchema } from "zod-to-json-schema";
 
@@ -45,11 +46,14 @@ export function createNodeFactory<T = any, R = any>(
  *
  * @returns A module node factory that also serializes as module.
  */
-export function lift<T, R>(
-  argumentSchema: JSONSchema,
-  resultSchema: JSONSchema,
-  implementation: (input: T) => R,
-): ModuleFactory<T, R>;
+export function lift<
+  T extends JSONSchema = JSONSchema,
+  R extends JSONSchema = JSONSchema,
+>(
+  argumentSchema: T,
+  resultSchema: R,
+  implementation: (input: Schema<T>) => Schema<R>,
+): ModuleFactory<SchemaWithoutCell<T>, SchemaWithoutCell<R>>;
 export function lift<T extends z.ZodTypeAny, R extends z.ZodTypeAny>(
   argumentSchema: T,
   resultSchema: R,
@@ -97,6 +101,14 @@ export function byRef<T, R>(ref: string): ModuleFactory<T, R> {
   });
 }
 
+export function handler<
+  E extends JSONSchema = JSONSchema,
+  T extends JSONSchema = JSONSchema,
+>(
+  eventSchema: E,
+  stateSchema: T,
+  handler: (event: Schema<E>, props: Schema<T>) => any,
+): HandlerFactory<SchemaWithoutCell<T>, SchemaWithoutCell<E>>;
 export function handler<E, T>(
   eventSchema: JSONSchema,
   stateSchema: JSONSchema,

typescript/packages/common-builder/src/recipe.ts

@@ -4,6 +4,7 @@ import {
   isOpaqueRef,
   isShadowRef,
   type JSONSchema,
+  type JSONSchemaWritable,
   makeOpaqueRef,
   type Module,
   type Node,
@@ -27,6 +28,7 @@ import {
   toJSONWithAliases,
   traverseValue,
 } from "./utils.ts";
+import { SchemaWithoutCell } from "./schema-to-ts.ts";
 import { z } from "zod";
 import { zodToJsonSchema } from "zod-to-json-schema";
 
@@ -71,6 +73,21 @@ export function recipe<T, R>(
   resultSchema: JSONSchema,
   fn: (input: OpaqueRef<Required<T>>) => Opaque<R>,
 ): RecipeFactory<T, R>;
+export function recipe<S extends JSONSchema>(
+  argumentSchema: S,
+  fn: (input: OpaqueRef<Required<SchemaWithoutCell<S>>>) => any,
+): RecipeFactory<SchemaWithoutCell<S>, ReturnType<typeof fn>>;
+export function recipe<S extends JSONSchema, R>(
+  argumentSchema: S,
+  fn: (input: OpaqueRef<Required<SchemaWithoutCell<S>>>) => Opaque<R>,
+): RecipeFactory<SchemaWithoutCell<S>, R>;
+export function recipe<S extends JSONSchema, RS extends JSONSchema>(
+  argumentSchema: S,
+  resultSchema: RS,
+  fn: (
+    input: OpaqueRef<Required<SchemaWithoutCell<S>>>,
+  ) => Opaque<SchemaWithoutCell<RS>>,
+): RecipeFactory<SchemaWithoutCell<S>, SchemaWithoutCell<RS>>;
 export function recipe<T, R>(
   argumentSchema: string | JSONSchema | z.ZodTypeAny,
   resultSchema:
@@ -246,15 +263,15 @@ function factoryFromRecipe<T, R>(
     if (external) setValueAtPath(initial, paths.get(cell)!, external);
   });
 
-  let argumentSchema: JSONSchema;
+  let argumentSchema: JSONSchemaWritable;
 
   if (typeof argumentSchemaArg === "string") {
     // TODO(seefeld): initial is likely not needed anymore
-    // TODO(seefeld): But we need a new one for the result
+    // TODO(seefeld): But we need a derived schema for the result
     argumentSchema = createJsonSchema(defaults, {});
     argumentSchema.description = argumentSchemaArg;
 
-    delete argumentSchema.properties?.[UI]; // TODO(seefeld): This should be a schema for views
+    delete (argumentSchema.properties as any)?.[UI]; // TODO(seefeld): This should be a schema for views
     if (argumentSchema.properties?.internal?.properties) {
       for (
         const key of Object.keys(
@@ -269,15 +286,15 @@ function factoryFromRecipe<T, R>(
   } else if (argumentSchemaArg instanceof z.ZodType) {
     argumentSchema = zodToJsonSchema(argumentSchemaArg) as JSONSchema;
   } else {
-    argumentSchema = argumentSchemaArg as unknown as JSONSchema;
+    argumentSchema = argumentSchemaArg;
   }
 
   const resultSchema: JSONSchema = resultSchemaArg instanceof z.ZodType
     ? (zodToJsonSchema(resultSchemaArg) as JSONSchema)
     : (resultSchemaArg ?? ({} as JSONSchema));
 
   const serializedNodes = Array.from(nodes).map((node) => {
-    const module = toJSONWithAliases(node.module, paths) as Module;
+    const module = toJSONWithAliases(node.module, paths) as unknown as Module;
     const inputs = toJSONWithAliases(node.inputs, paths)!;
     const outputs = toJSONWithAliases(node.outputs, paths)!;
     return { module, inputs, outputs } satisfies Node;