New subscription and wallet sync protocol

chia/_tests/core/full_node/stores/test_coin_store.py

@@ -523,7 +523,7 @@ def random_coin_records() -> RandomCoinRecords:
         coin = Coin(
             std_hash(b"Parent Coin Id " + i.to_bytes(4, byteorder="big")),
             puzzle_hash,
-            uint64(1000),
+            uint64(i),
         )
 
         if is_hinted:
@@ -552,12 +552,16 @@ def random_coin_records() -> RandomCoinRecords:
 @pytest.mark.parametrize("include_spent", [True, False])
 @pytest.mark.parametrize("include_unspent", [True, False])
 @pytest.mark.parametrize("include_hinted", [True, False])
+@pytest.mark.parametrize(
+    "min_amount", [uint64(0), uint64(30000), uint64(0xFFFF), uint64(0x7FFF), uint64(0x8000), uint64(0x8000000000000000)]
+)
 async def test_coin_state_batches(
     db_version: int,
     random_coin_records: RandomCoinRecords,
     include_spent: bool,
     include_unspent: bool,
     include_hinted: bool,
+    min_amount: uint64,
 ) -> None:
     async with DBConnection(db_version) as db_wrapper:
         # Initialize coin and hint stores.
@@ -577,6 +581,8 @@ async def test_coin_state_batches(
                 continue
             if cr.coin.puzzle_hash not in ph_set and not include_hinted:
                 continue
+            if cr.coin.amount < min_amount:
+                continue
             expected_crs.append(cr)
 
         height: Optional[uint32] = uint32(0)
@@ -586,25 +592,30 @@ async def test_coin_state_batches(
         def height_of(coin_state: CoinState) -> int:
             return max(coin_state.created_height or 0, coin_state.spent_height or 0)
 
-        while height is not None:
-            (coin_states, height) = await coin_store.batch_coin_states_by_puzzle_hashes(
-                remaining_phs[:15000],
-                min_height=height,
-                include_spent=include_spent,
-                include_unspent=include_unspent,
-                include_hinted=include_hinted,
-            )
+        while len(remaining_phs) > 0:
+            while height is not None:
+                (coin_states, height) = await coin_store.batch_coin_states_by_puzzle_hashes(
+                    remaining_phs[: CoinStore.MAX_PUZZLE_HASH_BATCH_SIZE],
+                    min_height=height,
+                    include_spent=include_spent,
+                    include_unspent=include_unspent,
+                    include_hinted=include_hinted,
+                    min_amount=min_amount,
+                    max_items=7000,
+                )
 
-            # Ensure that all of the returned coin states are in order.
-            assert all(height_of(coin_states[i]) <= height_of(coin_states[i + 1]) for i in range(len(coin_states) - 1))
+                # Ensure that all of the returned coin states are in order.
+                assert all(
+                    height_of(coin_states[i]) <= height_of(coin_states[i + 1]) for i in range(len(coin_states) - 1)
+                )
 
-            all_coin_states += coin_states
+                all_coin_states += coin_states
 
-            if height is None:
-                remaining_phs = remaining_phs[15000:]
+                if height is None:
+                    remaining_phs = remaining_phs[CoinStore.MAX_PUZZLE_HASH_BATCH_SIZE :]
 
-                if len(remaining_phs) > 0:
-                    height = uint32(0)
+                    if len(remaining_phs) > 0:
+                        height = uint32(0)
 
         assert len(all_coin_states) == len(expected_crs)
 
@@ -685,6 +696,18 @@ async def test_batch_many_coin_states(db_version: int, cut_off_middle: bool) ->
         assert len(all_coin_states) == (25001 if cut_off_middle else 50000)
 
 
+@pytest.mark.anyio
+async def test_batch_no_puzzle_hashes(db_version: int) -> None:
+    async with DBConnection(db_version) as db_wrapper:
+        # Initialize coin and hint stores.
+        coin_store = await CoinStore.create(db_wrapper)
+        await HintStore.create(db_wrapper)
+
+        coin_states, height = await coin_store.batch_coin_states_by_puzzle_hashes([])
+        assert coin_states == []
+        assert height is None
+
+
 @pytest.mark.anyio
 async def test_unsupported_version() -> None:
     with pytest.raises(RuntimeError, match="CoinStore does not support database schema v1"):

chia/_tests/core/full_node/stores/test_hint_store.py

@@ -108,6 +108,38 @@ async def test_duplicates(db_version: int) -> None:
         assert rows[0][0] == 1
 
 
+@pytest.mark.anyio
+async def test_coin_ids_multi(db_version: int) -> None:
+    async with DBConnection(db_version) as db_wrapper:
+        hint_store = await HintStore.create(db_wrapper)
+        hints = [32 * i.to_bytes(1, byteorder="big", signed=False) for i in range(256)]
+        coin_ids = [bytes32(32 * i.to_bytes(1, byteorder="big", signed=False)) for i in range(256)]
+
+        expected: dict[bytes, list[bytes32]] = {}
+        expected[hints[0]] = coin_ids[0:10]
+        expected[hints[1]] = coin_ids[10:15]
+        expected[hints[2]] = [coin_ids[15]]
+        expected[hints[3]] = [coin_ids[16]]
+        expected[hints[4]] = coin_ids[17:20]
+
+        items = [(coin_id, hint) for hint, coin_ids in expected.items() for coin_id in coin_ids]
+
+        await hint_store.add_hints(items)
+
+        # Try all at once
+        actual = await hint_store.get_coin_ids_multi(set(expected.keys()))
+
+        assert set(actual) == {coin_id for coin_ids in expected.values() for coin_id in coin_ids}
+
+        # Try empty set
+        assert await hint_store.get_coin_ids_multi(set()) == []
+
+        # Try one at a time
+        for hint, coin_ids in expected.items():
+            actual = await hint_store.get_coin_ids_multi({hint})
+            assert set(actual) == set(coin_ids)
+
+
 @pytest.mark.anyio
 async def test_hints_in_blockchain(
     wallet_nodes: Tuple[

chia/_tests/core/full_node/test_subscriptions.py

@@ -401,3 +401,22 @@ def test_subscription_list() -> None:
 
     assert sub.coin_subscriptions(peer1) == {coin1, coin2}
     assert sub.puzzle_subscriptions(peer1) == {ph1, ph2}
+
+
+def test_clear_subscriptions() -> None:
+    subs = PeerSubscriptions()
+
+    subs.add_puzzle_subscriptions(peer1, [ph1, ph2], 4)
+    subs.add_coin_subscriptions(peer1, [coin1, coin2], 4)
+
+    subs.clear_puzzle_subscriptions(peer1)
+    assert subs.coin_subscriptions(peer1) == {coin1, coin2}
+    assert subs.puzzle_subscription_count() == 0
+
+    subs.add_puzzle_subscriptions(peer1, [ph1, ph2], 4)
+    subs.clear_coin_subscriptions(peer1)
+    assert subs.puzzle_subscriptions(peer1) == {ph1, ph2}
+    assert subs.coin_subscription_count() == 0
+
+    subs.clear_puzzle_subscriptions(peer1)
+    assert subs.peer_subscription_count(peer1) == 0

chia/_tests/core/server/test_server.py

@@ -80,7 +80,7 @@ async def test_connection_versions(
     outgoing_connection = wallet_node.server.all_connections[full_node.server.node_id]
     incoming_connection = full_node.server.all_connections[wallet_node.server.node_id]
     for connection in [outgoing_connection, incoming_connection]:
-        assert connection.protocol_version == Version(protocol_version[NodeType.FULL_NODE])
+        assert connection.protocol_version == Version(protocol_version[NodeType.WALLET])
         assert connection.version == __version__
         assert connection.get_version() == connection.version
 

chia/_tests/util/build_network_protocol_files.py

@@ -87,7 +87,7 @@ def visit_wallet_protocol(visitor: Callable[[Any, str], None]) -> None:
     visitor(respond_header_blocks, "respond_header_blocks")
     visitor(coin_state, "coin_state")
     visitor(register_for_ph_updates, "register_for_ph_updates")
-    visitor(reject_block_headers, "reject_block_headers"),
+    visitor(reject_block_headers, "reject_block_headers")
     visitor(respond_to_ph_updates, "respond_to_ph_updates")
     visitor(register_for_coin_updates, "register_for_coin_updates")
     visitor(respond_to_coin_updates, "respond_to_coin_updates")
@@ -96,6 +96,17 @@ def visit_wallet_protocol(visitor: Callable[[Any, str], None]) -> None:
     visitor(respond_children, "respond_children")
     visitor(request_ses_info, "request_ses_info")
     visitor(respond_ses_info, "respond_ses_info")
+    visitor(coin_state_filters, "coin_state_filters")
+    visitor(request_remove_puzzle_subscriptions, "request_remove_puzzle_subscriptions")
+    visitor(respond_remove_puzzle_subscriptions, "respond_remove_puzzle_subscriptions")
+    visitor(request_remove_coin_subscriptions, "request_remove_coin_subscriptions")
+    visitor(respond_remove_coin_subscriptions, "respond_remove_coin_subscriptions")
+    visitor(request_puzzle_state, "request_puzzle_state")
+    visitor(reject_puzzle_state, "reject_puzzle_state")
+    visitor(respond_puzzle_state, "respond_puzzle_state")
+    visitor(request_coin_state, "request_coin_state")
+    visitor(respond_coin_state, "respond_coin_state")
+    visitor(reject_coin_state, "reject_coin_state")
 
 
 def visit_harvester_protocol(visitor: Callable[[Any, str], None]) -> None:

chia/_tests/util/network_protocol_data.py

@@ -686,6 +686,53 @@
     [[uint32(1), uint32(2), uint32(3)], [uint32(4), uint32(606340525)]],
 )
 
+coin_state_filters = wallet_protocol.CoinStateFilters(True, True, True, uint64(0))
+
+hashes = [
+    bytes32(bytes.fromhex("59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7")),
+    bytes32(bytes.fromhex("d4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca")),
+    bytes32(bytes.fromhex("0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1")),
+]
+
+request_remove_puzzle_subscriptions = wallet_protocol.RequestRemovePuzzleSubscriptions(hashes)
+
+respond_remove_puzzle_subscriptions = wallet_protocol.RespondRemovePuzzleSubscriptions(hashes)
+
+request_remove_coin_subscriptions = wallet_protocol.RequestRemoveCoinSubscriptions(hashes)
+
+respond_remove_coin_subscriptions = wallet_protocol.RespondRemoveCoinSubscriptions(hashes)
+
+request_puzzle_state = wallet_protocol.RequestPuzzleState(
+    hashes,
+    uint32(0),
+    bytes32(bytes.fromhex("9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411")),
+    coin_state_filters,
+    True,
+)
+
+respond_puzzle_state = wallet_protocol.RespondPuzzleState(
+    hashes,
+    uint32(432487),
+    bytes32(bytes.fromhex("9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411")),
+    True,
+    [coin_state],
+)
+
+reject_puzzle_state = wallet_protocol.RejectPuzzleState(uint8(wallet_protocol.RejectStateReason.REORG))
+
+request_coin_state = wallet_protocol.RequestCoinState(
+    hashes,
+    uint32(0),
+    bytes32(bytes.fromhex("9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411")),
+    False,
+)
+
+respond_coin_state = wallet_protocol.RespondCoinState(hashes, [coin_state])
+
+reject_coin_state = wallet_protocol.RejectCoinState(
+    uint8(wallet_protocol.RejectStateReason.EXCEEDED_SUBSCRIPTION_LIMIT)
+)
+
 
 ### HARVESTER PROTOCOL
 pool_difficulty = harvester_protocol.PoolDifficulty(

chia/_tests/util/protocol_messages_json.py

@@ -2012,6 +2012,113 @@
     "heights": [[1, 2, 3], [4, 606340525]],
 }
 
+coin_state_filters_json: Dict[str, Any] = {
+    "include_spent": True,
+    "include_unspent": True,
+    "include_hinted": True,
+    "min_amount": 0,
+}
+
+request_remove_puzzle_subscriptions_json: Dict[str, Any] = {
+    "puzzle_hashes": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ]
+}
+
+respond_remove_puzzle_subscriptions_json: Dict[str, Any] = {
+    "puzzle_hashes": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ]
+}
+
+request_remove_coin_subscriptions_json: Dict[str, Any] = {
+    "coin_ids": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ]
+}
+
+respond_remove_coin_subscriptions_json: Dict[str, Any] = {
+    "coin_ids": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ]
+}
+
+request_puzzle_state_json: Dict[str, Any] = {
+    "puzzle_hashes": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ],
+    "previous_height": 0,
+    "header_hash": "0x9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411",
+    "filters": {"include_spent": True, "include_unspent": True, "include_hinted": True, "min_amount": 0},
+    "subscribe_when_finished": True,
+}
+
+reject_puzzle_state_json: Dict[str, Any] = {"reason": 0}
+
+respond_puzzle_state_json: Dict[str, Any] = {
+    "puzzle_hashes": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ],
+    "height": 432487,
+    "header_hash": "0x9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411",
+    "is_finished": True,
+    "coin_states": [
+        {
+            "coin": {
+                "parent_coin_info": "0xd56f435d3382cb9aa5f50f51816e4c54487c66402339901450f3c810f1d77098",
+                "puzzle_hash": "0x9944f63fcc251719b2f04c47ab976a167f96510736dc6fdfa8e037d740f4b5f3",
+                "amount": 6602327684212801382,
+            },
+            "spent_height": 2287030048,
+            "created_height": 3361305811,
+        }
+    ],
+}
+
+request_coin_state_json: Dict[str, Any] = {
+    "coin_ids": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ],
+    "previous_height": 0,
+    "header_hash": "0x9620d602399252a8401a44669a9d7a6fc328358868a427e827d721c233e2b411",
+    "subscribe": False,
+}
+
+respond_coin_state_json: Dict[str, Any] = {
+    "coin_ids": [
+        "0x59710628755b6d7f7d0b5d84d5c980e7a1c52e55f5a43b531312402bd9045da7",
+        "0xd4a68c9dc42d625092c3e71a657cce469ae4180d1b0632256d2da8ffc0a9beca",
+        "0x0e03ce4c43d7d60886f27af7da0ea9749a46b977b3743f3fd2e97b169dc539c1",
+    ],
+    "coin_states": [
+        {
+            "coin": {
+                "parent_coin_info": "0xd56f435d3382cb9aa5f50f51816e4c54487c66402339901450f3c810f1d77098",
+                "puzzle_hash": "0x9944f63fcc251719b2f04c47ab976a167f96510736dc6fdfa8e037d740f4b5f3",
+                "amount": 6602327684212801382,
+            },
+            "spent_height": 2287030048,
+            "created_height": 3361305811,
+        }
+    ],
+}
+
+reject_coin_state_json: Dict[str, Any] = {"reason": 1}
+
 pool_difficulty_json: Dict[str, Any] = {
     "difficulty": 14819251421858580996,
     "sub_slot_iters": 12852879676624401630,