From aa3d3bda1d107ec6ed593ea7557787cf12a700ad Mon Sep 17 00:00:00 2001 From: Thunkar <5404052+Thunkar@users.noreply.github.com> Date: Sun, 18 Jan 2026 18:11:55 +0000 Subject: [PATCH] feat: wallet full batchability Allows ALL methods in the wallet to be batched (because why not). Initially I was going for allowing them on a case-by-case basis, but integrating the Wallet SDK has made apparent apps can have many interaction flows and this extra flexibility is nice. Updated the migration notes and in the process realized I missed exporting some artifacts in https://github.com/AztecProtocol/aztec-packages/pull/19457, so here they are (protocol contract wrappers) Co-authored-by: thunkar --- .../docs/resources/migration_notes.md | 79 ++++++++++++ yarn-project/aztec.js/src/api/protocol.ts | 7 ++ .../aztec.js/src/contract/batch_call.ts | 4 +- .../aztec.js/src/wallet/wallet.test.ts | 60 +++++++-- yarn-project/aztec.js/src/wallet/wallet.ts | 118 ++++++++++-------- .../wallet-sdk/src/base-wallet/base_wallet.ts | 5 +- 6 files changed, 202 insertions(+), 71 deletions(-) diff --git a/docs/docs-developers/docs/resources/migration_notes.md b/docs/docs-developers/docs/resources/migration_notes.md index 59c0a8c93b90..8da517cef429 100644 --- a/docs/docs-developers/docs/resources/migration_notes.md +++ b/docs/docs-developers/docs/resources/migration_notes.md @@ -9,6 +9,85 @@ Aztec is in full-speed development. Literally every version breaks compatibility ## TBD +### [Aztec.js] Wallet batching now supports all methods + +The `BatchedMethod` type is now a discriminated union that ensures type safety: the `args` must match the specific method `name`. This prevents runtime errors from mismatched arguments. + +```diff +- // Before: Only 5 methods could be batched +- const results = await wallet.batch([ +- { name: "registerSender", args: [address, "alias"] }, +- { name: "sendTx", args: [payload, options] }, +- ]); ++ // After: All methods can be batched ++ const results = await wallet.batch([ ++ { name: "getChainInfo", args: [] }, ++ { name: "getContractMetadata", args: [contractAddress] }, ++ { name: "registerSender", args: [address, "alias"] }, ++ { name: "simulateTx", args: [payload, options] }, ++ { name: "sendTx", args: [payload, options] }, ++ ]); +``` + +### [Aztec.js] Refactored `getContractMetadata` and `getContractClassMetadata` in Wallet + +The contract metadata methods in the `Wallet` interface have been refactored to provide more granular information and avoid expensive round-trips. + +**`ContractMetadata`:** + +```diff + { +- contractInstance?: ContractInstanceWithAddress, ++ instance?: ContractInstanceWithAddress; // Instance registered in the Wallet, if any + isContractInitialized: boolean; // Is the init nullifier onchain? (already there) + isContractPublished: boolean; // Has the contract been published? (already there) ++ isContractUpdated: boolean; // Has the contract been updated? ++ updatedContractClassId?: Fr; // If updated, the new class ID + } +``` + +**`ContractClassMetadata`:** + +This method loses the ability to request the contract artifact via the `includeArtifact` flag + +```diff + { +- contractClass?: ContractClassWithId; +- artifact?: ContractArtifact; + isContractClassPubliclyRegistered: boolean; // Is the class registered onchain? ++ isArtifactRegistered: boolean; // Does the Wallet know about this artifact? + } +``` + +- Removes expensive artifact/class transfers between wallet and app +- Separates PXE storage info (`instance`, `isArtifactRegistered`) from public chain info (`isContractPublished`, `isContractClassPubliclyRegistered`) +- Makes it easier to determine if actions like `registerContract` are needed + +### [Aztec.js] Removed `UnsafeContract` and protocol contract helper functions + +The `UnsafeContract` class and async helper functions (`getFeeJuice`, `getClassRegistryContract`, `getInstanceRegistryContract`) have been removed. Protocol contracts are now accessed via auto-generated type-safe wrappers with only the ABI (no bytecode). Since PXE always has protocol contract artifacts available, importing and using these contracts from `aztec.js` is very lightweight and follows the same pattern as regular user contracts. + +**Migration:** + +```diff +- import { getFeeJuice, getClassRegistryContract, getInstanceRegistryContract } from '@aztec/aztec.js/contracts'; ++ import { FeeJuiceContract, ContractClassRegistryContract, ContractInstanceRegistryContract } from '@aztec/aztec.js/protocol'; + +- const feeJuice = await getFeeJuice(wallet); ++ const feeJuice = FeeJuiceContract.at(wallet); + await feeJuice.methods.check_balance(feeLimit).send().wait(); + +- const classRegistry = await getClassRegistryContract(wallet); ++ const classRegistry = ContractClassRegistryContract.at(wallet); + await classRegistry.methods.publish(...).send().wait(); + +- const instanceRegistry = await getInstanceRegistryContract(wallet); ++ const instanceRegistry = ContractInstanceRegistryContract.at(wallet); + await instanceRegistry.methods.publish_for_public_execution(...).send().wait(); +``` + +**Note:** The higher-level utilities like `publishInstance`, `publishContractClass`, and `broadcastPrivateFunction` from `@aztec/aztec.js/deployment` are still available and unchanged. These utilities use the new wrappers internally. + ### [Aztec.nr] Renamed Router contract `Router` contract has been renamed as `PublicChecks` contract. diff --git a/yarn-project/aztec.js/src/api/protocol.ts b/yarn-project/aztec.js/src/api/protocol.ts index a511f4a20135..62be26a31981 100644 --- a/yarn-project/aztec.js/src/api/protocol.ts +++ b/yarn-project/aztec.js/src/api/protocol.ts @@ -1,2 +1,9 @@ export { ProtocolContractAddress } from '@aztec/protocol-contracts'; export { INITIAL_L2_BLOCK_NUM } from '@aztec/constants'; + +export { AuthRegistryContract } from '../contract/protocol_contracts/auth-registry.js'; +export { ContractClassRegistryContract } from '../contract/protocol_contracts/contract-class-registry.js'; +export { ContractInstanceRegistryContract } from '../contract/protocol_contracts/contract-instance-registry.js'; +export { FeeJuiceContract } from '../contract/protocol_contracts/fee-juice.js'; +export { MultiCallEntrypointContract } from '../contract/protocol_contracts/multi-call-entrypoint.js'; +export { PublicChecksContract } from '../contract/protocol_contracts/public-checks.js'; diff --git a/yarn-project/aztec.js/src/contract/batch_call.ts b/yarn-project/aztec.js/src/contract/batch_call.ts index 552ac96419af..963e1f4a0fc5 100644 --- a/yarn-project/aztec.js/src/contract/batch_call.ts +++ b/yarn-project/aztec.js/src/contract/batch_call.ts @@ -76,13 +76,13 @@ export class BatchCall extends BaseContractInteraction { { indexedExecutionPayloads: [], utility: [], publicIndex: 0, privateIndex: 0 }, ); - const batchRequests: Array | BatchedMethod<'simulateTx'>> = []; + const batchRequests: BatchedMethod[] = []; // Add utility calls to batch for (const [call] of utility) { batchRequests.push({ name: 'simulateUtility' as const, - args: [call, options?.authWitnesses] as const, + args: [call, options?.authWitnesses], }); } diff --git a/yarn-project/aztec.js/src/wallet/wallet.test.ts b/yarn-project/aztec.js/src/wallet/wallet.test.ts index 67dc7f27f836..c4e77dc5a442 100644 --- a/yarn-project/aztec.js/src/wallet/wallet.test.ts +++ b/yarn-project/aztec.js/src/wallet/wallet.test.ts @@ -21,7 +21,6 @@ import { import type { Aliased, BatchResults, - BatchableMethods, BatchedMethod, ContractClassMetadata, ContractMetadata, @@ -235,6 +234,10 @@ describe('WalletSchema', () => { const simulateOpts: SimulateOptions = { from: await AztecAddress.random(), }; + const profileOpts: ProfileOptions = { + from: await AztecAddress.random(), + profileMode: 'gates', + }; const call = { name: 'testFunction', @@ -267,24 +270,57 @@ describe('WalletSchema', () => { storageLayout: {}, }; - const methods: BatchedMethod[] = [ + const eventMetadata: EventMetadataDefinition = { + eventSelector: EventSelector.fromField(new Fr(1)), + abiType: { kind: 'field' }, + fieldNames: ['field1'], + }; + + const methods: BatchedMethod[] = [ + { name: 'getChainInfo', args: [] }, + { name: 'getTxReceipt', args: [TxHash.random()] }, + { name: 'getContractMetadata', args: [address1] }, + { name: 'getContractClassMetadata', args: [Fr.random()] }, + { + name: 'getPrivateEvents', + args: [eventMetadata, { contractAddress: address1, scopes: [address2], fromBlock: BlockNumber(1) }], + }, { name: 'registerSender', args: [address1, 'alias1'] }, + { name: 'getAddressBook', args: [] }, + { name: 'getAccounts', args: [] }, { name: 'registerContract', args: [mockInstance, mockArtifact, undefined] }, - { name: 'sendTx', args: [exec, opts] }, - { name: 'simulateUtility', args: [call, [AuthWitness.random()]] }, { name: 'simulateTx', args: [exec, simulateOpts] }, + { name: 'simulateUtility', args: [call, [AuthWitness.random()]] }, + { name: 'profileTx', args: [exec, profileOpts] }, + { name: 'sendTx', args: [exec, opts] }, + { name: 'createAuthWit', args: [address1, Fr.random()] }, ]; const results = await context.client.batch(methods); - expect(results).toHaveLength(5); - expect(results[0]).toEqual({ name: 'registerSender', result: expect.any(AztecAddress) }); - expect(results[1]).toEqual({ + expect(results).toHaveLength(14); + expect(results[0]).toEqual({ name: 'getChainInfo', result: { chainId: expect.any(Fr), version: expect.any(Fr) } }); + expect(results[1]).toEqual({ name: 'getTxReceipt', result: expect.any(TxReceipt) }); + expect(results[2]).toEqual({ + name: 'getContractMetadata', + result: expect.objectContaining({ isContractInitialized: expect.any(Boolean) }), + }); + expect(results[3]).toEqual({ + name: 'getContractClassMetadata', + result: expect.objectContaining({ isArtifactRegistered: expect.any(Boolean) }), + }); + expect(results[4]).toEqual({ name: 'getPrivateEvents', result: expect.any(Array) }); + expect(results[5]).toEqual({ name: 'registerSender', result: expect.any(AztecAddress) }); + expect(results[6]).toEqual({ name: 'getAddressBook', result: expect.any(Array) }); + expect(results[7]).toEqual({ name: 'getAccounts', result: expect.any(Array) }); + expect(results[8]).toEqual({ name: 'registerContract', result: expect.objectContaining({ address: expect.any(AztecAddress) }), }); - expect(results[2]).toEqual({ name: 'sendTx', result: expect.any(TxHash) }); - expect(results[3]).toEqual({ name: 'simulateUtility', result: expect.any(UtilitySimulationResult) }); - expect(results[4]).toEqual({ name: 'simulateTx', result: expect.any(TxSimulationResult) }); + expect(results[9]).toEqual({ name: 'simulateTx', result: expect.any(TxSimulationResult) }); + expect(results[10]).toEqual({ name: 'simulateUtility', result: expect.any(UtilitySimulationResult) }); + expect(results[11]).toEqual({ name: 'profileTx', result: expect.any(TxProfileResult) }); + expect(results[12]).toEqual({ name: 'sendTx', result: expect.any(TxHash) }); + expect(results[13]).toEqual({ name: 'createAuthWit', result: expect.any(AuthWitness) }); }); }); @@ -381,7 +417,7 @@ class MockWallet implements Wallet { return Promise.resolve(AuthWitness.random()); } - async batch[]>(methods: T): Promise> { + async batch(methods: T): Promise> { const results: any[] = []; for (const method of methods) { const { name, args } = method; @@ -390,7 +426,7 @@ class MockWallet implements Wallet { // 2. `args` matches the parameter types of that specific method // 3. The return type is correctly mapped in BatchResults // We use dynamic dispatch here for simplicity, but the types are enforced at the call site. - const fn = this[name] as (...args: any[]) => Promise; + const fn = (this as any)[name] as (...args: any[]) => Promise; const result = await fn.apply(this, args); // Wrap result with method name for discriminated union deserialization results.push({ name, result }); diff --git a/yarn-project/aztec.js/src/wallet/wallet.ts b/yarn-project/aztec.js/src/wallet/wallet.ts index 2aa5d89b6250..dcfbd5efc587 100644 --- a/yarn-project/aztec.js/src/wallet/wallet.ts +++ b/yarn-project/aztec.js/src/wallet/wallet.ts @@ -54,7 +54,7 @@ export type Aliased = { /** * Options for simulating interactions with the wallet. Overrides the fee settings of an interaction with - * a simplified version that only hints at the wallet wether the interaction contains a + * a simplified version that only hints at the wallet whether the interaction contains a * fee payment method or not */ export type SimulateOptions = Omit & { @@ -64,7 +64,7 @@ export type SimulateOptions = Omit & { /** * Options for profiling interactions with the wallet. Overrides the fee settings of an interaction with - * a simplified version that only hints at the wallet wether the interaction contains a + * a simplified version that only hints at the wallet whether the interaction contains a * fee payment method or not */ export type ProfileOptions = Omit & { @@ -74,7 +74,7 @@ export type ProfileOptions = Omit & { /** * Options for sending/proving interactions with the wallet. Overrides the fee settings of an interaction with - * a simplified version that only hints at the wallet wether the interaction contains a + * a simplified version that only hints at the wallet whether the interaction contains a * fee payment method or not */ export type SendOptions = Omit & { @@ -83,36 +83,40 @@ export type SendOptions = Omit & { }; /** - * Helper type that represents all methods that can be batched. + * Helper type that represents all methods that can be batched (all methods except batch itself). */ -export type BatchableMethods = Pick< - Wallet, - 'registerContract' | 'sendTx' | 'registerSender' | 'simulateUtility' | 'simulateTx' ->; +export type BatchableMethods = Omit; /** - * From the batchable methods, we create a type that represents a method call with its name and arguments. - * This is what the wallet will accept as arguments to the `batch` method. + * A method call with its name and arguments. */ -export type BatchedMethod = { +type BatchedMethodInternal = { /** The method name */ name: T; /** The method arguments */ args: Parameters; }; +/** + * Union of all possible batched method calls. + * This ensures type safety: the `args` must match the specific `name`. + */ +export type BatchedMethod = { + [K in keyof BatchableMethods]: BatchedMethodInternal; +}[keyof BatchableMethods]; + /** * Helper type to extract the return type of a batched method */ export type BatchedMethodResult = - T extends BatchedMethod ? Awaited> : never; + T extends BatchedMethodInternal ? Awaited> : never; /** * Wrapper type for batch results that includes the method name for discriminated union deserialization. * Each result is wrapped as \{ name: 'methodName', result: ActualResult \} to allow proper deserialization * when AztecAddress and TxHash would otherwise be ambiguous (both are hex strings). */ -export type BatchedMethodResultWrapper> = { +export type BatchedMethodResultWrapper = { /** The method name */ name: T['name']; /** The method result */ @@ -122,7 +126,7 @@ export type BatchedMethodResultWrapper[]> = { +export type BatchResults = { [K in keyof T]: BatchedMethodResultWrapper; }; @@ -209,7 +213,7 @@ export type Wallet = { profileTx(exec: ExecutionPayload, opts: ProfileOptions): Promise; sendTx(exec: ExecutionPayload, opts: SendOptions): Promise; createAuthWit(from: AztecAddress, messageHashOrIntent: Fr | IntentInnerHash | CallIntent): Promise; - batch[]>(methods: T): Promise>; + batch(methods: T): Promise>; }; export const FunctionCallSchema = z.object({ @@ -278,29 +282,6 @@ export const MessageHashOrIntentSchema = z.union([ }), ]); -export const BatchedMethodSchema = z.union([ - z.object({ - name: z.literal('registerSender'), - args: z.tuple([schemas.AztecAddress, optional(z.string())]), - }), - z.object({ - name: z.literal('registerContract'), - args: z.tuple([ContractInstanceWithAddressSchema, optional(ContractArtifactSchema), optional(schemas.Fr)]), - }), - z.object({ - name: z.literal('sendTx'), - args: z.tuple([ExecutionPayloadSchema, SendOptionsSchema]), - }), - z.object({ - name: z.literal('simulateUtility'), - args: z.tuple([FunctionCallSchema, optional(z.array(AuthWitness.schema))]), - }), - z.object({ - name: z.literal('simulateTx'), - args: z.tuple([ExecutionPayloadSchema, SimulateOptionsSchema]), - }), -]); - export const EventMetadataDefinitionSchema = z.object({ eventSelector: schemas.EventSelector, abiType: AbiTypeSchema, @@ -335,7 +316,11 @@ export const ContractClassMetadataSchema = z.object({ isContractClassPubliclyRegistered: z.boolean(), }); -export const WalletSchema: ApiSchemaFor = { +/** + * Record of all wallet method schemas (excluding batch). + * This is the single source of truth for method schemas - batch schemas are derived from this. + */ +const WalletMethodSchemas = { getChainInfo: z .function() .args() @@ -368,19 +353,46 @@ export const WalletSchema: ApiSchemaFor = { profileTx: z.function().args(ExecutionPayloadSchema, ProfileOptionsSchema).returns(TxProfileResult.schema), sendTx: z.function().args(ExecutionPayloadSchema, SendOptionsSchema).returns(TxHash.schema), createAuthWit: z.function().args(schemas.AztecAddress, MessageHashOrIntentSchema).returns(AuthWitness.schema), +}; + +/** + * Creates batch schemas from the individual wallet methods. + * This allows us to define them once and derive batch schemas automatically, + * reducing duplication and ensuring consistency. + */ +function createBatchSchemas, z.ZodTypeAny>>>( + methodSchemas: T, +) { + const names = Object.keys(methodSchemas) as (keyof T)[]; + + const namesAndArgs = names.map(name => + z.object({ + name: z.literal(name), + args: methodSchemas[name].parameters(), + }), + ); + + const namesAndReturns = names.map(name => + z.object({ + name: z.literal(name), + result: methodSchemas[name].returnType(), + }), + ); + + // Type assertion needed because discriminatedUnion expects a tuple type [T, T, ...T[]] + // but we're building the array dynamically. The runtime behavior is correct. + return { + input: z.discriminatedUnion('name', namesAndArgs as [(typeof namesAndArgs)[0], ...typeof namesAndArgs]), + output: z.discriminatedUnion('name', namesAndReturns as [(typeof namesAndReturns)[0], ...typeof namesAndReturns]), + }; +} + +const { input: BatchedMethodSchema, output: BatchedResultSchema } = createBatchSchemas(WalletMethodSchemas); + +export { BatchedMethodSchema, BatchedResultSchema }; + +export const WalletSchema: ApiSchemaFor = { + ...WalletMethodSchemas, // @ts-expect-error - ApiSchemaFor cannot properly type generic methods with readonly arrays - batch: z - .function() - .args(z.array(BatchedMethodSchema)) - .returns( - z.array( - z.discriminatedUnion('name', [ - z.object({ name: z.literal('registerSender'), result: schemas.AztecAddress }), - z.object({ name: z.literal('registerContract'), result: ContractInstanceWithAddressSchema }), - z.object({ name: z.literal('sendTx'), result: TxHash.schema }), - z.object({ name: z.literal('simulateUtility'), result: UtilitySimulationResult.schema }), - z.object({ name: z.literal('simulateTx'), result: TxSimulationResult.schema }), - ]), - ), - ), + batch: z.function().args(z.array(BatchedMethodSchema)).returns(z.array(BatchedResultSchema)), }; diff --git a/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts b/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts index c7411e26796a..b1593973c16d 100644 --- a/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts +++ b/yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts @@ -4,7 +4,6 @@ import type { FeePaymentMethod } from '@aztec/aztec.js/fee'; import type { Aliased, BatchResults, - BatchableMethods, BatchedMethod, PrivateEvent, PrivateEventFilter, @@ -131,9 +130,7 @@ export abstract class BaseWallet implements Wallet { return account.createAuthWit(messageHashOrIntent); } - public async batch[]>( - methods: T, - ): Promise> { + public async batch(methods: T): Promise> { const results: any[] = []; for (const method of methods) { const { name, args } = method;