starred/chia-blockchain
1
0
Fork 0
mirror of https://github.com/Chia-Network/chia-blockchain.git synced 2026-08-24 02:24:23 -05:00
Code Issues Packages Projects Releases Wiki Activity

[LABS-484] Extract clawback functionality into its own module (#21160)

Browse Source
This commit is contained in:
Matt Hauff
2026-08-05 15:33:42 -07:00
committed by GitHub
parent 7648ef7d87
commit fec92d5b2b
5 changed files with 236 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
chia/_tests/wallet/test_wallet.py
+9
View File
@@ -541,6 +541,15 @@ class TestWalletSimulator:
wallet_environments.tx_config,
)
assert len(resp.transaction_ids) == 1
# just test that an extra submission doesn't do anything
await env.rpc_client.spend_clawback_coins(
SpendClawbackCoins(
coin_ids=[merkle_coin.name()],
fee=uint64(test_fee),
push=True,
),
wallet_environments.tx_config,
)
await wallet_environments.process_pending_states(
[
chia/wallet/clawback_manager.py
+197
View File
@@ -0,0 +1,197 @@
from __future__ import annotations
import logging
import time
from collections.abc import Awaitable, Callable
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass, replace
from chia_rs import Coin, CoinSpend, G2Element
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint32, uint64
from chia.types.mempool_inclusion_status import MempoolInclusionStatus
from chia.util.hash import std_hash
from chia.util.streamable import UInt32Range, UInt64Range
from chia.wallet.conditions import (
AssertCoinAnnouncement,
Condition,
CreateCoin,
CreateCoinAnnouncement,
parse_timelock_info,
)
from chia.wallet.puzzles.clawback.drivers import generate_clawback_spend_bundle
from chia.wallet.puzzles.clawback.metadata import ClawbackMetadata
from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.util.compute_memos import compute_memos
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.tx_config import TXConfig
from chia.wallet.util.wallet_types import CoinType, WalletType
from chia.wallet.wallet import Wallet
from chia.wallet.wallet_action_scope import WalletActionScope
from chia.wallet.wallet_blockchain import WalletBlockchain
from chia.wallet.wallet_coin_store import WalletCoinStore
from chia.wallet.wallet_puzzle_store import WalletPuzzleStore
from chia.wallet.wallet_spend_bundle import WalletSpendBundle
from chia.wallet.wallet_transaction_store import WalletTransactionStore
@dataclass(frozen=True, kw_only=True)
class ClawbackManager:
log: logging.Logger
blockchain: WalletBlockchain
coin_store: WalletCoinStore
puzzle_store: WalletPuzzleStore
transaction_store: WalletTransactionStore
xch_wallet: Wallet
auto_claim_tx_fee: uint64
auto_claim_batch_size: int
timestamp_for_height: Callable[[uint32], Awaitable[uint64]]
puzzle_hash_encoder: Callable[[bytes32], str]
action_scope_sandbox: Callable[[TXConfig, bool], AbstractAsyncContextManager[WalletActionScope]]
async def auto_claim_coins(self, action_scope: WalletActionScope) -> None:
# Get unspent clawback coin
current_timestamp = self.blockchain.get_latest_timestamp()
clawback_coins: dict[Coin, ClawbackMetadata] = {}
unspent_coins = await self.coin_store.get_coin_records(
coin_type=CoinType.CLAWBACK,
wallet_type=WalletType.STANDARD_WALLET,
spent_range=UInt32Range(stop=uint32(0)),
amount_range=UInt64Range(
start=action_scope.config.tx_config.coin_selection_config.min_coin_amount,
stop=action_scope.config.tx_config.coin_selection_config.max_coin_amount,
),
)
for coin in unspent_coins.records:
try:
metadata = coin.parsed_metadata()
assert isinstance(metadata, ClawbackMetadata)
if await metadata.is_recipient(self.puzzle_store):
coin_timestamp = await self.timestamp_for_height(coin.confirmed_block_height)
if current_timestamp - coin_timestamp >= metadata.time_lock:
clawback_coins[coin.coin] = metadata
if len(clawback_coins) >= self.auto_claim_batch_size:
await self.spend_clawback_coins(clawback_coins, self.auto_claim_tx_fee, action_scope)
clawback_coins = {}
except Exception as e:
self.log.error(f"Failed to claim clawback coin {coin.coin.name().hex()}: %s", e)
if len(clawback_coins) > 0:
await self.spend_clawback_coins(clawback_coins, self.auto_claim_tx_fee, action_scope)
async def spend_clawback_coins(
self,
clawback_coins: dict[Coin, ClawbackMetadata],
fee: uint64,
action_scope: WalletActionScope,
force: bool = False,
extra_conditions: tuple[Condition, ...] = tuple(),
) -> None:
assert len(clawback_coins) > 0
coin_spends: list[CoinSpend] = []
message = std_hash(b"".join([c.name() for c in clawback_coins.keys()]))
derivation_record = None
amount = uint64(0)
for coin, metadata in clawback_coins.items():
try:
self.log.info(f"Claiming clawback coin {coin.name().hex()}")
# Get incoming tx
incoming_tx = await self.transaction_store.get_transaction_record(coin.name())
assert incoming_tx is not None, f"Cannot find incoming tx for clawback coin {coin.name().hex()}"
if incoming_tx.sent > 0 and not force:
self.log.error(
f"Clawback coin {coin.name().hex()} is already in a pending spend bundle. {incoming_tx}"
)
continue
recipient_puzhash = metadata.recipient_puzzle_hash
sender_puzhash = metadata.sender_puzzle_hash
is_recipient: bool = await metadata.is_recipient(self.puzzle_store)
if is_recipient:
derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(recipient_puzhash)
else:
derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(sender_puzhash)
assert derivation_record is not None
amount = uint64(amount + coin.amount)
# Remove the clawback hint since it is unnecessary for the XCH coin
memos: list[bytes] = [] if len(incoming_tx.memos) == 0 else next(iter(incoming_tx.memos.items()))[1][1:]
inner_puzzle = self.xch_wallet.puzzle_for_pk(derivation_record.pubkey)
inner_solution = self.xch_wallet.make_solution(
primaries=[
CreateCoin(
derivation_record.puzzle_hash,
uint64(coin.amount),
memos, # Forward memo of the first coin
)
],
conditions=(
extra_conditions
if len(coin_spends) > 0 or fee == 0
else (*extra_conditions, CreateCoinAnnouncement(message))
),
)
coin_spend: CoinSpend = generate_clawback_spend_bundle(coin, metadata, inner_puzzle, inner_solution)
coin_spends.append(coin_spend)
# Update incoming tx to prevent double spend and mark it is pending
await self.transaction_store.increment_sent(incoming_tx.name, "", MempoolInclusionStatus.PENDING, None)
except Exception as e:
self.log.error(f"Failed to create clawback spend bundle for {coin.name().hex()}: {e}")
if len(coin_spends) == 0:
return
spend_bundle = WalletSpendBundle(coin_spends, G2Element())
if fee > 0:
async with self.action_scope_sandbox(action_scope.config.tx_config, False) as inner_action_scope:
async with action_scope.use() as interface:
async with inner_action_scope.use() as inner_interface:
inner_interface.side_effects.selected_coins = interface.side_effects.selected_coins
await self.xch_wallet.create_tandem_xch_tx(
fee,
inner_action_scope,
extra_conditions=(
AssertCoinAnnouncement(asserted_id=coin_spends[0].coin.name(), asserted_msg=message),
),
)
async with inner_action_scope.use() as inner_interface:
# This should not be looked to for best practice.
# Ideally, the two spend bundles can exist separately on each tx record until they are pushed.
# This is not very supported behavior at the moment
# so to avoid any potential backwards compatibility issues,
# we're moving the spend bundle from this TX to the main
interface.side_effects.transactions.extend(
[replace(tx, spend_bundle=None) for tx in inner_interface.side_effects.transactions]
)
interface.side_effects.selected_coins.extend(inner_interface.side_effects.selected_coins)
spend_bundle = WalletSpendBundle.aggregate(
[
spend_bundle,
*(
tx.spend_bundle
for tx in inner_action_scope.side_effects.transactions
if tx.spend_bundle is not None
),
]
)
assert derivation_record is not None
tx_record = TransactionRecord(
confirmed_at_height=uint32(0),
created_at_time=uint64(time.time()),
to_puzzle_hash=derivation_record.puzzle_hash,
to_address=self.puzzle_hash_encoder(derivation_record.puzzle_hash),
amount=amount,
fee_amount=uint64(fee),
confirmed=False,
sent=uint32(0),
spend_bundle=spend_bundle,
additions=spend_bundle.additions(),
removals=spend_bundle.removals(),
wallet_id=uint32(1),
sent_to=[],
trade_id=None,
type=uint32(TransactionType.OUTGOING_CLAWBACK),
name=spend_bundle.name(),
memos=compute_memos(spend_bundle),
valid_times=parse_timelock_info(extra_conditions),
)
async with action_scope.use() as interface:
interface.side_effects.transactions.append(tx_record)
chia/wallet/wallet_node.py
+10 -1
View File
@@ -323,6 +323,12 @@ class WalletNode:
if "auto_claim" not in self.config or self.config["auto_claim"] != auto_claim_config_json:
# Update in memory config
self.config["auto_claim"] = auto_claim_config_json
if self._wallet_state_manager is not None:
self.wallet_state_manager.clawback_manager = dataclasses.replace(
self.wallet_state_manager.clawback_manager,
auto_claim_tx_fee=auto_claim_config.tx_fee,
auto_claim_batch_size=auto_claim_config.batch_size,
)
# Update config file
with lock_and_load_config(self.root_path, "config.yaml") as config:
config["wallet"]["auto_claim"] = self.config["auto_claim"]
@@ -1299,7 +1305,10 @@ class WalletNode:
# Check if any coin needs auto spending
if self.config.get("auto_claim", {}).get("enabled", False):
await self.wallet_state_manager.auto_claim_coins()
async with self.wallet_state_manager.new_action_scope(
self.wallet_state_manager.tx_config, push=True
) as action_scope:
await self.wallet_state_manager.clawback_manager.auto_claim_coins(action_scope)
if new_peak_hb.foliage_transaction_block is not None:
await self._retry_fee_failed_transactions()
chia/wallet/wallet_rpc_api.py
+2 -2
View File
@@ -1500,7 +1500,7 @@ class WalletRpcApi:
batch_size = (
request.batch_size
if request.batch_size is not None
else self.service.wallet_state_manager.config.get("auto_claim", {}).get("batch_size", 50)
else self.service.wallet_state_manager.clawback_manager.auto_claim_batch_size
)
records_list = list(coin_records.coin_id_to_record.values())
for i in range(0, len(records_list), batch_size):
@@ -1511,7 +1511,7 @@ class WalletRpcApi:
except WalletCoinRecordMetadataParsingError as e:
log.error("Failed to spend clawback coin: %s", e)
continue
await self.service.wallet_state_manager.spend_clawback_coins(
await self.service.wallet_state_manager.clawback_manager.spend_clawback_coins(
# Semantically, we're guaranteed the right type here, but the typing isn't there
coin_batch, # type: ignore[arg-type]
request.fee,
chia/wallet/wallet_state_manager.py
+18 -155
View File
@@ -48,12 +48,11 @@ from chia.wallet.cat_wallet.cat_info import CATCoinData, CATInfo, CRCATInfo
from chia.wallet.cat_wallet.cat_utils import CAT_MOD, CAT_MOD_HASH, construct_cat_puzzle, match_cat_puzzle
from chia.wallet.cat_wallet.cat_wallet import CATWallet
from chia.wallet.cat_wallet.r_cat_wallet import RCATWallet
from chia.wallet.clawback_manager import ClawbackManager
from chia.wallet.conditions import (
AssertCoinAnnouncement,
Condition,
ConditionValidTimes,
CreateCoin,
CreateCoinAnnouncement,
parse_timelock_info,
)
from chia.wallet.db_wallet.db_wallet_puzzles import MIRROR_PUZZLE_HASH
@@ -88,7 +87,7 @@ from chia.wallet.outer_puzzles import AssetType
from chia.wallet.plotnft_wallet.plotnft_store import PlotNFTStore
from chia.wallet.plotnft_wallet.plotnft_wallet import PlotNFT2Wallet
from chia.wallet.puzzle_drivers import PuzzleInfo
from chia.wallet.puzzles.clawback.drivers import generate_clawback_spend_bundle, match_clawback_puzzle
from chia.wallet.puzzles.clawback.drivers import match_clawback_puzzle
from chia.wallet.puzzles.clawback.metadata import ClawbackMetadata, ClawbackVersion
from chia.wallet.remote_wallet.remote_coin_store import RemoteCoinStore
from chia.wallet.remote_wallet.remote_wallet import RemoteWallet
@@ -122,7 +121,7 @@ from chia.wallet.vc_wallet.vc_wallet import VCWallet
from chia.wallet.wallet import Wallet
from chia.wallet.wallet_action_scope import PlotNFTTargetStateInfo, WalletActionScope, new_wallet_action_scope
from chia.wallet.wallet_blockchain import WalletBlockchain
from chia.wallet.wallet_coin_record import MetadataTypes, WalletCoinRecord
from chia.wallet.wallet_coin_record import WalletCoinRecord
from chia.wallet.wallet_coin_store import CoinRecordOrder, WalletCoinStore
from chia.wallet.wallet_info import WalletInfo
from chia.wallet.wallet_interested_store import WalletInterestedStore
@@ -206,6 +205,7 @@ class WalletStateManager:
initial_num_public_keys: int
decorator_manager: PuzzleDecoratorManager
signer: WalletSigner
clawback_manager: ClawbackManager
@staticmethod
async def create(
@@ -297,6 +297,20 @@ class WalletStateManager:
self.main_wallet = await Wallet.create(self, main_wallet_info)
self.clawback_manager = ClawbackManager(
log=self.log,
blockchain=self.blockchain,
coin_store=self.coin_store,
puzzle_store=self.puzzle_store,
transaction_store=self.tx_store,
xch_wallet=self.main_wallet,
auto_claim_tx_fee=uint64(self.config.get("auto_claim", {}).get("tx_fee", 0)),
auto_claim_batch_size=self.config.get("auto_claim", {}).get("batch_size", 50),
timestamp_for_height=self.wallet_node.get_timestamp_for_height,
puzzle_hash_encoder=self.encode_puzzle_hash,
action_scope_sandbox=self.new_action_scope,
)
self.wallets = {main_wallet_info.id: self.main_wallet}
self.asset_to_wallet_map = {
@@ -1036,157 +1050,6 @@ class WalletStateManager:
logged_in_fingerprint=self.root_pubkey.get_fingerprint(),
)
async def auto_claim_coins(self) -> None:
# Get unspent clawback coin
current_timestamp = self.blockchain.get_latest_timestamp()
clawback_coins: dict[Coin, ClawbackMetadata] = {}
tx_fee = uint64(self.config.get("auto_claim", {}).get("tx_fee", 0))
unspent_coins = await self.coin_store.get_coin_records(
coin_type=CoinType.CLAWBACK,
wallet_type=WalletType.STANDARD_WALLET,
spent_range=UInt32Range(stop=uint32(0)),
amount_range=UInt64Range(
start=self.tx_config.coin_selection_config.min_coin_amount,
stop=self.tx_config.coin_selection_config.max_coin_amount,
),
)
async with self.new_action_scope(self.tx_config, push=True) as action_scope:
for coin in unspent_coins.records:
try:
metadata: MetadataTypes = coin.parsed_metadata()
assert isinstance(metadata, ClawbackMetadata)
if await metadata.is_recipient(self.puzzle_store):
coin_timestamp = await self.wallet_node.get_timestamp_for_height(coin.confirmed_block_height)
if current_timestamp - coin_timestamp >= metadata.time_lock:
clawback_coins[coin.coin] = metadata
if len(clawback_coins) >= self.config.get("auto_claim", {}).get("batch_size", 50):
await self.spend_clawback_coins(clawback_coins, tx_fee, action_scope)
clawback_coins = {}
except Exception as e:
self.log.error(f"Failed to claim clawback coin {coin.coin.name().hex()}: %s", e)
if len(clawback_coins) > 0:
await self.spend_clawback_coins(clawback_coins, tx_fee, action_scope)
async def spend_clawback_coins(
self,
clawback_coins: dict[Coin, ClawbackMetadata],
fee: uint64,
action_scope: WalletActionScope,
force: bool = False,
extra_conditions: tuple[Condition, ...] = tuple(),
) -> None:
assert len(clawback_coins) > 0
coin_spends: list[CoinSpend] = []
message: bytes32 = std_hash(b"".join([c.name() for c in clawback_coins.keys()]))
now: uint64 = uint64(time.time())
derivation_record: DerivationRecord | None = None
amount: uint64 = uint64(0)
for coin, metadata in clawback_coins.items():
try:
self.log.info(f"Claiming clawback coin {coin.name().hex()}")
# Get incoming tx
incoming_tx = await self.tx_store.get_transaction_record(coin.name())
assert incoming_tx is not None, f"Cannot find incoming tx for clawback coin {coin.name().hex()}"
if incoming_tx.sent > 0 and not force:
self.log.error(
f"Clawback coin {coin.name().hex()} is already in a pending spend bundle. {incoming_tx}"
)
continue
recipient_puzhash: bytes32 = metadata.recipient_puzzle_hash
sender_puzhash: bytes32 = metadata.sender_puzzle_hash
is_recipient: bool = await metadata.is_recipient(self.puzzle_store)
if is_recipient:
derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(recipient_puzhash)
else:
derivation_record = await self.puzzle_store.get_derivation_record_for_puzzle_hash(sender_puzhash)
assert derivation_record is not None
amount = uint64(amount + coin.amount)
# Remove the clawback hint since it is unnecessary for the XCH coin
memos: list[bytes] = [] if len(incoming_tx.memos) == 0 else next(iter(incoming_tx.memos.items()))[1][1:]
inner_puzzle: Program = self.main_wallet.puzzle_for_pk(derivation_record.pubkey)
inner_solution: Program = self.main_wallet.make_solution(
primaries=[
CreateCoin(
derivation_record.puzzle_hash,
uint64(coin.amount),
memos, # Forward memo of the first coin
)
],
conditions=(
extra_conditions
if len(coin_spends) > 0 or fee == 0
else (*extra_conditions, CreateCoinAnnouncement(message))
),
)
coin_spend: CoinSpend = generate_clawback_spend_bundle(coin, metadata, inner_puzzle, inner_solution)
coin_spends.append(coin_spend)
# Update incoming tx to prevent double spend and mark it is pending
await self.tx_store.increment_sent(incoming_tx.name, "", MempoolInclusionStatus.PENDING, None)
except Exception as e:
self.log.error(f"Failed to create clawback spend bundle for {coin.name().hex()}: {e}")
if len(coin_spends) == 0:
return
spend_bundle = WalletSpendBundle(coin_spends, G2Element())
if fee > 0:
async with self.new_action_scope(action_scope.config.tx_config, push=False) as inner_action_scope:
async with action_scope.use() as interface:
async with inner_action_scope.use() as inner_interface:
inner_interface.side_effects.selected_coins = interface.side_effects.selected_coins
await self.main_wallet.create_tandem_xch_tx(
fee,
inner_action_scope,
extra_conditions=(
AssertCoinAnnouncement(asserted_id=coin_spends[0].coin.name(), asserted_msg=message),
),
)
async with inner_action_scope.use() as inner_interface:
# This should not be looked to for best practice.
# Ideally, the two spend bundles can exist separately on each tx record until they are pushed.
# This is not very supported behavior at the moment
# so to avoid any potential backwards compatibility issues,
# we're moving the spend bundle from this TX to the main
interface.side_effects.transactions.extend(
[
dataclasses.replace(tx, spend_bundle=None)
for tx in inner_interface.side_effects.transactions
]
)
interface.side_effects.selected_coins.extend(inner_interface.side_effects.selected_coins)
spend_bundle = WalletSpendBundle.aggregate(
[
spend_bundle,
*(
tx.spend_bundle
for tx in inner_action_scope.side_effects.transactions
if tx.spend_bundle is not None
),
]
)
assert derivation_record is not None
tx_record = TransactionRecord(
confirmed_at_height=uint32(0),
created_at_time=now,
to_puzzle_hash=derivation_record.puzzle_hash,
to_address=self.encode_puzzle_hash(derivation_record.puzzle_hash),
amount=amount,
fee_amount=uint64(fee),
confirmed=False,
sent=uint32(0),
spend_bundle=spend_bundle,
additions=spend_bundle.additions(),
removals=spend_bundle.removals(),
wallet_id=uint32(1),
sent_to=[],
trade_id=None,
type=uint32(TransactionType.OUTGOING_CLAWBACK),
name=spend_bundle.name(),
memos=compute_memos(spend_bundle),
valid_times=parse_timelock_info(extra_conditions),
)
async with action_scope.use() as interface:
interface.side_effects.transactions.append(tx_record)
async def filter_spam(self, new_coin_state: list[CoinState]) -> list[CoinState]:
xch_spam_amount = self.config.get("xch_spam_amount", 1000000)