feat(simulator): opt-in state-override stream for pAMM-aware settlement simulation

Cargo.toml

@@ -145,6 +145,7 @@ time = "0.3.47"
 token-info = { path = "crates/token-info" }
 tokio = { version = "1.38.0", features = ["tracing"] }
 tokio-stream = { version = "0.1.15", features = ["sync"] }
+tokio-tungstenite = { version = "0.28", default-features = false }
 toml = "0.8.14"
 tower = "0.5"
 tower-http = "0.6"

crates/configs/src/price_estimation.rs

@@ -109,6 +109,14 @@ pub struct PriceEstimation {
     #[serde(default)]
     pub tenderly: Option<crate::simulator::TenderlyConfig>,
 
+    /// Optional background websocket stream of eth_call-style state override
+    /// sets, applied during quote/trade verification simulation. Same wire
+    /// format and semantics as the driver's
+    /// `[simulator.state-override-stream]`. When absent, no stream task is
+    /// spawned and behavior is unchanged.
+    #[serde(default)]
+    pub state_override_stream: Option<crate::simulator::StateOverrideStream>,
+
     /// The CoinGecko native price configuration.
     pub coin_gecko: Option<CoinGeckoConfig>,
 
@@ -120,6 +128,7 @@ impl Default for PriceEstimation {
     fn default() -> Self {
         Self {
             tenderly: Default::default(),
+            state_override_stream: Default::default(),
             price_estimation_rate_limiter: None,
             amount_to_estimate_prices_with: None,
             one_inch: Default::default(),

crates/configs/src/simulator.rs

@@ -30,6 +30,31 @@ pub struct Config {
     /// - tenderly (using TenderlyConfig)
     #[serde(default)]
     pub kind: SimulatorKind,
+
+    /// Optional background websocket stream of eth_call-style state override
+    /// sets, applied on top of latest state during settlement gas estimation.
+    /// When absent, no stream task is spawned and behavior is unchanged.
+    #[serde(default)]
+    pub state_override_stream: Option<StateOverrideStream>,
+}
+
+/// Configuration for a background websocket stream delivering eth_call-style
+/// state override sets (the standard Ethereum State Override Set wire format),
+/// so live-quote venues (e.g. pAMMs) simulate against current in-memory state
+/// rather than stale previous-block state.
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(rename_all = "kebab-case", deny_unknown_fields)]
+pub struct StateOverrideStream {
+    /// Websocket URL pushing JSON frames of eth_call-style state overrides.
+    pub ws_url: Url,
+
+    /// Ignore the snapshot if the last frame is older than this.
+    #[serde(with = "humantime_serde", default = "default_override_max_age")]
+    pub max_age: Duration,
+}
+
+const fn default_override_max_age() -> Duration {
+    Duration::from_secs(5)
 }
 
 fn default_ethrpc_batch_delay() -> Duration {
@@ -207,4 +232,44 @@ mod tests {
             _ => panic!("Config should be of type Tenderly"),
         };
     }
+
+    #[test]
+    fn deserialize_state_override_stream_absent_by_default() {
+        let config: Config = toml::from_str("").unwrap();
+        assert!(config.state_override_stream.is_none());
+    }
+
+    #[test]
+    fn deserialize_state_override_stream_with_defaults() {
+        let toml = r#"
+        [state-override-stream]
+        ws-url = "wss://example.com/stream"
+        "#;
+        let config: Config = toml::from_str(toml).unwrap();
+        let stream = config.state_override_stream.unwrap();
+        assert_eq!(stream.ws_url.as_str(), "wss://example.com/stream");
+        assert_eq!(stream.max_age, Duration::from_secs(5));
+    }
+
+    #[test]
+    fn deserialize_state_override_stream_full() {
+        let toml = r#"
+        [state-override-stream]
+        ws-url = "wss://example.com/stream"
+        max-age = "5s"
+        "#;
+        let config: Config = toml::from_str(toml).unwrap();
+        let stream = config.state_override_stream.unwrap();
+        assert_eq!(stream.max_age, Duration::from_secs(5));
+    }
+
+    #[test]
+    fn deserialize_state_override_stream_rejects_unknown_fields() {
+        let toml = r#"
+        [state-override-stream]
+        ws-url = "wss://example.com/stream"
+        bogus = true
+        "#;
+        assert!(toml::from_str::<Config>(toml).is_err());
+    }
 }

crates/driver/example.toml

@@ -111,3 +111,12 @@ max-order-age = "1m"
 # base-url = "https://api.liquorice.tech/"
 # api-key = "..."
 # http_timeout = "10s"
+
+# [simulator.state-override-stream] # Optional background websocket stream of
+# # eth_call-style state override sets (standard Ethereum State Override Set wire
+# # format), so live-quote venues (e.g. pAMMs) simulate against current in-memory
+# # state. When absent, no stream task is spawned and behavior is unchanged.
+# # Requires a node supporting eth_estimateGas state/block overrides (geth
+# # >=1.13, reth, nethermind, recent erigon, anvil).
+# ws-url = "wss://eu.rpc.titanbuilder.xyz/ws/pamm_quote_stream"
+# max-age = "5s"        # ignore the snapshot if the last frame is older than this

crates/driver/src/run.rs

@@ -162,6 +162,7 @@ fn simulator(
     eth: simulator::Ethereum,
     http_factory: &HttpClientFactory,
 ) -> Simulator {
+    let block_stream = eth.current_block().clone();
     let mut simulator = match &config.simulator {
         configs::simulator::Config {
             kind: configs::simulator::SimulatorKind::Tenderly(config),
@@ -179,6 +180,10 @@ fn simulator(
     if let Some(gas) = config.disable_gas_simulation {
         simulator.disable_gas(gas);
     }
+    if let Some(cfg) = &config.simulator.state_override_stream {
+        simulator
+            .set_simulation_overrides(simulator::state_override_stream::spawn(cfg, block_stream));
+    }
 
     simulator
 }

crates/e2e/Cargo.toml

@@ -36,6 +36,7 @@ autopilot = { workspace = true, features = ["test-util"] }
 axum = { workspace = true }
 balance-overrides = { workspace = true }
 bigdecimal = { workspace = true }
+chain = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
 configs = { workspace = true, features = ["test-util"] }
@@ -47,6 +48,7 @@ driver = { workspace = true }
 eth-domain-types = { workspace = true }
 ethrpc = { workspace = true, features = ["test-util"] }
 futures = { workspace = true }
+gas-price-estimation = { workspace = true }
 hex-literal = { workspace = true }
 itertools = { workspace = true }
 liquidity-sources = { workspace = true }

crates/e2e/tests/e2e/main.rs

@@ -41,6 +41,7 @@ mod refunder;
 mod replace_order;
 mod smart_contract_orders;
 mod solver_competition;
+mod state_override;
 mod submission;
 mod token_metadata;
 mod tracking_insufficient_funds;

crates/e2e/tests/e2e/state_override.rs

@@ -0,0 +1,100 @@
+use {
+    alloy::{
+        primitives::{Address, B256, U256, address, hex, map::B256Map},
+        providers::{Provider, ProviderBuilder},
+        rpc::types::{
+            TransactionRequest,
+            state::{AccountOverride, StateOverride},
+        },
+    },
+    chain::Chain,
+    e2e::nodes::{NODE_HOST, Node},
+    ethrpc::Web3,
+    gas_price_estimation::FakeGasPriceEstimator,
+    simulator::Ethereum,
+    std::{str::FromStr, sync::Arc},
+};
+
+/// Runtime bytecode of a contract that reverts in `gate()` (selector
+/// `0x7a0ebc88`) unless storage slot 0 is nonzero. Compiled with forge.
+const SLOT_GATE_INIT: &str = "0x6080604052348015600e575f5ffd5b506101238061001c5f395ff3fe6080604052348015600e575f5ffd5b50600436106026575f3560e01c80637a0ebc8814602a575b5f5ffd5b60306032565b005b5f5f5490505f5f1b81036078576040517f08c379a0000000000000000000000000000000000000000000000000000000008152600401606f9060d1565b60405180910390fd5b50565b5f82825260208201905092915050565b7f736c6f74207a65726f00000000000000000000000000000000000000000000005f82015250565b5f60bd600983607b565b915060c682608b565b602082019050919050565b5f6020820190508181035f83015260e68160b3565b905091905056fea2646970667358221220538cd705c3b870d0412519ce867cae3b5484bd1b0acf57a8ec779199e45097ec64736f6c634300081c0033";
+const GATE_SELECTOR: &str = "0x7a0ebc88";
+const DEFAULT_SENDER: Address = address!("f39Fd6e51aad88F6F4ce6aB8827279cffFb92266");
+const SLOT_0: &str = "0x0000000000000000000000000000000000000000000000000000000000000000";
+
+async fn deploy() -> Address {
+    let provider = ProviderBuilder::new().connect_http(NODE_HOST.parse().unwrap());
+    let tx = TransactionRequest {
+        from: Some(DEFAULT_SENDER),
+        input: hex::decode(SLOT_GATE_INIT.trim_start_matches("0x"))
+            .unwrap()
+            .into(),
+        gas: Some(0x200000),
+        ..Default::default()
+    };
+    let pending = provider.send_transaction(tx).await.unwrap();
+    let receipt = pending.get_receipt().await.unwrap();
+    receipt.contract_address.unwrap()
+}
+
+fn gate_call(to: Address) -> TransactionRequest {
+    TransactionRequest {
+        from: Some(DEFAULT_SENDER),
+        to: Some(to.into()),
+        input: hex::decode(GATE_SELECTOR.trim_start_matches("0x"))
+            .unwrap()
+            .into(),
+        ..Default::default()
+    }
+}
+
+fn slot0_override(contract: Address) -> StateOverride {
+    let mut override_map = StateOverride::default();
+    let mut diff = B256Map::default();
+    diff.insert(B256::from_str(SLOT_0).unwrap(), B256::from(U256::from(1)));
+    override_map.insert(
+        contract,
+        AccountOverride {
+            state_diff: Some(diff),
+            ..Default::default()
+        },
+    );
+    override_map
+}
+
+async fn ethereum() -> Ethereum {
+    let web3 = Web3::new_from_url(NODE_HOST);
+    let block_stream = ethrpc::block_stream::mock_single_block(Default::default());
+    Ethereum::new(
+        web3,
+        Chain::Mainnet,
+        configs::simulator::Addresses::default(),
+        Arc::new(FakeGasPriceEstimator::default()),
+        block_stream,
+        U256::from(30_000_000),
+    )
+}
+
+#[tokio::test]
+#[ignore = "requires anvil; run via `just test-e2e local_node_estimate_gas_state_override`"]
+async fn local_node_estimate_gas_state_override() {
+    let mut node = Node::new().await;
+    let contract = deploy().await;
+    let eth = ethereum().await;
+
+    let without = eth.estimate_gas(gate_call(contract), None).await;
+    assert!(
+        without.is_err(),
+        "gate() should revert without the slot-0 override, got {without:?}",
+    );
+
+    let with = eth
+        .estimate_gas(gate_call(contract), Some(slot0_override(contract)))
+        .await;
+    assert!(
+        with.is_ok(),
+        "gate() should succeed with the slot-0 override, got {with:?}",
+    );
+
+    node.kill().await;
+}

crates/price-estimation/src/factory.rs

@@ -27,7 +27,7 @@ use {
     number::nonzero::NonZeroU256,
     rate_limit::RateLimiter,
     reqwest::Url,
-    simulator::{simulation_builder::SettlementSimulator, tenderly},
+    simulator::{self, simulation_builder::SettlementSimulator, tenderly},
     std::{collections::HashMap, num::NonZeroUsize, sync::Arc},
     token_info::TokenInfoFetching,
 };
@@ -119,6 +119,10 @@ impl<'a> PriceEstimatorFactory<'a> {
         });
         let settlement_contract =
             GPv2Settlement::GPv2Settlement::new(network.settlement, web3.provider.clone());
+        let simulation_overrides = args
+            .state_override_stream
+            .as_ref()
+            .map(|cfg| simulator::state_override_stream::spawn(cfg, network.block_stream.clone()));
         let simulator = SettlementSimulator::new(
             settlement_contract,
             network.flash_loan_router,
@@ -128,6 +132,7 @@ impl<'a> PriceEstimatorFactory<'a> {
             balance_overrides,
             network.block_stream.clone(),
             tenderly.clone(),
+            simulation_overrides,
         )
         .await?;
 

crates/simulator/Cargo.toml

@@ -43,11 +43,12 @@ serde_with = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true }
+tokio = { workspace = true, features = ["sync", "time", "rt", "macros"] }
+tokio-tungstenite = { workspace = true, features = ["connect", "rustls-tls-webpki-roots"] }
 
 [dev-dependencies]
 mockall = { workspace = true }
 testlib = { workspace = true }
-tokio = { workspace = true }
 
 [features]
 test-util = ["dep:mockall"]

crates/simulator/src/encoding.rs

@@ -490,6 +490,25 @@ pub(crate) async fn finish_simulation_builder(
         }
         None => return Err(BuildError::NoSolver),
     };
+
+    // Streamed overrides are appended as Custom requests so
+    // build_final_state_overrides' merge arbitrates against the verifier's own.
+    // Only latest-block simulations get overrides; historical-block sims replay
+    // against a pinned block whose pAMM state is not the live stream's.
+    if matches!(builder.block, Block::Latest)
+        && let Some(stream) = builder.simulator.0.simulation_overrides.as_ref()
+        && let Some(state_overrides) = stream.current()
+    {
+        for (account, state) in state_overrides.iter() {
+            builder
+                .account_override_requests
+                .push(AccountOverrideRequest::Custom {
+                    account: *account,
+                    state: state.clone(),
+                });
+        }
+    }
+
     let (state_overrides, state_override_errors) = build_final_state_overrides(
         builder.account_override_requests,
         builder.simulator.0.state_overrides.as_ref(),

crates/simulator/src/ethereum/mod.rs

@@ -2,7 +2,7 @@ use {
     crate::ethereum::contracts::Contracts,
     alloy_primitives::U256,
     alloy_provider::{Provider, network::TransactionBuilder},
-    alloy_rpc_types::TransactionRequest,
+    alloy_rpc_types::{TransactionRequest, state::StateOverride},
     anyhow::{Result, anyhow},
     chain::Chain,
     eth_domain_types::AccessList,
@@ -120,7 +120,11 @@ impl Ethereum {
             .into())
     }
 
-    pub async fn estimate_gas<T>(&self, tx: T) -> Result<eth_domain_types::Gas, Error>
+    pub async fn estimate_gas<T>(
+        &self,
+        tx: T,
+        overrides: Option<StateOverride>,
+    ) -> Result<eth_domain_types::Gas, Error>
     where
         T: Into<TransactionRequest>,
     {
@@ -130,10 +134,14 @@ impl Ethereum {
             _ => tx,
         };
 
+        // `None` omits the RPC param entirely (byte-identical to today).
+        // Requires a node supporting eth_estimateGas state overrides
+        // (geth >=1.13, reth, nethermind, recent erigon, anvil).
         let estimated_gas = self
             .web3
             .provider
             .estimate_gas(tx)
+            .overrides_opt(overrides)
             .pending()
             .await
             .map_err(Error::Rpc)?