Change outputs should be part of unilateral actions

This also removes tx creation code from the wallet entirely
for the multi party sessions. Coin, output and change selection
all happens when enumurating actions.

Author: arminsabouri

src/actions.rs

@@ -1,6 +1,7 @@
 use std::{iter::Sum, ops::Add};
 
 use bdk_coin_select::{Target, TargetFee, TargetOutputs};
+use bitcoin::Amount;
 use log::debug;
 
 use crate::{
@@ -44,12 +45,12 @@ fn piecewise_linear(x: f64, points: &[(f64, f64)]) -> f64 {
 /// An Action a wallet can perform
 #[derive(Debug)]
 pub(crate) enum Action {
-    /// Spend a payment obligation unilaterally with pre-selected inputs
-    UnilateralPayments(Vec<PaymentObligationId>, Vec<Outpoint>),
+    /// Spend a payment obligation unilaterally with pre-selected inputs and pre-computed change
+    UnilateralPayments(Vec<PaymentObligationId>, Vec<Outpoint>, Vec<Amount>),
     /// Accept a cospend invitation
     AcceptCospendProposal((MessageId, BulletinBoardId)),
-    /// Contribute outputs to a cospend session that is waiting for them
-    ContributeOutputsToSession(BulletinBoardId, Vec<PaymentObligationId>),
+    /// Contribute outputs to a cospend session that is waiting for them, with pre-computed change
+    ContributeOutputsToSession(BulletinBoardId, Vec<PaymentObligationId>, Vec<Amount>),
     /// Continue to participate in a multi-party payjoin
     ContinueParticipateInCospend(BulletinBoardId),
     /// Taker records non-committal interest in cospending with each orderbook UTXO
@@ -163,8 +164,8 @@ fn simulate_one_action(wallet_handle: &WalletHandleMut, action: &Action) -> Pred
 
     // POs handled: derived from action since confirmation is deferred to block
     let payment_obligations_handled: Vec<PaymentObligationId> = match action {
-        Action::UnilateralPayments(po_ids, _) => po_ids.clone(),
-        Action::ContributeOutputsToSession(_, po_ids) => po_ids.clone(),
+        Action::UnilateralPayments(po_ids, _, _) => po_ids.clone(),
+        Action::ContributeOutputsToSession(_, po_ids, _) => po_ids.clone(),
         _ => vec![],
     };
 
@@ -296,6 +297,36 @@ fn target_for_obligations(pos: &[PaymentObligationData], wallet: &WalletHandleMu
     }
 }
 
+/// Compute pre-selected change outputs for a `ContributeOutputsToSession` action.
+/// If the session has pre-selected inputs (from the aggregator), uses those exactly.
+/// Otherwise falls back to full BNB / spend-all selection over all wallet UTXOs.
+fn change_for_session_contribution(
+    bb_id: &BulletinBoardId,
+    pos: &[PaymentObligationData],
+    wallet: &WalletHandleMut,
+) -> Vec<Amount> {
+    let session = wallet
+        .info()
+        .active_multi_party_payjoins
+        .get(bb_id)
+        .unwrap();
+    let session_input_outpoints: Vec<Outpoint> =
+        session.inputs.iter().map(|i| i.outpoint).collect();
+    let target = target_for_obligations(pos, wallet);
+    if session_input_outpoints.is_empty() {
+        let candidates = wallet.handle().coin_candidates();
+        if let Some((_, change)) = select_bnb(&candidates, target) {
+            return change;
+        }
+        select_all(&candidates, target).1
+    } else {
+        let candidates = wallet
+            .handle()
+            .coin_candidates_for(&session_input_outpoints);
+        select_all(&candidates, target).1
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct UnilateralSpender;
 
@@ -315,12 +346,12 @@ impl Strategy for UnilateralSpender {
         let mut actions = vec![];
         for po in state.payment_obligations.iter() {
             let target = target_for_obligations(std::slice::from_ref(po), wallet);
-            if let Some(inputs) = select_bnb(&candidates, target) {
-                actions.push(Action::UnilateralPayments(vec![po.id], inputs));
+            if let Some((inputs, change)) = select_bnb(&candidates, target) {
+                actions.push(Action::UnilateralPayments(vec![po.id], inputs, change));
             }
-            let all_inputs = select_all(&candidates);
+            let (all_inputs, change) = select_all(&candidates, target);
             if !all_inputs.is_empty() {
-                actions.push(Action::UnilateralPayments(vec![po.id], all_inputs));
+                actions.push(Action::UnilateralPayments(vec![po.id], all_inputs, change));
             }
         }
         if actions.is_empty() {
@@ -346,11 +377,12 @@ impl Strategy for Consolidator {
         wallet: &WalletHandleMut,
     ) -> Vec<Action> {
         let candidates = wallet.handle().coin_candidates();
-        let all_inputs = select_all(&candidates);
         let mut actions = Vec::new();
         for po in state.payment_obligations.iter() {
+            let target = target_for_obligations(std::slice::from_ref(po), wallet);
+            let (all_inputs, change) = select_all(&candidates, target);
             if !all_inputs.is_empty() {
-                actions.push(Action::UnilateralPayments(vec![po.id], all_inputs.clone()));
+                actions.push(Action::UnilateralPayments(vec![po.id], all_inputs, change));
             }
         }
         actions.push(Action::Wait);
@@ -380,12 +412,12 @@ impl Strategy for BatchSpender {
         let target = target_for_obligations(&state.payment_obligations, wallet);
         let candidates = wallet.handle().coin_candidates();
         let mut actions = vec![];
-        if let Some(inputs) = select_bnb(&candidates, target) {
-            actions.push(Action::UnilateralPayments(po_ids.clone(), inputs));
+        if let Some((inputs, change)) = select_bnb(&candidates, target) {
+            actions.push(Action::UnilateralPayments(po_ids.clone(), inputs, change));
         }
-        let all_inputs = select_all(&candidates);
+        let (all_inputs, change) = select_all(&candidates, target);
         if !all_inputs.is_empty() {
-            actions.push(Action::UnilateralPayments(po_ids, all_inputs));
+            actions.push(Action::UnilateralPayments(po_ids, all_inputs, change));
         }
         if actions.is_empty() {
             actions.push(Action::Wait);
@@ -424,11 +456,17 @@ impl Strategy for MakerStrategy {
             }
         }
 
-        // Contribute outputs to sessions that are waiting for them (SentInputs state)
+        // Contribute outputs to sessions that are waiting for them (AcceptedProposal state)
         for (bb_id, session) in wallet.info().active_multi_party_payjoins.iter() {
             if session.state == TxConstructionState::AcceptedProposal {
                 for po in state.payment_obligations.iter() {
-                    actions.push(Action::ContributeOutputsToSession(*bb_id, vec![po.id]));
+                    let change =
+                        change_for_session_contribution(bb_id, std::slice::from_ref(po), wallet);
+                    actions.push(Action::ContributeOutputsToSession(
+                        *bb_id,
+                        vec![po.id],
+                        change,
+                    ));
                 }
             }
         }
@@ -446,11 +484,11 @@ impl Strategy for MakerStrategy {
         } else {
             vec![]
         };
-        // Selected inputs are already embedded in each action — no simulation needed.
+        // Selected inputs are already embedded in each action i.e no simulation needed.
         let per_action_inputs: Vec<std::collections::HashSet<Outpoint>> = unilateral_actions
             .iter()
             .filter_map(|a| match a {
-                Action::UnilateralPayments(_, inputs) => Some(inputs.iter().copied().collect()),
+                Action::UnilateralPayments(_, inputs, _) => Some(inputs.iter().copied().collect()),
                 _ => None,
             })
             .collect();
@@ -490,11 +528,17 @@ impl Strategy for TakerStrategy {
     ) -> Vec<Action> {
         let mut actions = vec![];
 
-        // Contribute outputs to sessions awaiting them (SentInputs state)
+        // Contribute outputs to sessions awaiting them (AcceptedProposal state)
         for (bb_id, session) in wallet.info().active_multi_party_payjoins.iter() {
             if session.state == TxConstructionState::AcceptedProposal {
                 for po in state.payment_obligations.iter() {
-                    actions.push(Action::ContributeOutputsToSession(*bb_id, vec![po.id]));
+                    let change =
+                        change_for_session_contribution(bb_id, std::slice::from_ref(po), wallet);
+                    actions.push(Action::ContributeOutputsToSession(
+                        *bb_id,
+                        vec![po.id],
+                        change,
+                    ));
                 }
             }
         }
@@ -724,7 +768,7 @@ mod tests {
         assert!(actions.iter().any(|a| matches!(a, Action::Wait)));
         assert!(!actions
             .iter()
-            .any(|a| matches!(a, Action::UnilateralPayments(_, _))));
+            .any(|a| matches!(a, Action::UnilateralPayments(_, _, _))));
     }
 
     #[test]
@@ -752,13 +796,16 @@ mod tests {
 
         let actions = strategy.enumerate_candidate_actions(&view, &wallet);
 
-        // No UTXOs — coin selection produces nothing, falls back to Wait.
+        // No UTXOs coin selection produces nothing, falls back to Wait.
         assert_eq!(actions.len(), 1);
         assert!(matches!(actions[0], Action::Wait));
     }
 
     #[test]
     fn test_composite_strategy_combines_actions() {
+        // TODO: this test is kinda useless, we need to add UTXOs to the sim and test the composite strategy.
+        // Otherwise we are just testing that both strategies fall back to Wait when there are no UTXOs.
+        // This is bc coin selection uses `wallet.handle().coin_candidates();` not `state.utxos`.
         let mut sim = test_sim();
         let wallet = WalletId(0).with_mut(&mut sim);
         let composite = CompositeStrategy {
@@ -785,24 +832,9 @@ mod tests {
 
         let actions = composite.enumerate_candidate_actions(&view, &wallet);
 
-        // Should include actions from both strategies
-        // UnilateralSpender: 2 actions (one per obligation, single-PO each)
-        // BatchSpender: 1 action (all obligations in one tx)
-        assert_eq!(actions.len(), 3);
-
-        let single_po_count = actions
-            .iter()
-            .filter(|a| matches!(a, Action::UnilateralPayments(ids, _) if ids.len() == 1))
-            .count();
-        assert_eq!(single_po_count, 2);
-
-        let batch_count = actions
-            .iter()
-            .filter(|a| matches!(a, Action::UnilateralPayments(ids, _) if ids.len() == 2))
-            .count();
-        assert_eq!(batch_count, 1);
-        // No UTXOs both strategies fall back to Wait, composite collects both.
-        assert_eq!(actions.len(), 3);
+        // Wallet has no UTXOs in the sim, both strategies fall back to Wait.
+        // Composite collects one Wait from each strategy.
+        assert_eq!(actions.len(), 2);
         assert!(actions.iter().all(|a| matches!(a, Action::Wait)));
     }
 
@@ -820,7 +852,7 @@ mod tests {
             from: WalletId(0),
             to: WalletId(1),
         };
-        // No orderbook UTXOs — taker has nothing to propose to
+        // No orderbook UTXOs, taker has nothing to propose to
         let view = create_test_wallet_view(vec![po]);
 
         let actions = strategy.enumerate_candidate_actions(&view, &wallet);
@@ -982,7 +1014,7 @@ mod tests {
 
         assert!(actions
             .iter()
-            .any(|a| matches!(a, Action::ContributeOutputsToSession(id, ids) if *id == bb_id && ids.len() == 1)));
+            .any(|a| matches!(a, Action::ContributeOutputsToSession(id, ids, _) if *id == bb_id && ids.len() == 1)));
         // Should NOT emit ContinueParticipateInCospend for this session
         assert!(!actions
             .iter()

src/coin_selection.rs

@@ -2,6 +2,7 @@ use bdk_coin_select::{
     metrics::LowestFee, Candidate, ChangePolicy, CoinSelector, DrainWeights, Target,
     TR_DUST_RELAY_MIN_VALUE,
 };
+use bitcoin::Amount;
 use log::warn;
 
 use crate::transaction::Outpoint;
@@ -18,30 +19,50 @@ pub(crate) fn long_term_feerate() -> bdk_coin_select::FeeRate {
     bdk_coin_select::FeeRate::from_sat_per_wu(2.5)
 }
 
-/// Run BNB coin selection over candidates for the given target.
-/// Falls back to greedy selection if BNB finds no solution.
-/// Returns None if no selection can meet the target.
-pub(crate) fn select_bnb(candidates: &[CoinCandidate], target: Target) -> Option<Vec<Outpoint>> {
-    let bdk_candidates: Vec<Candidate> = candidates
+fn change_policy_for(target: Target) -> ChangePolicy {
+    ChangePolicy::min_value_and_waste(
+        DrainWeights::default(),
+        TR_DUST_RELAY_MIN_VALUE,
+        target.fee.rate,
+        long_term_feerate(),
+    )
+}
+
+fn bdk_candidates(candidates: &[CoinCandidate]) -> Vec<Candidate> {
+    candidates
         .iter()
         .map(|c| Candidate {
             value: c.amount_sats,
             weight: c.weight_wu,
             input_count: 1,
             is_segwit: c.is_segwit,
         })
-        .collect();
+        .collect()
+}
 
-    let mut coin_selector = CoinSelector::new(&bdk_candidates);
+fn drain_to_change(drain: bdk_coin_select::Drain) -> Vec<Amount> {
+    if drain.value > 0 {
+        vec![Amount::from_sat(drain.value)]
+    } else {
+        vec![]
+    }
+}
+
+/// Run BNB coin selection over candidates for the given target.
+/// Falls back to greedy selection if BNB finds no solution.
+/// Returns None if no selection can meet the target.
+/// Returns (selected_inputs, change_outputs).
+pub(crate) fn select_bnb(
+    candidates: &[CoinCandidate],
+    target: Target,
+) -> Option<(Vec<Outpoint>, Vec<Amount>)> {
+    let bdk = bdk_candidates(candidates);
+    let mut coin_selector = CoinSelector::new(&bdk);
 
-    let drain_weights = DrainWeights::default();
-    let dust_limit = TR_DUST_RELAY_MIN_VALUE;
-    let ltfr = long_term_feerate();
-    let change_policy =
-        ChangePolicy::min_value_and_waste(drain_weights, dust_limit, target.fee.rate, ltfr);
+    let change_policy = change_policy_for(target);
     let metric = LowestFee {
         target,
-        long_term_feerate: ltfr,
+        long_term_feerate: long_term_feerate(),
         change_policy,
     };
 
@@ -52,15 +73,26 @@ pub(crate) fn select_bnb(candidates: &[CoinCandidate], target: Target) -> Option
         }
     }
 
-    Some(
-        coin_selector
-            .apply_selection(candidates)
-            .map(|c| c.outpoint)
-            .collect(),
-    )
+    let inputs = coin_selector
+        .apply_selection(candidates)
+        .map(|c| c.outpoint)
+        .collect();
+    let change = drain_to_change(coin_selector.drain(target, change_policy));
+    Some((inputs, change))
 }
 
-/// Return all candidate outpoints (consolidation / spend-all strategy).
-pub(crate) fn select_all(candidates: &[CoinCandidate]) -> Vec<Outpoint> {
-    candidates.iter().map(|c| c.outpoint).collect()
+/// Select all candidates (consolidation / spend-all strategy).
+/// Returns (selected_inputs, change_outputs).
+pub(crate) fn select_all(
+    candidates: &[CoinCandidate],
+    target: Target,
+) -> (Vec<Outpoint>, Vec<Amount>) {
+    let bdk = bdk_candidates(candidates);
+    let mut coin_selector = CoinSelector::new(&bdk);
+    coin_selector.select_all();
+
+    let change_policy = change_policy_for(target);
+    let inputs = candidates.iter().map(|c| c.outpoint).collect();
+    let change = drain_to_change(coin_selector.drain(target, change_policy));
+    (inputs, change)
 }