@aztec/foundation 0.0.1-commit.2eb6648a → 0.0.1-commit.2f68f620

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/bigint-utils.ts

@@ -0,0 +1,25 @@
+/**
+ * Converts bigint values to strings recursively in a log object to avoid serialization issues.
+ */
+export function convertBigintsToStrings(obj: unknown): unknown {
+  if (typeof obj === 'bigint') {
+    return String(obj);
+  }
+
+  if (Array.isArray(obj)) {
+    return obj.map(item => convertBigintsToStrings(item));
+  }
+
+  if (obj !== null && typeof obj === 'object') {
+    if (typeof (obj as any).toJSON === 'function') {
+      return convertBigintsToStrings((obj as any).toJSON());
+    }
+    const result: Record<string, unknown> = {};
+    for (const key in obj) {
+      result[key] = convertBigintsToStrings((obj as Record<string, unknown>)[key]);
+    }
+    return result;
+  }
+
+  return obj;
+}

package/src/log/gcloud-logger-config.ts

@@ -1,5 +1,7 @@
 import type { pino } from 'pino';
 
+import { convertBigintsToStrings } from './bigint-utils.js';
+
 /* eslint-disable camelcase */
 
 const GOOGLE_CLOUD_TRACE_ID = 'logging.googleapis.com/trace';
@@ -15,6 +17,9 @@ export const GoogleCloudLoggerConfig = {
   messageKey: 'message',
   formatters: {
     log(object: Record<string, unknown>): Record<string, unknown> {
+      // Convert bigints to strings recursively to avoid serialization issues
+      object = convertBigintsToStrings(object) as Record<string, unknown>;
+
       // Add trace context attributes following Cloud Logging structured log format described
       // in https://cloud.google.com/logging/docs/structured-logging#special-payload-fields
       const { trace_id, span_id, trace_flags, ...rest } = object;

package/src/log/log-filters.ts

@@ -19,22 +19,40 @@ export function getLogLevelFromFilters(filters: LogFilters, module: string): Log
   return undefined;
 }
 
-export function assertLogLevel(level: string): asserts level is LogLevel {
-  if (!LogLevels.includes(level as LogLevel)) {
-    throw new Error(`Invalid log level: ${level}`);
+/**
+ * Parses the LOG_LEVEL env string into a default level and per-module filter overrides.
+ *
+ * Format: `<default_level>;<level>:<module1>,<module2>;<level>:<module3>;...`
+ * - First segment (before the first `;`) is the default log level for all modules.
+ * - Remaining segments are `level:module` pairs: apply the given level to the listed modules (comma-separated).
+ * - Later filters override earlier ones for overlapping module matches.
+ * - The `aztec:` prefix is stripped from module names; spaces are trimmed.
+ *
+ * @example
+ * ```ts
+ * parseLogLevel('debug;warn:module1,module2;error:module3', 'info')
+ * // => ['debug', [['module3', 'error'], ['module2', 'warn'], ['module1', 'warn']]]
+ * ```
+ */
+export function parseLogLevelEnvVar(
+  logLevelEnvVar: string | undefined,
+  defaultLevel: LogLevel,
+): [LogLevel, LogFilters] {
+  if (!logLevelEnvVar) {
+    return [defaultLevel, []];
   }
+  const [level] = logLevelEnvVar.split(';', 1);
+  assertValidLogLevel(level);
+  return [level, parseFilters(logLevelEnvVar.slice(level.length + 1))];
 }
 
-export function parseEnv(env: string | undefined, defaultLevel: LogLevel): [LogLevel, LogFilters] {
-  if (!env) {
-    return [defaultLevel, []];
+function assertValidLogLevel(level: string): asserts level is LogLevel {
+  if (!LogLevels.includes(level as LogLevel)) {
+    throw new Error(`Invalid log level: ${level}`);
   }
-  const [level] = env.split(';', 1);
-  assertLogLevel(level);
-  return [level, parseFilters(env.slice(level.length + 1))];
 }
 
-export function parseFilters(definition: string | undefined): LogFilters {
+function parseFilters(definition: string | undefined): LogFilters {
   if (!definition) {
     return [];
   }
@@ -48,7 +66,7 @@ export function parseFilters(definition: string | undefined): LogFilters {
       throw new Error(`Invalid log filter statement: ${statement}`);
     }
     const sanitizedLevel = level.trim().toLowerCase();
-    assertLogLevel(sanitizedLevel);
+    assertValidLogLevel(sanitizedLevel);
     for (const module of modules.split(',')) {
       filters.push([
         module

package/src/log/pino-logger.ts

@@ -7,8 +7,9 @@ import { inspect } from 'util';
 import { compactArray } from '../collection/array.js';
 import type { EnvVar } from '../config/index.js';
 import { parseBooleanEnv } from '../config/parse-env.js';
+import { convertBigintsToStrings } from './bigint-utils.js';
 import { GoogleCloudLoggerConfig } from './gcloud-logger-config.js';
-import { getLogLevelFromFilters, parseEnv } from './log-filters.js';
+import { getLogLevelFromFilters, parseLogLevelEnvVar } from './log-filters.js';
 import type { LogLevel } from './log-levels.js';
 import type { LogData, LogFn } from './log_fn.js';
 
@@ -20,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;
@@ -47,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;
@@ -126,7 +129,7 @@ function isLevelEnabled(logger: pino.Logger<'verbose', boolean>, level: LogLevel
 
 // Load log levels from environment variables.
 const defaultLogLevel = process.env.NODE_ENV === 'test' ? 'silent' : 'info';
-export const [logLevel, logFilters] = parseEnv(process.env.LOG_LEVEL, defaultLogLevel);
+export const [logLevel, logFilters] = parseLogLevelEnvVar(process.env.LOG_LEVEL, defaultLogLevel);
 
 // Define custom logging levels for pino.
 const customLevels = { verbose: 25 };
@@ -165,6 +168,9 @@ const pinoOpts: pino.LoggerOptions<keyof typeof customLevels> = {
       ...redactedPaths.map(p => `opts.${p}`),
     ],
   },
+  formatters: {
+    log: obj => convertBigintsToStrings(obj) as Record<string, unknown>,
+  },
   ...(useGcloudLogging ? GoogleCloudLoggerConfig : {}),
 };
 
@@ -371,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/queue/base_memory_queue.ts

@@ -122,7 +122,7 @@ export abstract class BaseMemoryQueue<T> {
    * @param handler - A function that takes an item of type T and returns a Promise<void> after processing the item.
    * @returns A Promise<void> that resolves when the queue is finished processing.
    */
-  public async process(handler: (item: T) => Promise<void>) {
+  public async process(handler: (item: T) => Promise<void> | void) {
     try {
       while (true) {
         const item = await this.get();

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

@@ -104,6 +104,36 @@ export async function retryUntil<T>(
   }
 }
 
+/**
+ * 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;
 };
 
@@ -37,11 +24,24 @@ export type ApiSchema = {
 };
 
 /** Return whether an API schema defines a valid function schema for a given method name. */
-export function schemaHasMethod(schema: ApiSchema, methodName: string) {
+export function schemaHasMethod<T extends ApiSchema>(
+  schema: T,
+  methodName: string,
+): methodName is Extract<keyof T, string> {
   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));
 }
 

package/src/serialize/buffer_reader.ts

@@ -1,3 +1,4 @@
+import { toBigIntBE } from '../bigint-buffer/index.js';
 import type { Tuple } from './types.js';
 
 /**
@@ -130,6 +131,20 @@ export class BufferReader {
     return result;
   }
 
+  /**
+   * Reads a 256-bit signed integer (two's complement) from the buffer at the current index position.
+   * Updates the index position by 32 bytes after reading the number.
+   *
+   * @returns The read 256 bit signed value as a bigint.
+   */
+  public readInt256(): bigint {
+    this.#rangeCheck(32);
+    const unsigned = toBigIntBE(this.buffer.subarray(this.index, this.index + 32));
+    this.index += 32;
+    const signBit = 1n << 255n;
+    return unsigned >= signBit ? unsigned - (1n << 256n) : unsigned;
+  }
+
   /** Alias for readUInt256 */
   public readBigInt(): bigint {
     return this.readUInt256();