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

package/src/dispatcher_pool.ts

@@ -0,0 +1,316 @@
+import { getSchnorrAccountContractArtifact } from '@aztec/accounts/schnorr/lazy';
+import { BackendType, Barretenberg, BarretenbergSync } from '@aztec/bb.js';
+import type { Logger } from '@aztec/foundation/log';
+import { openEphemeralStore } from '@aztec/kv-store/lmdb-v2';
+import { LazyProtocolContractsProvider } from '@aztec/protocol-contracts/providers/lazy';
+import { ContractStore } from '@aztec/pxe/client/lazy';
+import { getStandardAuthRegistry } from '@aztec/standard-contracts/auth-registry';
+import { getContractClassFromArtifact } from '@aztec/stdlib/contract';
+
+import { existsSync } from 'node:fs';
+import { cpus } from 'node:os';
+import { fileURLToPath } from 'node:url';
+import { Worker } from 'node:worker_threads';
+
+import { TXE_REQUIRED_PROTOCOL_CONTRACTS } from './index.js';
+import type { TXEForeignCallInput } from './index.js';
+import type { ForeignCallResult } from './utils/encoding.js';
+
+void Barretenberg.initSingleton({ backend: BackendType.Wasm, skipSrsInit: true, threads: 1 });
+void BarretenbergSync.initSingleton({ backend: BackendType.Wasm });
+
+/**
+ * Opens a fresh LMDB in a tmp dir and writes the protocol contracts in
+ * {@link TXE_REQUIRED_PROTOCOL_CONTRACTS} plus the standard AuthRegistry and the SchnorrAccount artifact,
+ * returning the directory path and the SchnorrAccount class id (hex). The store handle is intentionally kept
+ * alive: closing it would trigger the ephemeral-store cleanup hook and remove the tmp
+ * directory, so any worker that has not yet cloned would find it missing.
+ */
+export async function buildSharedContractStore(): Promise<{ dataDir: string; schnorrClassId: string }> {
+  const kvStore = await openEphemeralStore('txe-shared-contracts', undefined, 2);
+  const dataDir = kvStore.dataDirectory;
+  const contractStore = new ContractStore(kvStore);
+  const provider = new LazyProtocolContractsProvider();
+  const [protocolContracts, standardContracts, schnorrArtifact] = await Promise.all([
+    Promise.all(TXE_REQUIRED_PROTOCOL_CONTRACTS.map(name => provider.getProtocolContractArtifact(name))),
+    Promise.all([getStandardAuthRegistry()]),
+    getSchnorrAccountContractArtifact(),
+  ]);
+  const schnorrClass = await getContractClassFromArtifact(schnorrArtifact);
+  await Promise.all([
+    ...[...protocolContracts, ...standardContracts].flatMap(({ instance, artifact, contractClass }) => [
+      contractStore.addContractArtifact(artifact, contractClass),
+      contractStore.addContractInstance(instance),
+    ]),
+    contractStore.addContractArtifact(schnorrArtifact, schnorrClass),
+  ]);
+  return { dataDir, schnorrClassId: schnorrClass.id.toString() };
+}
+
+/**
+ * Resolves `worker.bundle.js` whether this code is running unbundled (next to dispatcher_pool.js
+ * inside `dest/`) or bundled into `dest/bin/index.js` (one directory deeper). `import.meta.url`
+ * refers to whichever module the calling code actually lives in; we try both relative locations
+ * and use whichever exists.
+ */
+function resolveWorkerBundlePath(): URL {
+  const candidates = [new URL('./worker.bundle.js', import.meta.url), new URL('../worker.bundle.js', import.meta.url)];
+  return candidates.find(u => existsSync(fileURLToPath(u))) ?? candidates[0];
+}
+
+interface SerializedError {
+  message: string;
+  name?: string;
+  stack?: string;
+}
+
+type WorkerMessage =
+  | { type: 'result'; requestId: number; ok: true; value: ForeignCallResult }
+  | { type: 'result'; requestId: number; ok: false; error: SerializedError }
+  | {
+      type: 'memstat';
+      sessions: number;
+      rss: number;
+      heapTotal: number;
+      heapUsed: number;
+      external: number;
+      arrayBuffers: number;
+    };
+
+interface WorkerSlot {
+  worker: Worker;
+  sessions: Set<number>;
+  inFlightRequestIds: Set<number>;
+}
+
+interface PendingRequest {
+  resolve: (value: ForeignCallResult) => void;
+  reject: (err: Error) => void;
+  workerIdx: number;
+}
+
+function deserializeError(payload: SerializedError | undefined): Error {
+  if (!payload) {
+    return new Error('Worker returned an error with no payload');
+  }
+  const err = new Error(payload.message);
+  if (payload.name) {
+    err.name = payload.name;
+  }
+  if (payload.stack) {
+    err.stack = payload.stack;
+  }
+  return err;
+}
+
+export interface TXEDispatcherPoolOptions {
+  /** Number of worker threads */
+  workers?: number;
+}
+
+/**
+ * Main-thread router that owns a pool of TXE worker threads. Each worker runs its own
+ * {@link TXEDispatcher} and handles foreign-call requests for the sessions assigned to it.
+ *
+ * Routing is sticky by `session_id`: new sessions go to a freshly spawned worker (up to
+ * `maxWorkers`) or, once the cap is reached, to the existing worker with the fewest sessions.
+ * Each session's state — TXESession, native world state, KV stores — stays single-threaded
+ * within its worker; different sessions run in parallel across workers.
+ *
+ * The main thread builds a shared protocol-contracts LMDB once and passes its path via
+ * `workerData`. Workers clone the data file on demand instead of re-registering the contracts.
+ */
+export class TXEDispatcherPool {
+  private readonly workers: WorkerSlot[] = [];
+  private readonly sessionToWorker = new Map<number, number>();
+  private readonly pending = new Map<number, PendingRequest>();
+  private nextRequestId = 0;
+  private readonly maxWorkers: number;
+  private readonly readyPromise: Promise<void>;
+  private contractStoreSourceDir?: string;
+  private schnorrClassId?: string;
+  private workerPath?: URL;
+
+  constructor(
+    private readonly logger: Logger,
+    opts: TXEDispatcherPoolOptions = {},
+  ) {
+    // TXE still doesn't scale well beyond 16 workers due to contention
+    this.maxWorkers = Math.max(1, opts.workers ?? Math.min(cpus().length, 16));
+    this.readyPromise = this.init();
+  }
+
+  /** Resolves once the shared protocol-contracts store is built. Workers spawn lazily on demand. */
+  public ready(): Promise<void> {
+    return this.readyPromise;
+  }
+
+  private async init(): Promise<void> {
+    const t0 = Date.now();
+    const { dataDir, schnorrClassId } = await buildSharedContractStore();
+    this.contractStoreSourceDir = dataDir;
+    this.schnorrClassId = schnorrClassId;
+    this.workerPath = resolveWorkerBundlePath();
+    this.logger.debug(`TXE dispatcher pool ready (lazy spawn, cap=${this.maxWorkers})`, {
+      contractStoreSourceDir: dataDir,
+      schnorrClassId,
+      ms: Date.now() - t0,
+    });
+  }
+
+  /**
+   * Spawns a fresh worker and returns its index in `this.workers`. Caller must have awaited
+   * `init()` so the shared contract store path is set. Messages posted before the worker has
+   * finished loading are queued by Node's worker_threads transport.
+   */
+  private spawnWorker(): number {
+    const workerIdx = this.workers.length;
+    const w = new Worker(this.workerPath!, {
+      workerData: { contractStoreSourceDir: this.contractStoreSourceDir, schnorrClassId: this.schnorrClassId },
+    });
+    const slot: WorkerSlot = {
+      worker: w,
+      sessions: new Set(),
+      inFlightRequestIds: new Set(),
+    };
+    this.workers.push(slot);
+    w.on('message', (msg: WorkerMessage) => this.handleMessage(workerIdx, msg));
+    w.on('error', err => this.handleWorkerError(workerIdx, err));
+    w.on('exit', code => {
+      if (code !== 0) {
+        this.handleWorkerError(workerIdx, new Error(`Worker ${workerIdx} exited with code ${code}`));
+      }
+    });
+    this.logger.debug(`Spawning TXE worker ${workerIdx} on demand`, {
+      cap: this.maxWorkers,
+      poolSize: this.workers.length,
+    });
+    return workerIdx;
+  }
+
+  /** Routes a session-dispose request to the worker that owns the session. Fire-and-forget. */
+  disposeSession(sessionId: number): void {
+    const workerIdx = this.sessionToWorker.get(sessionId);
+    if (workerIdx === undefined) {
+      throw new Error(`disposeSession: no worker mapped for session ${sessionId}`);
+    }
+    this.sessionToWorker.delete(sessionId);
+    const slot = this.workers[workerIdx];
+    if (!slot) {
+      throw new Error(`disposeSession: worker ${workerIdx} (session ${sessionId}) missing from pool`);
+    }
+    slot.sessions.delete(sessionId);
+    slot.worker.postMessage({ type: 'dispose-session', sessionId });
+  }
+
+  // eslint-disable-next-line camelcase
+  async resolve_foreign_call(callData: TXEForeignCallInput): Promise<ForeignCallResult> {
+    // Make sure the shared contract store + worker bundle path are in place before we spawn.
+    await this.readyPromise;
+    const sessionId = callData.session_id;
+    let workerIdx = this.sessionToWorker.get(sessionId);
+    if (workerIdx === undefined) {
+      workerIdx = this.pickOrSpawnWorker();
+      this.sessionToWorker.set(sessionId, workerIdx);
+      this.workers[workerIdx].sessions.add(sessionId);
+      this.logger.debug(`Routing new session ${sessionId} to worker ${workerIdx}`);
+    }
+
+    const slot = this.workers[workerIdx];
+    const requestId = this.nextRequestId++;
+    return new Promise<ForeignCallResult>((resolve, reject) => {
+      this.pending.set(requestId, { resolve, reject, workerIdx: workerIdx! });
+      slot.inFlightRequestIds.add(requestId);
+      slot.worker.postMessage({ type: 'foreign-call', requestId, callData });
+    });
+  }
+
+  /** Terminates all workers. Used by tests; production TXE exits the process directly. */
+  async shutdown(): Promise<void> {
+    const slots = this.workers.splice(0, this.workers.length);
+    await Promise.all(slots.map(s => s.worker.terminate()));
+    for (const [requestId, pending] of this.pending) {
+      pending.reject(new Error('TXE dispatcher pool was shut down'));
+      this.pending.delete(requestId);
+    }
+  }
+
+  /**
+   * Returns the index of the worker that should handle a new session. Spawns a fresh worker if
+   * the pool hasn't reached its cap; otherwise picks the existing worker with the fewest
+   * assigned sessions. Always returns a valid index into `this.workers`.
+   */
+  private pickOrSpawnWorker(): number {
+    if (this.workers.length < this.maxWorkers) {
+      return this.spawnWorker();
+    }
+    let minIdx = 0;
+    let minCount = this.workers[0].sessions.size;
+    for (let i = 1; i < this.workers.length; i++) {
+      const count = this.workers[i].sessions.size;
+      if (count < minCount) {
+        minIdx = i;
+        minCount = count;
+      }
+    }
+    return minIdx;
+  }
+
+  private handleMessage(workerIdx: number, msg: WorkerMessage): void {
+    switch (msg.type) {
+      case 'result':
+        this.handleResult(workerIdx, msg);
+        return;
+      case 'memstat':
+        this.logger.debug(`worker ${workerIdx} memstat`, {
+          worker: workerIdx,
+          sessions: msg.sessions,
+          rssMiB: Math.round(msg.rss / 1024 / 1024),
+          heapTotalMiB: Math.round(msg.heapTotal / 1024 / 1024),
+          heapUsedMiB: Math.round(msg.heapUsed / 1024 / 1024),
+          externalMiB: Math.round(msg.external / 1024 / 1024),
+          arrayBuffersMiB: Math.round(msg.arrayBuffers / 1024 / 1024),
+        });
+        return;
+    }
+  }
+
+  private handleResult(workerIdx: number, msg: WorkerMessage & { type: 'result' }): void {
+    const pending = this.pending.get(msg.requestId);
+    if (!pending) {
+      throw new Error(`handleResult: request ${msg.requestId} (worker ${workerIdx}) not in pending map`);
+    }
+    this.pending.delete(msg.requestId);
+    const slot = this.workers[workerIdx];
+    if (!slot) {
+      throw new Error(`handleResult: worker ${workerIdx} (request ${msg.requestId}) missing from pool`);
+    }
+    slot.inFlightRequestIds.delete(msg.requestId);
+    if (msg.ok) {
+      pending.resolve(msg.value);
+    } else {
+      pending.reject(deserializeError(msg.error));
+    }
+  }
+
+  private handleWorkerError(workerIdx: number, err: Error): void {
+    this.logger.error(`TXE worker ${workerIdx} crashed; sessions assigned to it will fail`, err);
+    const slot = this.workers[workerIdx];
+    if (!slot) {
+      throw new Error(`handleWorkerError: worker ${workerIdx} missing from pool (orig err: ${err.message})`);
+    }
+    for (const requestId of slot.inFlightRequestIds) {
+      const pending = this.pending.get(requestId);
+      if (!pending) {
+        throw new Error(`handleWorkerError: in-flight request ${requestId} (worker ${workerIdx}) not in pending map`);
+      }
+      pending.reject(err);
+      this.pending.delete(requestId);
+    }
+    slot.inFlightRequestIds.clear();
+    for (const sessionId of slot.sessions) {
+      this.sessionToWorker.delete(sessionId);
+    }
+    slot.sessions.clear();
+  }
+}

package/src/index.ts

@@ -1,213 +1,151 @@
-import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
-import {
-  AztecAddress,
-  type ContractArtifact,
-  type ContractInstanceWithAddress,
-  Fr,
-  PublicKeys,
-  deriveKeys,
-  getContractInstanceFromDeployParams,
-  loadContractArtifact,
-} from '@aztec/aztec.js';
-import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
+import { Fr } from '@aztec/aztec.js/fields';
 import type { Logger } from '@aztec/foundation/log';
-import { type ProtocolContract, protocolContractNames } from '@aztec/protocol-contracts';
-import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import type { ApiSchemaFor, ZodFor } from '@aztec/stdlib/schemas';
+import { cloneEphemeralStoreFrom } from '@aztec/kv-store/lmdb-v2';
+import type { ProtocolContractName } from '@aztec/protocol-contracts';
+import { ContractStore } from '@aztec/pxe/client/lazy';
+import type { ApiSchemaFor } from '@aztec/stdlib/schemas';
+import { zodFor } from '@aztec/stdlib/schemas';
 
-import { readFile, readdir } from 'fs/promises';
 import { join } from 'path';
 import { z } from 'zod';
 
-import { TXEService } from './txe_service/txe_service.js';
+// 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,
   ForeignCallArgsSchema,
-  type ForeignCallArray,
   type ForeignCallResult,
   ForeignCallResultSchema,
-  type ForeignCallSingle,
-  fromArray,
-  fromSingle,
-  toForeignCallResult,
-  toSingle,
-} from './util/encoding.js';
-
-const TXESessions = new Map<number, TXEService>();
+} from './utils/encoding.js';
+import { TXEArtifactResolver } from './utils/txe_artifact_resolver.js';
 
-const TXEArtifactsCache = new Map<string, { artifact: ContractArtifact; instance: ContractInstanceWithAddress }>();
+// 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[] = [];
 
-type MethodNames<T> = {
-  [K in keyof T]: T[K] extends (...args: any[]) => any ? K : never;
-}[keyof T];
+const sessions = new Map<number, TXESession>();
 
-type TXEForeignCallInput = {
+export type TXEForeignCallInput = {
   session_id: number;
-  function: MethodNames<TXEService> | 'reset';
+  function: TXEOracleFunctionName;
   root_path: string;
   package_name: string;
   inputs: ForeignCallArgs;
 };
 
-const TXEForeignCallInputSchema = z.object({
-  // eslint-disable-next-line camelcase
-  session_id: z.number().int().nonnegative(),
-  function: z.string() as ZodFor<MethodNames<TXEService> | 'reset'>,
-  // eslint-disable-next-line camelcase
-  root_path: z.string(),
-  // eslint-disable-next-line camelcase
-  package_name: z.string(),
-  inputs: ForeignCallArgsSchema,
-}) satisfies ZodFor<TXEForeignCallInput>;
-
-class TXEDispatcher {
-  private protocolContracts!: ProtocolContract[];
-
-  constructor(private logger: Logger) {}
-
-  async #processDeployInputs({ inputs, root_path: rootPath, package_name: packageName }: TXEForeignCallInput) {
-    const [pathStr, contractName, initializer] = inputs.slice(0, 3).map(input =>
-      fromArray(input as ForeignCallArray)
-        .map(char => String.fromCharCode(char.toNumber()))
-        .join(''),
-    );
-
-    const decodedArgs = fromArray(inputs[4] as ForeignCallArray);
-    const secret = fromSingle(inputs[5] as ForeignCallSingle);
-    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 });
-    }
+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().nonnegative(),
+    function: z.string() as z.ZodType<TXEOracleFunctionName>,
+    // eslint-disable-next-line camelcase
+    root_path: z.string(),
+    // eslint-disable-next-line camelcase
+    package_name: z.string(),
+    inputs: ForeignCallArgsSchema,
+  }),
+);
+
+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. The
+   * {@link TXEArtifactResolver} looks the artifact up from the cloned store via this class id
+   * instead of recomputing it via `getSchnorrAccountContractArtifact()` + `computeArtifactHash()`.
+   */
+  schnorrClassId: string;
+}
 
-    inputs.splice(0, 2, artifact, instance, toSingle(secret));
+export class TXEDispatcher {
+  private contractStore!: ContractStore;
+  private artifactResolver!: TXEArtifactResolver;
+  private readonly contractStoreSourceDir: string;
+  private readonly schnorrClassId: Fr;
+
+  constructor(
+    private logger: Logger,
+    opts: TXEDispatcherOptions,
+  ) {
+    this.contractStoreSourceDir = opts.contractStoreSourceDir;
+    this.schnorrClassId = Fr.fromString(opts.schnorrClassId);
   }
 
-  async #processAddAccountInputs({ inputs }: TXEForeignCallInput) {
-    const secret = fromSingle(inputs[0] as ForeignCallSingle);
-
-    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 });
+  /**
+   * 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;
     }
-
-    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: TXEForeignCallInput): Promise<ForeignCallResult> {
-    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);
+  }
 
-    const txeService = TXESessions.get(sessionId);
-    const response = await (txeService as any)[functionName](...inputs);
-    return response;
+  /**
+   * 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();
   }
 }
 
-const TXEDispatcherApiSchema: ApiSchemaFor<TXEDispatcher> = {
+/** Diagnostic-only: number of sessions currently held by this worker. */
+export function activeSessionCount(): number {
+  return sessions.size;
+}
+
+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(),
+});