fix(tx): make NoteConsumptionChecker bundle-aware

CHANGELOG.md

@@ -59,6 +59,7 @@
 
 ### Fixes
 
+- `NoteConsumptionChecker` now groups input notes into bundles that must be consumed together and searches over those bundles, so a feature note and the `FEE_SPONSORSHIP` notes bound to it are kept as a unit instead of being dropped alongside an unrelated note whose fee is uncovered ([#3710](https://github.com/0xMiden/protocol/issues/3710)).
 - Fixed `AuthNetworkAccount` accepting empty fee-only transactions, which let callers drain the account's native fee-asset vault ([#3729](https://github.com/0xMiden/protocol/pull/3729)).
 - [BREAKING] AggLayer bridge token registration now rejects keys owned by another faucet, and token-key cleanup verifies ownership before clearing a mapping ([#3754](https://github.com/0xMiden/protocol/pull/3754)).
 - [BREAKING] AggLayer bridges now allow faucet deregistration while paused, so compromised faucets can be revoked without resuming claims and bridge-outs ([#3750](https://github.com/0xMiden/protocol/pull/3753)).

crates/miden-standards/src/note/fee_sponsorship.rs

@@ -149,6 +149,24 @@ impl FeeSponsorshipNote {
         FEE_SPONSORSHIP_SCRIPT.root()
     }
 
+    /// Returns the ID of the feature note that `note` sponsors, or `None` if `note` is not a
+    /// FEE_SPONSORSHIP note or its storage does not decode as [`FeeSponsorshipNoteStorage`].
+    ///
+    /// This is the binding that `collect_sponsored_fees` relies on: a sponsorship names its
+    /// feature note by ID, so the pair may appear at any positions in the input notes and several
+    /// sponsorships may be bound to the same feature note. Use this to group input notes into the
+    /// sets that must be consumed together, since neither half of a bound pair is consumable on
+    /// its own.
+    pub fn sponsored_feature_note_id(note: &Note) -> Option<NoteId> {
+        if note.script().root() != Self::script_root() {
+            return None;
+        }
+
+        FeeSponsorshipNoteStorage::try_from(note.storage().items())
+            .ok()
+            .map(|storage| storage.feature_note_id())
+    }
+
     /// Returns the account ID of the sponsor which created the note.
     pub fn sender(&self) -> AccountId {
         self.sender

crates/miden-testing/tests/scripts/fee_collection.rs

@@ -43,6 +43,8 @@ use miden_standards::note::{
 };
 use miden_standards::testing::note::NoteBuilder;
 use miden_testing::{Auth, MockChain, MockTransaction, assert_transaction_executor_error};
+use miden_tx::auth::UnreachableAuth;
+use miden_tx::{NoteConsumptionChecker, TransactionExecutor};
 use rand::SeedableRng;
 use rand::rngs::Xoshiro256PlusPlus;
 use rand::seq::SliceRandom;
@@ -1299,3 +1301,194 @@ async fn sponsoring_with_a_non_native_fee_asset_is_rejected() -> anyhow::Result<
 
     Ok(())
 }
+
+// NOTE CONSUMPTION CHECKER
+// ================================================================================================
+
+/// Runs [`NoteConsumptionChecker::check_notes_consumability`] over `notes` against the network
+/// account and returns the note ids it reports as `(successful, failed)`.
+async fn check_consumability(
+    mock_chain: &MockChain,
+    network_account_id: AccountId,
+    notes: Vec<Note>,
+) -> anyhow::Result<(BTreeSet<NoteId>, BTreeSet<NoteId>)> {
+    let mut builder = mock_chain.build_transaction(network_account_id);
+    for note in &notes {
+        builder = builder.authenticated_input_note(note.id());
+    }
+    let mock_tx = builder.build()?;
+
+    let executor = TransactionExecutor::<'_, '_, _, UnreachableAuth>::new(&mock_tx)
+        .with_source_manager(mock_tx.source_manager());
+    let info = NoteConsumptionChecker::new(&executor)
+        .check_notes_consumability(
+            network_account_id,
+            mock_tx.tx_inputs().block_header().block_num(),
+            notes,
+            mock_tx.tx_args().clone(),
+        )
+        .await?;
+
+    Ok((
+        info.successful().iter().map(|note| note.note().id()).collect(),
+        info.failed().iter().map(|note| note.note().id()).collect(),
+    ))
+}
+
+/// An uncovered feature note does not drag the intact (feature note, FEE_SPONSORSHIP) pairs
+/// sharing its batch down with it.
+///
+/// The uncovered note aborts fee collection in the auth procedure, which runs after note
+/// processing and so is reported as an epilogue failure. That routes the batch into a search for
+/// the largest executable set, which grows its candidate set one bundle at a time: a bound pair is
+/// a single bundle, so it is tested - and kept - as a unit even though neither of its halves is
+/// consumable alone.
+///
+/// Writing `F` for a feature note, `S` for the sponsorship bound to it and `S'` for an underfunded
+/// one, the cases below are:
+///
+/// ```text
+/// [F0, S0]           ->  successful {F0, S0}   failed {}
+/// [F0, S0, F1]       ->  successful {F0, S0}   failed {F1}
+/// [F0, S0, F1, S1']  ->  successful {F0, S0}   failed {F1, S1'}
+/// ```
+#[rstest]
+#[case::no_sponsorship(None)]
+#[case::underfunded_sponsorship(Some(FEE_AMOUNT - 1))]
+#[tokio::test]
+async fn note_checker_keeps_intact_pairs_alongside_an_uncovered_note(
+    #[case] uncovered_sponsored_amount: Option<u64>,
+) -> anyhow::Result<()> {
+    let mut test_builder = Test::builder()
+        .feature_note_fee(AssetAmount::new(FEE_AMOUNT)?)
+        .num_feature_notes(2)
+        .sponsorship(0, fee_asset(FEE_AMOUNT)?);
+    if let Some(amount) = uncovered_sponsored_amount {
+        test_builder = test_builder.sponsorship(1, fee_asset(amount)?);
+    }
+    let Test {
+        mock_chain,
+        network_account,
+        feature_notes,
+        sponsorship_notes,
+    } = test_builder.build()?;
+
+    // control: the intact pair is consumable on its own
+    let intact = vec![feature_notes[0].clone(), sponsorship_notes[0].clone()];
+    let intact_ids: BTreeSet<NoteId> = intact.iter().map(Note::id).collect();
+    let (successful, failed) =
+        check_consumability(&mock_chain, network_account.id(), intact.clone()).await?;
+    assert_eq!(successful, intact_ids, "the intact pair should be consumable on its own");
+    assert!(failed.is_empty(), "no note of an intact pair should fail");
+
+    // subject: the same pair, plus the uncovered feature note and its underfunded sponsorship
+    let mut poisoned = intact;
+    poisoned.push(feature_notes[1].clone());
+    let mut uncovered_ids = BTreeSet::from([feature_notes[1].id()]);
+    if uncovered_sponsored_amount.is_some() {
+        poisoned.push(sponsorship_notes[1].clone());
+        uncovered_ids.insert(sponsorship_notes[1].id());
+    }
+
+    let (successful, failed) =
+        check_consumability(&mock_chain, network_account.id(), poisoned).await?;
+    assert_eq!(successful, intact_ids, "the intact pair should survive the uncovered note");
+    assert_eq!(failed, uncovered_ids, "only the uncovered note and its sponsorship should fail");
+
+    Ok(())
+}
+
+/// A FEE_SPONSORSHIP note whose feature note is absent is not bundled with anything: it fails on
+/// its own - it can neither be collected, since fee collection rejects a sponsorship without its
+/// feature note, nor reclaimed by the network account - and leaves an intact pair untouched.
+#[tokio::test]
+async fn note_checker_fails_an_orphan_sponsorship_alone() -> anyhow::Result<()> {
+    let Test {
+        mock_chain,
+        network_account,
+        feature_notes,
+        sponsorship_notes,
+    } = Test::builder()
+        .feature_note_fee(AssetAmount::new(FEE_AMOUNT)?)
+        .num_feature_notes(2)
+        .sponsorship(0, fee_asset(FEE_AMOUNT)?)
+        .sponsorship(1, fee_asset(FEE_AMOUNT)?)
+        .build()?;
+
+    // the sponsorship for the second feature note is included, but that feature note is not
+    let notes = vec![
+        feature_notes[0].clone(),
+        sponsorship_notes[0].clone(),
+        sponsorship_notes[1].clone(),
+    ];
+
+    let (successful, failed) =
+        check_consumability(&mock_chain, network_account.id(), notes).await?;
+
+    assert_eq!(
+        successful,
+        BTreeSet::from([feature_notes[0].id(), sponsorship_notes[0].id()]),
+        "the intact pair should survive the orphan sponsorship"
+    );
+    assert_eq!(
+        failed,
+        BTreeSet::from([sponsorship_notes[1].id()]),
+        "only the orphan sponsorship should fail"
+    );
+
+    Ok(())
+}
+
+/// Bundling follows the note IDs the sponsorships name, not the order the notes are passed in, and
+/// several sponsorships topping up one feature note join the same bundle.
+#[rstest]
+#[tokio::test]
+async fn note_checker_bundles_by_note_id_regardless_of_order(
+    #[values(false, true)] sponsorships_first: bool,
+) -> anyhow::Result<()> {
+    let Test {
+        mock_chain,
+        network_account,
+        feature_notes,
+        sponsorship_notes,
+    } = Test::builder()
+        .feature_note_fee(AssetAmount::new(FEE_AMOUNT)?)
+        .num_feature_notes(2)
+        // the first feature note's fee is split across two sponsorships, so its bundle holds three
+        // notes; the second feature note stays uncovered
+        .sponsorship(0, fee_asset(FEE_AMOUNT / 2)?)
+        .sponsorship(0, fee_asset(FEE_AMOUNT - FEE_AMOUNT / 2)?)
+        .build()?;
+
+    let covered = [
+        feature_notes[0].clone(),
+        sponsorship_notes[0].clone(),
+        sponsorship_notes[1].clone(),
+    ];
+    let mut notes = if sponsorships_first {
+        vec![
+            sponsorship_notes[0].clone(),
+            sponsorship_notes[1].clone(),
+            feature_notes[0].clone(),
+        ]
+    } else {
+        covered.to_vec()
+    };
+    notes.push(feature_notes[1].clone());
+
+    let (successful, failed) =
+        check_consumability(&mock_chain, network_account.id(), notes).await?;
+
+    assert_eq!(
+        successful,
+        covered.iter().map(Note::id).collect::<BTreeSet<_>>(),
+        "the feature note and both of its sponsorships should be consumed together"
+    );
+    assert_eq!(
+        failed,
+        BTreeSet::from([feature_notes[1].id()]),
+        "only the uncovered feature note should fail"
+    );
+
+    Ok(())
+}