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

package/dest/index.js

@@ -1,14 +1,10 @@
-import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
 import { loadContractArtifact } from '@aztec/aztec.js/abi';
 import { AztecAddress } from '@aztec/aztec.js/addresses';
 import { getContractInstanceFromInstantiationParams } 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 { 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 { ContractStore } from '@aztec/pxe/client/lazy';
 import { computeArtifactHash } from '@aztec/stdlib/contract';
 import { zodFor } from '@aztec/stdlib/schemas';
 import { createHash } from 'crypto';
@@ -16,20 +12,55 @@ import { createReadStream } from 'fs';
 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 { TXESession } from './txe_session.js';
-import { ForeignCallArgsSchema, ForeignCallResultSchema, addressFromSingle, fromArray, fromSingle, toSingle } from './util/encoding.js';
+import { ForeignCallArgsSchema, ForeignCallResultSchema, addressFromSingle, fromArray, fromSingle, toSingle } 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 = [];
 const sessions = new Map();
-/*
- * 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.
- */ const TXEArtifactsCache = new Map();
-const TXEArtifactsCacheInFlight = new Map();
-const TXEForeignCallInputSchema = zodFor()(z.object({
+/**
+ * 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.
+ */ class AsyncCache {
+    cache = new Map();
+    inFlight = new Map();
+    getOrCompute(key, compute) {
+        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();
+// 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();
+export const TXEForeignCallInputSchema = zodFor()(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(),
     // eslint-disable-next-line camelcase
     root_path: z.string(),
@@ -37,11 +68,30 @@ const TXEForeignCallInputSchema = zodFor()(z.object({
     package_name: z.string(),
     inputs: ForeignCallArgsSchema
 }));
-class TXEDispatcher {
+export class TXEDispatcher {
     logger;
     contractStore;
-    constructor(logger){
+    contractStoreSourceDir;
+    schnorrClassId;
+    constructor(logger, opts){
         this.logger = logger;
+        this.contractStoreSourceDir = opts.contractStoreSourceDir;
+        this.schnorrClassId = Fr.fromString(opts.schnorrClassId);
+    }
+    /**
+   * 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.
+   */ async warmUp() {
+        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
+        });
     }
     fastHashFile(path) {
         return new Promise((resolve)=>{
@@ -87,88 +137,68 @@ class TXEDispatcher {
         }
         const fileHash = await this.fastHashFile(artifactPath);
         const cacheKey = `${contractDirectory ?? ''}-${contractFilename}-${initializer}-${decodedArgs.map((arg)=>arg.toString()).join('-')}-${publicKeysHash}-${salt}-${deployer}-${fileHash}`;
-        let instance;
-        let artifact;
-        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'));
-                    const artifactWithoutHash = loadContractArtifact(artifactJSON);
-                    const computedArtifact = {
-                        ...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;
+        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'));
+                const artifactWithoutHash = loadContractArtifact(artifactJSON);
+                return {
+                    ...artifactWithoutHash,
+                    artifactHash: await computeArtifactHash(artifactWithoutHash)
                 };
-                TXEArtifactsCacheInFlight.set(cacheKey, compute());
-            }
-            ({ artifact, instance } = await TXEArtifactsCacheInFlight.get(cacheKey));
-        }
+            });
+            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));
     }
     async #processAddAccountInputs({ inputs }) {
         const secret = fromSingle(inputs[0]);
         const cacheKey = `SchnorrAccountContract-${secret}`;
-        let artifact;
-        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 = {
-                        ...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);
     }
     // eslint-disable-next-line camelcase
@@ -177,17 +207,7 @@ class TXEDispatcher {
         this.logger.debug(`Calling ${functionName} on session ${sessionId}`);
         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));
         }
         switch(functionName){
@@ -204,17 +224,36 @@ 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) {
+        const session = sessions.get(sessionId);
+        if (!session) {
+            return;
+        }
+        sessions.delete(sessionId);
+        await session.dispose();
+    }
 }
-const TXEDispatcherApiSchema = {
+/** Diagnostic-only: number of sessions currently held by this worker. */ export function activeSessionCount() {
+    return sessions.size;
+}
+export const TXEDispatcherApiSchema = {
     // 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) {
-    return createSafeJsonRpcServer(new TXEDispatcher(logger), TXEDispatcherApiSchema, {
-        http200OnError: true
-    });
-}