Fix coldkey-wide root claims and stale staking relationships

docs/guides/root-reborn.mdx

@@ -239,8 +239,11 @@ root-stakes to.
   deadline and nothing expires.
 - **Claim fee.** The inclusion fee scales with how many ALPHA types the
   basket holds. Both root-claim calls reserve a conservative 256-unit work
-  envelope at inclusion, independent of how many networks currently exist,
-  and refund the unused part after the claim. The amount you actually spend
+  envelope on every chain and refund the unused part after the claim. Admission
+  counts only root-relevant validator hotkeys plus their stored basket rows;
+  unrelated subnet stakes are filtered instead of being multiplied by the total
+  network count. Classifying the staking-hotkey vector is separately capped at
+  256 relationships. The amount you actually spend
   follows the holdings scanned and redeemed (around τ0.057 on a full
   128-holding basket).
   `btcli root claim --dry-run` shows reserved versus spent, compares the

docs/guides/staking.mdx

@@ -190,8 +190,11 @@ btcli stake remove --netuid 0 --hotkey 5F... --amount all --claim  # claim all,
 ```
 
 The claim fee scales with how many ALPHA types are in the basket. Both
-root-claim calls reserve a conservative 256-unit work envelope at inclusion,
-independent of the current network count, and refund the unused part after.
+root-claim calls reserve a conservative 256-unit work envelope on every chain
+and refund the unused part after. Admission counts only root-relevant validator
+hotkeys plus their stored basket rows; unrelated subnet stakes are filtered
+instead of being multiplied by the total network count. Classifying the
+staking-hotkey vector is separately capped at 256 relationships.
 You actually spend according to the holdings scanned and redeemed (around
 τ0.057 on a full 128-holding basket). `--dry-run` and the confirm step show
 reserved versus spent, warn if that spent fee exceeds accrued yield, and

pallets/subtensor/src/lib.rs

@@ -72,10 +72,8 @@ pub const MIN_ALPHA_LOW: u16 = 1_639;
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
 /// Benchmark upper bound and admission envelope for `claim_root` /
-/// `claim_root_scan` (Linear<1, N>). Weight calculation cannot walk storage,
-/// so both claim paths reserve this many units and refuse work that would
-/// exceed the envelope. Post-dispatch weight is refunded to the work
-/// actually performed.
+/// `claim_root_scan` (`Linear<1, N>`). Both claim paths reserve this many units
+/// and refund unused weight after dispatch.
 pub const MAX_ROOT_CLAIM_WORK: u32 = 256;
 
 /// Minimum number of positive destination weights required by `set_root_weights`. Softened

pallets/subtensor/src/macros/dispatches.rs

@@ -1966,11 +1966,8 @@ mod dispatches {
         /// # Events
         /// * `RootClaimed`: On successfully claiming the root emissions for a coldkey.
         #[pallet::call_index(121)]
-        // Signer is not in the call data, so admission uses the conservative
-        // MAX_ROOT_CLAIM_WORK envelope. Execution refuses a fat coldkey that
-        // would exceed it — use claim_root_with_hotkey per validator.
         #[pallet::weight(
-            <T as crate::pallet::Config>::WeightInfo::claim_root(Pallet::<T>::root_claim_declared_work())
+            Pallet::<T>::root_claim_declared_weight()
         )]
         pub fn claim_root(
             origin: OriginFor<T>,
@@ -1979,7 +1976,13 @@ mod dispatches {
             let coldkey: T::AccountId = ensure_signed(origin)?;
             let _ = subnets; // ignored: basket claims are fund-level, not per-subnet
 
-            let hotkeys = StakingHotkeys::<T>::get(&coldkey);
+            let staking_hotkeys = StakingHotkeys::<T>::get(&coldkey);
+            let selection_scanned = u32::try_from(staking_hotkeys.len()).unwrap_or(u32::MAX);
+            ensure!(
+                selection_scanned <= Self::root_claim_declared_work(),
+                Error::<T>::RootClaimTooHeavy
+            );
+            let hotkeys = Self::root_claim_hotkeys(&coldkey, staking_hotkeys);
             ensure!(
                 Self::root_claim_fits_declared_budget(&hotkeys),
                 Error::<T>::RootClaimTooHeavy
@@ -1988,7 +1991,7 @@ mod dispatches {
             let outcome = Self::do_root_claim(coldkey.clone(), hotkeys)?;
             Self::maybe_add_coldkey_index(&coldkey);
 
-            let weight = Self::root_claim_actual_weight(hotkey_count, &outcome);
+            let weight = Self::root_claim_actual_weight(hotkey_count, selection_scanned, &outcome);
             Ok((Some(weight), Pays::Yes).into())
         }
 
@@ -2008,7 +2011,7 @@ mod dispatches {
         /// * `RootClaimed`: On successfully claiming the root emissions for this coldkey+hotkey.
         #[pallet::call_index(148)]
         #[pallet::weight(
-            <T as crate::pallet::Config>::WeightInfo::claim_root(crate::MAX_ROOT_CLAIM_WORK)
+            Pallet::<T>::root_claim_declared_weight()
         )]
         pub fn claim_root_with_hotkey(
             origin: OriginFor<T>,
@@ -2023,7 +2026,7 @@ mod dispatches {
             let outcome = Self::do_root_claim(coldkey.clone(), vec![hotkey])?;
             Self::maybe_add_coldkey_index(&coldkey);
 
-            let weight = Self::root_claim_actual_weight(1, &outcome);
+            let weight = Self::root_claim_actual_weight(1, 0, &outcome);
             Ok((Some(weight), Pays::Yes).into())
         }
 

pallets/subtensor/src/macros/errors.rs

@@ -361,8 +361,8 @@ mod errors {
         /// (dividends accumulate in place) until weight setting is switched on by
         /// governance or a later upgrade.
         RootWeightSettingDisabled,
-        /// Coldkey-wide `claim_root` would process more work units than the
-        /// pre-dispatch envelope ([`crate::MAX_ROOT_CLAIM_WORK`]). Use
+        /// A root claim would process more root hotkeys and basket rows than the
+        /// fixed admission envelope. Use
         /// `claim_root_with_hotkey` per validator so admission weight matches
         /// the holdings actually walked.
         RootClaimTooHeavy,

pallets/subtensor/src/migrations/migrate_cleanup_staking_hotkeys.rs

@@ -6,7 +6,9 @@ use scale_info::prelude::string::String;
 use sp_std::collections::btree_set::BTreeSet;
 use sp_std::vec::Vec;
 
-pub const MIGRATION_NAME: &[u8] = b"migrate_cleanup_staking_hotkeys";
+// Fresh marker reruns the same bounded cleanup after share-pool dust canonicalization. The
+// original `migrate_cleanup_staking_hotkeys` pass may already be marked complete on-chain.
+pub const MIGRATION_NAME: &[u8] = b"migrate_cleanup_staking_hotkeys_v2";
 
 /// Persistent progress for the bounded `StakingHotkeys` cleanup.
 ///
@@ -48,11 +50,10 @@ fn row_load_weight<T: Config>() -> Weight {
 
 /// A conservative keep predicate for one hotkey/coldkey relationship.
 ///
-/// Any stored share row is retained, including a zero-valued legacy row. The storage-bloat
-/// migration runs first and clears zero legacy rows, while treating an unexpected remaining row
-/// as live makes this cleanup fail safe. A basket watermark is also sufficient to retain the
-/// relationship because zero-root-stake claimants still need to be discoverable by claims and
-/// coldkey swaps.
+/// Any stored non-zero share row is retained. The storage-bloat migration runs first and clears
+/// exact-zero Alpha and AlphaV2 rows; treating an unexpected remaining row as live keeps this
+/// cleanup fail safe. A basket watermark is also sufficient to retain the relationship because
+/// zero-root-stake claimants still need to be discoverable by claims and coldkey swaps.
 fn relationship_must_remain<T: Config>(hotkey: &T::AccountId, coldkey: &T::AccountId) -> bool {
     AlphaV2::<T>::iter_prefix((hotkey, coldkey))
         .next()

pallets/subtensor/src/migrations/migrate_storage_bloat_v2.rs

@@ -6,7 +6,9 @@ use scale_info::prelude::string::String;
 use sp_io::{hashing::twox_128, storage};
 use sp_std::vec::Vec;
 
-const MIGRATION_NAME: &[u8] = b"migrate_storage_bloat_v2";
+// Fresh marker reruns the bounded GC with the AlphaV2 zero-row target added below. The original
+// v2 pass may already be marked complete on-chain.
+const MIGRATION_NAME: &[u8] = b"migrate_storage_bloat_v3";
 
 #[derive(Clone, Copy, PartialEq, Eq)]
 enum CleanupMode {
@@ -90,6 +92,12 @@ const TARGETS: &[CleanupTarget] = &[
         storage: "StakingHotkeys",
         mode: CleanupMode::ClearIfZero,
     },
+    // Appended so an interrupted v2 cursor keeps the same target indices after upgrade.
+    CleanupTarget {
+        pallet: "SubtensorModule",
+        storage: "AlphaV2",
+        mode: CleanupMode::ClearIfZero,
+    },
 ];
 
 /// Persistent progress for the bounded storage cleanup.

pallets/subtensor/src/staking/claim_root.rs

@@ -4,7 +4,7 @@ use frame_support::storage::{TransactionOutcome, with_transaction};
 use frame_support::weights::{Weight, WeightMeter};
 use sp_core::Get;
 use sp_runtime::DispatchError;
-use sp_runtime::traits::AccountIdConversion;
+use sp_runtime::traits::{AccountIdConversion, Zero};
 use sp_std::collections::btree_map::BTreeMap;
 use sp_std::collections::btree_set::BTreeSet;
 use substrate_fixed::types::I96F32;
@@ -636,21 +636,6 @@ impl<T: Config> Pallet<T> {
             return Ok(outcome);
         }
 
-        // A claim that pays on some holdings but floors a valuable holding to take == 0 must
-        // not settle: burning all owed shares would forfeit that holding to remaining
-        // shareholders (e.g. 99/100 of a 1-alpha position → floor 0, then the leftover
-        // 1-share holder owns the whole unit). Same posture as the all-zero-take rollback
-        // below — leave the watermark untouched. Worthless rows (realizable 0) are ignored.
-        for (_, slot_alpha, slot_value, _) in valued_holdings.iter() {
-            let alpha = slot_alpha.to_u64();
-            if alpha == 0 {
-                continue;
-            }
-            if Self::mul_div_u64(alpha, owed_shares, shares_total) == 0 && *slot_value > 0 {
-                return Ok(outcome);
-            }
-        }
-
         let escrow = Self::get_beta_escrow_account_id();
 
         // Redeemed slots are counted outside the transaction: a rolled-back redemption
@@ -666,8 +651,25 @@ impl<T: Config> Pallet<T> {
             let mut written_off: u32 = 0;
 
             for (netuid, slot_alpha, slot_value, terminal_garbage) in valued_holdings.iter() {
+                let slot_entitlement =
+                    Self::basket_payout_from(owed_shares, *slot_value, shares_total);
                 // This staker's pro-rata slice of the holding: slot_alpha * owed / P.
-                let take: u64 = Self::mul_div_u64(slot_alpha.to_u64(), owed_shares, shares_total);
+                let proportional_take =
+                    Self::mul_div_u64(slot_alpha.to_u64(), owed_shares, shares_total);
+                // A high-value alpha row can owe at least one rao even when its proportional
+                // alpha slice floors to zero. Sell one atomic alpha unit, pay no more than the
+                // marked entitlement below, and retain the sale surplus as fund root cash.
+                // Root is already denominated in rao, so take == entitlement there; terminal
+                // rows have no realizable entitlement and keep the ordinary floor.
+                let take = if proportional_take == 0
+                    && slot_entitlement > 0
+                    && !netuid.is_root()
+                    && !terminal_garbage
+                {
+                    1
+                } else {
+                    proportional_take
+                };
                 if take == 0 {
                     continue;
                 }
@@ -729,8 +731,6 @@ impl<T: Config> Pallet<T> {
                 // fraction, so pay only the priced entitlement and retain the surplus as
                 // fund cash. Otherwise a permissionless deposit followed by a claim can
                 // extract the difference from earlier holders.
-                let slot_entitlement =
-                    Self::basket_payout_from(owed_shares, *slot_value, shares_total);
                 let realized_tao = tao.to_u64();
                 // A final claimant has no remaining holders to retain a surplus for. Give
                 // them every realized rao so no root cash is stranded behind zero shares.
@@ -895,51 +895,82 @@ impl<T: Config> Pallet<T> {
         swept
     }
 
-    /// Live existing-network count (including root). Ghost `NetworksAdded=false`
-    /// leftovers are excluded so they cannot inflate the single-hotkey quote.
-    pub(crate) fn root_claim_existing_networks() -> u32 {
-        (Self::get_all_subnet_netuids().len() as u32).max(1)
-    }
-
-    /// Pre-dispatch work units for both claim paths. Weight calculation must
-    /// stay storage-independent (no `NetworksAdded` or basket walks here);
-    /// execution then refuses work that would exceed this envelope.
+    /// Fixed admission budget for both claim paths.
     pub(crate) fn root_claim_declared_work() -> u32 {
         crate::MAX_ROOT_CLAIM_WORK
     }
 
-    /// True when a coldkey-wide claim's reachable work (hotkeys × existing
-    /// networks, and the actual basket rows) fits the admission envelope.
+    /// Pre-dispatch weight for both independently bounded dimensions: full claim work and
+    /// scan-only work.
+    pub(crate) fn root_claim_declared_weight() -> Weight {
+        let limit = Self::root_claim_declared_work();
+        <T as crate::pallet::Config>::WeightInfo::claim_root(limit).saturating_add(
+            <T as crate::pallet::Config>::WeightInfo::claim_root_scan(limit),
+        )
+    }
+
+    /// Hotkeys relevant to a coldkey-wide root claim. Ordinary subnet-only staking hotkeys
+    /// are deliberately excluded. A negative basket watermark keeps an unstaked claimant
+    /// eligible because it encodes shares which still need to be redeemed.
+    pub(crate) fn root_claim_hotkeys(
+        coldkey: &T::AccountId,
+        staking_hotkeys: Vec<T::AccountId>,
+    ) -> Vec<T::AccountId> {
+        staking_hotkeys
+            .into_iter()
+            .filter(|hotkey| {
+                !Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, NetUid::ROOT)
+                    .is_zero()
+                    || BasketClaimed::<T>::get(hotkey, coldkey) < 0
+            })
+            .collect()
+    }
+
+    /// True when the hotkeys plus the basket storage rows the claim will scan fit the fixed
+    /// admission envelope. Count raw Alpha/AlphaV2 rows so legacy duplicates and malformed
+    /// zero rows are charged conservatively, and stop as soon as the bound is exceeded.
     pub(crate) fn root_claim_fits_declared_budget(hotkeys: &[T::AccountId]) -> bool {
         let budget = Self::root_claim_declared_work();
-        let networks = Self::root_claim_existing_networks();
-        let hotkey_count = hotkeys.len() as u32;
-        if hotkey_count.saturating_mul(networks) > budget {
+        let mut work = u32::try_from(hotkeys.len()).unwrap_or(u32::MAX);
+        if work > budget {
             return false;
         }
-        let mut rows: u32 = 0;
+
+        let escrow = Self::get_beta_escrow_account_id();
         for hotkey in hotkeys {
-            rows = rows.saturating_add(Self::get_basket_holdings(hotkey).len() as u32);
-            if rows > budget {
-                return false;
+            for _ in Alpha::<T>::iter_prefix((hotkey, &escrow)) {
+                work = work.saturating_add(1);
+                if work > budget {
+                    return false;
+                }
+            }
+            for _ in AlphaV2::<T>::iter_prefix((hotkey, &escrow)) {
+                work = work.saturating_add(1);
+                if work > budget {
+                    return false;
+                }
             }
         }
         true
     }
 
-    /// Actual post-dispatch weight of a root claim: full benchmark units for the slots
-    /// that did real work (redeemed or swept — a swap plus stake writes each, floored at
-    /// the hotkey count so walking empty hotkeys stays covered) plus the cheap per-row
-    /// scan cost for holdings that were only valued. This is what lets a fund's claim fee
-    /// decay as dust rows are consolidated, and makes a below-threshold no-op cost a scan
-    /// instead of a full claim. Work above [`crate::MAX_ROOT_CLAIM_WORK`] is
-    /// refused at dispatch (`RootClaimTooHeavy`) rather than admitted cheaply.
+    /// Actual post-dispatch weight of a root claim: full benchmark units for relationships
+    /// classified and slots that did real work (redeemed or swept — a swap plus stake writes
+    /// each, floored at the selected hotkey count) plus the cheap per-row scan cost for holdings
+    /// that were only valued. This is what lets a fund's claim fee decay as dust rows are
+    /// consolidated, and makes a below-threshold no-op cost a scan instead of a full claim.
+    /// Work above the fixed admission budget
+    /// is refused at dispatch (`RootClaimTooHeavy`) rather than admitted cheaply.
     pub(crate) fn root_claim_actual_weight(
         hotkey_count: u32,
+        selection_scanned: u32,
         outcome: &RootClaimOutcome,
     ) -> Weight {
         let active = hotkey_count
             .max(outcome.realized.saturating_add(outcome.swept))
+            // Classifying a StakingHotkeys relationship reads the position's share-pool state
+            // and basket watermark. Price it conservatively as a full hotkey unit.
+            .max(selection_scanned)
             .max(1);
         let scanned = outcome.rows.saturating_sub(outcome.realized);
         <T as crate::pallet::Config>::WeightInfo::claim_root(active).saturating_add(