@aztec/foundation 0.0.1-commit.9ef841308 → 0.0.1-commit.a5db02d

package/src/json-rpc/server/safe_json_rpc_server.ts

@@ -10,7 +10,13 @@ import { ZodError } from 'zod';
 
 import { type Logger, createLogger } from '../../log/index.js';
 import { promiseWithResolvers } from '../../promise/utils.js';
-import { type ApiSchema, type ApiSchemaFor, parseWithOptionals, schemaHasMethod } from '../../schemas/index.js';
+import {
+  type ApiSchema,
+  type ApiSchemaFor,
+  getSchemaParameters,
+  parseWithOptionals,
+  schemaHasMethod,
+} from '../../schemas/index.js';
 import { jsonStringify } from '../convert.js';
 import { assert } from '../js_utils.js';
 
@@ -53,8 +59,10 @@ export class SafeJsonRpcServer {
     config: Partial<SafeJsonRpcServerConfig> = {},
     /** Health check function */
     private readonly healthCheck: StatusCheckFn = () => true,
-    /** Additional middlewares */
+    /** Additional Koa middlewares */
     private extraMiddlewares: Application.Middleware[] = [],
+    /** Additional per-request diagnostics middlewares */
+    private diagnosticsMiddleware?: DiagnosticsMiddleware,
     /** Logger */
     private log = createLogger('json-rpc:server'),
   ) {
@@ -155,7 +163,7 @@ export class SafeJsonRpcServer {
           };
           return;
         }
-        const resp = await this.processBatch(ctx.request.body);
+        const resp = await this.processBatch(ctx.request.body, ctx.request.headers);
         if (Array.isArray(resp)) {
           ctx.status = 200;
           ctx.body = resp;
@@ -164,7 +172,7 @@ export class SafeJsonRpcServer {
           ctx.body = resp;
         }
       } else {
-        const resp = await this.processRequest(ctx.request.body);
+        const resp = await this.processRequest(ctx.request.body, ctx.request.headers);
         if ('error' in resp) {
           ctx.status = this.config.http200OnError ? 200 : 400;
         }
@@ -176,11 +184,11 @@ export class SafeJsonRpcServer {
     return router;
   }
 
-  private async processBatch(requests: any[]) {
+  private async processBatch(requests: any[], headers: http.IncomingHttpHeaders = {}) {
     if (requests.length === 0) {
       return { jsonrpc: '2.0', error: { code: -32600, message: 'Invalid Request' }, id: null };
     }
-    const results = await Promise.allSettled(requests.map(req => this.processRequest(req)));
+    const results = await Promise.allSettled(requests.map(req => this.processRequest(req, headers)));
     return results.map(res => {
       if (res.status === 'fulfilled') {
         return res.value;
@@ -191,7 +199,7 @@ export class SafeJsonRpcServer {
     });
   }
 
-  private async processRequest(request: any) {
+  private async processRequest(request: any, headers: http.IncomingHttpHeaders = {}) {
     if (!request || typeof request !== 'object') {
       return { jsonrpc: '2.0', error: { code: -32600, message: 'Invalid Request' }, id: null };
     }
@@ -206,7 +214,16 @@ export class SafeJsonRpcServer {
       return { jsonrpc, id, error: { code: -32601, message: `Method not found: ${method}` } };
     } else {
       try {
-        const result = await this.proxy.call(method, params);
+        let result: any;
+
+        if (this.diagnosticsMiddleware) {
+          await this.diagnosticsMiddleware({ id: id ?? null, method, params, headers }, async () => {
+            result = await this.proxy.call(method, params);
+          });
+        } else {
+          result = await this.proxy.call(method, params);
+        }
+
         return { jsonrpc, id, result };
       } catch (err: any) {
         if (err && err instanceof ZodError) {
@@ -309,7 +326,7 @@ export class SafeJsonProxy<T extends object = any> implements Proxy {
     assert(schemaHasMethod(this.schema, methodName), `Method ${methodName} not found in schema`);
     const method = this.handler[methodName as keyof T];
     assert(typeof method === 'function', `Method ${methodName} is not a function`);
-    const args = await parseWithOptionals(jsonParams, this.schema[methodName].parameters());
+    const args = await parseWithOptionals(jsonParams, getSchemaParameters(this.schema[methodName]));
     const ret = await method.apply(this.handler, args);
     this.log.debug(format('response', methodName, ret));
     return ret;
@@ -377,6 +394,7 @@ function makeAggregateHealthcheck(namedHandlers: NamespacedApiHandlers, log?: Lo
 export type SafeJsonRpcServerOptions = Partial<
   SafeJsonRpcServerConfig & {
     healthCheck: StatusCheckFn;
+    diagnostic: DiagnosticsMiddleware;
     log: Logger;
     middlewares: Application.Middleware[];
   }
@@ -391,10 +409,10 @@ export function createNamespacedSafeJsonRpcServer(
   handlers: NamespacedApiHandlers,
   options: Omit<SafeJsonRpcServerOptions, 'healthcheck'> = {},
 ): SafeJsonRpcServer {
-  const { middlewares, log } = options;
+  const { diagnostic, middlewares, log } = options;
   const proxy = new NamespacedSafeJsonProxy(handlers);
   const healthCheck = makeAggregateHealthcheck(handlers, log);
-  return new SafeJsonRpcServer(proxy, options, healthCheck, middlewares, log);
+  return new SafeJsonRpcServer(proxy, options, healthCheck, middlewares, diagnostic, log);
 }
 
 export function createSafeJsonRpcServer<T extends object = any>(
@@ -402,9 +420,9 @@ export function createSafeJsonRpcServer<T extends object = any>(
   schema: ApiSchemaFor<T>,
   options: SafeJsonRpcServerOptions = {},
 ) {
-  const { log, healthCheck, middlewares: extraMiddlewares } = options;
+  const { diagnostic, log, healthCheck, middlewares: extraMiddlewares } = options;
   const proxy = new SafeJsonProxy(handler, schema);
-  return new SafeJsonRpcServer(proxy, options, healthCheck, extraMiddlewares, log);
+  return new SafeJsonRpcServer(proxy, options, healthCheck, extraMiddlewares, diagnostic, log);
 }
 
 /**

package/src/log/pino-logger.ts

@@ -21,6 +21,8 @@ export type LoggerBindings = {
   instanceId?: string;
 };
 
+const MAX_MODULE_NAME_LENGTH = 256;
+
 // Allow global hooks for providing default bindings.
 // Used by withLoggerBindings in pino-logger-server to propagate bindings via AsyncLocalStorage.
 type LogBindingsHandler = () => LoggerBindings | undefined;
@@ -48,7 +50,7 @@ function getBindingsFromHandlers(): LoggerBindings | undefined {
 }
 
 export function createLogger(module: string, bindings?: LoggerBindings): Logger {
-  module = module.replace(/^aztec:/, '');
+  module = module.slice(0, MAX_MODULE_NAME_LENGTH).replace(/^aztec:/, '');
 
   const resolvedBindings = { ...getBindingsFromHandlers(), ...bindings };
   const actor = resolvedBindings?.actor;
@@ -375,5 +377,20 @@ export type Logger = { [K in LogLevel]: LogFn } & { /** Error log function */ er
  * @returns A string with both the log message and the error message.
  */
 function formatErr(msg: string, err?: unknown): string {
-  return err ? `${msg}: ${inspect(err)}` : msg;
+  if (!err) {
+    return msg;
+  }
+
+  try {
+    return `${msg}: ${inspect(err)}`;
+  } catch {
+    // inspect can crash on error objects with broken property descriptors
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-base-to-string
+      return `${msg}: ${err instanceof Error ? `${err.name}: ${err.message}` : String(err)}`;
+    } catch {
+      // if even String(err) fails, return the original message with a note about the error
+      return `${msg}: [unserializable error]`;
+    }
+  }
 }

package/src/noir/noir_package_config.ts

@@ -10,26 +10,6 @@ const noirLocalDependencySchema = z.object({
   path: z.string(),
 });
 
-const noirPackageConfigSchema = z.object({
-  package: z.object({
-    name: z.string().default(''),
-    type: z.enum(['lib', 'contract', 'bin']).default('bin'),
-    entry: z.string().optional(),
-    description: z.string().optional(),
-    authors: z.array(z.string()).optional(),
-    // eslint-disable-next-line camelcase
-    compiler_version: z.string().optional(),
-    backend: z.string().optional(),
-    license: z.string().optional(),
-  }),
-  dependencies: z.record(z.union([noirGitDependencySchema, noirLocalDependencySchema])).default({}),
-});
-
-/**
- * Noir package configuration.
- */
-export type NoirPackageConfig = z.infer<typeof noirPackageConfigSchema>;
-
 /**
  * A remote package dependency.
  */
@@ -45,6 +25,38 @@ export type NoirLocalDependencyConfig = z.infer<typeof noirLocalDependencySchema
  */
 export type NoirDependencyConfig = NoirGitDependencyConfig | NoirLocalDependencyConfig;
 
+/**
+ * Noir package configuration.
+ */
+export type NoirPackageConfig = {
+  package: {
+    name: string;
+    type: 'lib' | 'contract' | 'bin';
+    entry?: string | undefined;
+    description?: string | undefined;
+    authors?: string[] | undefined;
+    compiler_version?: string | undefined;
+    backend?: string | undefined;
+    license?: string | undefined;
+  };
+  dependencies: Record<string, NoirDependencyConfig>;
+};
+
+const noirPackageConfigSchema: z.ZodType<NoirPackageConfig> = z.object({
+  package: z.object({
+    name: z.string().default(''),
+    type: z.enum(['lib', 'contract', 'bin']).default('bin'),
+    entry: z.string().optional(),
+    description: z.string().optional(),
+    authors: z.array(z.string()).optional(),
+    // eslint-disable-next-line camelcase
+    compiler_version: z.string().optional(),
+    backend: z.string().optional(),
+    license: z.string().optional(),
+  }),
+  dependencies: z.record(z.string(), z.union([noirGitDependencySchema, noirLocalDependencySchema])).default({}),
+});
+
 /**
  * Checks that an object is a package configuration.
  * @param config - Config to check

package/src/promise/running-promise.ts

@@ -22,16 +22,17 @@ export function makeLoggingErrorHandler(
  * at a specified polling interval. It allows starting, stopping, and checking the status of the
  * internally managed promise. The class also supports interrupting the polling process when stopped.
  */
-export class RunningPromise {
+export class RunningPromise<T = void> {
   private running = false;
   private runningPromise = Promise.resolve();
   private interruptibleSleep = new InterruptibleSleep();
   private requested: PromiseWithResolvers<void> | undefined = undefined;
+  private requestedArg: T | undefined = undefined;
 
   public static readonly EXIT: typeof EXIT = EXIT;
 
   constructor(
-    private fn: () => void | Promise<void>,
+    private fn: (arg?: T) => void | Promise<void>,
     private logger = createLogger('running-promise'),
     private pollingIntervalMS = 10000,
     private handleError: ErrorHandler = makeLoggingErrorHandler(logger),
@@ -51,7 +52,7 @@ export class RunningPromise {
       while (this.running) {
         const hasRequested = this.requested !== undefined;
         try {
-          await this.fn();
+          await this.fn(this.requestedArg);
         } catch (err) {
           const code = await this.handleError(err);
           if (code === RunningPromise.EXIT) {
@@ -64,6 +65,7 @@ export class RunningPromise {
         if (hasRequested) {
           this.requested!.resolve();
           this.requested = undefined;
+          this.requestedArg = undefined;
         }
 
         // If no immediate run was requested, sleep for the polling interval.
@@ -101,15 +103,16 @@ export class RunningPromise {
    * Triggers an immediate run of the function, bypassing the polling interval.
    * If the function is currently running, it will be allowed to continue and then called again immediately.
    */
-  public async trigger() {
+  public async trigger(arg?: T): Promise<void> {
     if (!this.running) {
-      return this.fn();
+      return this.fn(arg);
     }
 
     let requested = this.requested;
     if (!requested) {
       requested = promiseWithResolvers<void>();
       this.requested = requested;
+      this.requestedArg = arg;
       this.interruptibleSleep.interrupt();
     }
     await requested!.promise;

package/src/queue/batch_queue.ts

@@ -105,6 +105,7 @@ export class BatchQueue<T, K extends string | number> {
       return Promise.resolve();
     }
 
+    this.flushCurrentBatch();
     this.container.end();
     return runningPromise;
   }

package/src/retry/index.ts

@@ -1,7 +1,7 @@
 import { TimeoutError } from '../error/index.js';
 import { type Logger, createLogger } from '../log/index.js';
 import { sleep } from '../sleep/index.js';
-import { Timer } from '../timer/index.js';
+import { type DateProvider, Timer } from '../timer/index.js';
 
 /** An error that indicates that the operation should not be retried. */
 export class NoRetryError extends Error {}
@@ -72,6 +72,34 @@ export async function retry<Result>(
   }
 }
 
+/**
+ * Timeout specification accepted by {@link retryUntil}. Either a plain number of seconds, an explicit
+ * `{ timeout }` in seconds, or an absolute `{ deadline }` with an optional {@link DateProvider} used to
+ * read the current time. A deadline is converted to the remaining seconds at call time; a deadline that
+ * has already passed yields a zero remaining budget, matching the immediate-timeout semantics of a
+ * non-positive numeric timeout.
+ */
+export type RetryUntilTimeout = number | { timeout: number } | { deadline: Date; dateProvider?: DateProvider };
+
+/**
+ * Resolves a {@link RetryUntilTimeout} to a numeric timeout in seconds for the legacy timer-based loop.
+ * A numeric/`{ timeout }` value passes through unchanged (0 means never time out, negative times out on the
+ * first interval). A `{ deadline }` is the remaining seconds until the deadline; when the deadline is already
+ * at or past `now` it resolves to a negative value so the loop times out immediately instead of being read as
+ * the never-timeout sentinel 0.
+ */
+function resolveRetryUntilTimeoutSeconds(timeout: RetryUntilTimeout): number {
+  if (typeof timeout === 'number') {
+    return timeout;
+  }
+  if ('deadline' in timeout) {
+    const now = timeout.dateProvider?.now() ?? Date.now();
+    const remainingSeconds = (timeout.deadline.getTime() - now) / 1000;
+    return remainingSeconds > 0 ? remainingSeconds : -1;
+  }
+  return timeout.timeout;
+}
+
 /**
  * Retry an asynchronous function until it returns a truthy value or the specified timeout is exceeded.
  * The function is retried periodically with a fixed interval between attempts. The operation can be named for better error messages.
@@ -79,16 +107,20 @@ export async function retry<Result>(
  *
  * @param fn - The asynchronous function to be retried, which should return a truthy value upon success or undefined otherwise.
  * @param name - The optional name of the operation, used for generating timeout error message.
- * @param timeout - The optional maximum time, in seconds, to keep retrying before throwing a timeout error. Defaults to 0 (never timeout).
+ * @param timeout - The maximum time to keep retrying before throwing a timeout error. Accepts a number of
+ * seconds (0 = never time out), an explicit `{ timeout }` in seconds, or an absolute `{ deadline }` with an
+ * optional `dateProvider`. A deadline already in the past times out on the first interval, the same as a
+ * zero/negative numeric timeout. Defaults to 0 (never timeout).
  * @param interval - The optional interval, in seconds, between retry attempts. Defaults to 1 second.
  * @returns A Promise that resolves with the successful (truthy) result of the provided function, or rejects if timeout is exceeded.
  */
 export async function retryUntil<T>(
   fn: () => (T | undefined) | Promise<T | undefined>,
   name = '',
-  timeout = 0,
+  timeout: RetryUntilTimeout = 0,
   interval = 1,
 ) {
+  const timeoutSeconds = resolveRetryUntilTimeoutSeconds(timeout);
   const timer = new Timer();
   while (true) {
     const result = await fn();
@@ -98,12 +130,42 @@ export async function retryUntil<T>(
 
     await sleep(interval * 1000);
 
-    if (timeout && timer.s() > timeout) {
+    if (timeoutSeconds && timer.s() > timeoutSeconds) {
       throw new TimeoutError(name ? `Timeout awaiting ${name}` : 'Timeout');
     }
   }
 }
 
+/**
+ * Retry an asynchronous function until it returns a truthy value or the maximum number of retries is exceeded.
+ * The function is retried periodically with a fixed interval between attempts.
+ *
+ * @param fn - The asynchronous function to be retried, which should return a truthy value upon success or undefined otherwise.
+ * @param name - The optional name of the operation, used for generating error messages.
+ * @param maxRetries - The maximum number of retry attempts before throwing an error.
+ * @param retryInterval - The optional interval, in seconds, between retry attempts. Defaults to 1 second.
+ * @returns A Promise that resolves with the successful (truthy) result of the provided function, or rejects if retries are exhausted.
+ */
+export async function retryTimes<T>(
+  fn: () => (T | undefined) | Promise<T | undefined>,
+  name = '',
+  maxRetries: number,
+  retryInterval = 1,
+) {
+  for (let attempt = 0; attempt <= maxRetries; attempt++) {
+    const result = await fn();
+    if (result) {
+      return result;
+    }
+
+    if (attempt < maxRetries) {
+      await sleep(retryInterval * 1000);
+    }
+  }
+
+  throw new Error(name ? `Retries exhausted awaiting ${name}` : 'Retries exhausted');
+}
+
 /**
  * Convenience wrapper around retryUntil with fast polling for tests.
  * Uses 10s timeout and 100ms polling interval by default.

package/src/schemas/api.ts

@@ -2,32 +2,19 @@ import type { z } from 'zod';
 
 import type { ZodNullableOptional } from './utils.js';
 
-// Forces usage of ZodNullableOptional in parameters schemas so we properly accept nulls for optional parameters.
-type ZodParameterTypeFor<T> = undefined extends T
-  ? ZodNullableOptional<z.ZodType<T, z.ZodTypeDef, any>>
-  : z.ZodType<T, z.ZodTypeDef, any>;
-
-type ZodReturnTypeFor<T> = z.ZodType<T, z.ZodTypeDef, any>;
-
-// This monstruosity is used for mapping function arguments to their schema representation.
-// The complexity is required to satisfy ZodTuple which requires a fixed length tuple and
-// has a very annoying type of [] | [ZodTypeAny, ...ZodTypeAny], and most types fail to match
-// the second option. While a purely recursive approach works, it fails when trying to deal
-// with optional arguments (ie optional items in the tuple), and ts does not really like them
-// during a recursion and fails with infinite stack depth.
-// This type appears to satisfy everyone. Everyone but me.
-type ZodMapParameterTypes<T> = T extends []
-  ? []
-  : T extends [item: infer Head, ...infer Rest]
-    ? [ZodParameterTypeFor<Head>, ...{ [K in keyof Rest]: ZodParameterTypeFor<Rest[K]> }]
-    : T extends [item?: infer Head, ...infer Rest]
-      ? [ZodNullableOptional<ZodParameterTypeFor<Head>>, ...{ [K in keyof Rest]: ZodParameterTypeFor<Rest[K]> }]
-      : never;
+// Forces usage of schemas that accept nulls for optional parameters.
+type ZodParameterTypeFor<T> = undefined extends T ? ZodNullableOptional<z.ZodType<T, any>> : z.ZodType<T, any>;
+
+type ZodReturnTypeFor<T> = z.ZodType<T, any>;
+
+type ZodMapParameterTypes<T extends readonly unknown[]> = {
+  [K in keyof T]-?: ZodParameterTypeFor<T[K]>;
+};
 
 /** Maps all functions in an interface to their schema representation. */
 export type ApiSchemaFor<T> = {
   [K in keyof T]: T[K] extends (...args: infer Args) => Promise<infer Ret>
-    ? z.ZodFunction<z.ZodTuple<ZodMapParameterTypes<Args>, z.ZodUnknown>, ZodReturnTypeFor<Ret>>
+    ? z.ZodFunction<z.ZodTuple<ZodMapParameterTypes<Args>, null>, ZodReturnTypeFor<Ret>>
     : never;
 };
 
@@ -44,7 +31,17 @@ export function schemaHasMethod<T extends ApiSchema>(
   return (
     typeof methodName === 'string' &&
     Object.hasOwn(schema, methodName) &&
-    typeof schema[methodName].parameters === 'function' &&
-    typeof schema[methodName].returnType === 'function'
+    schema[methodName].def.input !== undefined &&
+    schema[methodName].def.output !== undefined
   );
 }
+
+/** Returns the parameter tuple schema for an API method schema. */
+export function getSchemaParameters<T extends z.ZodFunction<z.ZodTuple<any, any>, z.ZodTypeAny>>(schema: T) {
+  return schema.def.input;
+}
+
+/** Returns the return value schema for an API method schema. */
+export function getSchemaReturnType<T extends z.ZodFunction<z.ZodTuple<any, any>, z.ZodTypeAny>>(schema: T) {
+  return schema.def.output;
+}

package/src/schemas/parse.ts

@@ -4,16 +4,19 @@ import { times } from '../collection/array.js';
 
 /** Parses the given arguments using a tuple from the provided schemas. */
 export function parse<T extends [] | [z.ZodTypeAny, ...z.ZodTypeAny[]]>(args: IArguments, ...schemas: T) {
-  return z.tuple(schemas).parse(args);
+  return schemas.length === 0
+    ? z.tuple([]).parse(args)
+    : z.tuple(schemas as [z.ZodTypeAny, ...z.ZodTypeAny[]]).parse(args);
 }
 
 /**
  * Parses the given arguments against a tuple, allowing empty for optional items.
- * @dev Zod doesn't like tuplues with optional items. See https://github.com/colinhacks/zod/discussions/949.
+ * @dev Zod doesn't like tuples with optional items. See https://github.com/colinhacks/zod/discussions/949.
  */
-export function parseWithOptionals<T extends z.AnyZodTuple>(args: any[], schema: T): Promise<T['_output']> {
-  const missingCount = schema.items.length - args.length;
-  const optionalCount = schema.items.filter(isOptional).length;
+export function parseWithOptionals<T extends z.ZodTuple<any, any>>(args: any[], schema: T): Promise<z.output<T>> {
+  const items = schema.def.items;
+  const missingCount = items.length - args.length;
+  const optionalCount = items.filter(isOptional).length;
   const toParse =
     missingCount > 0 && missingCount <= optionalCount ? args.concat(times(missingCount, () => undefined)) : args;
   return schema.parseAsync(toParse);
@@ -24,6 +27,6 @@ function isOptional(schema: z.ZodTypeAny) {
     return schema.isOptional();
   } catch {
     // See https://github.com/colinhacks/zod/issues/1911
-    return schema._def.typeName === 'ZodOptional';
+    return schema.def.type === 'optional';
   }
 }

package/src/schemas/schemas.ts

@@ -44,15 +44,15 @@ export const schemas = {
   ]),
 
   /** Coerces any input to bigint. */
-  BigInt: z.union([z.bigint(), z.number(), z.string()]).pipe(z.coerce.bigint()),
+  BigInt: z.union([z.bigint(), z.number(), z.string()]).pipe(z.coerce.bigint<string | number | bigint>()),
 
   /** Coerces any input to integer number. */
-  Integer: z.union([z.bigint(), z.number(), z.string()]).pipe(z.coerce.number().int()),
+  Integer: z.union([z.bigint(), z.number(), z.string()]).pipe(z.coerce.number<string | number | bigint>().int()),
 
   /** Coerces input to UInt32. */
   UInt32: z.union([z.bigint(), z.number(), z.string()]).pipe(
     z.coerce
-      .number()
+      .number<string | number | bigint>()
       .int()
       .min(0)
       .max(2 ** 32 - 1),
@@ -61,7 +61,7 @@ export const schemas = {
   /** Coerces input to UInt64. */
   UInt64: z.union([z.bigint(), z.number(), z.string()]).pipe(
     z.coerce
-      .bigint()
+      .bigint<string | number | bigint>()
       .min(0n)
       .max(2n ** 64n - 1n),
   ),

package/src/schemas/types.ts

@@ -1,6 +1,6 @@
 import type { ZodType } from 'zod';
 
-export type ZodFor<T> = ZodType<T, any, any>;
+export type ZodFor<T> = ZodType<T, any>;
 
 /**
  * Creates a schema validator that enforces all properties of type T are present in the schema.
@@ -26,7 +26,7 @@ export type ZodFor<T> = ZodType<T, any, any>;
  * ```
  */
 export function zodFor<T>() {
-  return (schema => schema) as <S extends ZodType<any, any, any>>(
+  return (schema => schema) as <S extends ZodType<any, any>>(
     schema: keyof T extends keyof S['_output']
       ? keyof S['_output'] extends keyof T
         ? S

package/src/schemas/utils.ts

@@ -1,16 +1,4 @@
-import {
-  OK,
-  type ParseInput,
-  type ParseReturnType,
-  ZodFirstPartyTypeKind,
-  type ZodObject,
-  ZodOptional,
-  ZodParsedType,
-  type ZodRawShape,
-  type ZodType,
-  type ZodTypeAny,
-  z,
-} from 'zod';
+import { type ZodObject, type ZodRawShape, type ZodType, type ZodTypeAny, z } from 'zod';
 
 import { pick } from '../collection/object.js';
 import { isHex, withoutHexPrefix } from '../string/index.js';
@@ -41,24 +29,7 @@ export const bufferSchema: ZodFor<Buffer> = z.union([
   z.instanceof(Buffer),
 ]);
 
-export class ZodNullableOptional<T extends ZodTypeAny> extends ZodOptional<T> {
-  _isNullableOptional = true;
-
-  override _parse(input: ParseInput): ParseReturnType<this['_output']> {
-    const parsedType = this._getType(input);
-    if (parsedType === ZodParsedType.undefined || parsedType === ZodParsedType.null) {
-      return OK(undefined);
-    }
-    return this._def.innerType._parse(input);
-  }
-
-  static override create<T extends ZodTypeAny>(type: T): ZodNullableOptional<T> {
-    return new ZodNullableOptional({
-      innerType: type,
-      typeName: ZodFirstPartyTypeKind.ZodOptional,
-    }) as any;
-  }
-}
+export type ZodNullableOptional<T extends ZodTypeAny> = ZodType<z.output<T> | undefined, z.input<T> | null | undefined>;
 
 /**
  * Declares a parameter as optional. Use this over z.optional in order to accept nulls as undefineds.
@@ -66,7 +37,7 @@ export class ZodNullableOptional<T extends ZodTypeAny> extends ZodOptional<T> {
  * need to convert nulls to undefineds as we parse.
  */
 export function optional<T extends ZodTypeAny>(schema: T) {
-  return ZodNullableOptional.create(schema);
+  return schema.nullish().transform(value => value ?? undefined);
 }
 
 type ToJsonIs<T, TRet> = T extends { toJSON(): TRet } ? T : never;
@@ -83,7 +54,6 @@ export function hexSchemaFor<TClass extends { fromString(str: string): any } | {
   TClass extends { fromString(str: string): infer TInstance } | { fromBuffer(buf: Buffer): infer TInstance }
     ? ToJsonIs<TInstance, string>
     : never,
-  any,
   string
 > {
   const stringSchema = refinement ? z.string().refine(refinement, `Not a valid instance`) : z.string();
@@ -101,11 +71,7 @@ export function hexSchemaFor<TClass extends { fromString(str: string): any } | {
 export function bufferSchemaFor<TClass extends { fromBuffer(buf: Buffer): any }>(
   klazz: TClass,
   refinement?: (buf: Buffer) => boolean,
-): ZodType<
-  TClass extends { fromBuffer(buf: Buffer): infer TInstance } ? ToJsonIs<TInstance, Buffer> : never,
-  any,
-  string
-> {
+): ZodType<TClass extends { fromBuffer(buf: Buffer): infer TInstance } ? ToJsonIs<TInstance, Buffer> : never, string> {
   return bufferSchema.refine(refinement ?? (() => true), 'Not a valid buffer').transform(klazz.fromBuffer.bind(klazz));
 }