@aztec/foundation 0.0.1-commit.b2a5d0dd1 → 0.0.1-commit.b3d3157a

package/src/json-rpc/fixtures/test_state.ts

@@ -130,14 +130,14 @@ export class TestState implements TestStateApi {
 }
 
 export const TestStateSchema: ApiSchemaFor<TestStateApi> = {
-  getNote: z.function().args(z.number()).returns(TestNote.schema.optional()),
-  getNotes: z.function().args(optional(schemas.Integer)).returns(z.array(TestNote.schema)),
-  getNotes2: z.function().args(optional(schemas.BigInt)).returns(z.array(TestNote.schema)),
-  getNotes3: z.function().args(optional(schemas.Integer)).returns(z.array(TestNote.schema)),
-  clear: z.function().returns(z.void()),
-  addNotes: z.function().args(z.array(TestNote.schema)).returns(z.array(TestNote.schema)),
-  fail: z.function().returns(z.void()),
-  count: z.function().returns(z.number()),
-  getStatus: z.function().returns(z.object({ status: z.string(), count: schemas.BigInt })),
-  getTuple: z.function().returns(z.tuple([z.string(), optional(z.string()), z.number()])),
+  getNote: z.function({ input: z.tuple([z.number()]), output: TestNote.schema.optional() }),
+  getNotes: z.function({ input: z.tuple([optional(schemas.Integer)]), output: z.array(TestNote.schema) }),
+  getNotes2: z.function({ input: z.tuple([optional(schemas.BigInt)]), output: z.array(TestNote.schema) }),
+  getNotes3: z.function({ input: z.tuple([optional(schemas.Integer)]), output: z.array(TestNote.schema) }),
+  clear: z.function({ input: z.tuple([]), output: z.void() }),
+  addNotes: z.function({ input: z.tuple([z.array(TestNote.schema)]), output: z.array(TestNote.schema) }),
+  fail: z.function({ input: z.tuple([]), output: z.void() }),
+  count: z.function({ input: z.tuple([]), output: z.number() }),
+  getStatus: z.function({ input: z.tuple([]), output: z.object({ status: z.string(), count: schemas.BigInt }) }),
+  getTuple: z.function({ input: z.tuple([]), output: z.tuple([z.string(), optional(z.string()), z.number()]) }),
 };

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/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/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));
 }