Skip to content

feat: implement RoundCostAwarenessHook for budget visibility #781 #1020

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jhawpetoss6-collab wants to merge 1 commit into
base: main
Choose a base branch
from
+19 −0
Open

feat: implement RoundCostAwarenessHook for budget visibility #781 #1020

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions massgen/hooks/RoundCostAwarenessHook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
class RoundCostAwarenessHook:
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify whether RoundCostAwarenessHook is registered and follows existing hook base patterns.

set -euo pipefail

echo "== Hook class declarations =="
rg -nP --type=py -C2 'class\s+\w+Hook\s*(\(|:)' massgen

echo
echo "== RoundCostAwarenessHook references =="
rg -nP --type=py -C3 '\bRoundCostAwarenessHook\b|hook_type\s*==\s*"RoundCostAwareness"|register_hooks_from_config\s*\(' massgen

echo
echo "== Existing hook inheritance patterns =="
rg -nP --type=py -C2 'class\s+\w+\(PatternHook\)|class\s+\w+\(FunctionHook\)' massgen

Repository: massgen/MassGen

Length of output: 30750

🏁 Script executed:

cat -n massgen/hooks/RoundCostAwarenessHook.py

Repository: massgen/MassGen

Length of output: 908

🏁 Script executed:

ls -la massgen/hooks/ && echo "---" && cat massgen/hooks/__init__.py 2>/dev/null || echo "No __init__.py in hooks/"

Repository: massgen/MassGen

Length of output: 276

🏁 Script executed:

echo "=== Searching for any imports of RoundCostAwarenessHook ===" && \
rg -nP "from.*RoundCostAwarenessHook|import.*RoundCostAwarenessHook|massgen\.hooks" --type=py massgen && \
echo "=== Checking if register_hooks_from_config dynamically loads hooks ===" && \
rg -nP -A20 "def register_hooks_from_config" massgen/mcp_tools/hooks.py | head -60

Repository: massgen/MassGen

Length of output: 1090

🏁 Script executed:

rg -nP -A60 "def register_hooks_from_config" massgen/mcp_tools/hooks.py | head -120

Repository: massgen/MassGen

Length of output: 2617

🏁 Script executed:

rg -nP -A100 "def _create_hook_from_config" massgen/mcp_tools/hooks.py | head -150

Repository: massgen/MassGen

Length of output: 3930

Inherit from PatternHook and implement required execute() method—hook will not execute in framework otherwise.

RoundCostAwarenessHook cannot function as a hook in the MassGen framework. The framework expects all hooks to inherit from PatternHook and implement an async def execute(function_name, arguments, context, **kwargs) -> HookResult method. This class has neither. If someone tries to register it via configuration, it will fail at runtime.

Additionally:

  • Line 7-19: Add Google-style docstrings to new methods (__init__, update_cost, get_status_string)
  • Line 13: if self.budget and fails for valid budget=0. Use explicit is not None check instead.
  • Line 14: Use structured logging (logger.warning) instead of print()
  • Line 12: Validate that cost is a valid numeric type before accumulating
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@massgen/hooks/RoundCostAwarenessHook.py` at line 1, RoundCostAwarenessHook
must inherit from PatternHook and implement the required async
execute(function_name, arguments, context, **kwargs) -> HookResult so the
framework can call it; update the class signature to subclass PatternHook and
add an async def execute(...) that returns a HookResult and delegates to
existing methods. In __init__, update_cost, and get_status_string add
Google-style docstrings, ensure update_cost validates that cost is numeric
(int/float) before accumulating, change the budget check from "if self.budget
and" to "if self.budget is not None and ..." to allow budget=0, and replace
print(...) with structured logging via logger.warning (or logger.error as
appropriate). Ensure execute uses the class methods
(update_cost/get_status_string) and imports/returns the correct HookResult type
expected by the framework.

"""
Hook for per-round cost awareness in MassGen.
Injects cumulative cost metadata with tool results.
Addresses issue #781.
"""
def __init__(self, budget=None):
self.cumulative_cost = 0.0
self.budget = budget
Comment on lines +7 to +9
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Budget boundary handling is incorrect for 0 and invalid for negatives.

Using if self.budget treats 0 as “no budget,” disabling warnings/status budget formatting. Also, negative budgets are accepted silently. Use explicit None checks and validate non-negative budget in __init__.

Proposed fix 
 class RoundCostAwarenessHook:
@@
     def __init__(self, budget=None):
+        if budget is not None and budget < 0:
+            raise ValueError("budget must be >= 0")
         self.cumulative_cost = 0.0
         self.budget = budget
@@
-        if self.budget and self.cumulative_cost > self.budget * 0.75:
+        if self.budget is not None and self.cumulative_cost > self.budget * 0.75:
             print(f"Warning: Nearing budget limit [${self.cumulative_cost:.2f} / ${self.budget:.2f}]")
@@
-        if self.budget:
+        if self.budget is not None:
             return f"Cost: ${self.cumulative_cost:.2f} / ${self.budget:.2f} budget"

Also applies to: 13-18

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@massgen/hooks/RoundCostAwarenessHook.py` around lines 7 - 9, In
RoundCostAwarenessHook.__init__, validate the incoming budget explicitly: if
budget is None leave as-is, but if budget is negative raise a ValueError to
reject invalid budgets; set self.budget = budget (allow 0). Then update any
places in this class (e.g., methods that currently do checks like "if
self.budget") to use an explicit None check ("if self.budget is not None") so a
budget of 0 is handled correctly when formatting warnings/status or comparing
cumulative_cost. Ensure all references to budget use the same None-check logic.


def update_cost(self, cost):
self.cumulative_cost += cost
if self.budget and self.cumulative_cost > self.budget * 0.75:
Comment on lines +11 to +13
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

update_cost should reject negative/non-numeric costs.

Without validation, cumulative_cost can be decreased or crash on invalid input types.

Proposed fix 
     def update_cost(self, cost):
+        if not isinstance(cost, (int, float)):
+            raise TypeError("cost must be numeric")
+        if cost < 0:
+            raise ValueError("cost must be >= 0")
         self.cumulative_cost += cost
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@massgen/hooks/RoundCostAwarenessHook.py` around lines 11 - 13, In
RoundCostAwarenessHook.update_cost validate the incoming cost before mutating
self.cumulative_cost: ensure cost is a numeric (e.g., int or float), not NaN,
and non-negative; if the check fails raise a ValueError (or TypeError for wrong
type) and do not change cumulative_cost or trigger budget logic; add this guard
at the top of update_cost to reject negative/non-numeric inputs and only then
increment self.cumulative_cost and evaluate the budget threshold.

print(f"Warning: Nearing budget limit [${self.cumulative_cost:.2f} / ${self.budget:.2f}]")

def get_status_string(self):
if self.budget:
return f"Cost: ${self.cumulative_cost:.2f} / ${self.budget:.2f} budget"
return f"Cost: ${self.cumulative_cost:.2f}"