feat: expose wallet.registerContractClass (prep for simulating contract upgrades)

yarn-project/aztec.js/src/wallet/capabilities.ts

@@ -138,23 +138,25 @@ export interface ContractsCapability {
 export interface GrantedContractsCapability extends ContractsCapability {}
 
 /**
- * Contract class capability - for querying contract class metadata.
+ * Contract class capability - for querying contract class meatadata and registering contract classes.
  *
  * Maps to wallet methods:
- * - getContractClassMetadata
+ * - getContractClassMetadata (when canGetMetadata: true)
+ * - registerContractClass (when canRegister: true)
  *
  * Contract classes are identified by their class ID (Fr), not by contract address.
  * Multiple contract instances can share the same class. This capability grants
- * permission to query metadata for specific contract classes.
+ * permission to query metadata for, and register, specific contract classes.
  *
  * Apps typically acquire this permission automatically when registering a contract
  * with an artifact (the wallet auto-grants permission for that contract's class ID).
  *
  * @example
- * // Query specific contract classes
+ * // Register and query a specific contract class
  * \{
  *   type: 'contractClasses',
- *   classes: [classId1, classId2],
+ *   classes: [classId1],
+ *   canRegister: true,
  *   canGetMetadata: true
  * \}
  *
@@ -177,6 +179,9 @@ export interface ContractClassesCapability {
    */
   classes: '*' | Fr[];
 
+  /** Can register a contract class artifact in the local PXE. Maps to: registerContractClass */
+  canRegister?: boolean;
+
   /** Can query contract class metadata. Maps to: getContractClassMetadata */
   canGetMetadata: boolean;
 }

yarn-project/aztec.js/src/wallet/wallet.test.ts

@@ -157,6 +157,19 @@ describe('WalletSchema', () => {
     });
   });
 
+  it('registerContractClass', async () => {
+    const mockArtifact: ContractArtifact = {
+      name: 'TestContract',
+      aztecVersion: DEV_VERSION,
+      functions: [],
+      nonDispatchPublicFunctions: [],
+      outputs: { structs: {}, globals: {} },
+      fileMap: {},
+      storageLayout: {},
+    };
+    await context.client.registerContractClass(mockArtifact);
+  });
+
   it('simulateTx', async () => {
     const exec: ExecutionPayload = {
       calls: [],
@@ -451,6 +464,8 @@ class MockWallet implements Wallet {
     };
   }
 
+  async registerContractClass(_artifact: any): Promise<void> {}
+
   async simulateTx(_exec: ExecutionPayload, _opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return TxSimulationResultWithAppOffset.fromResultAndOffset(await TxSimulationResult.random(), 0);
   }

yarn-project/aztec.js/src/wallet/wallet.ts

@@ -272,6 +272,12 @@ export type Wallet = {
     artifact?: ContractArtifact,
     secretKey?: Fr,
   ): Promise<ContractInstanceWithAddress>;
+  /**
+   * Registers a contract class artifact in the local PXE without binding it to any instance.
+   * Useful for simulation flows that need the artifact available locally before any on-chain
+   * upgrade has taken effect. No chain check.
+   */
+  registerContractClass(artifact: ContractArtifact): Promise<void>;
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset>;
   executeUtility(call: FunctionCall, opts: ExecuteUtilityOptions): Promise<UtilityExecutionResult>;
   profileTx(exec: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult>;
@@ -429,6 +435,7 @@ export const GrantedContractsCapabilitySchema = ContractsCapabilitySchema;
 export const ContractClassesCapabilitySchema = z.object({
   type: z.literal('contractClasses'),
   classes: z.union([z.literal('*'), z.array(schemas.Fr)]),
+  canRegister: optional(z.boolean()),
   canGetMetadata: z.boolean(),
 });
 
@@ -559,6 +566,7 @@ const WalletMethodSchemas = {
     .function()
     .args(ContractInstanceWithAddressSchema, optional(ContractArtifactSchema), optional(schemas.Fr))
     .returns(ContractInstanceWithAddressSchema),
+  registerContractClass: z.function().args(ContractArtifactSchema).returns(z.void()),
   simulateTx: z
     .function()
     .args(ExecutionPayloadSchema, SimulateOptionsSchema)

yarn-project/end-to-end/src/test-wallet/worker_wallet.ts

@@ -166,6 +166,10 @@ export class WorkerWallet implements Wallet {
     return this.call('registerContract', instance, artifact, secretKey);
   }
 
+  registerContractClass(artifact: ContractArtifact): Promise<void> {
+    return this.call('registerContractClass', artifact);
+  }
+
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return this.call('simulateTx', exec, opts);
   }

yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts

@@ -348,6 +348,10 @@ export abstract class BaseWallet implements Wallet {
     return instance;
   }
 
+  registerContractClass(artifact: ContractArtifact): Promise<void> {
+    return this.pxe.registerContractClass(artifact);
+  }
+
   /**
    * Simulates calls through the standard PXE path (account entrypoint).
    * @param executionPayload - The execution payload to simulate.