refactor(store): get account vault details from account state forest

CHANGELOG.md

@@ -2,8 +2,9 @@
 
 ## v0.14.10 (TBD)
 
+- Optimize `GetAccount` implementation to serve vault assets from `AccountStateForest`.
 - Trace additional RPC request properties e.g. `account.id` in `GetAccount` ([#1983](https://github.com/0xMiden/node/pull/1983)).
-- Added `accept`, `origin`, `user-agent`, `forwarded`, `x-forwarded-for` and `x-real-ip` headers to telemetry for gRPC requests ([#1982](https://github.com/0xMiden/node/pull/1982).
+- Added `accept`, `origin`, `user-agent`, `forwarded`, `x-forwarded-for` and `x-real-ip` headers to telemetry for gRPC requests ([#1982](https://github.com/0xMiden/node/pull/1982)).
 
 ## v0.14.9 (2026-04-21)
 

crates/store/src/account_state_forest/mod.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeSet;
 
 use miden_crypto::hash::rpo::Rpo256;
 use miden_crypto::merkle::smt::ForestInMemoryBackend;
-use miden_node_proto::domain::account::AccountStorageMapDetails;
+use miden_node_proto::domain::account::{AccountStorageMapDetails, AccountVaultDetails};
 use miden_node_utils::ErrorReport;
 use miden_protocol::account::delta::{AccountDelta, AccountStorageDelta, AccountVaultDelta};
 use miden_protocol::account::{
@@ -12,7 +12,7 @@ use miden_protocol::account::{
     StorageMapWitness,
     StorageSlotName,
 };
-use miden_protocol::asset::{AssetVaultKey, AssetWitness, FungibleAsset};
+use miden_protocol::asset::{Asset, AssetVaultKey, AssetWitness, FungibleAsset};
 use miden_protocol::block::BlockNumber;
 use miden_protocol::crypto::merkle::smt::{
     ForestOperation,
@@ -270,6 +270,27 @@ impl AccountStateForest {
         witnessees
     }
 
+    /// Enumerates vault contents for the specified account at the requested block.
+    #[instrument(target = COMPONENT, skip_all)]
+    pub(crate) fn get_vault_details(
+        &self,
+        account_id: AccountId,
+        block_num: BlockNumber,
+    ) -> Result<AccountVaultDetails, WitnessError> {
+        let lineage = Self::vault_lineage_id(account_id);
+        let tree = self.get_tree_id(lineage, block_num).ok_or(WitnessError::RootNotFound)?;
+        // TODO: we should be checking `.entry_count()` instead of pulling entries from the tree
+        // once the optimization making `.entry_count()` cheap once `miden-crypto` is upgraded to
+        // > 0.23.
+        let entries = self.forest.entries(tree).map_err(Self::map_forest_error_to_witness)?;
+        let assets = entries
+            .take(AccountVaultDetails::MAX_RETURN_ENTRIES + 1)
+            .map(|entry| Asset::from_key_value_words(entry.key, entry.value))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(AccountVaultDetails::from_assets(assets))
+    }
+
     /// Opens a storage map and returns storage map details with SMT proofs for the given keys.
     ///
     /// Returns `None` if no storage root is tracked for this account/slot/block combination.

crates/store/src/account_state_forest/tests.rs

@@ -1,13 +1,20 @@
 use assert_matches::assert_matches;
-use miden_node_proto::domain::account::StorageMapEntries;
+use miden_node_proto::domain::account::{AccountVaultDetails, StorageMapEntries};
 use miden_protocol::Felt;
-use miden_protocol::account::{AccountCode, StorageMapKey};
-use miden_protocol::asset::{Asset, AssetVault, FungibleAsset};
+use miden_protocol::account::{AccountCode, AccountStorageMode, AccountType, StorageMapKey};
+use miden_protocol::asset::{
+    Asset,
+    AssetVault,
+    FungibleAsset,
+    NonFungibleAsset,
+    NonFungibleAssetDetails,
+};
 use miden_protocol::crypto::merkle::smt::SmtProof;
 use miden_protocol::testing::account_id::{
     ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
     ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE,
     ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE_2,
+    AccountIdBuilder,
 };
 
 use super::*;
@@ -161,6 +168,57 @@ fn vault_incremental_updates_with_add_and_remove() {
     assert_eq!(root_after_120, root_full_state_120);
 }
 
+#[test]
+fn vault_details_returns_latest_and_historical_assets() {
+    let mut forest = AccountStateForest::new();
+    let account_id = dummy_account();
+    let faucet_id = dummy_faucet();
+
+    let block_1 = BlockNumber::GENESIS.child();
+    let asset_100 = dummy_fungible_asset(faucet_id, 100);
+    let full_delta = dummy_full_state_delta(account_id, &[asset_100]);
+    forest.update_account(block_1, &full_delta).unwrap();
+
+    let block_2 = block_1.child();
+    let mut vault_delta_2 = AccountVaultDelta::default();
+    vault_delta_2.add_asset(dummy_fungible_asset(faucet_id, 50)).unwrap();
+    let delta_2 = dummy_partial_delta(account_id, vault_delta_2, AccountStorageDelta::default());
+    forest.update_account(block_2, &delta_2).unwrap();
+
+    let historical = forest.get_vault_details(account_id, block_1).unwrap();
+    assert_eq!(historical, AccountVaultDetails::Assets(vec![asset_100]));
+
+    let latest = forest.get_vault_details(account_id, block_2).unwrap();
+    assert_eq!(latest, AccountVaultDetails::Assets(vec![dummy_fungible_asset(faucet_id, 150)]));
+}
+
+#[test]
+fn vault_details_limit_exceeded_for_large_vault() {
+    let mut forest = AccountStateForest::new();
+    let account_id = dummy_account();
+    let block_num = BlockNumber::GENESIS.child();
+
+    let faucet_id = AccountIdBuilder::new()
+        .account_type(AccountType::NonFungibleFaucet)
+        .storage_mode(AccountStorageMode::Public)
+        .build_with_seed([7; 32]);
+    let assets = (0..=AccountVaultDetails::MAX_RETURN_ENTRIES)
+        .map(|i| {
+            let details =
+                NonFungibleAssetDetails::new(faucet_id, vec![i as u8, (i >> 8) as u8]).unwrap();
+            Asset::NonFungible(NonFungibleAsset::new(&details).unwrap())
+        })
+        .collect::<Vec<_>>();
+
+    let full_delta = dummy_full_state_delta(account_id, &assets);
+    forest.update_account(block_num, &full_delta).unwrap();
+
+    assert_eq!(
+        forest.get_vault_details(account_id, block_num).unwrap(),
+        AccountVaultDetails::LimitExceeded
+    );
+}
+
 #[test]
 fn forest_versions_are_continuous_for_sequential_updates() {
     use std::collections::BTreeMap;

crates/store/src/db/mod.rs

@@ -453,19 +453,6 @@ impl Db {
         .await
     }
 
-    /// Queries vault assets at a specific block
-    #[instrument(target = COMPONENT, skip_all, ret(level = "debug"), err)]
-    pub async fn select_account_vault_at_block(
-        &self,
-        account_id: AccountId,
-        block_num: BlockNumber,
-    ) -> Result<Vec<Asset>> {
-        self.transact("Get account vault at block", move |conn| {
-            queries::select_account_vault_at_block(conn, account_id, block_num)
-        })
-        .await
-    }
-
     /// Queries the account code by its commitment hash.
     ///
     /// Returns `None` if no code exists with that commitment.

crates/store/src/db/models/queries/accounts.rs

@@ -47,10 +47,7 @@ use crate::db::{AccountVaultValue, schema};
 use crate::errors::DatabaseError;
 
 mod at_block;
-pub(crate) use at_block::{
-    select_account_header_with_storage_header_at_block,
-    select_account_vault_at_block,
-};
+pub(crate) use at_block::select_account_header_with_storage_header_at_block;
 
 mod delta;
 use delta::{

crates/store/src/db/models/queries/accounts/at_block.rs

@@ -1,8 +1,7 @@
-use diesel::prelude::{Queryable, QueryableByName};
+use diesel::prelude::Queryable;
 use diesel::query_dsl::methods::SelectDsl;
 use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection};
 use miden_protocol::account::{AccountHeader, AccountId, AccountStorageHeader};
-use miden_protocol::asset::Asset;
 use miden_protocol::block::BlockNumber;
 use miden_protocol::utils::serde::{Deserializable, Serializable};
 use miden_protocol::{Felt, Word};
@@ -100,65 +99,3 @@ pub(crate) fn select_account_header_with_storage_header_at_block(
 
     Ok(Some((account_header, storage_header)))
 }
-
-// ACCOUNT VAULT
-// ================================================================================================
-
-/// Query vault assets at a specific block by finding the most recent update for each `vault_key`.
-///
-/// Uses a single raw SQL query with a subquery join:
-/// ```sql
-/// SELECT a.asset FROM account_vault_assets a
-/// INNER JOIN (
-///     SELECT vault_key, MAX(block_num) as max_block
-///     FROM account_vault_assets
-///     WHERE account_id = ? AND block_num <= ?
-///     GROUP BY vault_key
-/// ) latest ON a.vault_key = latest.vault_key AND a.block_num = latest.max_block
-/// WHERE a.account_id = ?
-/// ```
-pub(crate) fn select_account_vault_at_block(
-    conn: &mut SqliteConnection,
-    account_id: AccountId,
-    block_num: BlockNumber,
-) -> Result<Vec<Asset>, DatabaseError> {
-    use diesel::sql_types::{BigInt, Binary};
-
-    let account_id_bytes = account_id.to_bytes();
-    let block_num_sql = block_num.to_raw_sql();
-
-    let entries: Vec<Option<Vec<u8>>> = diesel::sql_query(
-        r"
-        SELECT a.asset FROM account_vault_assets a
-        INNER JOIN (
-            SELECT vault_key, MAX(block_num) as max_block
-            FROM account_vault_assets
-            WHERE account_id = ? AND block_num <= ?
-            GROUP BY vault_key
-        ) latest ON a.vault_key = latest.vault_key AND a.block_num = latest.max_block
-        WHERE a.account_id = ?
-        ",
-    )
-    .bind::<Binary, _>(&account_id_bytes)
-    .bind::<BigInt, _>(block_num_sql)
-    .bind::<Binary, _>(&account_id_bytes)
-    .load::<AssetRow>(conn)?
-    .into_iter()
-    .map(|row| row.asset)
-    .collect();
-
-    // Convert to assets, filtering out deletions (None values)
-    let mut assets = Vec::new();
-    for asset_bytes in entries.into_iter().flatten() {
-        let asset = Asset::read_from_bytes(&asset_bytes)?;
-        assets.push(asset);
-    }
-
-    Ok(assets)
-}
-
-#[derive(QueryableByName)]
-struct AssetRow {
-    #[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Binary>)]
-    asset: Option<Vec<u8>>,
-}

crates/store/src/db/models/queries/accounts/delta/tests.rs

@@ -41,9 +41,9 @@ use miden_standards::account::auth::AuthSingleSig;
 use miden_standards::code_builder::CodeBuilder;
 
 use crate::db::migrations::MIGRATIONS;
+use crate::db::models::queries::accounts::tests::select_account_vault_at_block;
 use crate::db::models::queries::accounts::{
     select_account_header_with_storage_header_at_block,
-    select_account_vault_at_block,
     select_full_account,
     upsert_accounts,
 };

crates/store/src/db/models/queries/accounts/tests.rs

@@ -252,6 +252,66 @@ fn assert_storage_map_slot_entries(
     assert_eq!(&entries, expected, "map entries mismatch");
 }
 
+/// Test helper: query vault assets at a specific block by finding the most recent
+/// update for each `vault_key`.
+///
+/// Uses a single raw SQL query with a subquery join:
+/// ```sql
+/// SELECT a.asset FROM account_vault_assets a
+/// INNER JOIN (
+///     SELECT vault_key, MAX(block_num) as max_block
+///     FROM account_vault_assets
+///     WHERE account_id = ? AND block_num <= ?
+///     GROUP BY vault_key
+/// ) latest ON a.vault_key = latest.vault_key AND a.block_num = latest.max_block
+/// WHERE a.account_id = ?
+/// ```
+pub(super) fn select_account_vault_at_block(
+    conn: &mut SqliteConnection,
+    account_id: AccountId,
+    block_num: BlockNumber,
+) -> Result<Vec<Asset>, DatabaseError> {
+    use diesel::sql_types::{BigInt, Binary};
+
+    let account_id_bytes = account_id.to_bytes();
+    let block_num_sql = block_num.to_raw_sql();
+
+    let entries: Vec<Option<Vec<u8>>> = diesel::sql_query(
+        r"
+        SELECT a.asset FROM account_vault_assets a
+        INNER JOIN (
+            SELECT vault_key, MAX(block_num) as max_block
+            FROM account_vault_assets
+            WHERE account_id = ? AND block_num <= ?
+            GROUP BY vault_key
+        ) latest ON a.vault_key = latest.vault_key AND a.block_num = latest.max_block
+        WHERE a.account_id = ?
+        ",
+    )
+    .bind::<Binary, _>(&account_id_bytes)
+    .bind::<BigInt, _>(block_num_sql)
+    .bind::<Binary, _>(&account_id_bytes)
+    .load::<AssetRow>(conn)?
+    .into_iter()
+    .map(|row| row.asset)
+    .collect();
+
+    // Convert to assets, filtering out deletions (None values)
+    let mut assets = Vec::new();
+    for asset_bytes in entries.into_iter().flatten() {
+        let asset = Asset::read_from_bytes(&asset_bytes)?;
+        assets.push(asset);
+    }
+
+    Ok(assets)
+}
+
+#[derive(QueryableByName)]
+struct AssetRow {
+    #[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Binary>)]
+    asset: Option<Vec<u8>>,
+}
+
 // ACCOUNT HEADER AT BLOCK TESTS
 // ================================================================================================
 

crates/store/src/state/mod.rs

@@ -744,11 +744,17 @@ impl State {
             Some(commitment) if commitment == account_header.vault_root() => {
                 AccountVaultDetails::empty()
             },
-            Some(_) => {
-                let vault_assets =
-                    self.db.select_account_vault_at_block(account_id, block_num).await?;
-                AccountVaultDetails::from_assets(vault_assets)
-            },
+            Some(_) => self
+                .forest
+                .read()
+                .instrument(tracing::info_span!("acquire_forest_for_vault"))
+                .await
+                .get_vault_details(account_id, block_num)
+                .map_err(|err| {
+                    DatabaseError::DataCorrupted(format!(
+                        "failed to reconstruct vault for account {account_id} at block {block_num}: {err}"
+                    ))
+                })?,
             None => AccountVaultDetails::empty(),
         };
 
@@ -775,8 +781,11 @@ impl State {
         let mut storage_map_details_by_index = vec![None; storage_request_slots.len()];
 
         if !map_keys_requests.is_empty() {
-            let forest_guard =
-                self.forest.read().instrument(tracing::info_span!("acquire_forest")).await;
+            let forest_guard = self
+                .forest
+                .read()
+                .instrument(tracing::info_span!("acquire_forest_for_storage_map"))
+                .await;
             for (index, slot_name, keys) in map_keys_requests {
                 let details = forest_guard
                     .get_storage_map_details_for_keys(