Automatically stake PNK rewards instead of transferring them #2099
Conversation
WalkthroughImplements conditional auto-staking of PNK rewards via a new SortitionModule.setStakeReward API. KlerosCore reward execution now attempts to stake juror PNK rewards per court, falling back to direct token transfer if staking fails. Interfaces and university/base modules are updated accordingly, and tests adjust expected balances. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Core as KlerosCore(_executeRewards)
participant Sortition as SortitionModule
participant PNK as Pinakion Token
Core->>PNK: Transfer fee reward (ETH/Token path unchanged)
Note over Core: Compute pnkReward
Core->>Sortition: setStakeReward(account, courtID, pnkReward)
alt setStakeReward succeeds
Note over Sortition,Core: PNK reward credited to juror stake
else setStakeReward fails
Core->>PNK: safeTransfer(account, pnkReward)
Note over Core: Fallback direct PNK payout
end
Note over Core: sumFeeRewardPaid / sumPnkRewardPaid updated
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Suggested labels
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (9)
contracts/src/arbitration/interfaces/ISortitionModule.sol (2)
32-33: Add NatSpec to document new reward-staking API and its return semanticsThe new function is clear, but please document:
- When it returns false (e.g., juror not staked in that court, or other preconditions).
- That it must not create a new stake path (only top-up existing court stake).
- Whether phase constraints apply (see my comment on SortitionModuleBase).
This helps downstream implementers and keeps interface behavior unambiguous.
32-33: Optional: encode phase behavior in the interface commentIf you adopt phase gating in setStakeReward, please reflect that in the interface comments so callers (Core) understand that return false may simply mean “not Staking phase,” and should fallback to transfer.
Would you like me to update the interface NatSpec if we merge the phase gate?
contracts/src/arbitration/SortitionModuleBase.sol (1)
309-318: Consider emitting a dedicated event for reward-based stake increasesToday, reward credits emit StakeSet only, making it hard to distinguish user-initiated staking from reward compounding in off-chain analytics. An optional, additive event (e.g., RewardStaked(address juror, uint96 courtID, uint256 reward)) helps indexers and dashboards.
contracts/src/arbitration/KlerosCoreBase.sol (1)
858-862: Avoid external call when pnkReward is zero, and rely on phase-gated setStakeRewardTwo small tweaks:
- Skip calling sortitionModule when pnkReward == 0 to save gas.
- If you accept gating setStakeReward to Staking phase (see SortitionModuleBase note), this code remains correct: it will fall back to transfer when false.
Diff within this hunk:
- // Stake the PNK reward if possible, by-passes delayed stakes and other checks usually done by validateStake() - if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) { - pinakion.safeTransfer(account, pnkReward); - } + // Stake the PNK reward if possible (phase-gated in SortitionModule); otherwise transfer. + if (pnkReward != 0) { + if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) { + pinakion.safeTransfer(account, pnkReward); + } + }If you don’t adopt the phase gate, please confirm you’re comfortable with sum-tree mutations during Generating/Drawing potentially affecting ongoing draws in other disputes.
contracts/src/arbitration/university/KlerosCoreUniversity.sol (2)
850-857: Prefer a safe, unified fee payout instead of raw send()University version still uses payable(account).send(feeReward) (unchecked boolean). Consider mirroring KlerosCoreBase’s _transferFeeToken pattern (SafeSend) to avoid silent failures and to standardize behavior across variants. If you want to keep University minimal, at least check the boolean and revert on failure.
Example minimal change:
- if (round.feeToken == NATIVE_CURRENCY) { - // The dispute fees were paid in ETH - payable(account).send(feeReward); - } else { - // The dispute fees were paid in ERC20 - round.feeToken.safeTransfer(account, feeReward); - } + if (round.feeToken == NATIVE_CURRENCY) { + require(payable(account).send(feeReward), "ETH fee payout failed"); + } else { + round.feeToken.safeTransfer(account, feeReward); + }
858-862: Skip SortitionModule call for zero reward and rely on phase-gate behaviorSame gas/tidiness nit as base contract: avoid calling setStakeReward when pnkReward == 0; if setStakeReward is gated to Staking phase, fallback to transfer preserves behavior during Drawing/Generating.
- if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) { - pinakion.safeTransfer(account, pnkReward); - } + if (pnkReward != 0) { + if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) { + pinakion.safeTransfer(account, pnkReward); + } + }contracts/test/foundry/KlerosCore.t.sol (1)
2400-2403: Fix comments to reflect the new autostake behaviorThe assertions are correct, but the inline comments contradict them:
- Core balance comment says “Was 21500. 1000 was transferred to staker2” — here 1000 was autostaked (into staker2’s stake via core custody), not transferred to wallet.
- staker2 comment says “thus -19k from the default” — the wallet is down by 20k (stake), the 1k reward is staked, not in wallet.
Proposed diff (comments only):
- assertEq(pinakion.balanceOf(address(core)), 21500, "Wrong token balance of the core"); // Was 21500. 1000 was transferred to staker2 + assertEq(pinakion.balanceOf(address(core)), 21500, "Wrong token balance of the core"); // 1000 PNK reward was autostaked (held by core) ... - assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2"); // 20k stake and 1k added as a reward, thus -19k from the default + assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2"); // 20k deposited as stake; 1k reward autostaked (not in wallet)contracts/src/arbitration/university/SortitionModuleUniversity.sol (2)
220-226: _setStake entry point looks good; note minor gas/readability tweaks.No functional issues spotted. Optional micro-tweaks you may consider:
- Make the deposit/withdraw branches explicit:
if (_pnkDeposit > 0) { ... } else if (_pnkWithdrawal > 0) { ... }to skip the no-op subtraction when both are zero.- In the cleanup loop that walks courtIDs in reverse, use unchecked decrement to save gas.
- In the ancestor update, consider applying a single delta per court (
±(new - current)) instead of two consecutive updates to reduce SSTOREs.If helpful, I can draft a focused gas-optimized variant of the loop/update block.
206-218: setStakeReward behavior confirmed; optional O(1) lookup refactor and event emission recommended
- Verified that in SortitionModuleUniversity.sol (lines 206–218)
setStakeReward(address,uint96,uint256)is declared as
external override onlyByCore returns (bool success)and uses the intended early-return logic:
● returnstruewhen_reward == 0
● returnsfalsewhen the juror’s current stake is zero, allowing the caller to fallback to a direct transfer- Checked KlerosCoreUniversity.sol (line 859) and KlerosCoreBase.sol (line 859) and confirmed that both core implementations handle a
falsereturn by callingpinakion.safeTransfer(account, pnkReward)as a fallback—ensuring no funds are lost when auto-staking isn’t possible- The method signature and caller expectations align correctly; no breaking changes detected
Optional improvements to consider:
- Replace the O(n) scan in
_stakeOf(_account,_courtID)with a direct storage lookup (e.g.jurors[_account].stakesByCourtID[_courtID]) to reduce gas in hot reward loops- Emit a dedicated
RewardAutoStaked(address indexed account, uint96 indexed courtID, uint256 reward, uint256 newStake)event immediately after_setStakefor easier off-chain indexing and auditing
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (6)
contracts/src/arbitration/KlerosCoreBase.sol(1 hunks)contracts/src/arbitration/SortitionModuleBase.sol(1 hunks)contracts/src/arbitration/interfaces/ISortitionModule.sol(1 hunks)contracts/src/arbitration/university/KlerosCoreUniversity.sol(1 hunks)contracts/src/arbitration/university/SortitionModuleUniversity.sol(1 hunks)contracts/test/foundry/KlerosCore.t.sol(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: contracts-testing
- GitHub Check: Analyze (javascript)
🔇 Additional comments (4)
contracts/src/arbitration/KlerosCoreBase.sol (2)
849-854: Micro: reuse helper to compute rewards (already done) – looks goodThe switch to _applyCoherence improves readability and reduces duplication. No action needed.
855-857: Uniform fee transfer helper: good callUsing _transferFeeToken unifies ETH/ERC20 payout logic and inherits SafeSend behavior. Nice.
contracts/src/arbitration/university/KlerosCoreUniversity.sol (1)
843-848: LGTM on reward computationThe reward math mirrors the base contract and maintains sumPnkRewardPaid/sumFeeRewardPaid accounting. Good.
contracts/src/arbitration/university/SortitionModuleUniversity.sol (1)
193-195: Good extraction: single source of truth for stake mutations via _setStake.Routing setStake through the internal helper removes duplication and ensures the rewards path shares identical accounting side-effects and events. LGTM.
| address _account, | ||
| uint96 _courtID, | ||
| uint256 _reward | ||
| ) external override onlyByCore returns (bool success) { | ||
| if (_reward == 0) return true; // No reward to add. | ||
|
|
||
| uint256 currentStake = stakeOf(_account, _courtID); | ||
| if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake. | ||
|
|
||
| uint256 newStake = currentStake + _reward; | ||
| _setStake(_account, _courtID, _reward, 0, newStake); | ||
| return true; | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Do not mutate sortition trees during Generating/Drawing: gate setStakeReward by phase
setStakeReward calls _setStake directly, bypassing validateStake() and the phase guard that defers writes during non-staking phases. That diverges from the module’s invariant “sum trees are only updated in Staking phase,” and can alter draw probabilities mid-Drawing for unrelated disputes. This can lead to inconsistent sortition outcomes and hard-to-reason temporal coupling between execution and drawing.
Recommend returning false when phase != Phase.staking so KlerosCore falls back to transfer (as designed), preserving invariants and avoiding mid-phase mutations.
Diff within this hunk:
function setStakeReward(
address _account,
uint96 _courtID,
uint256 _reward
) external override onlyByCore returns (bool success) {
if (_reward == 0) return true; // No reward to add.
+ // Preserve invariant: do not modify trees outside Staking phase.
+ if (phase != Phase.staking) {
+ return false;
+ }
+
uint256 currentStake = stakeOf(_account, _courtID);
if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake.
uint256 newStake = currentStake + _reward;
_setStake(_account, _courtID, _reward, 0, newStake);
return true;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| address _account, | |
| uint96 _courtID, | |
| uint256 _reward | |
| ) external override onlyByCore returns (bool success) { | |
| if (_reward == 0) return true; // No reward to add. | |
| uint256 currentStake = stakeOf(_account, _courtID); | |
| if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake. | |
| uint256 newStake = currentStake + _reward; | |
| _setStake(_account, _courtID, _reward, 0, newStake); | |
| return true; | |
| } | |
| function setStakeReward( | |
| address _account, | |
| uint96 _courtID, | |
| uint256 _reward | |
| ) external override onlyByCore returns (bool success) { | |
| if (_reward == 0) return true; // No reward to add. | |
| // Preserve invariant: do not modify trees outside Staking phase. | |
| if (phase != Phase.staking) { | |
| return false; | |
| } | |
| uint256 currentStake = stakeOf(_account, _courtID); | |
| if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake. | |
| uint256 newStake = currentStake + _reward; | |
| _setStake(_account, _courtID, _reward, 0, newStake); | |
| return true; | |
| } |
🤖 Prompt for AI Agents
In contracts/src/arbitration/SortitionModuleBase.sol around lines 319 to 331,
setStakeReward currently calls _setStake unconditionally which mutates sortition
trees during non-staking phases; change the function to check the current phase
and return false unless Phase.staking (or call validateStake() that enforces the
same guard) so that writes are deferred and KlerosCore will fall back to
transfer, avoiding mid-Drawing mutations and preserving the invariant that sum
trees are only updated in Staking phase.
✅ Deploy Preview for kleros-v2-testnet-devtools ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
❌ Deploy Preview for kleros-v2-university failed. Why did it fail? →
|
✅ Deploy Preview for kleros-v2-neo ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
✅ Deploy Preview for kleros-v2-testnet ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
PR-Codex overview
This PR introduces a new
setStakeRewardfunction to manage PNK rewards for jurors and updates the reward calculation process in multiple contracts. It also corrects token balance assertions in tests.Detailed summary
setStakeRewardfunction inISortitionModule.solto deposit PNK rewards.KlerosCoreBase.solandKlerosCoreUniversity.sol.KlerosCore.t.sol.KlerosCoreBase.solandKlerosCoreUniversity.solto transfer rewards conditionally.setStakeRewardfunction inSortitionModuleBase.solandSortitionModuleUniversity.sol.Summary by CodeRabbit