From 255b64df0dde2aec30696f3b1e576edaf6fea0d1 Mon Sep 17 00:00:00 2001
From: dancoombs <dan.coombs@alchemy.com>
Date: Fri, 14 Jun 2024 12:23:53 -0500
Subject: [PATCH] feat(pool): allow staked senders to have multiple UOs in
 best_operations

---
 crates/pool/src/mempool/uo_pool.rs      | 31 +++++++++++++++++++------
 test/spec-tests/local/launcher.sh       |  2 +-
 test/spec-tests/v0_6/bundler-spec-tests |  2 +-
 test/spec-tests/v0_7/bundler-spec-tests |  2 +-
 4 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/crates/pool/src/mempool/uo_pool.rs b/crates/pool/src/mempool/uo_pool.rs
index 9ee07f13f..925de0bef 100644
--- a/crates/pool/src/mempool/uo_pool.rs
+++ b/crates/pool/src/mempool/uo_pool.rs
@@ -641,8 +641,12 @@ where
                         .div_mod(self.config.num_shards.into())
                         .1
                         == shard_index.into())) &&
-                // filter out ops from senders we've already seen
-                senders.insert(op.uo.sender())
+                // filter out ops from unstaked senders we've already seen
+                if !op.account_is_staked {
+                    senders.insert(op.uo.sender())
+                } else {
+                    true
+                }
             })
             .take(max)
             .map(Into::into)
@@ -1031,13 +1035,13 @@ mod tests {
     async fn test_account_reputation() {
         let address = Address::random();
         let (pool, uos) = create_pool_insert_ops(vec![
-            create_op_with_errors(address, 0, 2, None, None, true),
-            create_op_with_errors(address, 1, 2, None, None, true),
-            create_op_with_errors(address, 1, 2, None, None, true),
+            create_op_with_errors(address, 0, 2, None, None, true), // accept
+            create_op_with_errors(address, 1, 2, None, None, true), // accept
+            create_op_with_errors(address, 1, 2, None, None, true), // reject
         ])
         .await;
-        // Only return 1 op per sender
-        check_ops(pool.best_operations(3, 0).unwrap(), vec![uos[0].clone()]);
+        // staked, so include all ops
+        check_ops(pool.best_operations(3, 0).unwrap(), uos[0..2].to_vec());
 
         let rep = pool.dump_reputation();
         assert_eq!(rep.len(), 1);
@@ -1489,6 +1493,19 @@ mod tests {
             .is_err());
     }
 
+    #[tokio::test]
+    async fn test_best_staked() {
+        let address = Address::random();
+        let (pool, uos) = create_pool_insert_ops(vec![
+            create_op_with_errors(address, 0, 2, None, None, true),
+            create_op_with_errors(address, 1, 2, None, None, true),
+            create_op_with_errors(address, 2, 2, None, None, true),
+        ])
+        .await;
+        // staked, so include all ops
+        check_ops(pool.best_operations(3, 0).unwrap(), uos);
+    }
+
     #[derive(Clone, Debug)]
     struct OpWithErrors {
         op: UserOperationVariant,
diff --git a/test/spec-tests/local/launcher.sh b/test/spec-tests/local/launcher.sh
index ca440b695..fb967e5e7 100755
--- a/test/spec-tests/local/launcher.sh
+++ b/test/spec-tests/local/launcher.sh
@@ -7,7 +7,7 @@ case $1 in
  start)
 	docker-compose up -d
 	sleep 10
-	cast send --unlocked --from $(cast rpc eth_accounts | tail -n 1 | tr -d '[]"') --value 1ether 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266 > /dev/null
+	cast send --unlocked --from $(cast rpc eth_accounts | tail -n 1 | tr -d '[]"') --value 100ether 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266 > /dev/null
 	(cd ../$2/bundler-spec-tests/@account-abstraction && yarn deploy --network localhost)
 	../../../target/debug/rundler node --log.file out.log &
 	while [[ "$(curl -s -o /dev/null -w ''%{http_code}'' localhost:3000/health)" != "200" ]]; do sleep 1 ; done
diff --git a/test/spec-tests/v0_6/bundler-spec-tests b/test/spec-tests/v0_6/bundler-spec-tests
index 6bed71394..e6e2b1adf 160000
--- a/test/spec-tests/v0_6/bundler-spec-tests
+++ b/test/spec-tests/v0_6/bundler-spec-tests
@@ -1 +1 @@
-Subproject commit 6bed71394894ed514fc3c35585a3217a40c2e2b3
+Subproject commit e6e2b1adf5b592a45b96829bcb41a4966c566a55
diff --git a/test/spec-tests/v0_7/bundler-spec-tests b/test/spec-tests/v0_7/bundler-spec-tests
index 56888c610..032e69d59 160000
--- a/test/spec-tests/v0_7/bundler-spec-tests
+++ b/test/spec-tests/v0_7/bundler-spec-tests
@@ -1 +1 @@
-Subproject commit 56888c610d00cc8ea1cdefbae0209d158603a8bb
+Subproject commit 032e69d598c3cdf09b27b58eb6ee294e4d1580be