@aztec/txe 0.0.1-commit.96dac018d → 0.0.1-commit.993d240

package/src/oracle/txe_oracle_top_level_context.ts

@@ -1,9 +1,6 @@
 import {
   CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
-  DEFAULT_DA_GAS_LIMIT,
-  DEFAULT_L2_GAS_LIMIT,
-  DEFAULT_TEARDOWN_DA_GAS_LIMIT,
-  DEFAULT_TEARDOWN_L2_GAS_LIMIT,
+  MAX_PRIVATE_LOGS_PER_TX,
   NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
 } from '@aztec/constants';
 import { BlockNumber } from '@aztec/foundation/branded-types';
@@ -12,17 +9,19 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { LogLevels, type Logger, applyStringFormatting, createLogger } from '@aztec/foundation/log';
 import { TestDateProvider } from '@aztec/foundation/timer';
 import type { KeyStore } from '@aztec/key-store';
-import type { AccessScopes } from '@aztec/pxe/client/lazy';
 import {
   AddressStore,
+  CapsuleService,
   CapsuleStore,
   type ContractStore,
+  type ExecutionHooks,
   NoteStore,
-  ORACLE_VERSION,
+  ORACLE_VERSION_MAJOR,
   PrivateEventStore,
   RecipientTaggingStore,
   SenderAddressBookStore,
   SenderTaggingStore,
+  composeHooks,
   enrichPublicSimulationError,
 } from '@aztec/pxe/server';
 import {
@@ -64,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';
@@ -80,13 +80,13 @@ 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 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 { TXEPublicContractDataSource } from '../utils/txe_public_contract_data_source.js';
 import type { ITxeExecutionOracle } from './interfaces.js';
 
 export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracle {
@@ -116,22 +116,34 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     this.logger.debug('Entering Top Level Context');
   }
 
-  utilityAssertCompatibleOracleVersion(version: number): void {
-    if (version !== ORACLE_VERSION) {
+  private contractOracleVersion: { major: number; minor: number } | undefined;
+
+  assertCompatibleOracleVersion(major: number, minor: number): void {
+    if (major !== ORACLE_VERSION_MAJOR) {
+      const hint =
+        major > ORACLE_VERSION_MAJOR
+          ? 'The contract was compiled with a newer version of Aztec.nr than this aztec cli version supports. Upgrade your aztec cli version to a compatible version.'
+          : 'The contract was compiled with an older version of Aztec.nr than this aztec cli version supports. Recompile the contract with a compatible version of Aztec.nr.';
       throw new Error(
-        `Incompatible oracle version. TXE is using version '${ORACLE_VERSION}', but got a request for '${version}'.`,
+        `Incompatible aztec cli version: ${hint} See https://docs.aztec.network/errors/8 (expected oracle major version ${ORACLE_VERSION_MAJOR}, got ${major})`,
       );
     }
+    this.contractOracleVersion = { major, minor };
+  }
+
+  // Prefixed with "nonOracleFunction" as it is not used as an oracle handler.
+  nonOracleFunctionGetContractOracleVersion(): { major: number; minor: number } | undefined {
+    return this.contractOracleVersion;
   }
 
   // This is typically only invoked in private contexts, but it is convenient to also have it in top-level for testing
   // setup.
-  utilityGetRandomField(): Fr {
+  getRandomField(): Fr {
     return Fr.random();
   }
 
   // We instruct users to debug contracts via this oracle, so it makes sense that they'd expect it to also work in tests
-  utilityLog(level: number, message: string, fields: Fr[]): Promise<void> {
+  log(level: number, message: string, fields: Fr[]): Promise<void> {
     if (!LogLevels[level]) {
       throw new Error(`Invalid log level: ${level}`);
     }
@@ -141,25 +153,25 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     return Promise.resolve();
   }
 
-  txeGetDefaultAddress(): AztecAddress {
+  getDefaultAddress(): AztecAddress {
     return DEFAULT_ADDRESS;
   }
 
-  async txeGetNextBlockNumber(): Promise<BlockNumber> {
+  async getNextBlockNumber(): Promise<BlockNumber> {
     return BlockNumber((await this.getLastBlockNumber()) + 1);
   }
 
-  txeGetNextBlockTimestamp(): Promise<bigint> {
+  getNextBlockTimestamp(): Promise<bigint> {
     return Promise.resolve(this.nextBlockTimestamp);
   }
 
-  async txeGetLastBlockTimestamp() {
-    return (await this.stateMachine.node.getBlockHeader('latest'))!.globalVariables.timestamp;
+  async getLastBlockTimestamp() {
+    return (await this.stateMachine.node.getBlockData('latest'))!.header.globalVariables.timestamp;
   }
 
-  async txeGetLastTxEffects() {
+  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
@@ -168,12 +180,22 @@ 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) {
     if (contractAddress.equals(DEFAULT_ADDRESS)) {
-      this.logger.debug(`Skipping sync in txeGetPrivateEvents because the events correspond to the default address.`);
+      this.logger.debug(`Skipping sync in getPrivateEvents because the events correspond to the default address.`);
       return;
     }
 
@@ -190,7 +212,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     );
   }
 
-  async txeGetPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress) {
+  async getPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress) {
     return (
       await this.privateEventStore.getPrivateEvents(selector, {
         contractAddress,
@@ -201,7 +223,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     ).map(e => e.packedEvent);
   }
 
-  async txeAdvanceBlocksBy(blocks: number) {
+  async advanceBlocksBy(blocks: number) {
     this.logger.debug(`time traveling ${blocks} blocks`);
 
     for (let i = 0; i < blocks; i++) {
@@ -209,24 +231,23 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     }
   }
 
-  txeAdvanceTimestampBy(duration: UInt64) {
+  advanceTimestampBy(duration: UInt64) {
     this.logger.debug(`time traveling ${duration} seconds`);
     this.nextBlockTimestamp += duration;
   }
 
-  async txeDeploy(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(artifact: ContractArtifact, instance: ContractInstanceWithAddress, secret: Fr) {
     // 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.txeAddAccount(artifact, instance, secret);
+      await this.addAccount(artifact, instance, secret);
     } else {
       await this.contractStore.addContractInstance(instance);
       await this.contractStore.addContractArtifact(artifact);
@@ -234,7 +255,16 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     }
   }
 
-  async txeAddAccount(artifact: ContractArtifact, instance: ContractInstanceWithAddress, secret: Fr) {
+  /**
+   * 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) {
     const partialAddress = await computePartialAddress(instance);
 
     this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`);
@@ -249,7 +279,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     return completeAddress;
   }
 
-  async txeCreateAccount(secret: Fr) {
+  async createAccount(secret: Fr) {
     // This is a foot gun !
     const completeAddress = await this.keyStore.addAccount(secret, secret);
     await this.accountStore.setAccount(completeAddress.address, completeAddress);
@@ -259,20 +289,21 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     return completeAddress;
   }
 
-  async txeAddAuthWitness(address: AztecAddress, messageHash: Fr) {
+  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);
   }
 
   async mineBlock(options: { nullifiers?: Fr[] } = {}) {
-    const blockNumber = await this.txeGetNextBlockNumber();
+    const blockNumber = await this.getNextBlockNumber();
 
     const txEffect = TxEffect.empty();
     txEffect.nullifiers = [getSingleTxBlockRequestHash(blockNumber), ...(options.nullifiers ?? [])];
@@ -296,14 +327,17 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     await this.stateMachine.handleL2Block(block);
   }
 
-  async txePrivateCallNewFlow(
-    from: AztecAddress,
+  async privateCallNewFlow(
+    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}`,
@@ -317,12 +351,10 @@ 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: AccessScopes) => {
+    const utilityExecutor = async (call: FunctionCall, execScopes: AztecAddress[]) => {
       await this.executeUtilityCall(call, execScopes, jobId);
     };
 
@@ -333,16 +365,13 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       utilityExecutor,
       blockHeader,
       jobId,
-      effectiveScopes,
+      scopes,
     );
 
-    const blockNumber = await this.txeGetNextBlockNumber();
+    const blockNumber = await this.getNextBlockNumber();
 
-    const callContext = new CallContext(from, targetContractAddress, functionSelector, isStaticCall);
-
-    const gasLimits = new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
-    const teardownGasLimits = new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_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);
 
@@ -376,17 +405,25 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       senderTaggingStore: this.senderTaggingStore,
       recipientTaggingStore: this.recipientTaggingStore,
       senderAddressBookStore: this.senderAddressBookStore,
-      capsuleStore: this.capsuleStore,
+      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.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.
@@ -409,7 +446,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       );
       const publicFunctionsCalldata = await Promise.all(
         publicCallRequests.map(async r => {
-          const calldata = await privateExecutionOracle.privateLoadFromExecutionCache(r.calldataHash);
+          const calldata = await privateExecutionOracle.getHashPreimage(r.calldataHash);
           return new HashedValues(calldata, r.calldataHash);
         }),
       );
@@ -494,11 +531,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();
@@ -520,26 +563,21 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     await forkedWorldTrees.close();
 
-    return executionResult.returnValues ?? [];
+    return { returnValues: executionResult.returnValues ?? [], offchainEffects };
   }
 
-  async txePublicCallNewFlow(
-    from: AztecAddress,
+  async publicCallNewFlow(
+    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}`,
     );
 
-    const blockNumber = await this.txeGetNextBlockNumber();
-
-    const gasLimits = new Gas(DEFAULT_DA_GAS_LIMIT, DEFAULT_L2_GAS_LIMIT);
-
-    const teardownGasLimits = new Gas(DEFAULT_TEARDOWN_DA_GAS_LIMIT, DEFAULT_TEARDOWN_L2_GAS_LIMIT);
-
-    const gasSettings = new GasSettings(gasLimits, teardownGasLimits, GasFees.empty(), GasFees.empty());
+    const blockNumber = await this.getNextBlockNumber();
 
     const txContext = new TxContext(this.chainId, this.version, gasSettings);
 
@@ -592,7 +630,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,
@@ -678,11 +716,12 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
     return returnValues ?? [];
   }
 
-  async txeExecuteUtilityFunction(
+  async executeUtilityFunction(
     targetContractAddress: AztecAddress,
     functionSelector: FunctionSelector,
     args: Fr[],
     jobId: string,
+    authorizedUtilityCallTargets: AztecAddress[],
   ) {
     const artifact = await this.contractStore.getFunctionArtifact(targetContractAddress, functionSelector);
     if (!artifact) {
@@ -699,7 +738,7 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       },
       blockHeader,
       jobId,
-      'ALL_SCOPES',
+      await this.keyStore.getAccounts(),
     );
 
     const call = FunctionCall.from({
@@ -713,10 +752,15 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
       returnTypes: [],
     });
 
-    return this.executeUtilityCall(call, 'ALL_SCOPES', jobId);
+    return this.executeUtilityCall(call, await this.keyStore.getAccounts(), jobId, authorizedUtilityCallTargets);
   }
 
-  private async executeUtilityCall(call: FunctionCall, scopes: AccessScopes, 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`);
@@ -729,6 +773,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: [],
@@ -741,12 +789,20 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
         aztecNode: this.stateMachine.node,
         recipientTaggingStore: this.recipientTaggingStore,
         senderAddressBookStore: this.senderAddressBookStore,
-        capsuleStore: this.capsuleStore,
+        capsuleService: new CapsuleService(this.capsuleStore, scopes),
         privateEventStore: this.privateEventStore,
+        messageContextService: this.stateMachine.messageContextService,
+        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);
@@ -774,7 +830,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_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;
+}