chore(PE-8188): add patch balances work

handlers-that-modify-balances.md

@@ -0,0 +1,141 @@
+# Handlers That Modify Balances
+
+## Direct Balance Modification Handlers
+
+### 1. Transfer
+
+**Handler**: `ActionMap.Transfer`  
+**Action**: Transfers tokens from sender to recipient  
+**Balance Changes**: Sender balance ↓, Recipient balance ↑
+
+### 2. CreateVault
+
+**Handler**: `ActionMap.CreateVault`  
+**Action**: Locks tokens in a time-locked vault for the sender  
+**Balance Changes**: Sender balance ↓
+
+### 3. VaultedTransfer
+
+**Handler**: `ActionMap.VaultedTransfer`  
+**Action**: Creates a time-locked vault for a recipient  
+**Balance Changes**: Sender balance ↓
+
+### 4. RevokeVault
+
+**Handler**: `ActionMap.RevokeVault`  
+**Action**: Controller revokes a vaultedTransfer before it matures  
+**Balance Changes**: Controller balance ↑
+
+### 5. IncreaseVault
+
+**Handler**: `ActionMap.IncreaseVault`  
+**Action**: Adds more tokens to an existing vault  
+**Balance Changes**: Sender balance ↓
+
+## Gateway/Staking Handlers
+
+### 6. JoinNetwork
+
+**Handler**: `ActionMap.JoinNetwork`  
+**Action**: Gateway joins the network with operator stake  
+**Balance Changes**: Operator balance ↓
+
+### 7. IncreaseOperatorStake
+
+**Handler**: `ActionMap.IncreaseOperatorStake`  
+**Action**: Increases gateway operator stake  
+**Balance Changes**: Operator balance ↓
+
+### 8. DecreaseOperatorStake
+
+**Handler**: `ActionMap.DecreaseOperatorStake`  
+**Action**: Decreases operator stake (with optional instant withdrawal)  
+**Balance Changes**: Operator balance ↑ (if instant withdrawal with penalty)
+
+### 9. DelegateStake
+
+**Handler**: `ActionMap.DelegateStake`  
+**Action**: Delegates stake to a gateway  
+**Balance Changes**: Delegator balance ↓
+
+### 10. DecreaseDelegateStake
+
+**Handler**: `ActionMap.DecreaseDelegateStake`  
+**Action**: Decreases delegated stake (with optional instant withdrawal)  
+**Balance Changes**: Delegator balance ↑ (if instant withdrawal with penalty)
+
+### 11. InstantWithdrawal
+
+**Handler**: `ActionMap.InstantWithdrawal`  
+**Action**: Instantly withdraws from a pending withdrawal vault with penalty  
+**Balance Changes**: Withdrawer balance ↑, Protocol balance ↑ (penalty fee)
+
+### 12. CancelWithdrawal
+
+**Handler**: `ActionMap.CancelWithdrawal`  
+**Action**: Cancels a withdrawal and returns to staked status  
+**Balance Changes**: No direct balance change (stake → vault → stake)
+
+## ArNS Name Handlers
+
+### 13. BuyName
+
+**Handler**: `ActionMap.BuyName`  
+**Action**: Purchases an ArNS name (lease or permabuy)  
+**Balance Changes**: Buyer balance ↓ (via funding plan), Protocol balance ↑, Returned name initiator balance ↑ (if applicable)
+
+### 14. ExtendLease
+
+**Handler**: `ActionMap.ExtendLease`  
+**Action**: Extends an ArNS name lease  
+**Balance Changes**: Owner balance ↓ (via funding plan), Protocol balance ↑
+
+### 15. IncreaseUndernameLimit
+
+**Handler**: `ActionMap.IncreaseUndernameLimit`  
+**Action**: Increases undername limit for an ArNS name  
+**Balance Changes**: Owner balance ↓ (via funding plan), Protocol balance ↑
+
+### 16. UpgradeName
+
+**Handler**: `ActionMap.UpgradeName`  
+**Action**: Upgrades a leased name to permabuy  
+**Balance Changes**: Owner balance ↓ (via funding plan), Protocol balance ↑
+
+## Primary Name Handlers
+
+### 17. RequestPrimaryName
+
+**Handler**: `ActionMap.RequestPrimaryName`  
+**Action**: Creates a primary name request  
+**Balance Changes**: Requester balance ↓ (via funding plan), Protocol balance ↑
+
+## System/Automatic Operations
+
+### 18. Epoch Distribution
+
+**Handler**: Automatic via `distributeEpoch`  
+**Action**: Distributes rewards to gateways and delegates  
+**Balance Changes**: Protocol balance ↓, Gateway operator balances ↑ (or stakes ↑ if auto-stake), Delegate stakes ↑
+
+### 19. Vault Pruning
+
+**Handler**: Automatic via `vaults.pruneVaults`  
+**Action**: Releases matured balance vaults  
+**Balance Changes**: Vault owner balance ↑
+
+### 20. Gateway Vault Pruning
+
+**Handler**: Automatic via gateway pruning  
+**Action**: Releases matured gateway exit/withdrawal vaults  
+**Balance Changes**: Gateway operator balance ↑, Delegate balances ↑
+
+### 21. Slashing
+
+**Handler**: Via `gar.slashOperatorStake`  
+**Action**: Slashes operator stake for repeated failures  
+**Balance Changes**: Gateway operator stake ↓, Protocol balance ↑
+
+---
+
+**Note**: Handlers with "funding plan" support can draw from multiple sources (balance, delegated stakes, vaults) using `gar.applyFundingPlan()`, which orchestrates the necessary balance reductions across different sources.

src/balances.lua

@@ -96,4 +96,17 @@ function balances.walletHasSufficientBalance(wallet, quantity)
 	return Balances[wallet] ~= nil and Balances[wallet] >= quantity
 end
 
+---@param addresses table<string> A table of addresses and their balances
+function balances.patchBalances(addresses)
+	assert(type(addresses) == "table", "Addresses must be a table")
+	--- For simplicity we always include the protocol balance in the patch message
+	local patchMessage = { device = "[email protected]", balances = { [ao.id] = Balances[ao.id] or 0 } }
+	for _, address in pairs(addresses) do
+		assert(type(address) == "string", "Address must be a string")
+		patchMessage.balances[address] = Balances[address] or 0
+	end
+
+	ao.send(patchMessage)
+end
+
 return balances

src/gar.lua

@@ -934,6 +934,8 @@ end
 --- @param msgId string The message ID
 --- @return PruneGatewaysResult # The result containing the pruned gateways, slashed gateways, and other stats
 function gar.pruneGateways(currentTimestamp, msgId)
+	local affectedAddresses = {}
+
 	--- @type PruneGatewaysResult
 	local result = {
 		prunedGateways = {},
@@ -1067,6 +1069,11 @@ function gar.pruneGateways(currentTimestamp, msgId)
 
 	result.gatewayObjectTallies = gatewayObjectTallies
 
+	for _, gateway in pairs(result.prunedGateways) do
+		table.insert(affectedAddresses, gateway)
+	end
+	balances.patchBalances(affectedAddresses)
+
 	return result
 end
 

src/main.lua

@@ -625,6 +625,7 @@ addEventingHandler(ActionMap.Transfer, utils.hasMatchingTag("Action", ActionMap.
 	msg.ioEvent:addField("RecipientFormatted", recipient)
 
 	local result = balances.transfer(recipient, msg.From, quantity, allowUnsafeAddresses)
+
 	if result ~= nil then
 		local senderNewBalance = result[msg.From]
 		local recipientNewBalance = result[recipient]
@@ -680,6 +681,7 @@ addEventingHandler(ActionMap.Transfer, utils.hasMatchingTag("Action", ActionMap.
 		Send(msg, debitNotice)
 		Send(msg, creditNotice)
 	end
+	balances.patchBalances({ msg.From, recipient })
 end)
 
 addEventingHandler(ActionMap.CreateVault, utils.hasMatchingTag("Action", ActionMap.CreateVault), function(msg)
@@ -715,6 +717,7 @@ addEventingHandler(ActionMap.CreateVault, utils.hasMatchingTag("Action", ActionM
 		},
 		Data = json.encode(vault),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(ActionMap.VaultedTransfer, utils.hasMatchingTag("Action", ActionMap.VaultedTransfer), function(msg)
@@ -782,6 +785,8 @@ addEventingHandler(ActionMap.VaultedTransfer, utils.hasMatchingTag("Action", Act
 		},
 		Data = json.encode(vault),
 	})
+	-- might not need recipient here
+	balances.patchBalances({ msg.From, recipient })
 end)
 
 addEventingHandler(ActionMap.RevokeVault, utils.hasMatchingTag("Action", ActionMap.RevokeVault), function(msg)
@@ -820,6 +825,7 @@ addEventingHandler(ActionMap.RevokeVault, utils.hasMatchingTag("Action", ActionM
 		Tags = { Action = ActionMap.RevokeVault .. "-Notice", ["Vault-Id"] = vaultId },
 		Data = json.encode(vault),
 	})
+	balances.patchBalances({ msg.From, recipient })
 end)
 
 addEventingHandler(ActionMap.ExtendVault, utils.hasMatchingTag("Action", ActionMap.ExtendVault), function(msg)
@@ -873,6 +879,7 @@ addEventingHandler(ActionMap.IncreaseVault, utils.hasMatchingTag("Action", Actio
 		Tags = { Action = ActionMap.IncreaseVault .. "-Notice" },
 		Data = json.encode(vault),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(ActionMap.BuyName, utils.hasMatchingTag("Action", ActionMap.BuyName), function(msg)
@@ -946,6 +953,7 @@ addEventingHandler(ActionMap.BuyName, utils.hasMatchingTag("Action", ActionMap.B
 			}),
 		})
 	end
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler("upgradeName", utils.hasMatchingTag("Action", ActionMap.UpgradeName), function(msg)
@@ -976,6 +984,7 @@ addEventingHandler("upgradeName", utils.hasMatchingTag("Action", ActionMap.Upgra
 			type = record.type,
 		}),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(ActionMap.ExtendLease, utils.hasMatchingTag("Action", ActionMap.ExtendLease), function(msg)
@@ -1001,6 +1010,7 @@ addEventingHandler(ActionMap.ExtendLease, utils.hasMatchingTag("Action", ActionM
 		Tags = { Action = ActionMap.ExtendLease .. "-Notice", Name = name },
 		Data = json.encode(fundFrom and result or recordResult),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(
@@ -1034,6 +1044,8 @@ addEventingHandler(
 			},
 			Data = json.encode(fundFrom and result or recordResult),
 		})
+		-- TODO patch other ARIO locations based on FUND
+		balances.patchBalances({ msg.From })
 	end
 )
 
@@ -1185,6 +1197,7 @@ addEventingHandler(ActionMap.JoinNetwork, utils.hasMatchingTag("Action", ActionM
 		Tags = { Action = ActionMap.JoinNetwork .. "-Notice" },
 		Data = json.encode(gateway),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(ActionMap.LeaveNetwork, utils.hasMatchingTag("Action", ActionMap.LeaveNetwork), function(msg)
@@ -1272,6 +1285,7 @@ addEventingHandler(
 			Tags = { Action = ActionMap.IncreaseOperatorStake .. "-Notice" },
 			Data = json.encode(gateway),
 		})
+		balances.patchBalances({ msg.From })
 	end
 )
 
@@ -1343,6 +1357,7 @@ addEventingHandler(
 			},
 			Data = json.encode(decreaseOperatorStakeResult.gateway),
 		})
+		balances.patchBalances({ msg.From })
 	end
 )
 
@@ -1376,6 +1391,7 @@ addEventingHandler(ActionMap.DelegateStake, utils.hasMatchingTag("Action", Actio
 		Tags = { Action = ActionMap.DelegateStake .. "-Notice", Gateway = gatewayTarget },
 		Data = json.encode(delegateResult),
 	})
+	balances.patchBalances({ msg.From })
 end)
 
 addEventingHandler(ActionMap.CancelWithdrawal, utils.hasMatchingTag("Action", ActionMap.CancelWithdrawal), function(msg)
@@ -1455,6 +1471,7 @@ addEventingHandler(
 				Data = json.encode(result),
 			})
 		end
+		balances.patchBalances({ msg.From })
 	end
 )
 
@@ -1526,6 +1543,7 @@ addEventingHandler(
 			},
 			Data = json.encode(result and result.updatedDelegate or {}),
 		})
+		balances.patchBalances({ msg.From })
 	end
 )
 
@@ -1782,6 +1800,7 @@ end, function(msg)
 	]]
 	--
 	print("Ticking from " .. lastCreatedEpochIndex .. " to " .. targetCurrentEpochIndex)
+
 	for epochIndexToTick = lastCreatedEpochIndex, targetCurrentEpochIndex do
 		local tickResult = tick.tickEpoch(msg.Timestamp, blockHeight, hashchain, msgId, epochIndexToTick)
 		if tickResult.pruneGatewaysResult ~= nil then
@@ -2137,17 +2156,12 @@ end)
 -- Pagination handlers
 
 addEventingHandler("paginatedRecords", function(msg)
-        return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
+	return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
 end, function(msg)
-        local page = utils.parsePaginationTags(msg)
-        local result = arns.getPaginatedRecords(
-                page.cursor,
-                page.limit,
-                page.sortBy or "startTimestamp",
-                page.sortOrder,
-                page.filters
-        )
-        Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
+	local page = utils.parsePaginationTags(msg)
+	local result =
+		arns.getPaginatedRecords(page.cursor, page.limit, page.sortBy or "startTimestamp", page.sortOrder, page.filters)
+	Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
 end)
 
 addEventingHandler("paginatedGateways", function(msg)
@@ -2454,6 +2468,8 @@ addEventingHandler(ActionMap.RedelegateStake, utils.hasMatchingTag("Action", Act
 		},
 		Data = json.encode(redelegationResult),
 	})
+
+	balances.patchBalances({ ao.id })
 end)
 
 addEventingHandler(ActionMap.RedelegationFee, utils.hasMatchingTag("Action", ActionMap.RedelegationFee), function(msg)
@@ -2521,7 +2537,7 @@ addEventingHandler("requestPrimaryName", utils.hasMatchingTag("Action", ActionMa
 	assertValidFundFrom(fundFrom)
 
 	local primaryNameResult = primaryNames.createPrimaryNameRequest(name, initiator, msg.Timestamp, msg.Id, fundFrom)
-
+	balances.patchBalances({ msg.From })
 	addPrimaryNameRequestData(msg.ioEvent, primaryNameResult)
 
 	--- if the from is the new owner, then send an approved notice to the from

src/vaults.lua

@@ -233,6 +233,8 @@ end
 --- @param currentTimestamp number The current timestamp
 --- @return Vault[] The pruned vaults
 function vaults.pruneVaults(currentTimestamp)
+	local affectedAddresses = {}
+
 	if not NextBalanceVaultsPruneTimestamp or currentTimestamp < NextBalanceVaultsPruneTimestamp then
 		-- No known pruning work to do
 		return {}
@@ -248,12 +250,14 @@ function vaults.pruneVaults(currentTimestamp)
 		for id, nestedVault in pairs(ownersVaults) do
 			if currentTimestamp >= nestedVault.endTimestamp then
 				balances.increaseBalance(owner, nestedVault.balance)
+				table.insert(affectedAddresses, owner)
 				prunedVaults[id] = vaults.removeVault(owner, id)
 			else
 				vaults.scheduleNextVaultsPruning(nestedVault.endTimestamp)
 			end
 		end
 	end
+	balances.patchBalances(affectedAddresses)
 	return prunedVaults
 end