starknet_os_runner: allow fee token address overide

Author: AvivYossef-starkware

crates/blockifier_reexecution/src/state_reader/offline_state_reader.rs

@@ -118,7 +118,7 @@ impl From<SerializableOfflineReexecutionData> for OfflineReexecutionData {
             offline_state_reader_prev_block,
             block_context_next_block: BlockContext::new(
                 block_info_next_block,
-                get_chain_info(&chain_id),
+                get_chain_info(&chain_id, None),
                 versioned_constants,
                 BouncerConfig::max(),
             ),

crates/blockifier_reexecution/src/state_reader/rpc_state_reader.rs

@@ -120,6 +120,8 @@ pub struct RpcStateReader {
     pub block_id: BlockId,
     pub(crate) retry_config: RetryConfig,
     pub chain_id: ChainId,
+    /// Optional override for the STRK fee token address.
+    pub strk_fee_token_address: Option<ContractAddress>,
     #[allow(dead_code)]
     pub(crate) contract_class_mapping_dumper: Arc<Mutex<Option<StarknetContractClassMapping>>>,
 }
@@ -132,6 +134,7 @@ impl Default for RpcStateReader {
             block_id: BlockId::Latest,
             retry_config: RetryConfig::default(),
             chain_id: ChainId::Mainnet,
+            strk_fee_token_address: None,
             contract_class_mapping_dumper: Arc::new(Mutex::new(None)),
         }
     }
@@ -153,15 +156,18 @@ impl RpcStateReader {
             block_id: BlockId::Number(block_number),
             retry_config: RetryConfig::default(),
             chain_id,
+            strk_fee_token_address: None,
             contract_class_mapping_dumper,
         }
     }
 
     /// Creates an RpcStateReader from a node URL, chain ID, and block ID.
+    /// Optionally accepts a STRK fee token address override for custom environments.
     pub fn new_with_config_from_url(
         node_url: String,
         chain_id: ChainId,
         block_id: BlockId,
+        strk_fee_token_address: Option<ContractAddress>,
     ) -> Self {
         let config = RpcStateReaderConfig::from_url(node_url);
         let contract_class_mapping_dumper = Arc::new(Mutex::new(None));
@@ -170,6 +176,7 @@ impl RpcStateReader {
             block_id,
             retry_config: RetryConfig::default(),
             chain_id,
+            strk_fee_token_address,
             contract_class_mapping_dumper,
         }
     }
@@ -301,7 +308,7 @@ impl RpcStateReader {
     pub fn get_block_context(&self) -> ReexecutionResult<BlockContext> {
         Ok(BlockContext::new(
             self.get_block_info()?,
-            get_chain_info(&self.chain_id),
+            get_chain_info(&self.chain_id, self.strk_fee_token_address),
             self.get_versioned_constants()?,
             BouncerConfig::max(),
         ))

crates/blockifier_reexecution/src/utils.rs

@@ -46,12 +46,18 @@ pub fn contract_class_to_compiled_classes(
     }
 }
 
-/// Returns the fee token addresses of mainnet.
-pub fn get_fee_token_addresses(chain_id: &ChainId) -> FeeTokenAddresses {
+/// Returns the fee token addresses for the given chain.
+/// If `strk_fee_token_address_override` is provided, it overrides the default STRK fee token
+/// address.
+pub fn get_fee_token_addresses(
+    chain_id: &ChainId,
+    strk_fee_token_address_override: Option<ContractAddress>,
+) -> FeeTokenAddresses {
     match chain_id {
         // Mainnet, testnet and integration systems have the same fee token addresses.
         ChainId::Mainnet | ChainId::Sepolia | ChainId::IntegrationSepolia => FeeTokenAddresses {
-            strk_fee_token_address: *STRK_FEE_CONTRACT_ADDRESS,
+            strk_fee_token_address: strk_fee_token_address_override
+                .unwrap_or(*STRK_FEE_CONTRACT_ADDRESS),
             eth_fee_token_address: *ETH_FEE_CONTRACT_ADDRESS,
         },
         unknown_chain => unimplemented!("Unknown chain ID {unknown_chain}."),
@@ -63,11 +69,16 @@ pub fn get_rpc_state_reader_config() -> RpcStateReaderConfig {
     RpcStateReaderConfig::from_url(RPC_NODE_URL.clone())
 }
 
-/// Returns the chain info of mainnet.
-pub fn get_chain_info(chain_id: &ChainId) -> ChainInfo {
+/// Returns the chain info for the given chain.
+/// If `strk_fee_token_address_override` is provided, it overrides the default STRK fee token
+/// address.
+pub fn get_chain_info(
+    chain_id: &ChainId,
+    strk_fee_token_address_override: Option<ContractAddress>,
+) -> ChainInfo {
     ChainInfo {
         chain_id: chain_id.clone(),
-        fee_token_addresses: get_fee_token_addresses(chain_id),
+        fee_token_addresses: get_fee_token_addresses(chain_id, strk_fee_token_address_override),
         is_l3: false,
     }
 }

crates/starknet_os_runner/src/config.rs

@@ -2,7 +2,7 @@
 
 use blockifier::blockifier::config::ContractClassManagerConfig;
 use serde::{Deserialize, Serialize};
-use starknet_api::core::ChainId;
+use starknet_api::core::{ChainId, ContractAddress};
 
 use crate::runner::RunnerConfig;
 
@@ -18,6 +18,8 @@ pub struct ProverConfig {
     pub rpc_node_url: String,
     /// Configuration for the runner.
     pub runner_config: RunnerConfig,
+    /// Optional override for the STRK fee token address (e.g., for custom environments).
+    pub strk_fee_token_address: Option<ContractAddress>,
 }
 
 impl Default for ProverConfig {
@@ -27,6 +29,7 @@ impl Default for ProverConfig {
             chain_id: ChainId::Mainnet,
             rpc_node_url: String::new(),
             runner_config: RunnerConfig::default(),
+            strk_fee_token_address: None,
         }
     }
 }

crates/starknet_os_runner/src/runner.rs

@@ -384,6 +384,8 @@ pub(crate) struct RpcRunnerFactory {
     contract_class_manager: ContractClassManager,
     /// Configuration for the runner.
     runner_config: RunnerConfig,
+    /// Optional override for the STRK fee token address (e.g., for custom environments).
+    strk_fee_token_address: Option<ContractAddress>,
 }
 
 impl RpcRunnerFactory {
@@ -393,8 +395,9 @@ impl RpcRunnerFactory {
         chain_id: ChainId,
         contract_class_manager: ContractClassManager,
         runner_config: RunnerConfig,
+        strk_fee_token_address: Option<ContractAddress>,
     ) -> Self {
-        Self { node_url, chain_id, contract_class_manager, runner_config }
+        Self { node_url, chain_id, contract_class_manager, runner_config, strk_fee_token_address }
     }
 
     /// Creates a runner configured for the given block ID.
@@ -404,6 +407,7 @@ impl RpcRunnerFactory {
             self.node_url.to_string(),
             self.chain_id.clone(),
             block_id,
+            self.strk_fee_token_address,
         );
 
         // Create the storage proofs provider.
@@ -414,6 +418,7 @@ impl RpcRunnerFactory {
             self.node_url.to_string(),
             self.chain_id.clone(),
             block_id,
+            self.strk_fee_token_address,
         );
 
         // Wrap in StateReaderAndContractManager for class resolution.

crates/starknet_os_runner/src/server/config.rs

@@ -2,10 +2,11 @@
 
 use std::net::{IpAddr, Ipv4Addr};
 use std::path::PathBuf;
+use std::str::FromStr;
 
 use clap::Parser;
 use serde::{Deserialize, Serialize};
-use starknet_api::core::ChainId;
+use starknet_api::core::{ChainId, ContractAddress};
 use tracing::info;
 
 use crate::config::ProverConfig;
@@ -117,6 +118,19 @@ impl ServiceConfig {
             }
         }
 
+        if let Some(hex_str) = args.strk_fee_token_address {
+            let strk_fee_token_address = ContractAddress::from_str(&hex_str).map_err(|e| {
+                ConfigError::InvalidArgument(format!("Invalid strk_fee_token_address: {}", e))
+            })?;
+            if Some(strk_fee_token_address) != config.prover_config.strk_fee_token_address {
+                info!(
+                    "CLI override: strk_fee_token_address: {:?} -> {:?}",
+                    config.prover_config.strk_fee_token_address, strk_fee_token_address
+                );
+                config.prover_config.strk_fee_token_address = Some(strk_fee_token_address);
+            }
+        }
+
         // Validate required fields.
         if config.prover_config.rpc_node_url.is_empty() {
             return Err(ConfigError::MissingRequiredField(
@@ -170,6 +184,10 @@ pub struct CliArgs {
     /// Maximum number of simultaneous JSON-RPC connections (default: 10).
     #[arg(long, value_name = "N")]
     pub max_connections: Option<u32>,
+
+    /// Override STRK fee token address (hex, e.g. for custom environments that share a chain ID).
+    #[arg(long, value_name = "ADDRESS")]
+    pub strk_fee_token_address: Option<String>,
 }
 
 /// Errors that can occur during configuration.

crates/starknet_os_runner/src/test_utils.rs

@@ -35,6 +35,7 @@ pub fn rpc_state_reader() -> RpcStateReader {
         node_url,
         ChainId::Mainnet,
         BlockId::Number(BlockNumber(TEST_BLOCK_NUMBER)),
+        None,
     )
 }
 

crates/starknet_os_runner/src/virtual_block_executor.rs

@@ -22,7 +22,7 @@ use blockifier_reexecution::utils::get_chain_info;
 use starknet_api::block::{BlockHash, BlockInfo};
 use starknet_api::block_hash::block_hash_calculator::{concat_counts, BlockHeaderCommitments};
 use starknet_api::contract_class::SierraVersion;
-use starknet_api::core::{ChainId, ClassHash};
+use starknet_api::core::{ChainId, ClassHash, ContractAddress};
 use starknet_api::transaction::fields::Fee;
 use starknet_api::transaction::{InvokeTransaction, MessageToL1, Transaction, TransactionHash};
 use starknet_api::versioned_constants_logic::VersionedConstantsTrait;
@@ -46,10 +46,12 @@ pub(crate) struct BaseBlockInfo {
     pub(crate) prev_base_block_hash: BlockHash,
 }
 
-impl TryFrom<(BlockHeader, ChainId)> for BaseBlockInfo {
+impl TryFrom<(BlockHeader, ChainId, Option<ContractAddress>)> for BaseBlockInfo {
     type Error = VirtualBlockExecutorError;
 
-    fn try_from((header, chain_id): (BlockHeader, ChainId)) -> Result<Self, Self::Error> {
+    fn try_from(
+        (header, chain_id, strk_fee_token_address): (BlockHeader, ChainId, Option<ContractAddress>),
+    ) -> Result<Self, Self::Error> {
         let base_block_hash = header.block_hash;
         let prev_base_block_hash = header.parent_hash;
         let base_block_header_commitments = BlockHeaderCommitments {
@@ -70,7 +72,7 @@ impl TryFrom<(BlockHeader, ChainId)> for BaseBlockInfo {
                 "Failed to convert block header to block info: {e}"
             ))
         })?;
-        let chain_info = get_chain_info(&chain_id);
+        let chain_info = get_chain_info(&chain_id, strk_fee_token_address);
         let mut versioned_constants = VersionedConstants::get(&block_info.starknet_version)
             .map_err(|e| {
                 VirtualBlockExecutorError::TransactionExecutionError(format!(
@@ -316,10 +318,18 @@ pub(crate) struct RpcVirtualBlockExecutor {
 
 #[allow(dead_code)]
 impl RpcVirtualBlockExecutor {
-    pub(crate) fn new(node_url: String, chain_id: ChainId, block_id: BlockId) -> Self {
+    pub(crate) fn new(
+        node_url: String,
+        chain_id: ChainId,
+        block_id: BlockId,
+        strk_fee_token_address: Option<ContractAddress>,
+    ) -> Self {
         Self {
             rpc_state_reader: RpcStateReader::new_with_config_from_url(
-                node_url, chain_id, block_id,
+                node_url,
+                chain_id,
+                block_id,
+                strk_fee_token_address,
             ),
             validate_txs: true,
         }
@@ -340,7 +350,11 @@ impl VirtualBlockExecutor for RpcVirtualBlockExecutor {
             .rpc_state_reader
             .get_block_header()
             .map_err(|e| VirtualBlockExecutorError::ReexecutionError(Box::new(e)))?;
-        BaseBlockInfo::try_from((block_header, self.rpc_state_reader.chain_id.clone()))
+        BaseBlockInfo::try_from((
+            block_header,
+            self.rpc_state_reader.chain_id.clone(),
+            self.rpc_state_reader.strk_fee_token_address,
+        ))
     }
 
     fn state_reader(

crates/starknet_os_runner/src/virtual_snos_prover.rs

@@ -97,6 +97,7 @@ impl VirtualSnosProver<RpcRunnerFactory> {
             prover_config.chain_id.clone(),
             contract_class_manager,
             prover_config.runner_config.clone(),
+            prover_config.strk_fee_token_address,
         );
         Self { runner }
     }