fix(avm): Add cancellation token to prevent C++ simulation race condition on timeout (#19219)

# fix(avm): Add cancellation token to prevent C++ simulation race
condition on timeout

## Summary

Fix race condition where C++ AVM simulation corrupts WorldState after
TypeScript timeout. This was done as follows:
- Add `CancellationToken` mechanism to safely stop C++ simulation before
reverting checkpoints
- **Primary Fix**: PublicProcessor cancels public simulation and
**waits** for it to die before proceeding and reverting checkpoints
- C++ simulation polls this cancellation token regularly so that we know
it will die quickly after cancellation
- **Secondary Fix**: Add mutex locking to checkpoint operations for
thread safety
- Add paired tests proving both the bug exists (without fix) and the fix
works (with cancellation)

## The Bug

When a C++ AVM simulation times out via `Promise.race()` in
`PublicProcessor`, a race condition occurs:

1. TypeScript's timeout fires and calls `revertCheckpoint()` on
WorldState
2. **But C++ simulation continues running** on a libuv worker thread
3. C++ holds a native handle to WorldState obtained before the timeout
4. C++ makes writes to WorldState **after** the checkpoint was reverted
5. This corrupts WorldState, causing `checkWorldStateUnchanged()` to
fail

The root causes are:
1. `GuardedMerkleTreeOperations` only guards TS merkle operations. The
C++ code interacts with the same WorldState without guards.
2. Nothing stops C++ from finishing simulation after PublicProcessor
reaches deadline.

```
Timeline of the race (BEFORE fix):

TypeScript                          C++ (libuv worker thread)
----------                          -------------------------
Start simulation -----------------> Begins executing opcodes
     |                                    |
     v                                    v
Promise.race timeout fires          Still running (e.g., in pad_trees())
     |                                    |
     v                                    |
     |                                    |
     v                                    |
IMMEDIATELY revertCheckpoint()            |
     |                                    v
     |                              Finishes pad_trees()
     |                              Makes write AFTER revert! (CORRUPTION)
     v                                    |
checkWorldStateUnchanged() FAILS!   
```

## The Fix

The key insight: **signaling cancellation is not enough** - we must
**wait** for C++ to actually stop.

```
Timeline (AFTER fix):

TypeScript                          C++ (libuv worker thread)
----------                          -------------------------
Start simulation -----------------> Begins executing opcodes
     |                                    |
     v                                    v
Promise.race timeout fires          Still running (e.g., in pad_trees())
     |                                    |
     v                                    |
cancel(100) sets atomic flag              | (doesn't check flag yet)
     |                                    |
     v                                    v
WAIT for simulation promise         Finishes pad_trees(), checks flag
     |                              Throws CancelledException
     |<-----------------------------Promise rejects
     v
NOW revertCheckpoint() (C++ is done)
     |
     v
checkWorldStateUnchanged() ✓ clean state
```

### C++ Side

1. **`CancellationToken` class** (`cancellation_token.hpp`):
   - Thread-safe `std::atomic<bool>` flag
   - `cancel()` - signals cancellation (called from TS thread)
- `check()` - throws `CancelledException` if cancelled (called from
worker thread)

3. **Check at each opcode** (`execution.cpp`, `hybrid_execution.cpp`):
   ```cpp
   while (!external_call_stack.empty()) {
       if (cancellation_token_) {
           cancellation_token_->check();
       }
       // ... execute opcode
   }
   ```

4. **Check before every WorldState write** (`raw_data_dbs.cpp`):
   ```cpp
   void PureRawMerkleDB::pad_tree(...) {
       check_cancellation();  // Throws if cancelled
       // ... proceed with write
   }
   ```

5. **Mutex locking for checkpoint operations**
(`cached_content_addressed_tree_store.hpp`):
   ```cpp
   void ContentAddressedCachedTreeStore::checkpoint() {
       std::unique_lock lock(mtx_);  // Prevent races with C++ writes
       cache_.checkpoint();
   }
   ```

### TypeScript Side

1. **`CppPublicTxSimulator.cancel(waitTimeoutMs)`** - Signal AND wait:
   ```typescript
   public async cancel(waitTimeoutMs?: number): Promise<void> {
       if (this.cancellationToken) {
           cancelSimulation(this.cancellationToken);
       }
       // Wait for simulation to actually complete
       if (waitTimeoutMs !== undefined && this.simulationPromise) {
           await Promise.race([
               this.simulationPromise.catch(() => {}),
               sleep(waitTimeoutMs),
           ]);
       }
   }
   ```
> Note: ideally we'd like to have no timeout here since we really want
to wait for C++ to recognize cancellation. The timeout is really just a
safeguard against some cancellation bug.

2. **`PublicProcessor`** timeout handler:
   ```typescript
   if (err?.name === 'PublicProcessorTimeoutError') {
       // Signal cancellation AND WAIT for C++ to stop (up to 100ms)
       await this.publicTxSimulator.cancel?.();
       // NOW safe to stop the guarded fork
       await this.guardedMerkleTree.stop();
   }
   ```

## Test Plan

Four paired tests at two levels verify the bug and fix:

### Replace bug and prove fix at TxSimulator level

Tests using `CppPublicTxSimulator` directly - **identical code**, only
difference is whether `cancel()` is called:

```typescript
async function runRaceConditionTest(useCancellation: boolean): Promise<number> {
    // ... setup ...
    const simulationPromise = simulator.simulate(tx);

    if (useCancellation) {
        await simulator.cancel(100);  // FIX: Signal AND wait
    }
    // BUG: No cancel, C++ continues during reverts

    await merkleTrees.revertCheckpoint();
    // Check for corruption...
}

it('BUG PROOF: race condition exists WITHOUT cancellation');   // Expects >0 corruptions
it('FIX PROOF: no race condition WITH cancellation');          // Expects 0 corruptions
```

### Replicate bug and prove fix at PublicProcessor level

Tests using full `PublicProcessor.process()` with deadline timeout:

```typescript
it('PublicProcessor BUG PROOF: state corruption occurs WITHOUT cancellation');
it('PublicProcessor FIX PROOF: no state corruption WITH cancellation');
```

### Running the tests

```bash
cd yarn-project/simulator
yarn test src/public/public_processor/apps_tests/timeout_race.test.ts
```

## Files Changed

### C++ (barretenberg)

| File | Change |
|------|--------|
| `vm2/simulation/lib/cancellation_token.hpp` | **New**: Thread-safe
cancellation token class |
| `vm2/simulation/lib/raw_data_dbs.hpp` | Add token to `PureRawMerkleDB`
|
| `vm2/simulation/lib/raw_data_dbs.cpp` | Check cancellation before all
writes |
| `vm2/simulation/gadgets/execution.hpp` | Add token to `Execution` |
| `vm2/simulation/gadgets/execution.cpp` | Check cancellation at each
opcode |
| `vm2/simulation/standalone/hybrid_execution.cpp` | Check cancellation
at each opcode |
| `vm2/avm_sim_api.hpp` | Thread token through API |
| `vm2/avm_sim_api.cpp` | Thread token through API |
| `vm2/simulation_helper.hpp` | Thread token to execution |
| `vm2/simulation_helper.cpp` | Thread token to execution |
| `nodejs_module/avm_simulate/avm_simulate_napi.hpp` | NAPI bindings for
token |
| `nodejs_module/avm_simulate/avm_simulate_napi.cpp` | NAPI bindings for
token |
| `nodejs_module/init_module.cpp` | Register NAPI functions |
| `crypto/merkle_tree/.../cached_content_addressed_tree_store.hpp` | Add
mutex to checkpoint ops |

### TypeScript (yarn-project)

| File | Change |
|------|--------|
| `native/src/native_module.ts` | Export `createCancellationToken`,
`cancelSimulation` |
| `simulator/.../public_tx_simulator_interface.ts` | Add
`cancel?(waitTimeoutMs?)` method |
| `simulator/.../cpp_public_tx_simulator.ts` | Track promise, implement
`cancel()` with wait |
| `simulator/.../public_processor.ts` | `await cancel(100)` before
reverts |
| `simulator/.../public_tx_simulation_tester.ts` | Add `cancel()` and
`getSimulator()` |
| `simulator/.../timeout_race.test.ts` | **New**: Paired proof tests |

## Why This Approach

1. **Wait, don't just signal**: The key fix is awaiting the simulation
promise after signaling
2. **Bounded wait**: 100ms timeout prevents indefinite blocking if C++
is stuck
3. **Check before writes**: Prevents partial/corrupted state from
cancelled operations
4. **Mutex on checkpoints**: Prevents races between C++ writes and TS
checkpoint ops
6. **Minimal overhead**: `std::atomic<bool>` check is very cheap (single
memory read)
7. **Backward compatible**: Token is optional, existing code works
unchanged
8. **Testable**: Paired tests definitively prove both the bug and the
fix

## Future Work

### **Move merkle tree guarding into C++** (out of scope for this PR)

Currently, `GuardedMerkleTreeOperations` in TypeScript wraps merkle tree
operations to prevent access after `stop()` is called. However, this
guard is ineffective for C++ because:

1. C++ obtains the native WorldState and bypasses TS guarding
2. C++ can still make writes after TS calls `stop()`

### Guard getRevision

I think this is lower-impact, but ideally getRevision should fail if
called on a stopped instance.

---

🤖 Description generated with [Claude
Code](https://claude.com/claude-code)

Author: dbanks12

barretenberg/cpp/src/barretenberg/crypto/merkle_tree/node_store/cached_content_addressed_tree_store.hpp

@@ -273,31 +273,36 @@ ContentAddressedCachedTreeStore<LeafValueType>::ContentAddressedCachedTreeStore(
 }
 
 // Much Like the commit/rollback/set finalized/remove historic blocks apis
-// These 3 apis (checkpoint/revert_checkpoint/commit_checkpoint) all assume they are not called
-// during the process of reading/writing uncommitted state
-// This is reasonable, they intended for use by forks at the point of starting/ending a function call
+// These checkpoint apis modify the cache's internal state.
+// They acquire the mutex to prevent races with concurrent read/write operations (e.g., when C++ AVM simulation
+// runs on a worker thread while TypeScript calls revert_checkpoint from a timeout handler).
 template <typename LeafValueType> void ContentAddressedCachedTreeStore<LeafValueType>::checkpoint()
 {
+    std::unique_lock lock(mtx_);
     cache_.checkpoint();
 }
 
 template <typename LeafValueType> void ContentAddressedCachedTreeStore<LeafValueType>::revert_checkpoint()
 {
+    std::unique_lock lock(mtx_);
     cache_.revert();
 }
 
 template <typename LeafValueType> void ContentAddressedCachedTreeStore<LeafValueType>::commit_checkpoint()
 {
+    std::unique_lock lock(mtx_);
     cache_.commit();
 }
 
 template <typename LeafValueType> void ContentAddressedCachedTreeStore<LeafValueType>::revert_all_checkpoints()
 {
+    std::unique_lock lock(mtx_);
     cache_.revert_all();
 }
 
 template <typename LeafValueType> void ContentAddressedCachedTreeStore<LeafValueType>::commit_all_checkpoints()
 {
+    std::unique_lock lock(mtx_);
     cache_.commit_all();
 }
 

barretenberg/cpp/src/barretenberg/nodejs_module/avm_simulate/avm_simulate_napi.cpp

@@ -10,6 +10,7 @@
 #include "barretenberg/serialize/msgpack_impl/msgpack_impl.hpp"
 #include "barretenberg/vm2/avm_sim_api.hpp"
 #include "barretenberg/vm2/common/avm_io.hpp"
+#include "barretenberg/vm2/simulation/lib/cancellation_token.hpp"
 
 namespace bb::nodejs {
 
@@ -116,15 +117,16 @@ Napi::Value AvmSimulateNapi::simulate(const Napi::CallbackInfo& cb_info)
 {
     Napi::Env env = cb_info.Env();
 
-    // Validate arguments - expects 4 arguments
+    // Validate arguments - expects 4-5 arguments
     // arg[0]: inputs Buffer (required)
     // arg[1]: contractProvider object (required)
     // arg[2]: worldStateHandle external (required)
     // arg[3]: logLevel number (required) - index into TS LogLevels array
+    // arg[4]: cancellationToken external (optional)
     if (cb_info.Length() < 4) {
         throw Napi::TypeError::New(env,
-                                   "Wrong number of arguments. Expected 4 arguments: inputs Buffer, contractProvider "
-                                   "object, worldStateHandle, and logLevel.");
+                                   "Wrong number of arguments. Expected 4-5 arguments: inputs Buffer, contractProvider "
+                                   "object, worldStateHandle, logLevel, and optional cancellationToken.");
     }
 
     if (!cb_info[0].IsBuffer()) {
@@ -144,6 +146,19 @@ Napi::Value AvmSimulateNapi::simulate(const Napi::CallbackInfo& cb_info)
         throw Napi::TypeError::New(env, "Fourth argument must be a log level number (0-7)");
     }
 
+    // Extract optional cancellation token (5th argument)
+    avm2::simulation::CancellationTokenPtr cancellation_token = nullptr;
+    if (cb_info.Length() > 4 && cb_info[4].IsExternal()) {
+        auto token_external = cb_info[4].As<Napi::External<avm2::simulation::CancellationToken>>();
+        // Wrap the raw pointer in a shared_ptr that does NOT delete (since the External owns it)
+        cancellation_token = std::shared_ptr<avm2::simulation::CancellationToken>(
+            token_external.Data(), [](avm2::simulation::CancellationToken*) {
+                // No-op deleter: the External
+                // (via shared_ptr destructor
+                // callback) owns the token
+            });
+    }
+
     // Extract log level and set logging flags
     int log_level = cb_info[3].As<Napi::Number>().Int32Value();
     set_logging_from_level(log_level);
@@ -189,42 +204,46 @@ Napi::Value AvmSimulateNapi::simulate(const Napi::CallbackInfo& cb_info)
     auto deferred = std::make_shared<Napi::Promise::Deferred>(env);
 
     // Create async operation that will run on a worker thread
-    auto* op = new AsyncOperation(env, deferred, [data, tsfns, ws_ptr](msgpack::sbuffer& result_buffer) {
-        // Ensure all thread-safe functions are released in all code paths
-        TsfnReleaser releaser = TsfnReleaser(tsfns.to_vector());
-
-        try {
-            // Deserialize inputs from msgpack
-            avm2::AvmFastSimulationInputs inputs;
-            msgpack::object_handle obj_handle =
-                msgpack::unpack(reinterpret_cast<const char*>(data->data()), data->size());
-            msgpack::object obj = obj_handle.get();
-            obj.convert(inputs);
-
-            // Create TsCallbackContractDB with TypeScript callbacks
-            TsCallbackContractDB contract_db(*tsfns.instance,
-                                             *tsfns.class_,
-                                             *tsfns.add_contracts,
-                                             *tsfns.bytecode,
-                                             *tsfns.debug_name,
-                                             *tsfns.create_checkpoint,
-                                             *tsfns.commit_checkpoint,
-                                             *tsfns.revert_checkpoint);
-
-            // Create AVM API and run simulation with the callback-based contracts DB and
-            // WorldState reference
-            avm2::AvmSimAPI avm;
-            avm2::TxSimulationResult result = avm.simulate(inputs, contract_db, *ws_ptr);
-
-            // Serialize the simulation result with msgpack into the return buffer to TS.
-            msgpack::pack(result_buffer, result);
-        } catch (const std::exception& e) {
-            // Rethrow with context (RAII wrappers will clean up automatically)
-            throw std::runtime_error(std::string("AVM simulation failed: ") + e.what());
-        } catch (...) {
-            throw std::runtime_error("AVM simulation failed with unknown exception");
-        }
-    });
+    auto* op =
+        new AsyncOperation(env, deferred, [data, tsfns, ws_ptr, cancellation_token](msgpack::sbuffer& result_buffer) {
+            // Ensure all thread-safe functions are released in all code paths
+            TsfnReleaser releaser = TsfnReleaser(tsfns.to_vector());
+
+            try {
+                // Deserialize inputs from msgpack
+                avm2::AvmFastSimulationInputs inputs;
+                msgpack::object_handle obj_handle =
+                    msgpack::unpack(reinterpret_cast<const char*>(data->data()), data->size());
+                msgpack::object obj = obj_handle.get();
+                obj.convert(inputs);
+
+                // Create TsCallbackContractDB with TypeScript callbacks
+                TsCallbackContractDB contract_db(*tsfns.instance,
+                                                 *tsfns.class_,
+                                                 *tsfns.add_contracts,
+                                                 *tsfns.bytecode,
+                                                 *tsfns.debug_name,
+                                                 *tsfns.create_checkpoint,
+                                                 *tsfns.commit_checkpoint,
+                                                 *tsfns.revert_checkpoint);
+
+                // Create AVM API and run simulation with the callback-based contracts DB,
+                // WorldState reference, and optional cancellation token
+                avm2::AvmSimAPI avm;
+                avm2::TxSimulationResult result = avm.simulate(inputs, contract_db, *ws_ptr, cancellation_token);
+
+                // Serialize the simulation result with msgpack into the return buffer to TS.
+                msgpack::pack(result_buffer, result);
+            } catch (const avm2::simulation::CancelledException& e) {
+                // Cancellation is an expected condition, rethrow with context
+                throw std::runtime_error("Simulation cancelled");
+            } catch (const std::exception& e) {
+                // Rethrow with context (RAII wrappers will clean up automatically)
+                throw std::runtime_error(std::string("AVM simulation failed: ") + e.what());
+            } catch (...) {
+                throw std::runtime_error("AVM simulation failed with unknown exception");
+            }
+        });
 
     // Napi is now responsible for destroying this object
     op->Queue();
@@ -299,4 +318,34 @@ Napi::Value AvmSimulateNapi::simulateWithHintedDbs(const Napi::CallbackInfo& cb_
     return deferred->Promise();
 }
 
+Napi::Value AvmSimulateNapi::createCancellationToken(const Napi::CallbackInfo& cb_info)
+{
+    Napi::Env env = cb_info.Env();
+
+    // Create a new CancellationToken. We use a shared_ptr to manage the lifetime,
+    // and the destructor callback in the External will clean it up when GC runs.
+    auto* token = new avm2::simulation::CancellationToken();
+
+    // Create an External with a destructor callback that deletes the token
+    return Napi::External<avm2::simulation::CancellationToken>::New(
+        env, token, [](Napi::Env /*env*/, avm2::simulation::CancellationToken* t) { delete t; });
+}
+
+Napi::Value AvmSimulateNapi::cancelSimulation(const Napi::CallbackInfo& cb_info)
+{
+    Napi::Env env = cb_info.Env();
+
+    if (cb_info.Length() < 1 || !cb_info[0].IsExternal()) {
+        throw Napi::TypeError::New(env, "Expected a CancellationToken External as argument");
+    }
+
+    auto token_external = cb_info[0].As<Napi::External<avm2::simulation::CancellationToken>>();
+    avm2::simulation::CancellationToken* token = token_external.Data();
+
+    // Signal cancellation - this is thread-safe (atomic store)
+    token->cancel();
+
+    return env.Undefined();
+}
+
 } // namespace bb::nodejs

barretenberg/cpp/src/barretenberg/nodejs_module/avm_simulate/avm_simulate_napi.hpp

@@ -26,13 +26,16 @@ class AvmSimulateNapi {
      *   - getContractInstance(address: string): Promise<Buffer | undefined>
      *   - getContractClass(classId: string): Promise<Buffer | undefined>
      * - info[2]: External WorldState handle (pointer to world_state::WorldState)
+     * - info[3]: Log level number (0-7)
+     * - info[4]: External CancellationToken handle (optional)
      *
      * Returns: Promise<Buffer> containing serialized simulation results
      *
      * @param info NAPI callback info containing arguments
      * @return Napi::Value Promise that resolves with simulation results
      */
     static Napi::Value simulate(const Napi::CallbackInfo& info);
+
     /**
      * @brief NAPI function to simulate AVM execution with pre-collected hints
      *
@@ -43,6 +46,27 @@ class AvmSimulateNapi {
      * @return Napi::Value Promise that resolves with simulation results
      */
     static Napi::Value simulateWithHintedDbs(const Napi::CallbackInfo& info);
+
+    /**
+     * @brief Create a cancellation token that can be used to cancel a simulation.
+     *
+     * Returns: External<CancellationToken> - a handle to a new cancellation token
+     *
+     * @param info NAPI callback info (no arguments expected)
+     * @return Napi::Value External handle to the cancellation token
+     */
+    static Napi::Value createCancellationToken(const Napi::CallbackInfo& info);
+
+    /**
+     * @brief Cancel a simulation by signaling the provided cancellation token.
+     *
+     * Expected arguments:
+     * - info[0]: External CancellationToken handle
+     *
+     * @param info NAPI callback info containing the token
+     * @return Napi::Value undefined
+     */
+    static Napi::Value cancelSimulation(const Napi::CallbackInfo& info);
 };
 
 } // namespace bb::nodejs

barretenberg/cpp/src/barretenberg/nodejs_module/init_module.cpp

@@ -16,6 +16,10 @@ Napi::Object Init(Napi::Env env, Napi::Object exports)
     exports.Set(Napi::String::New(env, "avmSimulate"), Napi::Function::New(env, bb::nodejs::AvmSimulateNapi::simulate));
     exports.Set(Napi::String::New(env, "avmSimulateWithHintedDbs"),
                 Napi::Function::New(env, bb::nodejs::AvmSimulateNapi::simulateWithHintedDbs));
+    exports.Set(Napi::String::New(env, "createCancellationToken"),
+                Napi::Function::New(env, bb::nodejs::AvmSimulateNapi::createCancellationToken));
+    exports.Set(Napi::String::New(env, "cancelSimulation"),
+                Napi::Function::New(env, bb::nodejs::AvmSimulateNapi::cancelSimulation));
     return exports;
 }
 

barretenberg/cpp/src/barretenberg/vm2/avm_sim_api.cpp

@@ -10,7 +10,8 @@ using namespace bb::avm2::simulation;
 
 TxSimulationResult AvmSimAPI::simulate(const FastSimulationInputs& inputs,
                                        simulation::ContractDBInterface& contract_db,
-                                       world_state::WorldState& ws)
+                                       world_state::WorldState& ws,
+                                       simulation::CancellationTokenPtr cancellation_token)
 {
     vinfo("Simulating...");
     AvmSimulationHelper simulation_helper;
@@ -23,7 +24,8 @@ TxSimulationResult AvmSimAPI::simulate(const FastSimulationInputs& inputs,
                                                                                inputs.config,
                                                                                inputs.tx,
                                                                                inputs.global_variables,
-                                                                               inputs.protocol_contracts));
+                                                                               inputs.protocol_contracts,
+                                                                               cancellation_token));
     } else {
         return AVM_TRACK_TIME_V("simulation/all",
                                 simulation_helper.simulate_fast_with_existing_ws(contract_db,
@@ -32,7 +34,8 @@ TxSimulationResult AvmSimAPI::simulate(const FastSimulationInputs& inputs,
                                                                                  inputs.config,
                                                                                  inputs.tx,
                                                                                  inputs.global_variables,
-                                                                                 inputs.protocol_contracts));
+                                                                                 inputs.protocol_contracts,
+                                                                                 cancellation_token));
     }
 }
 

barretenberg/cpp/src/barretenberg/vm2/avm_sim_api.hpp

@@ -2,6 +2,7 @@
 
 #include "barretenberg/vm2/common/avm_io.hpp"
 #include "barretenberg/vm2/simulation/interfaces/db.hpp"
+#include "barretenberg/vm2/simulation/lib/cancellation_token.hpp"
 
 namespace bb::avm2 {
 
@@ -14,7 +15,8 @@ class AvmSimAPI {
 
     TxSimulationResult simulate(const FastSimulationInputs& inputs,
                                 simulation::ContractDBInterface& contract_db,
-                                world_state::WorldState& ws);
+                                world_state::WorldState& ws,
+                                simulation::CancellationTokenPtr cancellation_token = nullptr);
     TxSimulationResult simulate_with_hinted_dbs(const AvmProvingInputs& inputs);
 };
 

barretenberg/cpp/src/barretenberg/vm2/simulation/gadgets/execution.cpp

@@ -1724,6 +1724,11 @@ EnqueuedCallResult Execution::execute(std::unique_ptr<ContextInterface> enqueued
     external_call_stack.push(std::move(enqueued_call_context));
 
     while (!external_call_stack.empty()) {
+        // Throws CancelledException if cancelled. No-op when cancellation_token_ is nullptr (non-NAPI paths).
+        if (cancellation_token_) {
+            cancellation_token_->check_and_throw();
+        }
+
         // We fix the context at this point. Even if the opcode changes the stack
         // we'll always use this in the loop.
         auto& context = *external_call_stack.top();

barretenberg/cpp/src/barretenberg/vm2/simulation/gadgets/execution.hpp

@@ -30,6 +30,7 @@
 #include "barretenberg/vm2/simulation/interfaces/poseidon2.hpp"
 #include "barretenberg/vm2/simulation/interfaces/sha256.hpp"
 #include "barretenberg/vm2/simulation/interfaces/to_radix.hpp"
+#include "barretenberg/vm2/simulation/lib/cancellation_token.hpp"
 #include "barretenberg/vm2/simulation/lib/execution_id_manager.hpp"
 #include "barretenberg/vm2/simulation/lib/instruction_info.hpp"
 #include "barretenberg/vm2/simulation/lib/serialization.hpp"
@@ -79,7 +80,8 @@ class Execution : public ExecutionInterface {
               EmitUnencryptedLogInterface& emit_unencrypted_log_component,
               DebugLoggerInterface& debug_log_component,
               HighLevelMerkleDBInterface& merkle_db,
-              CallStackMetadataCollectorInterface& call_stack_metadata_collector)
+              CallStackMetadataCollectorInterface& call_stack_metadata_collector,
+              CancellationTokenPtr cancellation_token = nullptr)
         : execution_components(execution_components)
         , instruction_info_db(instruction_info_db)
         , alu(alu)
@@ -100,6 +102,7 @@ class Execution : public ExecutionInterface {
         , events(event_emitter)
         , ctx_stack_events(ctx_stack_emitter)
         , call_stack_metadata_collector(call_stack_metadata_collector)
+        , cancellation_token_(std::move(cancellation_token))
     {}
 
     EnqueuedCallResult execute(std::unique_ptr<ContextInterface> enqueued_call_context) override;
@@ -265,6 +268,10 @@ class Execution : public ExecutionInterface {
     std::vector<MemoryValue> inputs;
     MemoryValue output;
     std::unique_ptr<GasTrackerInterface> gas_tracker;
+
+    // Optional cancellation token for stopping simulation on timeout.
+    // When nullptr, cancellation checks are skipped (no overhead for non-NAPI paths).
+    CancellationTokenPtr cancellation_token_;
 };
 
 } // namespace bb::avm2::simulation