@aztec/txe 0.0.1-commit.04852196a → 0.0.1-commit.04d373f

package/src/index.ts

@@ -1,4 +1,3 @@
-import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
 import { type NoirCompiledContract, loadContractArtifact } from '@aztec/aztec.js/abi';
 import { AztecAddress } from '@aztec/aztec.js/addresses';
 import {
@@ -7,12 +6,10 @@ import {
 } from '@aztec/aztec.js/contracts';
 import { Fr } from '@aztec/aztec.js/fields';
 import { PublicKeys, deriveKeys } from '@aztec/aztec.js/keys';
-import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 import type { Logger } from '@aztec/foundation/log';
-import { openTmpStore } from '@aztec/kv-store/lmdb-v2';
-import { protocolContractNames } from '@aztec/protocol-contracts';
-import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import { ContractStore } from '@aztec/pxe/server';
+import { cloneEphemeralStoreFrom } from '@aztec/kv-store/lmdb-v2';
+import type { ProtocolContractName } from '@aztec/protocol-contracts';
+import { ContractStore } from '@aztec/pxe/client/lazy';
 import { computeArtifactHash } from '@aztec/stdlib/contract';
 import type { ContractArtifactWithHash } from '@aztec/stdlib/contract';
 import type { ApiSchemaFor } from '@aztec/stdlib/schemas';
@@ -24,6 +21,9 @@ import { readFile, readdir } from 'fs/promises';
 import { join, parse } from 'path';
 import { z } from 'zod';
 
+// Side-effect import: registers the msgpackr Fr extension for the bundled `Fr` class. Must
+// be loaded before any `sendMessage` call. See msgpackr_fr_extension.ts for the why.
+import './msgpackr_fr_extension.js';
 import { type TXEOracleFunctionName, TXESession } from './txe_session.js';
 import {
   type ForeignCallArgs,
@@ -36,27 +36,57 @@ import {
   fromArray,
   fromSingle,
   toSingle,
-} from './util/encoding.js';
+} from './utils/encoding.js';
+
+// Protocol contracts TXE registers in its contract store. Only AuthRegistry is needed for the
+// current test suites; add a contract here if a lookup against a `0x000…00X` address fails.
+export const TXE_REQUIRED_PROTOCOL_CONTRACTS: ProtocolContractName[] = [];
 
 const sessions = new Map<number, TXESession>();
 
-/*
- * TXE typically has to load the same contract artifacts over and over again for multiple tests,
- * so we cache them here to avoid loading from disk repeatedly.
- *
- * The in-flight map coalesces concurrent requests for the same cache key so that
- * computeArtifactHash (very expensive) is only run once even under parallelism.
+/**
+ * Cache + in-flight map pair. Lookup hits the cache, then awaits an in-flight `compute()` if one
+ * exists, otherwise starts one and stores it. Guarantees `compute()` runs at most once per `key`
+ * across concurrent callers, which matters because `computeArtifactHash` is expensive.
  */
-const TXEArtifactsCache = new Map<
+class AsyncCache<K, V> {
+  private readonly cache = new Map<K, V>();
+  private readonly inFlight = new Map<K, Promise<V>>();
+
+  getOrCompute(key: K, compute: () => Promise<V>): Promise<V> {
+    const cached = this.cache.get(key);
+    if (cached !== undefined) {
+      return Promise.resolve(cached);
+    }
+    let pending = this.inFlight.get(key);
+    if (!pending) {
+      pending = (async () => {
+        try {
+          const value = await compute();
+          this.cache.set(key, value);
+          return value;
+        } finally {
+          this.inFlight.delete(key);
+        }
+      })();
+      this.inFlight.set(key, pending);
+    }
+    return pending;
+  }
+}
+
+// Full deploys (artifact + computed instance), keyed by the full deploy context (contract +
+// constructor args + publicKeys + salt + deployer). Hits on repeated identical deploys.
+const TXEDeploymentsCache = new AsyncCache<
   string,
   { artifact: ContractArtifactWithHash; instance: ContractInstanceWithAddress }
 >();
-const TXEArtifactsCacheInFlight = new Map<
-  string,
-  Promise<{ artifact: ContractArtifactWithHash; instance: ContractInstanceWithAddress }>
->();
 
-type TXEForeignCallInput = {
+// Loaded + hashed contract artifact, keyed by compiled-bytecode hash. Hits across deploys of the
+// same contract when constructor args / salt / deployer differ.
+const TXEArtifactsCache = new AsyncCache<string, ContractArtifactWithHash>();
+
+export type TXEForeignCallInput = {
   session_id: number;
   function: TXEOracleFunctionName;
   root_path: string;
@@ -64,10 +94,13 @@ type TXEForeignCallInput = {
   inputs: ForeignCallArgs;
 };
 
-const TXEForeignCallInputSchema = zodFor<TXEForeignCallInput>()(
+export const TXEForeignCallInputSchema = zodFor<TXEForeignCallInput>()(
   z.object({
+    // Nargo generates session_id as a u64, which may exceed Number.MAX_SAFE_INTEGER.
+    // Zod 4's `.int()` enforces the safe-integer bound, so we drop it here and only require
+    // the value to be a non-negative number (it is used solely as a Map key).
     // eslint-disable-next-line camelcase
-    session_id: z.number().int().nonnegative(),
+    session_id: z.number().nonnegative(),
     function: z.string() as z.ZodType<TXEOracleFunctionName>,
     // eslint-disable-next-line camelcase
     root_path: z.string(),
@@ -77,12 +110,56 @@ const TXEForeignCallInputSchema = zodFor<TXEForeignCallInput>()(
   }),
 );
 
-class TXEDispatcher {
+export interface TXEDispatcherOptions {
+  /**
+   * Path to an LMDB directory holding the required protocol contracts (see
+   * {@link TXE_REQUIRED_PROTOCOL_CONTRACTS}) and the SchnorrAccount artifact. When set, the
+   * dispatcher clones this directory into a fresh tmpdir on first use instead of registering
+   * the contracts itself.
+   */
+  contractStoreSourceDir: string;
+  /**
+   * Class id (hex) of the SchnorrAccount artifact pre-registered in the shared LMDB. When set,
+   * `#processAddAccountInputs` looks the artifact up from the cloned store instead of
+   * recomputing it via `getSchnorrAccountContractArtifact()` + `computeArtifactHash()`.
+   */
+  schnorrClassId: string;
+}
+
+export class TXEDispatcher {
   private contractStore!: ContractStore;
+  private readonly contractStoreSourceDir: string;
+  private readonly schnorrClassId: Fr;
+
+  constructor(
+    private logger: Logger,
+    opts: TXEDispatcherOptions,
+  ) {
+    this.contractStoreSourceDir = opts.contractStoreSourceDir;
+    this.schnorrClassId = Fr.fromString(opts.schnorrClassId);
+  }
 
-  constructor(private logger: Logger) {}
+  /**
+   * Clones the pre-populated LMDB at `contractStoreSourceDir` into a fresh per-instance tmpdir
+   * on first use, so this dispatcher has a writable store already containing the required
+   * protocol contracts + SchnorrAccount. Idempotent — subsequent calls are no-ops.
+   */
+  private async warmUp(): Promise<void> {
+    if (this.contractStore) {
+      return;
+    }
+    const t0 = Date.now();
+    const kvStore = await cloneEphemeralStoreFrom(
+      join(this.contractStoreSourceDir, 'data.mdb'),
+      'txe-contracts',
+      undefined,
+      2,
+    );
+    this.contractStore = new ContractStore(kvStore);
+    this.logger.debug('Cloned shared protocol-contracts store', { totalMs: Date.now() - t0 });
+  }
 
-  private fastHashFile(path: string) {
+  private fastHashFile(path: string): Promise<string> {
     return new Promise(resolve => {
       const fd = createReadStream(path);
       const hash = createHash('sha1');
@@ -90,7 +167,7 @@ class TXEDispatcher {
 
       fd.on('end', function () {
         hash.end();
-        resolve(hash.read());
+        resolve(hash.read() as string);
       });
 
       fd.pipe(hash);
@@ -140,44 +217,29 @@ class TXEDispatcher {
       .map(arg => arg.toString())
       .join('-')}-${publicKeysHash}-${salt}-${deployer}-${fileHash}`;
 
-    let instance;
-    let artifact: ContractArtifactWithHash;
-
-    if (TXEArtifactsCache.has(cacheKey)) {
-      this.logger.debug(`Using cached artifact for ${cacheKey}`);
-      ({ artifact, instance } = TXEArtifactsCache.get(cacheKey)!);
-    } else {
-      if (!TXEArtifactsCacheInFlight.has(cacheKey)) {
-        this.logger.debug(`Loading compiled artifact ${artifactPath}`);
-        const compute = async () => {
-          const artifactJSON = JSON.parse(await readFile(artifactPath, 'utf-8')) as NoirCompiledContract;
-          const artifactWithoutHash = loadContractArtifact(artifactJSON);
-          const computedArtifact: ContractArtifactWithHash = {
-            ...artifactWithoutHash,
-            // Artifact hash is *very* expensive to compute, so we do it here once
-            // and the TXE contract data provider can cache it
-            artifactHash: await computeArtifactHash(artifactWithoutHash),
-          };
-          this.logger.debug(
-            `Deploy ${computedArtifact.name} with initializer ${initializer}(${decodedArgs}) and public keys hash ${publicKeysHash.toString()}`,
-          );
-          const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
-            constructorArgs: decodedArgs,
-            skipArgsDecoding: true,
-            salt,
-            publicKeys,
-            constructorArtifact: initializer ? initializer : undefined,
-            deployer,
-          });
-          const result = { artifact: computedArtifact, instance: computedInstance };
-          TXEArtifactsCache.set(cacheKey, result);
-          TXEArtifactsCacheInFlight.delete(cacheKey);
-          return result;
-        };
-        TXEArtifactsCacheInFlight.set(cacheKey, compute());
-      }
-      ({ artifact, instance } = await TXEArtifactsCacheInFlight.get(cacheKey)!);
-    }
+    const { artifact, instance } = await TXEDeploymentsCache.getOrCompute(cacheKey, async () => {
+      this.logger.debug(`Loading compiled artifact ${artifactPath}`);
+      // Inner cache: artifact load + hash depends only on the compiled bytecode (`fileHash`), so
+      // subsequent deploys of the same contract — regardless of constructor args / deployer /
+      // salt — reuse the same `ContractArtifactWithHash`.
+      const computedArtifact = await TXEArtifactsCache.getOrCompute(fileHash, async () => {
+        const artifactJSON = JSON.parse(await readFile(artifactPath, 'utf-8')) as NoirCompiledContract;
+        const artifactWithoutHash = loadContractArtifact(artifactJSON);
+        return { ...artifactWithoutHash, artifactHash: await computeArtifactHash(artifactWithoutHash) };
+      });
+      this.logger.debug(
+        `Deploy ${computedArtifact.name} with initializer ${initializer}(${decodedArgs}) and public keys hash ${publicKeysHash.toString()}`,
+      );
+      const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
+        constructorArgs: decodedArgs,
+        skipArgsDecoding: true,
+        salt,
+        publicKeys,
+        constructorArtifact: initializer ? initializer : undefined,
+        deployer,
+      });
+      return { artifact: computedArtifact, instance: computedInstance };
+    });
 
     inputs.splice(0, 1, artifact, instance, toSingle(secret));
   }
@@ -187,43 +249,29 @@ class TXEDispatcher {
 
     const cacheKey = `SchnorrAccountContract-${secret}`;
 
-    let artifact: ContractArtifactWithHash;
-    let instance;
-
-    if (TXEArtifactsCache.has(cacheKey)) {
-      this.logger.debug(`Using cached artifact for ${cacheKey}`);
-      ({ artifact, instance } = TXEArtifactsCache.get(cacheKey)!);
-    } else {
-      if (!TXEArtifactsCacheInFlight.has(cacheKey)) {
-        const compute = async () => {
-          const keys = await deriveKeys(secret);
-          const args = [
-            keys.publicKeys.masterIncomingViewingPublicKey.x,
-            keys.publicKeys.masterIncomingViewingPublicKey.y,
-          ];
-          const computedArtifact: ContractArtifactWithHash = {
-            ...SchnorrAccountContractArtifact,
-            // Artifact hash is *very* expensive to compute, so we do it here once
-            // and the TXE contract data provider can cache it
-            artifactHash: await computeArtifactHash(SchnorrAccountContractArtifact),
-          };
-          const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
-            constructorArgs: args,
-            skipArgsDecoding: true,
-            salt: Fr.ONE,
-            publicKeys: keys.publicKeys,
-            constructorArtifact: 'constructor',
-            deployer: AztecAddress.ZERO,
-          });
-          const result = { artifact: computedArtifact, instance: computedInstance };
-          TXEArtifactsCache.set(cacheKey, result);
-          TXEArtifactsCacheInFlight.delete(cacheKey);
-          return result;
-        };
-        TXEArtifactsCacheInFlight.set(cacheKey, compute());
+    const { artifact, instance } = await TXEDeploymentsCache.getOrCompute(cacheKey, async () => {
+      const [artifactFromStore, classWithPreimage] = await Promise.all([
+        this.contractStore.getContractArtifact(this.schnorrClassId),
+        this.contractStore.getContractClassWithPreimage(this.schnorrClassId),
+      ]);
+      if (!artifactFromStore || !classWithPreimage) {
+        throw new Error(
+          `SchnorrAccount not found in shared contract store at class id ${this.schnorrClassId.toString()}`,
+        );
       }
-      ({ artifact, instance } = await TXEArtifactsCacheInFlight.get(cacheKey)!);
-    }
+      const computedArtifact = { ...artifactFromStore, artifactHash: classWithPreimage.artifactHash };
+      const keys = await deriveKeys(secret);
+      const args = [keys.publicKeys.ivpkM.x, keys.publicKeys.ivpkM.y];
+      const computedInstance = await getContractInstanceFromInstantiationParams(computedArtifact, {
+        constructorArgs: args,
+        skipArgsDecoding: true,
+        salt: Fr.ONE,
+        publicKeys: keys.publicKeys,
+        constructorArtifact: 'constructor',
+        deployer: AztecAddress.ZERO,
+      });
+      return { artifact: computedArtifact, instance: computedInstance };
+    });
 
     inputs.splice(0, 0, artifact, instance);
   }
@@ -235,17 +283,7 @@ class TXEDispatcher {
 
     if (!sessions.has(sessionId)) {
       this.logger.debug(`Creating new session ${sessionId}`);
-      if (!this.contractStore) {
-        const kvStore = await openTmpStore('txe-contracts');
-        this.contractStore = new ContractStore(kvStore);
-        const provider = new BundledProtocolContractsProvider();
-        for (const name of protocolContractNames) {
-          const { instance, artifact } = await provider.getProtocolContractArtifact(name);
-          await this.contractStore.addContractArtifact(artifact);
-          await this.contractStore.addContractInstance(instance);
-        }
-        this.logger.debug('Registered protocol contracts in shared contract store');
-      }
+      await this.warmUp();
       sessions.set(sessionId, await TXESession.init(this.contractStore));
     }
 
@@ -262,20 +300,31 @@ class TXEDispatcher {
 
     return await sessions.get(sessionId)!.processFunction(functionName, inputs);
   }
+
+  /**
+   * Releases a session and its resources (per-session LMDB + `NativeWorldStateService`).
+   * Called by the dispatcher pool when nargo closes its TCP connection for a test (see
+   * `rpc_server.ts`'s socket tracker). No-op if the session was never created — that happens
+   * when nargo opens a connection but errors before sending a request.
+   */
+  async disposeSession(sessionId: number): Promise<void> {
+    const session = sessions.get(sessionId);
+    if (!session) {
+      return;
+    }
+    sessions.delete(sessionId);
+    await session.dispose();
+  }
+}
+
+/** Diagnostic-only: number of sessions currently held by this worker. */
+export function activeSessionCount(): number {
+  return sessions.size;
 }
 
-const TXEDispatcherApiSchema: ApiSchemaFor<TXEDispatcher> = {
+export const TXEDispatcherApiSchema: ApiSchemaFor<TXEDispatcher> = {
   // eslint-disable-next-line camelcase
-  resolve_foreign_call: z.function().args(TXEForeignCallInputSchema).returns(ForeignCallResultSchema),
+  resolve_foreign_call: z.function({ input: z.tuple([TXEForeignCallInputSchema]), output: ForeignCallResultSchema }),
+  // disposeSession is invoked over IPC from the worker, not via RPC; required by ApiSchemaFor.
+  disposeSession: z.function({ input: z.tuple([z.number().nonnegative()]), output: z.void() }),
 };
-
-/**
- * Creates an RPC server that forwards calls to the TXE.
- * @param logger - Logger to output to
- * @returns A TXE RPC server.
- */
-export function createTXERpcServer(logger: Logger) {
-  return createSafeJsonRpcServer(new TXEDispatcher(logger), TXEDispatcherApiSchema, {
-    http200OnError: true,
-  });
-}

package/src/msgpackr_fr_extension.ts

@@ -0,0 +1,23 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+
+import { addExtension } from 'msgpackr';
+
+// `@aztec/native/msgpack_channel` registers a msgpackr extension that packs `Fr` instances to
+// their raw 32-byte buffer for the C++ world-state side to deserialize. Because `@aztec/native`
+// is externalized in TXE's esbuild config (it loads a `.node` binary that can't be bundled),
+// the `Fr` class the extension is keyed on belongs to the *external* `@aztec/foundation`
+// module instance — a different class identity than the `Fr` bundled into the TXE worker /
+// bin. msgpackr's extension match uses `instanceof`, so bundled `Fr` instances slip through
+// and msgpackr falls back to `Fr.toJSON()` (returns `Fr.toString()` — a `0x...` hex string).
+// The C++ side then sees a string where a 32-byte binary was expected and throws
+// `msgpack::type_error` whose `what()` is literally `"std::bad_cast"`.
+//
+// Fix: register the same extension under the BUNDLED `Fr` class identity. msgpackr's extension
+// table holds an entry per `Class`, so both Fr classes now resolve to the same `write` callback
+// and produce the same wire format. Imported as a side-effect from `./index.js` so both the
+// worker (via `./worker.ts`) and the main-thread RPC server (via `./rpc_server.ts`) get the
+// registration before any `sendMessage` runs.
+addExtension({
+  Class: Fr,
+  write: (fr: Fr) => fr.toBuffer(),
+});

package/src/oracle/interfaces.ts

@@ -6,6 +6,8 @@ import { BlockNumber } from '@aztec/foundation/branded-types';
 import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EventSelector, FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { GasSettings } from '@aztec/stdlib/gas';
+import type { PrivateLog } from '@aztec/stdlib/logs';
 import type { UInt64 } from '@aztec/stdlib/types';
 
 // These interfaces complement the ones defined in PXE, and combined with those contain the full list of oracles used by
@@ -36,6 +38,11 @@ export interface IAvmExecutionOracle {
   nullifierExists(siloedNullifier: Fr): Promise<boolean>;
   storageWrite(slot: Fr, value: Fr): Promise<void>;
   storageRead(slot: Fr, contractAddress: AztecAddress): Promise<Fr>;
+  returndataSize(): Promise<Fr>;
+  returndataCopy(rdOffset: number, copySize: number): Promise<Fr[]>;
+  call(l2Gas: number, daGas: number, address: AztecAddress, argsLength: number, args: Fr[]): Promise<Fr[]>;
+  staticCall(l2Gas: number, daGas: number, address: AztecAddress, argsLength: number, args: Fr[]): Promise<Fr[]>;
+  successCopy(): Promise<Fr>;
 }
 
 /**
@@ -58,28 +65,34 @@ export interface ITxeExecutionOracle {
     txHash: TxHash;
     noteHashes: Fr[];
     nullifiers: Fr[];
+    privateLogs: PrivateLog[];
   }>;
   getPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress): Promise<Fr[][]>;
   privateCallNewFlow(
-    from: AztecAddress,
+    from: AztecAddress | undefined,
     targetContractAddress: AztecAddress,
     functionSelector: FunctionSelector,
     args: Fr[],
     argsHash: Fr,
     isStaticCall: boolean,
+    additionalScopes: AztecAddress[],
     jobId: string,
-  ): Promise<Fr[]>;
+    authorizedUtilityCallTargets: AztecAddress[],
+    gasSettings: GasSettings,
+  ): Promise<{ returnValues: Fr[]; offchainEffects: Fr[][] }>;
   executeUtilityFunction(
     targetContractAddress: AztecAddress,
     functionSelector: FunctionSelector,
     args: Fr[],
     jobId: string,
+    authorizedUtilityCallTargets: AztecAddress[],
   ): Promise<Fr[]>;
   publicCallNewFlow(
-    from: AztecAddress,
+    from: AztecAddress | undefined,
     targetContractAddress: AztecAddress,
     calldata: Fr[],
     isStaticCall: boolean,
+    gasSettings: GasSettings,
   ): Promise<Fr[]>;
   // TODO(F-335): Drop this from here as it's not a real oracle handler - it's only called from
   // RPCTranslator::txeGetPrivateEvents and never from Noir.

package/src/oracle/txe_oracle_public_context.ts

@@ -122,6 +122,36 @@ export class TXEOraclePublicContext implements IAvmExecutionOracle {
     return value;
   }
 
+  returndataSize(): Promise<Fr> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  returndataCopy(_rdOffset: number, _copySize: number): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  call(_l2Gas: number, _daGas: number, _address: AztecAddress, _argsLength: number, _args: Fr[]): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  staticCall(_l2Gas: number, _daGas: number, _address: AztecAddress, _argsLength: number, _args: Fr[]): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  successCopy(): Promise<Fr> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
   async close(): Promise<L2Block> {
     this.logger.debug('Exiting Public Context, building block with collected side effects', {
       blockNumber: this.globalVariables.blockNumber,