Repository: stellar/rs-soroban-env
Pull request #1412 (2024-04-29 17:59:46)

Improved TTL testing UX and cleaned up storage diagnostics.
### What

Improved TTL testing UX and cleaned up storage diagnostics.

Both improvements rely on passing host into storage methods. This also allows us to use the ledger info and to encapsulate detailed diagnostic errors in the storage module.  We have to still keep the old storage methods around. There are a few users of storage getters, which probably don't have to use it, but the clean up will have some time and can only be performed in v22 due to semver.

There is a subtle change in observations because we now only allocate a new object in diagnostic mode (charged to shadow budget). However, this shouldn't cause divergence (because this version of env won't run in p20) and is highly unlikely to affect replay of p20 with the new env. If necessary, we can still emulate the p20 behavior.

The TTL testing support includes:

- Provide testutil getters for the `live_until_ledger` of all the ledger entries
- Treat expired temp entries as non-existent
- Emit special test-only errors for accessing expired persistent entries

### Why

Improving testing UX

### Known limitations

N/A



Patch:
diff --git a/soroban-env-host/src/auth.rs b/soroban-env-host/src/auth.rs
index a50bd379..67766a6b 100644
--- a/soroban-env-host/src/auth.rs
+++ b/soroban-env-host/src/auth.rs
@@ -2169,19 +2169,22 @@ impl Host {
         let live_until_ledger = live_until_ledger
             .max(self.get_min_live_until_ledger(xdr::ContractDataDurability::Temporary)?);
         self.with_mut_storage(|storage| {
-            if storage.has(&nonce_key, self.budget_ref()).map_err(|err| {
-                if err.error.is_type(ScErrorType::Storage)
-                    && err.error.is_code(ScErrorCode::ExceededLimit)
-                {
-                    return self.err(
-                        ScErrorType::Storage,
-                        ScErrorCode::ExceededLimit,
-                        "trying to access nonce outside of footprint for address",
-                        &[address.to_val()],
-                    );
-                }
-                err
-            })? {
+            if storage
+                .has_with_host(&nonce_key, self, None)
+                .map_err(|err| {
+                    if err.error.is_type(ScErrorType::Storage)
+                        && err.error.is_code(ScErrorCode::ExceededLimit)
+                    {
+                        return self.err(
+                            ScErrorType::Storage,
+                            ScErrorCode::ExceededLimit,
+                            "trying to access nonce outside of footprint for address",
+                            &[address.to_val()],
+                        );
+                    }
+                    err
+                })?
+            {
                 return Err(self.err(
                     ScErrorType::Auth,
                     ScErrorCode::ExistingValue,
@@ -2201,11 +2204,12 @@ impl Host {
                 data,
                 ext: LedgerEntryExt::V0,
             };
-            storage.put(
+            storage.put_with_host(
                 &nonce_key,
                 &Rc::metered_new(entry, self)?,
                 Some(live_until_ledger),
-                self.budget_ref(),
+                self,
+                None,
             )
         })
     }
diff --git a/soroban-env-host/src/builtin_contracts/stellar_asset_contract/balance.rs b/soroban-env-host/src/builtin_contracts/stellar_asset_contract/balance.rs
index 75a088b8..c13c0ebb 100644
--- a/soroban-env-host/src/builtin_contracts/stellar_asset_contract/balance.rs
+++ b/soroban-env-host/src/builtin_contracts/stellar_asset_contract/balance.rs
@@ -1,7 +1,6 @@
 use std::rc::Rc;
 
 use crate::{
-    budget::AsBudget,
     builtin_contracts::{
         base_types::{Address, BytesN},
         contract_error::ContractError,
@@ -418,7 +417,7 @@ fn transfer_account_balance(
         if new_balance >= min_balance && new_balance <= max_balance {
             ae.balance = new_balance;
             le = Host::modify_ledger_entry_data(host, &le, LedgerEntryData::Account(ae))?;
-            storage.put(&lk, &le, None, host.as_budget())
+            storage.put_with_host(&lk, &le, None, &host, None)
         } else {
             Err(err!(
                 host,
@@ -437,26 +436,17 @@ fn read_account_entry(
     storage: &mut Storage,
     lk: &Rc<LedgerKey>,
 ) -> Result<Rc<LedgerEntry>, HostError> {
-    storage
-        .try_get(&lk, host.as_budget())
-        .map_err(|e| {
-            host.decorate_account_footprint_error(
-                e,
-                lk,
-                "trying to access account entry outside of the footprint",
-            )
-        })?
-        .ok_or_else(|| {
-            let account_address = host.account_address_from_key(lk);
-            match account_address {
-                Ok(account_address) => host.error(
-                    ContractError::AccountMissingError.into(),
-                    "account entry is missing",
-                    &[account_address],
-                ),
-                Err(e) => e,
-            }
-        })
+    storage.try_get(&lk, &host, None)?.ok_or_else(|| {
+        let account_address = host.account_address_from_key(lk);
+        match account_address {
+            Ok(account_address) => host.error(
+                ContractError::AccountMissingError.into(),
+                "account entry is missing",
+                &[account_address],
+            ),
+            Err(e) => e,
+        }
+    })
 }
 
 fn read_trustline_entry(
@@ -464,26 +454,17 @@ fn read_trustline_entry(
     storage: &mut Storage,
     lk: &Rc<LedgerKey>,
 ) -> Result<Rc<LedgerEntry>, HostError> {
-    storage
-        .try_get(&lk, host.as_budget())
-        .map_err(|e| {
-            host.decorate_account_footprint_error(
-                e,
-                lk,
-                "trying to access account trustline entry outside of the footprint",
-            )
-        })?
-        .ok_or_else(|| {
-            let account_address = host.account_address_from_key(lk);
-            match account_address {
-                Ok(account_address) => host.error(
-                    ContractError::TrustlineMissingError.into(),
-                    "trustline entry is missing for account",
-                    &[account_address],
-                ),
-                Err(e) => e,
-            }
-        })
+    storage.try_get(&lk, &host, None)?.ok_or_else(|| {
+        let account_address = host.account_address_from_key(lk);
+        match account_address {
+            Ok(account_address) => host.error(
+                ContractError::TrustlineMissingError.into(),
+                "trustline entry is missing for account",
+                &[account_address],
+            ),
+            Err(e) => e,
+        }
+    })
 }
 
 // Metering: covered by components.
@@ -518,7 +499,7 @@ fn transfer_trustline_balance(
         if new_balance >= min_balance && new_balance <= max_balance {
             tl.balance = new_balance;
             le = Host::modify_ledger_entry_data(host, &le, LedgerEntryData::Trustline(tl))?;
-            storage.put(&lk, &le, None, host.as_budget())
+            storage.put_with_host(&lk, &le, None, &host, None)
         } else {
             Err(err!(
                 host,
@@ -788,7 +769,7 @@ fn set_trustline_authorization(
             tl.flags |= TrustLineFlags::AuthorizedToMaintainLiabilitiesFlag as u32;
         }
         le = Host::modify_ledger_entry_data(host, &le, LedgerEntryData::Trustline(tl))?;
-        storage.put(&lk, &le, None, host.as_budget())
+        storage.put_with_host(&lk, &le, None, &host, None)
     })
 }
 
diff --git a/soroban-env-host/src/host.rs b/soroban-env-host/src/host.rs
index 10d905fc..4bed4c7e 100644
--- a/soroban-env-host/src/host.rs
+++ b/soroban-env-host/src/host.rs
@@ -54,6 +54,7 @@ use self::{
     prng::Prng,
 };
 
+use crate::host::error::TryBorrowOrErr;
 #[cfg(any(test, feature = "testutils"))]
 pub use frame::ContractFunctionSet;
 pub(crate) use frame::Frame;
@@ -577,7 +578,6 @@ impl Host {
     }
 
     pub fn set_diagnostic_level(&self, diagnostic_level: DiagnosticLevel) -> Result<(), HostError> {
-        use crate::host::error::TryBorrowOrErr;
         *self.0.diagnostic_level.try_borrow_mut_or_err()? = diagnostic_level;
         Ok(())
     }
@@ -610,7 +610,6 @@ impl Host {
     where
         F: FnOnce() -> Result<(), HostError>,
     {
-        use crate::host::error::TryBorrowOrErr;
         if let Ok(cell) = self.0.diagnostic_level.try_borrow_or_err() {
             if matches!(*cell, DiagnosticLevel::Debug) {
                 return self.budget_ref().with_shadow_mode(f);
@@ -2022,8 +2021,7 @@ impl VmCallerEnv for Host {
             StorageType::Temporary | StorageType::Persistent => {
                 let key = self.storage_key_from_val(k, t.try_into()?)?;
                 self.try_borrow_storage_mut()?
-                    .has(&key, self.as_budget())
-                    .map_err(|e| self.decorate_contract_data_storage_error(e, k))?
+                    .has_with_host(&key, self, Some(k))?
             }
             StorageType::Instance => {
                 self.with_instance_storage(|s| Ok(s.map.get(&k, self)?.is_some()))?
@@ -2045,8 +2043,7 @@ impl VmCallerEnv for Host {
                 let key = self.storage_key_from_val(k, t.try_into()?)?;
                 let entry = self
                     .try_borrow_storage_mut()?
-                    .get(&key, self.as_budget())
-                    .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+                    .get_with_host(&key, self, Some(k))?;
                 match &entry.data {
                     LedgerEntryData::ContractData(e) => Ok(self.to_valid_host_val(&e.val)?),
                     _ => Err(self.err(
@@ -2084,8 +2081,7 @@ impl VmCallerEnv for Host {
             StorageType::Temporary | StorageType::Persistent => {
                 let key = self.storage_key_from_val(k, t.try_into()?)?;
                 self.try_borrow_storage_mut()?
-                    .del(&key, self.as_budget())
-                    .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+                    .del_with_host(&key, self, Some(k))?;
             }
             StorageType::Instance => {
                 self.with_mut_instance_storage(|s| {
@@ -2118,9 +2114,13 @@ impl VmCallerEnv for Host {
             ))?;
         }
         let key = self.storage_key_from_val(k, t.try_into()?)?;
-        self.try_borrow_storage_mut()?
-            .extend_ttl(self, key, threshold.into(), extend_to.into())
-            .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+        self.try_borrow_storage_mut()?.extend_ttl(
+            self,
+            key,
+            threshold.into(),
+            extend_to.into(),
+            Some(k),
+        )?;
         Ok(Val::VOID)
     }
 
@@ -2133,7 +2133,6 @@ impl VmCallerEnv for Host {
         let contract_id = self.get_current_contract_id_internal()?;
         let key = self.contract_instance_ledger_key(&contract_id)?;
         self.extend_contract_instance_ttl_from_contract_id(
-            &contract_id,
             key.clone(),
             threshold.into(),
             extend_to.into(),
@@ -2153,7 +2152,6 @@ impl VmCallerEnv for Host {
         let key = self.contract_instance_ledger_key(&contract_id)?;
 
         self.extend_contract_instance_ttl_from_contract_id(
-            &contract_id,
             key,
             threshold.into(),
             extend_to.into(),
@@ -2172,7 +2170,6 @@ impl VmCallerEnv for Host {
         let contract_id = self.contract_id_from_address(contract)?;
         let key = self.contract_instance_ledger_key(&contract_id)?;
         self.extend_contract_instance_ttl_from_contract_id(
-            &contract_id,
             key.clone(),
             threshold.into(),
             extend_to.into(),
@@ -3122,6 +3119,99 @@ impl Host {
     {
         f(self.0.budget.clone())
     }
+
+    /// Returns the ledger number until a contract with given address lives
+    /// (inclusive).
+    pub fn get_contract_instance_live_until_ledger(
+        &self,
+        contract: AddressObject,
+    ) -> Result<u32, HostError> {
+        let contract_id = self.contract_id_from_address(contract)?;
+        let key = self.contract_instance_ledger_key(&contract_id)?;
+        let (_, live_until) = self
+            .try_borrow_storage_mut()?
+            .get_with_live_until_ledger(&key, self, None)?;
+        live_until.ok_or_else(|| {
+            self.err(
+                ScErrorType::Storage,
+                ScErrorCode::InternalError,
+                "unexpected contract instance without TTL",
+                &[contract.into()],
+            )
+        })
+    }
+
+    /// Returns the ledger number until contract code entry for contract with
+    /// given address lives (inclusive).
+    pub fn get_contract_code_live_until_ledger(
+        &self,
+        contract: AddressObject,
+    ) -> Result<u32, HostError> {
+        let contract_id = self.contract_id_from_address(contract)?;
+        let key = self.contract_instance_ledger_key(&contract_id)?;
+        match self
+            .retrieve_contract_instance_from_storage(&key)?
+            .executable
+        {
+            ContractExecutable::Wasm(wasm_hash) => {
+                let key = self.contract_code_ledger_key(&wasm_hash)?;
+                let (_, live_until) = self
+                    .try_borrow_storage_mut()?
+                    .get_with_live_until_ledger(&key, self, None)?;
+                live_until.ok_or_else(|| {
+                    self.err(
+                        ScErrorType::Storage,
+                        ScErrorCode::InternalError,
+                        "unexpected contract code without TTL for a contract",
+                        &[contract.into()],
+                    )
+                })
+            }
+            ContractExecutable::StellarAsset => Err(self.err(
+                ScErrorType::Storage,
+                ScErrorCode::InvalidInput,
+                "Stellar Asset Contracts don't have contract code",
+                &[],
+            )),
+        }
+    }
+
+    /// Returns the ledger number until a current contract's data entry
+    /// with given key and storage type lives (inclusive).
+    /// Instance storage type is not supported by this function, use
+    /// `get_contract_instance_live_until_ledger` instead.
+    pub fn get_contract_data_live_until_ledger(
+        &self,
+        key: Val,
+        storage_type: StorageType,
+    ) -> Result<u32, HostError> {
+        let ledger_key = match storage_type {
+            StorageType::Temporary | StorageType::Persistent => {
+                self.storage_key_from_val(key, storage_type.try_into()?)?
+            }
+            StorageType::Instance => {
+                return Err(self.err(
+                    ScErrorType::Storage,
+                    ScErrorCode::InvalidAction,
+                    "`get_contract_data_live_until_ledger` doesn't support instance storage, use `get_contract_instance_live_until_ledger` instead.",
+                    &[],
+                ));
+            }
+        };
+        let (_, live_until) = self.try_borrow_storage_mut()?.get_with_live_until_ledger(
+            &ledger_key,
+            self,
+            Some(key),
+        )?;
+        live_until.ok_or_else(|| {
+            self.err(
+                ScErrorType::Storage,
+                ScErrorCode::InternalError,
+                "unexpected contract data without TTL",
+                &[key],
+            )
+        })
+    }
 }
 
 impl Host {
diff --git a/soroban-env-host/src/host/data_helper.rs b/soroban-env-host/src/host/data_helper.rs
index a474c299..4d687062 100644
--- a/soroban-env-host/src/host/data_helper.rs
+++ b/soroban-env-host/src/host/data_helper.rs
@@ -92,7 +92,9 @@ impl Host {
         &self,
         key: &Rc<LedgerKey>,
     ) -> Result<ScContractInstance, HostError> {
-        let entry = self.try_borrow_storage_mut()?.get(key, self.as_budget())?;
+        let entry = self
+            .try_borrow_storage_mut()?
+            .get_with_host(key, self, None)?;
         match &entry.data {
             LedgerEntryData::ContractData(e) => match &e.val {
                 ScVal::ContractInstance(instance) => instance.metered_clone(self),
@@ -130,8 +132,7 @@ impl Host {
         let key = self.contract_code_ledger_key(wasm_hash)?;
         match &self
             .try_borrow_storage_mut()?
-            .get(&key, self.as_budget())
-            .map_err(|e| self.decorate_contract_code_storage_error(e, wasm_hash))?
+            .get_with_host(&key, self, None)?
             .data
         {
             LedgerEntryData::ContractCode(e) => {
@@ -160,8 +161,7 @@ impl Host {
     pub(crate) fn wasm_exists(&self, wasm_hash: &Hash) -> Result<bool, HostError> {
         let key = self.contract_code_ledger_key(wasm_hash)?;
         self.try_borrow_storage_mut()?
-            .has(&key, self.as_budget())
-            .map_err(|e| self.decorate_contract_code_storage_error(e, wasm_hash))
+            .has_with_host(&key, self, None)
     }
 
     // Stores the contract instance specified with its parts (executable and
@@ -179,12 +179,11 @@ impl Host {
     ) -> Result<(), HostError> {
         if self
             .try_borrow_storage_mut()?
-            .has(key, self.as_budget())
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?
+            .has_with_host(key, self, None)?
         {
             let (current, live_until_ledger) = self
                 .try_borrow_storage_mut()?
-                .get_with_live_until_ledger(key, self.as_budget())?;
+                .get_with_live_until_ledger(key, self, None)?;
             let mut current = (*current).metered_clone(self)?;
 
             if let LedgerEntryData::ContractData(ref mut entry) = current.data {
@@ -212,14 +211,13 @@ impl Host {
                 ));
             }
 
-            self.try_borrow_storage_mut()?
-                .put(
-                    &key,
-                    &Rc::metered_new(current, self)?,
-                    live_until_ledger,
-                    self.as_budget(),
-                )
-                .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
+            self.try_borrow_storage_mut()?.put_with_host(
+                &key,
+                &Rc::metered_new(current, self)?,
+                live_until_ledger,
+                self,
+                None,
+            )?;
         } else {
             let data = ContractDataEntry {
                 contract: ScAddress::Contract(contract_id.metered_clone(self)?),
@@ -238,14 +236,13 @@ impl Host {
                 durability: ContractDataDurability::Persistent,
                 ext: ExtensionPoint::V0,
             };
-            self.try_borrow_storage_mut()?
-                .put(
-                    key,
-                    &Host::new_contract_data(self, data)?,
-                    Some(self.get_min_live_until_ledger(ContractDataDurability::Persistent)?),
-                    self.as_budget(),
-                )
-                .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
+            self.try_borrow_storage_mut()?.put_with_host(
+                key,
+                &Host::new_contract_data(self, data)?,
+                Some(self.get_min_live_until_ledger(ContractDataDurability::Persistent)?),
+                self,
+                None,
+            )?;
         }
         Ok(())
     }
@@ -263,8 +260,7 @@ impl Host {
             ContractExecutable::Wasm(wasm_hash) => {
                 let key = self.contract_code_ledger_key(&wasm_hash)?;
                 self.try_borrow_storage_mut()?
-                    .extend_ttl(self, key, threshold, extend_to)
-                    .map_err(|e| self.decorate_contract_code_storage_error(e, &wasm_hash))?;
+                    .extend_ttl(self, key, threshold, extend_to, None)?;
             }
             ContractExecutable::StellarAsset => {}
         }
@@ -273,19 +269,17 @@ impl Host {
 
     pub(crate) fn extend_contract_instance_ttl_from_contract_id(
         &self,
-        contract_id: &Hash,
         instance_key: Rc<LedgerKey>,
         threshold: u32,
         extend_to: u32,
     ) -> Result<(), HostError> {
-        self.try_borrow_storage_mut()?
-            .extend_ttl(
-                self,
-                instance_key.metered_clone(self)?,
-                threshold,
-                extend_to,
-            )
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?;
+        self.try_borrow_storage_mut()?.extend_ttl(
+            self,
+            instance_key.metered_clone(self)?,
+            threshold,
+            extend_to,
+            None,
+        )?;
         Ok(())
     }
 
@@ -304,18 +298,8 @@ impl Host {
     // notes on metering: `get` from storage is covered. Rest are free.
     pub(crate) fn load_account(&self, account_id: AccountId) -> Result<AccountEntry, HostError> {
         let acc = self.to_account_key(account_id)?;
-        self.with_mut_storage(|storage| {
-            match &storage
-                .get(&acc, self.as_budget())
-                .map_err(|e| {
-                    self.decorate_account_footprint_error(
-                        e,
-                        &acc,
-                        "trying to access account entry outside of the footprint",
-                    )
-                })?
-                .data
-            {
+        self.with_mut_storage(
+            |storage| match &storage.get_with_host(&acc, self, None)?.data {
                 LedgerEntryData::Account(ae) => ae.metered_clone(self),
                 e => Err(err!(
                     self,
@@ -323,8 +307,8 @@ impl Host {
                     "ledger entry is not account",
                     e.name()
                 )),
-            }
-        })
+            },
+        )
     }
 
     pub(crate) fn to_account_key(&self, account_id: AccountId) -> Result<Rc<LedgerKey>, HostError> {
@@ -488,18 +472,16 @@ impl Host {
         let durability: ContractDataDurability = t.try_into()?;
         let key = self.storage_key_from_val(k, durability)?;
         // Currently the storage stores the whole ledger entries, while this
-        // operation might only modify only the internal `ScVal` value. Thus we
+        // operation might only modify the internal `ScVal` value. Thus we
         // need to only overwrite the value in case if there is already an
         // existing ledger entry value for the key in the storage.
         if self
             .try_borrow_storage_mut()?
-            .has(&key, self.as_budget())
-            .map_err(|e| self.decorate_contract_data_storage_error(e, k))?
+            .has_with_host(&key, self, Some(k))?
         {
             let (current, live_until_ledger) = self
                 .try_borrow_storage_mut()?
-                .get_with_live_until_ledger(&key, self.as_budget())
-                .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+                .get_with_live_until_ledger(&key, self, Some(k))?;
             let mut current = (*current).metered_clone(self)?;
             match current.data {
                 LedgerEntryData::ContractData(ref mut entry) => {
@@ -514,14 +496,13 @@ impl Host {
                     ));
                 }
             }
-            self.try_borrow_storage_mut()?
-                .put(
-                    &key,
-                    &Rc::metered_new(current, self)?,
-                    live_until_ledger,
-                    self.as_budget(),
-                )
-                .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+            self.try_borrow_storage_mut()?.put_with_host(
+                &key,
+                &Rc::metered_new(current, self)?,
+                live_until_ledger,
+                self,
+                Some(k),
+            )?;
         } else {
             let data = ContractDataEntry {
                 contract: ScAddress::Contract(self.get_current_contract_id_internal()?),
@@ -530,14 +511,13 @@ impl Host {
                 durability,
                 ext: ExtensionPoint::V0,
             };
-            self.try_borrow_storage_mut()?
-                .put(
-                    &key,
-                    &Host::new_contract_data(self, data)?,
-                    Some(self.get_min_live_until_ledger(durability)?),
-                    self.as_budget(),
-                )
-                .map_err(|e| self.decorate_contract_data_storage_error(e, k))?;
+            self.try_borrow_storage_mut()?.put_with_host(
+                &key,
+                &Host::new_contract_data(self, data)?,
+                Some(self.get_min_live_until_ledger(durability)?),
+                self,
+                Some(k),
+            )?;
         }
 
         Ok(())
@@ -558,7 +538,9 @@ impl Host {
         val: &Rc<soroban_env_common::xdr::LedgerEntry>,
         live_until_ledger: Option<u32>,
     ) -> Result<(), HostError> {
-        self.with_mut_storage(|storage| storage.put(key, val, live_until_ledger, self.as_budget()))
+        self.with_mut_storage(|storage| {
+            storage.put_with_host(key, val, live_until_ledger, self, None)
+        })
     }
 
     // Performs the necessary setup to access the provided ledger key/entry in
diff --git a/soroban-env-host/src/host/error.rs b/soroban-env-host/src/host/error.rs
index 57d0bb49..9ee5ea54 100644
--- a/soroban-env-host/src/host/error.rs
+++ b/soroban-env-host/src/host/error.rs
@@ -1,7 +1,7 @@
 use crate::{
     budget::AsBudget,
     events::Events,
-    xdr::{self, Hash, LedgerKey, ScAddress, ScError, ScErrorCode, ScErrorType},
+    xdr::{self, LedgerKey, ScAddress, ScError, ScErrorCode, ScErrorType},
     ConversionError, EnvBase, Error, Host, TryFromVal, U32Val, Val,
 };
 
@@ -11,7 +11,6 @@ use core::fmt::Debug;
 use std::{
     cell::{Ref, RefCell, RefMut},
     ops::DerefMut,
-    rc::Rc,
 };
 
 use super::metered_clone::MeteredClone;
@@ -369,8 +368,8 @@ impl Host {
     // Extracts the account id from the given ledger key as address object `Val`.
     // Returns Void for unsupported entries.
     // Useful as a helper for error reporting.
-    pub(crate) fn account_address_from_key(&self, lk: &Rc<LedgerKey>) -> Result<Val, HostError> {
-        let account_id = match lk.as_ref() {
+    pub(crate) fn account_address_from_key(&self, lk: &LedgerKey) -> Result<Val, HostError> {
+        let account_id = match lk {
             LedgerKey::Account(e) => &e.account_id,
             LedgerKey::Trustline(e) => &e.account_id,
             _ => {
@@ -382,103 +381,6 @@ impl Host {
         ))
         .map(|a| a.to_val())
     }
-
-    pub(crate) fn decorate_account_footprint_error(
-        &self,
-        err: HostError,
-        lk: &Rc<LedgerKey>,
-        msg: &str,
-    ) -> HostError {
-        if err.error.is_type(ScErrorType::Storage) && err.error.is_code(ScErrorCode::ExceededLimit)
-        {
-            let account_address = self.account_address_from_key(lk);
-            match account_address {
-                Ok(account_address) => {
-                    return self.err(
-                        ScErrorType::Storage,
-                        ScErrorCode::ExceededLimit,
-                        msg,
-                        &[account_address],
-                    );
-                }
-                Err(e) => {
-                    return e;
-                }
-            }
-        }
-        err
-    }
-
-    pub(crate) fn decorate_contract_data_storage_error(
-        &self,
-        err: HostError,
-        key: Val,
-    ) -> HostError {
-        if !err.error.is_type(ScErrorType::Storage) {
-            return err;
-        }
-        if err.error.is_code(ScErrorCode::ExceededLimit) {
-            return self.err(
-                ScErrorType::Storage,
-                ScErrorCode::ExceededLimit,
-                "trying to access contract storage key outside of the footprint",
-                &[key],
-            );
-        }
-        if err.error.is_code(ScErrorCode::MissingValue) {
-            return self.err(
-                ScErrorType::Storage,
-                ScErrorCode::MissingValue,
-                "trying to get non-existing value for contract storage key",
-                &[key],
-            );
-        }
-        err
-    }
-
-    pub(crate) fn decorate_contract_instance_storage_error(
-        &self,
-        err: HostError,
-        contract_id: &Hash,
-    ) -> HostError {
-        if err.error.is_type(ScErrorType::Storage) && err.error.is_code(ScErrorCode::ExceededLimit)
-        {
-            return self.err(
-                ScErrorType::Storage,
-                ScErrorCode::ExceededLimit,
-                "trying to access contract instance key outside of the footprint",
-                // No need for metered clone here as we are on the unrecoverable
-                // error path.
-                &[self
-                    .add_host_object(ScAddress::Contract(contract_id.clone()))
-                    .map(|a| a.into())
-                    .unwrap_or(Val::VOID.into())],
-            );
-        }
-        err
-    }
-
-    pub(crate) fn decorate_contract_code_storage_error(
-        &self,
-        err: HostError,
-        wasm_hash: &Hash,
-    ) -> HostError {
-        if err.error.is_type(ScErrorType::Storage) && err.error.is_code(ScErrorCode::ExceededLimit)
-        {
-            return self.err(
-                ScErrorType::Storage,
-                ScErrorCode::ExceededLimit,
-                "trying to access contract code key outside of the footprint",
-                // No need for metered clone here as we are on the unrecoverable
-                // error path.
-                &[self
-                    .add_host_object(self.scbytes_from_hash(wasm_hash).unwrap_or_default())
-                    .map(|a| a.into())
-                    .unwrap_or(Val::VOID.into())],
-            );
-        }
-        err
-    }
 }
 
 pub(crate) trait DebugArg {
diff --git a/soroban-env-host/src/host/frame.rs b/soroban-env-host/src/host/frame.rs
index 024a2725..46a2186c 100644
--- a/soroban-env-host/src/host/frame.rs
+++ b/soroban-env-host/src/host/frame.rs
@@ -633,21 +633,14 @@ impl Host {
     ) -> Result<TestContractFrame, HostError> {
         let instance_key = self.contract_instance_ledger_key(&id)?;
         let instance = self.retrieve_contract_instance_from_storage(&instance_key)?;
-        Ok(TestContractFrame::new(
-            id.clone(),
-            func,
-            args.to_vec(),
-            instance,
-        ))
+        Ok(TestContractFrame::new(id, func, args.to_vec(), instance))
     }
 
     // Notes on metering: this is covered by the called components.
     fn call_contract_fn(&self, id: &Hash, func: &Symbol, args: &[Val]) -> Result<Val, HostError> {
         // Create key for storage
         let storage_key = self.contract_instance_ledger_key(id)?;
-        let instance = self
-            .retrieve_contract_instance_from_storage(&storage_key)
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &id))?;
+        let instance = self.retrieve_contract_instance_from_storage(&storage_key)?;
         Vec::<Val>::charge_bulk_init_cpy(args.len() as u64, self.as_budget())?;
         let args_vec = args.to_vec();
         match &instance.executable {
@@ -675,7 +668,7 @@ impl Host {
                     let wasm_key = self.contract_code_ledger_key(wasm_hash)?;
                     if self
                         .try_borrow_storage_mut()?
-                        .has(&wasm_key, self.budget_ref())?
+                        .has_with_host(&wasm_key, self, None)?
                     {
                         cache.get_module(self, wasm_hash)?
                     } else {
diff --git a/soroban-env-host/src/host/lifecycle.rs b/soroban-env-host/src/host/lifecycle.rs
index a01bac36..6b1ec8fc 100644
--- a/soroban-env-host/src/host/lifecycle.rs
+++ b/soroban-env-host/src/host/lifecycle.rs
@@ -1,5 +1,4 @@
 use crate::{
-    budget::AsBudget,
     err,
     host::{
         metered_clone::{MeteredAlloc, MeteredClone},
@@ -25,8 +24,7 @@ impl Host {
         let storage_key = self.contract_instance_ledger_key(&contract_id)?;
         if self
             .try_borrow_storage_mut()?
-            .has(&storage_key, self.as_budget())
-            .map_err(|e| self.decorate_contract_instance_storage_error(e, &contract_id))?
+            .has_with_host(&storage_key, self, None)?
         {
             return Err(self.err(
                 ScErrorType::Storage,
@@ -212,17 +210,13 @@ impl Host {
 
         // We will definitely put the contract in the ledger if it isn't there yet.
         #[allow(unused_mut)]
-        let mut should_put_contract = !storage
-            .has(&code_key, self.as_budget())
-            .map_err(|e| self.decorate_contract_code_storage_error(e, &Hash(hash_bytes)))?;
+        let mut should_put_contract = !storage.has_with_host(&code_key, self, None)?;
 
         // We may also, in the cache-supporting protocol, overwrite the contract if its ext field changed.
         if !should_put_contract
             && self.get_ledger_protocol_version()? >= super::ModuleCache::MIN_LEDGER_VERSION
         {
-            let entry = storage
-                .get(&code_key, self.as_budget())
-                .map_err(|e| self.decorate_contract_code_storage_error(e, &Hash(hash_bytes)))?;
+            let entry = storage.get_with_host(&code_key, self, None)?;
             if let crate::xdr::LedgerEntryData::ContractCode(ContractCodeEntry {
                 ext: old_ext,
                 ..
@@ -238,14 +232,13 @@ impl Host {
                 ext,
                 code: wasm_bytes_m,
             };
-            storage
-                .put(
-                    &code_key,
-                    &Host::new_contract_code(self, data)?,
-                    Some(self.get_min_live_until_ledger(ContractDataDurability::Persistent)?),
-                    self.as_budget(),
-                )
-                .map_err(|e| self.decorate_contract_code_storage_error(e, &Hash(hash_bytes)))?;
+            storage.put_with_host(
+                &code_key,
+                &Host::new_contract_code(self, data)?,
+                Some(self.get_min_live_until_ledger(ContractDataDurability::Persistent)?),
+                self,
+                None,
+            )?;
         }
         Ok(hash_obj)
     }
diff --git a/soroban-env-host/src/storage.rs b/soroban-env-host/src/storage.rs
index 9f3b37f2..6115d0a9 100644
--- a/soroban-env-host/src/storage.rs
+++ b/soroban-env-host/src/storage.rs
@@ -9,11 +9,13 @@
 
 use std::rc::Rc;
 
+use crate::budget::AsBudget;
+use crate::host::metered_clone::MeteredClone;
 use crate::{
     budget::Budget,
     host::metered_map::MeteredOrdMap,
     ledger_info::get_key_durability,
-    xdr::{ContractDataDurability, LedgerEntry, LedgerKey, ScErrorCode, ScErrorType},
+    xdr::{ContractDataDurability, LedgerEntry, LedgerKey, ScErrorCode, ScErrorType, ScVal},
     Env, Error, Host, HostError, Val,
 };
 
@@ -221,6 +223,24 @@ impl Storage {
         }
     }
 
+    fn try_get_full_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<Option<EntryWithLiveUntil>, HostError> {
+        let res = self
+            .try_get_full(key, host.as_budget())
+            .map_err(|e| host.decorate_storage_error(e, key.as_ref(), key_val))?;
+
+        #[cfg(any(test, feature = "testutils"))]
+        if !host.check_if_entry_is_live(key.as_ref(), &res, key_val)? {
+            return Ok(None);
+        }
+
+        Ok(res)
+    }
+
     /// Attempts to retrieve the [LedgerEntry] associated with a given
     /// [LedgerKey] in the [Storage], returning an error if the key is not
     /// found.
@@ -237,18 +257,31 @@ impl Storage {
         key: &Rc<LedgerKey>,
         budget: &Budget,
     ) -> Result<Rc<LedgerEntry>, HostError> {
-        self.try_get(key, budget)?
+        self.try_get_full(key, budget)?
             .ok_or_else(|| (ScErrorType::Storage, ScErrorCode::MissingValue).into())
+            .map(|e| e.0)
+    }
+    pub(crate) fn get_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<Rc<LedgerEntry>, HostError> {
+        self.try_get_full_with_host(key, host, key_val)?
+            .ok_or_else(|| (ScErrorType::Storage, ScErrorCode::MissingValue).into())
+            .map(|e| e.0)
+            .map_err(|e| host.decorate_storage_error(e, key.as_ref(), key_val))
     }
 
-    // Like `get`, but distinguishes between missing values (return `Ok(None)`)
+    // Like `get_with_host`, but distinguishes between missing values (return `Ok(None)`)
     // and out-of-footprint values or errors (`Err(...)`).
     pub(crate) fn try_get(
         &mut self,
         key: &Rc<LedgerKey>,
-        budget: &Budget,
+        host: &Host,
+        key_val: Option<Val>,
     ) -> Result<Option<Rc<LedgerEntry>>, HostError> {
-        self.try_get_full(key, budget)
+        self.try_get_full_with_host(key, host, key_val)
             .map(|ok| ok.map(|pair| pair.0))
     }
 
@@ -269,10 +302,14 @@ impl Storage {
     pub(crate) fn get_with_live_until_ledger(
         &mut self,
         key: &Rc<LedgerKey>,
-        budget: &Budget,
+        host: &Host,
+        key_val: Option<Val>,
     ) -> Result<EntryWithLiveUntil, HostError> {
-        self.try_get_full(key, budget)?
-            .ok_or_else(|| (ScErrorType::Storage, ScErrorCode::MissingValue).into())
+        self.try_get_full_with_host(key, host, key_val)
+            .and_then(|maybe_entry| {
+                maybe_entry.ok_or_else(|| (ScErrorType::Storage, ScErrorCode::MissingValue).into())
+            })
+            .map_err(|e| host.decorate_storage_error(e, key.as_ref(), key_val))
     }
 
     // Helper function `put` and `del` funnel into.
@@ -299,6 +336,30 @@ impl Storage {
         Ok(())
     }
 
+    fn put_opt_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        val: Option<EntryWithLiveUntil>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<(), HostError> {
+        let prev = host.as_budget().get_cpu_insns_consumed().unwrap();
+        #[cfg(any(test, feature = "testutils"))]
+        let _ = host.as_budget().with_observable_shadow_mode(|| {
+            let Ok(Some(entry_with_live_until)) =
+                self.map.get::<Rc<LedgerKey>>(key, host.as_budget())
+            else {
+                return Ok(());
+            };
+            let _ = host.check_if_entry_is_live(key.as_ref(), &entry_with_live_until, key_val)?;
+            Ok(())
+        })?;
+        let curr = host.as_budget().get_cpu_insns_consumed().unwrap();
+        assert_eq!(prev, curr);
+        self.put_opt(key, val, host.as_budget())
+            .map_err(|e| host.decorate_storage_error(e, key.as_ref(), key_val))
+    }
+
     /// Attempts to write to the [LedgerEntry] associated with a given
     /// [LedgerKey] in the [Storage].
     ///
@@ -319,6 +380,18 @@ impl Storage {
         self.put_opt(key, Some((val.clone(), live_until_ledger)), budget)
     }
 
+    pub(crate) fn put_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        val: &Rc<LedgerEntry>,
+        live_until_ledger: Option<u32>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<(), HostError> {
+        let _span = tracy_span!("storage put");
+        self.put_opt_with_host(key, Some((val.clone(), live_until_ledger)), host, key_val)
+    }
+
     /// Attempts to delete the [LedgerEntry] associated with a given [LedgerKey]
     /// in the [Storage].
     ///
@@ -333,6 +406,17 @@ impl Storage {
         self.put_opt(key, None, budget)
     }
 
+    pub(crate) fn del_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<(), HostError> {
+        let _span = tracy_span!("storage del");
+        self.put_opt_with_host(key, None, host, key_val)
+            .map_err(|e| host.decorate_storage_error(e, key.as_ref(), key_val))
+    }
+
     /// Attempts to determine the presence of a [LedgerEntry] associated with a
     /// given [LedgerKey] in the [Storage], returning `Ok(true)` if an entry
     /// with the key exists and `Ok(false)` if it does not.
@@ -355,6 +439,16 @@ impl Storage {
             .is_some())
     }
 
+    pub(crate) fn has_with_host(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        host: &Host,
+        key_val: Option<Val>,
+    ) -> Result<bool, HostError> {
+        let _span = tracy_span!("storage has");
+        Ok(self.try_get_full_with_host(key, host, key_val)?.is_some())
+    }
+
     /// Extends `key` to live `extend_to` ledgers from now (not counting the
     /// current ledger) if the current live_until_ledger for the entry is
     /// `threshold` ledgers or less away from the current ledger.
@@ -370,6 +464,7 @@ impl Storage {
         key: Rc<LedgerKey>,
         threshold: u32,
         extend_to: u32,
+        key_val: Option<Val>,
     ) -> Result<(), HostError> {
         let _span = tracy_span!("extend key");
         Self::check_supported_ledger_key_type(&key)?;
@@ -385,7 +480,7 @@ impl Storage {
 
         // Extending deleted/non-existing/out-of-footprint entries will result in
         // an error.
-        let (entry, old_live_until) = self.get_with_live_until_ledger(&key, host.budget_ref())?;
+        let (entry, old_live_until) = self.get_with_live_until_ledger(&key, &host, key_val)?;
         let old_live_until = old_live_until.ok_or_else(|| {
             host.err(
                 ScErrorType::Storage,
@@ -491,3 +586,157 @@ impl Storage {
         Ok(())
     }
 }
+
+fn get_key_type_string_for_error(lk: &LedgerKey) -> &str {
+    match lk {
+        LedgerKey::ContractData(cd) => match cd.key {
+            ScVal::LedgerKeyContractInstance => "contract instance",
+            ScVal::LedgerKeyNonce(_) => "nonce",
+            _ => "contract data key",
+        },
+        LedgerKey::ContractCode(_) => "contract code",
+        LedgerKey::Account(_) => "account",
+        LedgerKey::Trustline(_) => "account trustline",
+        // This shouldn't normally trigger, but it's safer to just return
+        // a safe default instead of an error in case if new key types are
+        // accessed.
+        _ => "ledger key",
+    }
+}
+
+impl Host {
+    fn decorate_storage_error(
+        &self,
+        err: HostError,
+        lk: &LedgerKey,
+        key_val: Option<Val>,
+    ) -> HostError {
+        let mut err = err;
+        self.with_debug_mode(|| {
+            if !err.error.is_type(ScErrorType::Storage) {
+                return Ok(());
+            }
+            if !err.error.is_code(ScErrorCode::ExceededLimit)
+                && !err.error.is_code(ScErrorCode::MissingValue)
+            {
+                return Ok(());
+            }
+
+            let key_type_str = get_key_type_string_for_error(lk);
+            // Accessing an entry outside of the footprint is a non-recoverable error, thus
+            // there is no way to observe the object pool being changed (host will continue
+            // propagating an error until there are no frames left and control is never
+            // returned to guest). This allows us to build a nicer error message.
+            // For the missing values we unfortunately can only safely use the existing `Val`s
+            // to enhance errors.
+            let can_create_new_objects = err.error.is_code(ScErrorCode::ExceededLimit);
+            let args = self
+                .get_args_for_error(lk, key_val, can_create_new_objects)
+                .unwrap_or_else(|_| vec![]);
+            if err.error.is_code(ScErrorCode::ExceededLimit) {
+                err = self.err(
+                    ScErrorType::Storage,
+                    ScErrorCode::ExceededLimit,
+                    format!("trying to access {} outside of the footprint", key_type_str).as_str(),
+                    args.as_slice(),
+                );
+            } else if err.error.is_code(ScErrorCode::MissingValue) {
+                err = self.err(
+                    ScErrorType::Storage,
+                    ScErrorCode::MissingValue,
+                    format!("trying to get non-existing value for {}", key_type_str).as_str(),
+                    args.as_slice(),
+                );
+            }
+
+            Ok(())
+        });
+        err
+    }
+
+    fn get_args_for_error(
+        &self,
+        lk: &LedgerKey,
+        key_val: Option<Val>,
+        can_create_new_objects: bool,
+    ) -> Result<Vec<Val>, HostError> {
+        let mut res = vec![];
+        match lk {
+            LedgerKey::ContractData(cd) => {
+                if can_create_new_objects {
+                    let address_val = self
+                        .add_host_object(cd.contract.metered_clone(self.as_budget())?)?
+                        .into();
+                    res.push(address_val);
+                }
+                match &cd.key {
+                    ScVal::LedgerKeyContractInstance => (),
+                    ScVal::LedgerKeyNonce(n) => {
+                        if can_create_new_objects {
+                            res.push(self.add_host_object(n.nonce)?.into());
+                        }
+                    }
+                    _ => {
+                        if let Some(key) = key_val {
+                            res.push(key);
+                        }
+                    }
+                }
+            }
+            LedgerKey::ContractCode(c) => {
+                if can_create_new_objects {
+                    res.push(
+                        self.add_host_object(self.scbytes_from_hash(&c.hash)?)?
+                            .into(),
+                    );
+                }
+            }
+            LedgerKey::Account(_) | LedgerKey::Trustline(_) => {
+                if can_create_new_objects {
+                    res.push(self.account_address_from_key(lk)?)
+                }
+            }
+            // This shouldn't normally trigger, but it's safer to just return
+            // a safe default instead of an error in case if new key types are
+            // accessed.
+            _ => (),
+        };
+        Ok(res)
+    }
+
+    #[cfg(any(test, feature = "testutils"))]
+    fn check_if_entry_is_live(
+        &self,
+        key: &LedgerKey,
+        entry_with_live_until: &Option<EntryWithLiveUntil>,
+        key_val: Option<Val>,
+    ) -> Result<bool, HostError> {
+        let Some((_, Some(live_until_ledger))) = &entry_with_live_until else {
+            return Ok(true);
+        };
+        let ledger_seq = self
+            .with_ledger_info(|li| Ok(li.sequence_number))
+            .unwrap_or_default();
+        if *live_until_ledger >= ledger_seq {
+            return Ok(true);
+        }
+        match get_key_durability(key) {
+            Some(ContractDataDurability::Temporary) => Ok(false),
+            Some(ContractDataDurability::Persistent) => {
+                let key_type_str = get_key_type_string_for_error(key);
+                let args = self
+                    .as_budget()
+                    .with_observable_shadow_mode(|| self.get_args_for_error(key, key_val, true))
+                    .unwrap_or_else(|_| vec![]);
+                let msg = format!("[testing-only] Accessed {} key that has been archived. Important: this error may only appear in tests; in the real network contracts aren't called at all if any archived entry is accessed.", key_type_str);
+                Err(self.err(
+                    ScErrorType::Storage,
+                    ScErrorCode::InternalError,
+                    msg.as_str(),
+                    args.as_slice(),
+                ))
+            }
+            None => Ok(true),
+        }
+    }
+}


Test patch:
diff --git a/soroban-env-host/observations/20/test_v_new_call_nonexistent_wasm_call.json b/soroban-env-host/observations/20/test_v_new_call_nonexistent_wasm_call.json
index 443db9a8..26495b34 100644
--- a/soroban-env-host/observations/20/test_v_new_call_nonexistent_wasm_call.json
+++ b/soroban-env-host/observations/20/test_v_new_call_nonexistent_wasm_call.json
@@ -3,6 +3,6 @@
   "   1 call vec_new()": "cpu:42436, mem:373, objs:-/1@10432200, store:-/1@6f750601, foot:1@2784223c",
   "   2 ret vec_new -> Ok(Vec(obj#3))": "cpu:42937, mem:437, objs:-/2@acf0be8f",
   "   3 call call(Address(obj#1), Symbol(test), Vec(obj#3))": "",
-  "   4 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:51600, mem:941",
-  "   5 end": "cpu:51600, mem:941, prngs:-/-, objs:-/2@acf0be8f, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
+  "   4 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:51600, mem:941, objs:-/3@69a02720",
+  "   5 end": "cpu:51600, mem:941, prngs:-/-, objs:-/3@69a02720, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/observations/20/test_v_new_call_runtime_deleted_wasm_call.json b/soroban-env-host/observations/20/test_v_new_call_runtime_deleted_wasm_call.json
index 07084c0d..b6600c93 100644
--- a/soroban-env-host/observations/20/test_v_new_call_runtime_deleted_wasm_call.json
+++ b/soroban-env-host/observations/20/test_v_new_call_runtime_deleted_wasm_call.json
@@ -7,6 +7,6 @@
   "   5 pop VM:3fa21985:test -> Ok(Symbol(pass))": "cpu:219364, mem:169803",
   "   6 ret call -> Ok(Symbol(pass))": " vm:-/-, stk:-, auth:-/-",
   "   7 call call(Address(obj#1), Symbol(test), Vec(obj#3))": "cpu:261360, mem:170112, store:-/1@6f750601, foot:1@2784223c",
-  "   8 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:269410, mem:170600",
-  "   9 end": "cpu:269410, mem:170600, prngs:-/-, objs:-/3@2988da9d, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
+  "   8 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:269410, mem:170600, objs:-/4@d193fe73",
+  "   9 end": "cpu:269410, mem:170600, prngs:-/-, objs:-/4@d193fe73, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/observations/20/test_v_new_rewrite_call_fail.json b/soroban-env-host/observations/20/test_v_new_rewrite_call_fail.json
index e3a38c3b..230c14d0 100644
--- a/soroban-env-host/observations/20/test_v_new_rewrite_call_fail.json
+++ b/soroban-env-host/observations/20/test_v_new_rewrite_call_fail.json
@@ -1,4 +1,4 @@
 {
   "   0 begin": "cpu:2746, mem:343, prngs:-/-, objs:-/-, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-",
-  "   1 end": "cpu:2106672, mem:534384, prngs:-/-, objs:-/2@846c64a, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-"
+  "   1 end": "cpu:2105703, mem:534272, prngs:-/-, objs:-/2@846c64a, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/observations/21/test_v_new_call_nonexistent_wasm_call.json b/soroban-env-host/observations/21/test_v_new_call_nonexistent_wasm_call.json
index 443db9a8..26495b34 100644
--- a/soroban-env-host/observations/21/test_v_new_call_nonexistent_wasm_call.json
+++ b/soroban-env-host/observations/21/test_v_new_call_nonexistent_wasm_call.json
@@ -3,6 +3,6 @@
   "   1 call vec_new()": "cpu:42436, mem:373, objs:-/1@10432200, store:-/1@6f750601, foot:1@2784223c",
   "   2 ret vec_new -> Ok(Vec(obj#3))": "cpu:42937, mem:437, objs:-/2@acf0be8f",
   "   3 call call(Address(obj#1), Symbol(test), Vec(obj#3))": "",
-  "   4 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:51600, mem:941",
-  "   5 end": "cpu:51600, mem:941, prngs:-/-, objs:-/2@acf0be8f, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
+  "   4 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:51600, mem:941, objs:-/3@69a02720",
+  "   5 end": "cpu:51600, mem:941, prngs:-/-, objs:-/3@69a02720, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/observations/21/test_v_new_call_runtime_deleted_wasm_call.json b/soroban-env-host/observations/21/test_v_new_call_runtime_deleted_wasm_call.json
index 07084c0d..b6600c93 100644
--- a/soroban-env-host/observations/21/test_v_new_call_runtime_deleted_wasm_call.json
+++ b/soroban-env-host/observations/21/test_v_new_call_runtime_deleted_wasm_call.json
@@ -7,6 +7,6 @@
   "   5 pop VM:3fa21985:test -> Ok(Symbol(pass))": "cpu:219364, mem:169803",
   "   6 ret call -> Ok(Symbol(pass))": " vm:-/-, stk:-, auth:-/-",
   "   7 call call(Address(obj#1), Symbol(test), Vec(obj#3))": "cpu:261360, mem:170112, store:-/1@6f750601, foot:1@2784223c",
-  "   8 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:269410, mem:170600",
-  "   9 end": "cpu:269410, mem:170600, prngs:-/-, objs:-/3@2988da9d, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
+  "   8 ret call -> Err(Error(Storage, ExceededLimit))": "cpu:269410, mem:170600, objs:-/4@d193fe73",
+  "   9 end": "cpu:269410, mem:170600, prngs:-/-, objs:-/4@d193fe73, vm:-/-, evt:-, store:-/1@6f750601, foot:1@2784223c, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/observations/21/test_v_new_rewrite_call_fail.json b/soroban-env-host/observations/21/test_v_new_rewrite_call_fail.json
index efdb8967..f9f929f2 100644
--- a/soroban-env-host/observations/21/test_v_new_rewrite_call_fail.json
+++ b/soroban-env-host/observations/21/test_v_new_rewrite_call_fail.json
@@ -1,4 +1,4 @@
 {
   "   0 begin": "cpu:2746, mem:343, prngs:-/-, objs:-/-, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-",
-  "   1 end": "cpu:1467786, mem:454165, prngs:-/-, objs:-/2@846c64a, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-"
+  "   1 end": "cpu:1466817, mem:454053, prngs:-/-, objs:-/2@846c64a, vm:-/-, evt:-, store:-/2@82fc3f1e, foot:2@5287c081, stk:-, auth:-/-"
 }
\ No newline at end of file
diff --git a/soroban-env-host/src/test/auth.rs b/soroban-env-host/src/test/auth.rs
index a69c36e5..f3497b70 100644
--- a/soroban-env-host/src/test/auth.rs
+++ b/soroban-env-host/src/test/auth.rs
@@ -302,11 +302,11 @@ impl AuthTest {
             .unwrap();
         self.host
             .with_mut_storage(|storage| {
-                if !storage.has(&nonce_key, self.host.budget_ref())? {
+                if !storage.has_with_host(&nonce_key, &self.host, None)? {
                     return Ok(None);
                 }
                 let (_, live_until_ledger) =
-                    storage.get_with_live_until_ledger(&nonce_key, self.host.budget_ref())?;
+                    storage.get_with_live_until_ledger(&nonce_key, &self.host, None)?;
                 Ok(live_until_ledger)
             })
             .unwrap()
@@ -553,7 +553,7 @@ fn test_single_authorized_call() {
             let key = test.host.to_account_key(account_id)?;
             // Note, that this represents 'correct footprint, missing value' scenario.
             // Incorrect footprint scenario is not covered (it's not auth specific).
-            storage.del(&key, test.host.budget_ref())
+            storage.del_with_host(&key, &test.host, None)
         })
         .unwrap();
 
diff --git a/soroban-env-host/src/test/lifecycle.rs b/soroban-env-host/src/test/lifecycle.rs
index 52517fa0..dc2cca07 100644
--- a/soroban-env-host/src/test/lifecycle.rs
+++ b/soroban-env-host/src/test/lifecycle.rs
@@ -24,9 +24,9 @@ use crate::testutils::{generate_account_id, generate_bytes_array};
 fn get_contract_wasm_ref(host: &Host, contract_id: Hash) -> Hash {
     let storage_key = host.contract_instance_ledger_key(&contract_id).unwrap();
     host.with_mut_storage(|s: &mut Storage| {
-        assert!(s.has(&storage_key, host.as_budget()).unwrap());
+        assert!(s.has_with_host(&storage_key, &host, None).unwrap());
 
-        match &s.get(&storage_key, host.as_budget()).unwrap().data {
+        match &s.get_with_host(&storage_key, host, None).unwrap().data {
             LedgerEntryData::ContractData(e) => match &e.val {
                 ScVal::ContractInstance(i) => match &i.executable {
                     ContractExecutable::Wasm(h) => Ok(h.clone()),
@@ -43,9 +43,9 @@ fn get_contract_wasm_ref(host: &Host, contract_id: Hash) -> Hash {
 fn get_contract_wasm(host: &Host, wasm_hash: Hash) -> Vec<u8> {
     let storage_key = host.contract_code_ledger_key(&wasm_hash).unwrap();
     host.with_mut_storage(|s: &mut Storage| {
-        assert!(s.has(&storage_key, host.as_budget()).unwrap());
+        assert!(s.has_with_host(&storage_key, &host, None).unwrap());
 
-        match &s.get(&storage_key, host.as_budget()).unwrap().data {
+        match &s.get_with_host(&storage_key, host, None).unwrap().data {
             LedgerEntryData::ContractCode(code_entry) => Ok(code_entry.code.to_vec()),
             _ => panic!("expected contract WASM code"),
         }
@@ -658,9 +658,8 @@ mod cap_54_55_56 {
         }
         fn reload(self, host: &Host) -> Result<Self, HostError> {
             host.with_mut_storage(|storage| {
-                let budget = host.budget_cloned();
-                let contract_entry = storage.get(&self.contract_key, &budget)?;
-                let wasm_entry = storage.get(&self.wasm_key, &budget)?;
+                let contract_entry = storage.get_with_host(&self.contract_key, host, None)?;
+                let wasm_entry = storage.get_with_host(&self.wasm_key, host, None)?;
                 Ok(ContractAndWasmEntries {
                     contract_key: self.contract_key,
                     contract_entry,
@@ -675,9 +674,8 @@ mod cap_54_55_56 {
             let wasm_key = host.contract_code_ledger_key(&wasm_hash)?;
 
             host.with_mut_storage(|storage| {
-                let budget = host.budget_cloned();
-                let contract_entry = storage.get(&contract_key, &budget)?;
-                let wasm_entry = storage.get(&wasm_key, &budget)?;
+                let contract_entry = storage.get_with_host(&contract_key, host, None)?;
+                let wasm_entry = storage.get_with_host(&wasm_key, host, None)?;
                 Ok(ContractAndWasmEntries {
                     contract_key,
                     contract_entry,
diff --git a/soroban-env-host/src/test/lifetime_extension.rs b/soroban-env-host/src/test/lifetime_extension.rs
index 0863d86f..31d2fa0f 100644
--- a/soroban-env-host/src/test/lifetime_extension.rs
+++ b/soroban-env-host/src/test/lifetime_extension.rs
@@ -49,8 +49,6 @@ impl InstanceCodeTest {
 }
 
 mod separate_instance_code_extension {
-    use crate::budget::AsBudget;
-
     use super::*;
 
     const PROTOCOL_SUPPORT_FOR_SEPARATE_EXTENSIONS: u32 = 21;
@@ -76,7 +74,8 @@ mod separate_instance_code_extension {
             .unwrap()
             .get_with_live_until_ledger(
                 &host.contract_instance_ledger_key(&contract).unwrap(),
-                host.as_budget(),
+                &host,
+                None,
             )
             .unwrap();
 
@@ -102,10 +101,7 @@ mod separate_instance_code_extension {
         let entry_with_live_until = host
             .try_borrow_storage_mut()
             .unwrap()
-            .get_with_live_until_ledger(
-                &host.contract_code_ledger_key(&code).unwrap(),
-                host.as_budget(),
-            )
+            .get_with_live_until_ledger(&host.contract_code_ledger_key(&code).unwrap(), &host, None)
             .unwrap();
 
         assert_eq!(entry_with_live_until.1, Some(9090));
@@ -126,10 +122,7 @@ mod separate_instance_code_extension {
         let code_entry_with_live_until = host
             .try_borrow_storage_mut()
             .unwrap()
-            .get_with_live_until_ledger(
-                &host.contract_code_ledger_key(&code).unwrap(),
-                host.as_budget(),
-            )
+            .get_with_live_until_ledger(&host.contract_code_ledger_key(&code).unwrap(), &host, None)
             .unwrap();
 
         assert_eq!(code_entry_with_live_until.1, Some(9090));
@@ -139,7 +132,8 @@ mod separate_instance_code_extension {
             .unwrap()
             .get_with_live_until_ledger(
                 &host.contract_instance_ledger_key(&contract).unwrap(),
-                host.as_budget(),
+                &host,
+                None,
             )
             .unwrap();
 
diff --git a/soroban-env-host/src/test/stellar_asset_contract.rs b/soroban-env-host/src/test/stellar_asset_contract.rs
index 0add85be..c7f6d15f 100644
--- a/soroban-env-host/src/test/stellar_asset_contract.rs
+++ b/soroban-env-host/src/test/stellar_asset_contract.rs
@@ -4,7 +4,6 @@ use crate::builtin_contracts::base_types::BytesN;
 use crate::testutils::simple_account_sign_fn;
 use crate::{
     auth::RecordedAuthPayload,
-    budget::AsBudget,
     builtin_contracts::{
         base_types::Address,
         contract_error::ContractError,
@@ -146,10 +145,12 @@ impl StellarAssetContractTest {
 
     fn get_trustline_balance(&self, key: &Rc<LedgerKey>) -> i64 {
         self.host
-            .with_mut_storage(|s| match &s.get(key, self.host.as_budget()).unwrap().data {
-                LedgerEntryData::Trustline(trustline) => Ok(trustline.balance),
-                _ => unreachable!(),
-            })
+            .with_mut_storage(
+                |s| match &s.get_with_host(key, &self.host, None).unwrap().data {
+                    LedgerEntryData::Trustline(trustline) => Ok(trustline.balance),
+                    _ => unreachable!(),
+                },
+            )
             .unwrap()
     }
     #[allow(clippy::too_many_arguments)]
@@ -182,7 +183,7 @@ impl StellarAssetContractTest {
     fn update_account_flags(&self, key: &Rc<LedgerKey>, new_flags: u32) {
         self.host
             .with_mut_storage(|s| {
-                let entry = s.get(key, self.host.as_budget()).unwrap();
+                let entry = s.get_with_host(key, &self.host, None).unwrap();
                 match entry.data.clone() {
                     LedgerEntryData::Account(mut account) => {
                         account.flags = new_flags;
@@ -191,7 +192,7 @@ impl StellarAssetContractTest {
                             &entry,
                             LedgerEntryData::Account(account),
                         )?;
-                        s.put(key, &update, None, self.host.as_budget())
+                        s.put_with_host(key, &update, None, &self.host, None)
                     }
                     _ => unreachable!(),
                 }
@@ -260,7 +261,7 @@ impl StellarAssetContractTest {
     fn update_trustline_flags(&self, key: &Rc<LedgerKey>, new_flags: u32) {
         self.host
             .with_mut_storage(|s| {
-                let entry = s.get(key, self.host.as_budget()).unwrap();
+                let entry = s.get_with_host(key, &self.host, None).unwrap();
                 match entry.data.clone() {
                     LedgerEntryData::Trustline(mut trustline) => {
                         trustline.flags = new_flags;
@@ -269,7 +270,7 @@ impl StellarAssetContractTest {
                             &entry,
                             LedgerEntryData::Trustline(trustline),
                         )?;
-                        s.put(key, &update, None, self.host.as_budget())
+                        s.put_with_host(key, &update, None, &self.host, None)
                     }
                     _ => unreachable!(),
                 }
@@ -3042,7 +3043,7 @@ fn test_custom_account_auth() {
             let key = test.host.contract_instance_ledger_key(&contract_id)?;
             // Note, that this represents 'correct footprint, missing value' scenario.
             // Incorrect footprint scenario is not covered (it's not auth specific).
-            storage.del(&key, test.host.budget_ref())
+            storage.del_with_host(&key, &test.host, None)
         })
         .unwrap();