@aztec/txe 0.0.1-commit.2448fdb → 0.0.1-commit.2606882

package/src/oracle/txe_oracle_top_level_context.ts

@@ -1,7 +1,6 @@
 import {
   CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
-  MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT,
-  MAX_PROCESSABLE_L2_GAS,
+  MAX_PRIVATE_LOGS_PER_TX,
   NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
 } from '@aztec/constants';
 import { BlockNumber } from '@aztec/foundation/branded-types';
@@ -15,13 +14,14 @@ import {
   CapsuleService,
   CapsuleStore,
   type ContractStore,
-  type ContractSyncService,
+  type ExecutionHooks,
   NoteStore,
   ORACLE_VERSION_MAJOR,
   PrivateEventStore,
   RecipientTaggingStore,
   SenderAddressBookStore,
   SenderTaggingStore,
+  composeHooks,
   enrichPublicSimulationError,
 } from '@aztec/pxe/server';
 import {
@@ -55,13 +55,7 @@ import { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { PublicSimulatorConfig } from '@aztec/stdlib/avm';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { type ContractInstanceWithAddress, computePartialAddress } from '@aztec/stdlib/contract';
-import {
-  FALLBACK_TEARDOWN_DA_GAS_LIMIT,
-  FALLBACK_TEARDOWN_L2_GAS_LIMIT,
-  Gas,
-  GasFees,
-  GasSettings,
-} from '@aztec/stdlib/gas';
+import { Gas, GasFees, GasSettings } from '@aztec/stdlib/gas';
 import { computeCalldataHash, computeProtocolNullifier, siloNullifier } from '@aztec/stdlib/hash';
 import {
   PartialPrivateTailPublicInputsForPublic,
@@ -69,6 +63,7 @@ import {
   PrivateToPublicAccumulatedData,
   PublicCallRequest,
 } from '@aztec/stdlib/kernel';
+import { hashPublicKey } from '@aztec/stdlib/keys';
 import { ChonkProof } from '@aztec/stdlib/proofs';
 import { makeGlobalVariables } from '@aztec/stdlib/testing';
 import { MerkleTreeId } from '@aztec/stdlib/trees';
@@ -85,13 +80,14 @@ import {
   collectNested,
 } from '@aztec/stdlib/tx';
 import type { UInt64 } from '@aztec/stdlib/types';
-import { ForkCheckpoint } from '@aztec/world-state';
+import { ForkCheckpoint } from '@aztec/world-state/native';
 
-import { DEFAULT_ADDRESS } from '../constants.js';
+import { DEFAULT_ADDRESS, MAX_PRIVATE_EVENTS_PER_TXE_QUERY, MAX_PRIVATE_EVENT_LEN } from '../constants.js';
 import type { TXEStateMachine } from '../state_machine/index.js';
-import type { TXEAccountStore } from '../util/txe_account_store.js';
-import { TXEPublicContractDataSource } from '../util/txe_public_contract_data_source.js';
 import { getSingleTxBlockRequestHash, insertTxEffectIntoWorldTrees, makeTXEBlock } from '../utils/block_creation.js';
+import type { TXEAccountStore } from '../utils/txe_account_store.js';
+import type { TXEArtifactResolver } from '../utils/txe_artifact_resolver.js';
+import { TXEPublicContractDataSource } from '../utils/txe_public_contract_data_source.js';
 import type { ITxeExecutionOracle } from './interfaces.js';
 
 export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracle {
@@ -116,7 +112,9 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     private version: Fr,
     private chainId: Fr,
     private authwits: Map<string, AuthWitness>,
-    private readonly contractSyncService: ContractSyncService,
+    private readonly artifactResolver: TXEArtifactResolver,
+    private readonly rootPath: string,
+    private readonly packageName: string,
   ) {
     this.logger = createLogger('txe:top_level_context');
     this.logger.debug('Entering Top Level Context');
@@ -149,7 +147,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
   }
 
   // We instruct users to debug contracts via this oracle, so it makes sense that they'd expect it to also work in tests
-  log(level: number, message: string, fields: Fr[]): Promise<void> {
+  log(level: number, message: string, _fieldsSize: number, fields: Fr[]): Promise<void> {
     if (!LogLevels[level]) {
       throw new Error(`Invalid log level: ${level}`);
     }
@@ -172,12 +170,12 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
   }
 
   async getLastBlockTimestamp() {
-    return (await this.stateMachine.node.getBlockHeader('latest'))!.globalVariables.timestamp;
+    return (await this.stateMachine.node.getBlockData('latest'))!.header.globalVariables.timestamp;
   }
 
   async getLastTxEffects() {
     const latestBlockNumber = await this.stateMachine.archiver.getBlockNumber();
-    const block = await this.stateMachine.archiver.getBlock(latestBlockNumber);
+    const block = await this.stateMachine.archiver.getBlock({ number: latestBlockNumber });
 
     if (block!.body.txEffects.length != 1) {
       // Note that calls like env.mine() will result in blocks with no transactions, hitting this
@@ -186,7 +184,17 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     const txEffects = block!.body.txEffects[0];
 
-    return { txHash: txEffects.txHash, noteHashes: txEffects.noteHashes, nullifiers: txEffects.nullifiers };
+    const privateLogs = txEffects.privateLogs;
+    if (privateLogs.length > MAX_PRIVATE_LOGS_PER_TX) {
+      throw new Error(`${privateLogs.length} private logs exceed max ${MAX_PRIVATE_LOGS_PER_TX}`);
+    }
+
+    return {
+      txHash: txEffects.txHash,
+      noteHashes: txEffects.noteHashes,
+      nullifiers: txEffects.nullifiers,
+      privateLogs,
+    };
   }
 
   async syncContractNonOracleMethod(contractAddress: AztecAddress, scope: AztecAddress, jobId: string) {
@@ -209,7 +217,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
   }
 
   async getPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress) {
-    return (
+    const events = (
       await this.privateEventStore.getPrivateEvents(selector, {
         contractAddress,
         scopes: [scope],
@@ -217,6 +225,15 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
         toBlock: (await this.getLastBlockNumber()) + 1,
       })
     ).map(e => e.packedEvent);
+
+    if (events.length > MAX_PRIVATE_EVENTS_PER_TXE_QUERY) {
+      throw new Error(`Array of length ${events.length} larger than maxLen ${MAX_PRIVATE_EVENTS_PER_TXE_QUERY}`);
+    }
+    if (events.some(e => e.length > MAX_PRIVATE_EVENT_LEN)) {
+      throw new Error(`Some private event has length larger than maxLen ${MAX_PRIVATE_EVENT_LEN}`);
+    }
+
+    return events;
   }
 
   async advanceBlocksBy(blocks: number) {
@@ -232,27 +249,71 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     this.nextBlockTimestamp += duration;
   }
 
-  async deploy(artifact: ContractArtifact, instance: ContractInstanceWithAddress, secret: Fr) {
+  private deploymentNullifier(address: AztecAddress): Promise<Fr> {
+    return siloNullifier(AztecAddress.fromNumber(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS), address.toField());
+  }
+
+  async deploy(
+    contractPath: string,
+    initializer: string,
+    args: Fr[],
+    secret: Fr,
+    salt: Fr,
+    deployer: AztecAddress,
+  ): Promise<Fr[]> {
+    const { artifact, instance } = await this.artifactResolver.resolveDeployArtifact({
+      rootPath: this.rootPath,
+      packageName: this.packageName,
+      contractPath,
+      initializer,
+      args,
+      secret,
+      salt,
+      deployer,
+    });
+
     // Emit deployment nullifier
     await this.mineBlock({
-      nullifiers: [
-        await siloNullifier(
-          AztecAddress.fromNumber(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS),
-          instance.address.toField(),
-        ),
-      ],
+      nullifiers: [await this.deploymentNullifier(instance.address)],
     });
 
     if (!secret.equals(Fr.ZERO)) {
-      await this.addAccount(artifact, instance, secret);
+      await this.registerContractAndAddAccount(artifact, instance, secret);
     } else {
       await this.contractStore.addContractInstance(instance);
       await this.contractStore.addContractArtifact(artifact);
       this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`);
     }
+
+    return [
+      instance.salt,
+      instance.deployer.toField(),
+      instance.currentContractClassId,
+      instance.initializationHash,
+      instance.immutablesHash,
+      ...instance.publicKeys.toFields(),
+    ];
+  }
+
+  /**
+   * Mines a single block containing only the deployment nullifiers for the contracts at the given addresses.
+   */
+  async mineDeploymentNullifiers(addresses: AztecAddress[]) {
+    await this.mineBlock({
+      nullifiers: await Promise.all(addresses.map(address => this.deploymentNullifier(address))),
+    });
   }
 
-  async addAccount(artifact: ContractArtifact, instance: ContractInstanceWithAddress, secret: Fr) {
+  async addAccount(secret: Fr) {
+    const { artifact, instance } = await this.artifactResolver.resolveAccountArtifact(secret);
+    return this.registerContractAndAddAccount(artifact, instance, secret);
+  }
+
+  private async registerContractAndAddAccount(
+    artifact: ContractArtifact,
+    instance: ContractInstanceWithAddress,
+    secret: Fr,
+  ) {
     const partialAddress = await computePartialAddress(instance);
 
     this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`);
@@ -279,12 +340,13 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
   async addAuthWitness(address: AztecAddress, messageHash: Fr) {
     const account = await this.accountStore.getAccount(address);
-    const privateKey = await this.keyStore.getMasterSecretKey(account.publicKeys.masterIncomingViewingPublicKey);
+    const ivpkMHash = await hashPublicKey(account.publicKeys.ivpkM);
+    const privateKey = await this.keyStore.getMasterSecretKey(ivpkMHash);
 
     const schnorr = new Schnorr();
-    const signature = await schnorr.constructSignature(messageHash.toBuffer(), privateKey);
+    const signature = await schnorr.constructSignature(messageHash, privateKey);
 
-    const authWitness = new AuthWitness(messageHash, [...signature.toBuffer()]);
+    const authWitness = new AuthWitness(messageHash, signature.toLimbFields());
 
     this.authwits.set(authWitness.requestHash.toString(), authWitness);
   }
@@ -315,13 +377,16 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
   }
 
   async privateCallNewFlow(
-    from: AztecAddress,
+    from: AztecAddress | undefined,
     targetContractAddress: AztecAddress = AztecAddress.zero(),
     functionSelector: FunctionSelector = FunctionSelector.empty(),
     args: Fr[],
     argsHash: Fr = Fr.zero(),
     isStaticCall: boolean = false,
+    additionalScopes: AztecAddress[] = [],
     jobId: string,
+    authorizedUtilityCallTargets: AztecAddress[],
+    gasSettings: GasSettings,
   ) {
     this.logger.verbose(
       `Executing external function ${await this.contractStore.getDebugFunctionName(targetContractAddress, functionSelector)}@${targetContractAddress} isStaticCall=${isStaticCall}`,
@@ -335,9 +400,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       throw new Error(message);
     }
 
-    // When `from` is the zero address (e.g. when deploying a new account contract), we return an
-    // empty scope list which acts as deny-all: no notes are visible and no keys are accessible.
-    const effectiveScopes = from.isZero() ? [] : [from];
+    const scopes = from === undefined ? additionalScopes : [from, ...additionalScopes];
 
     // Sync notes before executing private function to discover notes from previous transactions
     const utilityExecutor = async (call: FunctionCall, execScopes: AztecAddress[]) => {
@@ -351,16 +414,13 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       utilityExecutor,
       blockHeader,
       jobId,
-      effectiveScopes,
+      scopes,
     );
 
     const blockNumber = await this.getNextBlockNumber();
 
-    const callContext = new CallContext(from, targetContractAddress, functionSelector, isStaticCall);
-
-    const gasLimits = new Gas(MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT, MAX_PROCESSABLE_L2_GAS);
-    const teardownGasLimits = new Gas(FALLBACK_TEARDOWN_DA_GAS_LIMIT, FALLBACK_TEARDOWN_L2_GAS_LIMIT);
-    const gasSettings = new GasSettings(gasLimits, teardownGasLimits, GasFees.empty(), GasFees.empty());
+    const msgSender = from ?? AztecAddress.NULL_MSG_SENDER;
+    const callContext = new CallContext(msgSender, targetContractAddress, functionSelector, isStaticCall);
 
     const txContext = new TxContext(this.chainId, this.version, gasSettings);
 
@@ -394,19 +454,25 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       senderTaggingStore: this.senderTaggingStore,
       recipientTaggingStore: this.recipientTaggingStore,
       senderAddressBookStore: this.senderAddressBookStore,
-      capsuleService: new CapsuleService(this.capsuleStore, effectiveScopes),
+      capsuleService: new CapsuleService(this.capsuleStore, scopes),
       privateEventStore: this.privateEventStore,
       contractSyncService: this.stateMachine.contractSyncService,
       jobId,
       totalPublicCalldataCount: 0,
       sideEffectCounter: minRevertibleSideEffectCounter,
-      scopes: effectiveScopes,
+      scopes,
       // In TXE, the typical transaction entrypoint is skipped, so we need to simulate the actions that such a
       // contract would perform, including setting senderForTags.
       senderForTags: from,
       simulator,
       messageContextService: this.stateMachine.messageContextService,
-      l2TipsStore: this.stateMachine.node,
+      l2TipsStore: this.stateMachine.l2TipsProvider,
+      hooks: composeHooks({
+        authorizeUtilityCall: this.buildAuthorizeUtilityCallHook(
+          isStaticCall ? 'private view' : 'private',
+          authorizedUtilityCallTargets,
+        ),
+      }),
     });
 
     // Note: This is a slight modification of simulator.run without any of the checks. Maybe we should modify simulator.run with a boolean value to skip checks.
@@ -514,11 +580,17 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       }
     }
 
+    // Walk the nested private-call tree and collect every offchain effect the transaction emitted.
+    // PXE stores these on each `PrivateCallExecutionResult` and they never reach TXE via the
+    // `aztec_utl_emitOffchainEffect` foreign-call path (that path only fires at the top-level), so
+    // we pull them out here and the RPC wrapper will hand them to `TXESession` for buffering.
+    const offchainEffects = collectNested([executionResult], r => r.offchainEffects.map(e => e.data));
+
     if (isStaticCall) {
       await checkpoint!.revert();
 
       await forkedWorldTrees.close();
-      return executionResult.returnValues ?? [];
+      return { returnValues: executionResult.returnValues ?? [], offchainEffects };
     }
 
     const txEffect = TxEffect.empty();
@@ -540,14 +612,15 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     await forkedWorldTrees.close();
 
-    return executionResult.returnValues ?? [];
+    return { returnValues: executionResult.returnValues ?? [], offchainEffects };
   }
 
   async publicCallNewFlow(
-    from: AztecAddress,
+    from: AztecAddress | undefined,
     targetContractAddress: AztecAddress,
     calldata: Fr[],
     isStaticCall: boolean,
+    gasSettings: GasSettings,
   ) {
     this.logger.verbose(
       `Executing public function ${await this.contractStore.getDebugFunctionName(targetContractAddress, FunctionSelector.fromField(calldata[0]))}@${targetContractAddress} isStaticCall=${isStaticCall}`,
@@ -555,12 +628,6 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     const blockNumber = await this.getNextBlockNumber();
 
-    const gasLimits = new Gas(MAX_PROCESSABLE_DA_GAS_PER_CHECKPOINT, MAX_PROCESSABLE_L2_GAS);
-
-    const teardownGasLimits = new Gas(FALLBACK_TEARDOWN_DA_GAS_LIMIT, FALLBACK_TEARDOWN_L2_GAS_LIMIT);
-
-    const gasSettings = new GasSettings(gasLimits, teardownGasLimits, GasFees.empty(), GasFees.empty());
-
     const txContext = new TxContext(this.chainId, this.version, gasSettings);
 
     const anchorBlockHeader = await this.stateMachine.anchorBlockStore.getBlockHeader();
@@ -612,7 +679,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     // may require producing reverts.
     const revertibleAccumulatedData = PrivateToPublicAccumulatedData.empty();
     revertibleAccumulatedData.publicCallRequests[0] = new PublicCallRequest(
-      from,
+      from ?? AztecAddress.NULL_MSG_SENDER,
       targetContractAddress,
       isStaticCall,
       calldataHash,
@@ -703,6 +770,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     functionSelector: FunctionSelector,
     args: Fr[],
     jobId: string,
+    authorizedUtilityCallTargets: AztecAddress[],
   ) {
     const artifact = await this.contractStore.getFunctionArtifact(targetContractAddress, functionSelector);
     if (!artifact) {
@@ -733,10 +801,15 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       returnTypes: [],
     });
 
-    return this.executeUtilityCall(call, await this.keyStore.getAccounts(), jobId);
+    return this.executeUtilityCall(call, await this.keyStore.getAccounts(), jobId, authorizedUtilityCallTargets);
   }
 
-  private async executeUtilityCall(call: FunctionCall, scopes: AztecAddress[], jobId: string): Promise<Fr[]> {
+  private async executeUtilityCall(
+    call: FunctionCall,
+    scopes: AztecAddress[],
+    jobId: string,
+    authorizedUtilityCallTargets: AztecAddress[] = [],
+  ): Promise<Fr[]> {
     const entryPointArtifact = await this.contractStore.getFunctionArtifactWithDebugMetadata(call.to, call.selector);
     if (entryPointArtifact.functionType !== FunctionType.UTILITY) {
       throw new Error(`Cannot run ${entryPointArtifact.functionType} function as utility`);
@@ -749,6 +822,10 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     try {
       const anchorBlockHeader = await this.stateMachine.anchorBlockStore.getBlockHeader();
+      const simulator = new WASMSimulator();
+      const utilityExecutor = async (syncCall: FunctionCall, execScopes: AztecAddress[]) => {
+        await this.executeUtilityCall(syncCall, execScopes, jobId);
+      };
       const oracle = new UtilityExecutionOracle({
         contractAddress: call.to,
         authWitnesses: [],
@@ -764,12 +841,17 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
         capsuleService: new CapsuleService(this.capsuleStore, scopes),
         privateEventStore: this.privateEventStore,
         messageContextService: this.stateMachine.messageContextService,
-        contractSyncService: this.contractSyncService,
-        l2TipsStore: this.stateMachine.node,
+        contractSyncService: this.stateMachine.contractSyncService,
+        l2TipsStore: this.stateMachine.l2TipsProvider,
         jobId,
         scopes,
+        simulator,
+        hooks: composeHooks({
+          authorizeUtilityCall: this.buildAuthorizeUtilityCallHook('utility', authorizedUtilityCallTargets),
+        }),
+        utilityExecutor,
       });
-      const acirExecutionResult = await new WASMSimulator()
+      const acirExecutionResult = await simulator
         .executeUserCircuit(toACVMWitness(0, call.args), entryPointArtifact, new Oracle(oracle).toACIRCallback())
         .catch((err: Error) => {
           err.message = resolveAssertionMessageFromError(err, entryPointArtifact);
@@ -797,7 +879,20 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
   }
 
   private async getLastBlockNumber(): Promise<BlockNumber> {
-    const header = await this.stateMachine.node.getBlockHeader('latest');
-    return header ? header.globalVariables.blockNumber : BlockNumber.ZERO;
+    const block = await this.stateMachine.node.getBlock('latest');
+    return block ? block.header.globalVariables.blockNumber : BlockNumber.ZERO;
+  }
+
+  private buildAuthorizeUtilityCallHook(
+    callerContext: 'private' | 'private view' | 'utility',
+    authorizedTargets: AztecAddress[],
+  ): ExecutionHooks['authorizeUtilityCall'] | undefined {
+    if (authorizedTargets.length === 0) {
+      return undefined;
+    }
+    return req =>
+      Promise.resolve({
+        authorized: req.callerContext === callerContext && authorizedTargets.some(t => t.equals(req.target)),
+      });
   }
 }

package/src/oracle/txe_oracle_version.ts

@@ -0,0 +1,17 @@
+/**
+ * The TXE oracle version constants are used to check that the oracle interface used for tests is in sync between
+ * TXE and Aztec.nr. This is separate from the contract oracle version in `pxe/src/oracle_version.ts`, which covers
+ * oracles used during contract execution by PXE.
+ *
+ * The Noir counterparts are in `noir-projects/aztec-nr/aztec/src/test/helpers/txe_oracles.nr`.
+ */
+export const TXE_ORACLE_VERSION_MAJOR = 1;
+export const TXE_ORACLE_VERSION_MINOR = 2;
+
+/**
+ * This hash is computed from the TXE oracle interfaces (IAvmExecutionOracle and ITxeExecutionOracle) and is used to
+ * detect when those interfaces change. When it does, bump:
+ *   - TXE_ORACLE_VERSION_MAJOR (and reset MINOR to 0) for breaking changes, or
+ *   - TXE_ORACLE_VERSION_MINOR for additive changes (new oracle method added).
+ */
+export const TXE_ORACLE_INTERFACE_HASH = '9e5f6ad5fd170d1de5ddd417f19cce47b17382567e08360dc6a783154828e218';

package/src/oracle/txe_private_execution_oracle.ts

@@ -0,0 +1,30 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import { PrivateExecutionOracle } from '@aztec/pxe/simulator';
+import type { FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/**
+ * TXE-specific subclass of PrivateExecutionOracle that forbids operations not supported in
+ * TestEnvironment::private_context. TXE uses dedicated oracle flows (e.g. private_call) instead.
+ */
+export class TXEPrivateExecutionOracle extends PrivateExecutionOracle {
+  override callPrivateFunction(
+    _targetContractAddress: AztecAddress,
+    _functionSelector: FunctionSelector,
+    _argsHash: Fr,
+    _sideEffectCounter: number,
+    _isStaticCall: boolean,
+  ): Promise<{ endSideEffectCounter: Fr; returnsHash: Fr }> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::private_context`, use `private_call` instead',
+    );
+  }
+
+  override assertValidPublicCalldata(_calldataHash: Fr): Promise<void> {
+    throw new Error('Enqueueing public calls is not supported in TestEnvironment::private_context');
+  }
+
+  override notifyRevertiblePhaseStart(_minRevertibleSideEffectCounter: number): Promise<void> {
+    throw new Error('Enqueueing public calls is not supported in TestEnvironment::private_context');
+  }
+}

package/src/rpc_server.ts

@@ -0,0 +1,87 @@
+import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
+import type { Logger } from '@aztec/foundation/log';
+
+import type { Socket } from 'node:net';
+
+import { TXEDispatcherPool, buildSharedContractStore } from './dispatcher_pool.js';
+import { TXEDispatcher, TXEDispatcherApiSchema } from './index.js';
+
+/**
+ * Symbol used to tag an incoming TCP socket with the `session_id` it has been associated with.
+ * Hidden under a Symbol so we don't risk colliding with anything koa or http core adds.
+ */
+const SESSION_SYMBOL = Symbol('txeSessionId');
+
+type TaggedSocket = Socket & { [SESSION_SYMBOL]?: number };
+
+/**
+ * Creates the TXE RPC server. With `TXE_WORKERS=1` oracle calls run on the main thread (no
+ * worker_threads, no IPC overhead). With any other value oracle calls are
+ * routed to a pool of worker threads sized to that value, sticky by `session_id`.
+ *
+ * Each incoming TCP socket is tagged with the `session_id` of the first oracle call it carries —
+ * nargo uses one HTTP client per test, so the socket-to-session mapping is 1:1. When the socket
+ * closes (end of test), the dispatcher disposes the session and frees its world state + LMDB.
+ *
+ * Lives in its own module so the worker bundle does not pull in the HTTP server stack.
+ */
+export async function createTXERpcServer(logger: Logger) {
+  const workerCount = Number(process.env.TXE_WORKERS);
+  let dispatcher: TXEDispatcher | TXEDispatcherPool;
+  if (workerCount === 1) {
+    const { dataDir, schnorrClassId } = await buildSharedContractStore();
+    dispatcher = new TXEDispatcher(logger, { contractStoreSourceDir: dataDir, schnorrClassId });
+  } else {
+    dispatcher = new TXEDispatcherPool(logger, {
+      workers: Number.isFinite(workerCount) && workerCount > 1 ? workerCount : undefined,
+    });
+  }
+  const server = createSafeJsonRpcServer(dispatcher, TXEDispatcherApiSchema, {
+    http200OnError: true,
+    middlewares: [
+      async (ctx, next) => {
+        // The body parser runs further down the chain, so `ctx.request.body` is populated only
+        // after `next()` resolves.
+        await next();
+        const socket = ctx.req.socket as TaggedSocket;
+        if (socket[SESSION_SYMBOL] !== undefined) {
+          return;
+        }
+        const body = (ctx.request as { body?: unknown }).body;
+        const sessionId =
+          body && typeof body === 'object' && 'params' in body
+            ? extractSessionId((body as { params: unknown }).params)
+            : undefined;
+        if (sessionId === undefined) {
+          return;
+        }
+        socket[SESSION_SYMBOL] = sessionId;
+        logger.debug(`Tagged socket with session`, {
+          sessionId,
+          remoteAddress: socket.remoteAddress,
+          remotePort: socket.remotePort,
+        });
+        socket.once('close', () => {
+          logger.debug(`Disposing session on socket close`, { sessionId });
+          void dispatcher.disposeSession(sessionId);
+        });
+      },
+    ],
+  });
+  return server;
+}
+
+// Extracts `session_id` from a JSON-RPC `params` array. Always a `number` because session_id
+// comes off `JSON.parse`, which never produces BigInt; values above MAX_SAFE_INTEGER lose
+// precision but still work as a Map key.
+function extractSessionId(params: unknown): number | undefined {
+  if (!Array.isArray(params) || params.length === 0) {
+    return undefined;
+  }
+  const first = params[0];
+  if (first && typeof first === 'object' && 'session_id' in first) {
+    const sid = (first as { session_id: unknown }).session_id;
+    return typeof sid === 'number' ? sid : undefined;
+  }
+  return undefined;
+}