@aztec/txe 0.0.0-test.1 → 0.0.1-commit.017a351

package/dest/index.js

@@ -1,156 +1,100 @@
-import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
-import { AztecAddress, Fr, PublicKeys, deriveKeys, getContractInstanceFromDeployParams, loadContractArtifact } from '@aztec/aztec.js';
-import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
-import { protocolContractNames } from '@aztec/protocol-contracts';
-import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import { readFile, readdir } from 'fs/promises';
+import { Fr } from '@aztec/aztec.js/fields';
+import { cloneEphemeralStoreFrom } from '@aztec/kv-store/lmdb-v2';
+import { ContractStore } from '@aztec/pxe/client/lazy';
+import { zodFor } from '@aztec/stdlib/schemas';
 import { join } from 'path';
 import { z } from 'zod';
-import { TXEService } from './txe_service/txe_service.js';
-import { ForeignCallArgsSchema, ForeignCallResultSchema, fromArray, fromSingle, toForeignCallResult, toSingle } from './util/encoding.js';
-const TXESessions = new Map();
-const TXEArtifactsCache = new Map();
-const TXEForeignCallInputSchema = z.object({
+// 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 } from './utils/encoding.js';
+import { TXEArtifactResolver } from './utils/txe_artifact_resolver.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();
+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(),
     // eslint-disable-next-line camelcase
     package_name: z.string(),
     inputs: ForeignCallArgsSchema
-});
-class TXEDispatcher {
+}));
+export class TXEDispatcher {
     logger;
-    protocolContracts;
-    constructor(logger){
+    contractStore;
+    artifactResolver;
+    contractStoreSourceDir;
+    schnorrClassId;
+    constructor(logger, opts){
         this.logger = logger;
+        this.contractStoreSourceDir = opts.contractStoreSourceDir;
+        this.schnorrClassId = Fr.fromString(opts.schnorrClassId);
     }
-    async #processDeployInputs({ inputs, root_path: rootPath, package_name: packageName }) {
-        const [pathStr, contractName, initializer] = inputs.slice(0, 3).map((input)=>fromArray(input).map((char)=>String.fromCharCode(char.toNumber())).join(''));
-        const decodedArgs = fromArray(inputs[4]);
-        const secret = fromSingle(inputs[5]);
-        const publicKeys = secret.equals(Fr.ZERO) ? PublicKeys.default() : (await deriveKeys(secret)).publicKeys;
-        const publicKeysHash = await publicKeys.hash();
-        const cacheKey = `${pathStr}-${contractName}-${initializer}-${decodedArgs.map((arg)=>arg.toString()).join('-')}-${publicKeysHash.toString()}`;
-        let artifact;
-        let instance;
-        if (TXEArtifactsCache.has(cacheKey)) {
-            this.logger.debug(`Using cached artifact for ${cacheKey}`);
-            ({ artifact, instance } = TXEArtifactsCache.get(cacheKey));
-        } else {
-            let artifactPath = '';
-            // We're deploying the contract under test
-            // env.deploy_self("contractName")
-            if (!pathStr) {
-                artifactPath = join(rootPath, './target', `${packageName}-${contractName}.json`);
-            } else {
-                // We're deploying a contract that belongs in a workspace
-                // env.deploy("../path/to/workspace/root@packageName", "contractName")
-                if (pathStr.includes('@')) {
-                    const [workspace, pkg] = pathStr.split('@');
-                    const targetPath = join(rootPath, workspace, './target');
-                    this.logger.debug(`Looking for compiled artifact in workspace ${targetPath}`);
-                    artifactPath = join(targetPath, `${pkg}-${contractName}.json`);
-                } else {
-                    // We're deploying a standalone contract
-                    // env.deploy("../path/to/contract/root", "contractName")
-                    const targetPath = join(rootPath, pathStr, './target');
-                    this.logger.debug(`Looking for compiled artifact in ${targetPath}`);
-                    [artifactPath] = (await readdir(targetPath)).filter((file)=>file.endsWith(`-${contractName}.json`));
-                }
-            }
-            this.logger.debug(`Loading compiled artifact ${artifactPath}`);
-            artifact = loadContractArtifact(JSON.parse(await readFile(artifactPath, 'utf-8')));
-            this.logger.debug(`Deploy ${artifact.name} with initializer ${initializer}(${decodedArgs}) and public keys hash ${publicKeysHash.toString()}`);
-            instance = await getContractInstanceFromDeployParams(artifact, {
-                constructorArgs: decodedArgs,
-                skipArgsDecoding: true,
-                salt: Fr.ONE,
-                publicKeys,
-                constructorArtifact: initializer ? initializer : undefined,
-                deployer: AztecAddress.ZERO
-            });
-            TXEArtifactsCache.set(cacheKey, {
-                artifact,
-                instance
-            });
+    /**
+   * 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;
         }
-        inputs.splice(0, 2, 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 {
-            const keys = await deriveKeys(secret);
-            const args = [
-                keys.publicKeys.masterIncomingViewingPublicKey.x,
-                keys.publicKeys.masterIncomingViewingPublicKey.y
-            ];
-            artifact = SchnorrAccountContractArtifact;
-            instance = await getContractInstanceFromDeployParams(artifact, {
-                constructorArgs: args,
-                skipArgsDecoding: true,
-                salt: Fr.ONE,
-                publicKeys: keys.publicKeys,
-                constructorArtifact: 'constructor',
-                deployer: AztecAddress.ZERO
-            });
-            TXEArtifactsCache.set(cacheKey, {
-                artifact,
-                instance
-            });
-        }
-        inputs.splice(0, 0, artifact, instance);
+        const t0 = Date.now();
+        const kvStore = await cloneEphemeralStoreFrom(join(this.contractStoreSourceDir, 'data.mdb'), 'txe-contracts', undefined, 2);
+        this.contractStore = new ContractStore(kvStore);
+        this.artifactResolver = new TXEArtifactResolver(this.contractStore, this.schnorrClassId);
+        this.logger.debug('Cloned shared protocol-contracts store', {
+            totalMs: Date.now() - t0
+        });
     }
     // eslint-disable-next-line camelcase
     async resolve_foreign_call(callData) {
-        const { session_id: sessionId, function: functionName, inputs } = callData;
+        const { session_id: sessionId, function: functionName, inputs, root_path: rootPath, package_name: packageName } = callData;
         this.logger.debug(`Calling ${functionName} on session ${sessionId}`);
-        if (!TXESessions.has(sessionId) && functionName != 'reset') {
+        if (!sessions.has(sessionId)) {
             this.logger.debug(`Creating new session ${sessionId}`);
-            if (!this.protocolContracts) {
-                this.protocolContracts = await Promise.all(protocolContractNames.map((name)=>new BundledProtocolContractsProvider().getProtocolContractArtifact(name)));
-            }
-            TXESessions.set(sessionId, await TXEService.init(this.logger, this.protocolContracts));
+            await this.warmUp();
+            sessions.set(sessionId, await TXESession.init(this.contractStore, this.artifactResolver, rootPath, packageName));
         }
-        switch(functionName){
-            case 'reset':
-                {
-                    TXESessions.delete(sessionId) && this.logger.debug(`Called reset on session ${sessionId}, yeeting it out of existence`);
-                    return toForeignCallResult([]);
-                }
-            case 'deploy':
-                {
-                    await this.#processDeployInputs(callData);
-                    break;
-                }
-            case 'addAccount':
-                {
-                    await this.#processAddAccountInputs(callData);
-                    break;
-                }
+        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;
         }
-        const txeService = TXESessions.get(sessionId);
-        const response = await txeService[functionName](...inputs);
-        return response;
+        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
-    });
-}