chore: import upstream snapshot with attribution

libs/python/agent/cua_agent/callbacks/budget_manager.py

@@ -0,0 +1,56 @@
+from typing import Any, Dict, List
+
+from .base import AsyncCallbackHandler
+
+
+class BudgetExceededError(Exception):
+    """Exception raised when budget is exceeded."""
+
+    pass
+
+
+class BudgetManagerCallback(AsyncCallbackHandler):
+    """Budget manager callback that tracks usage costs and can stop execution when budget is exceeded."""
+
+    def __init__(
+        self, max_budget: float, reset_after_each_run: bool = True, raise_error: bool = False
+    ):
+        """
+        Initialize BudgetManagerCallback.
+
+        Args:
+            max_budget: Maximum budget allowed
+            reset_after_each_run: Whether to reset budget after each run
+            raise_error: Whether to raise an error when budget is exceeded
+        """
+        self.max_budget = max_budget
+        self.reset_after_each_run = reset_after_each_run
+        self.raise_error = raise_error
+        self.total_cost = 0.0
+
+    async def on_run_start(self, kwargs: Dict[str, Any], old_items: List[Dict[str, Any]]) -> None:
+        """Reset budget if configured to do so."""
+        if self.reset_after_each_run:
+            self.total_cost = 0.0
+
+    async def on_usage(self, usage: Dict[str, Any]) -> None:
+        """Track usage costs."""
+        if "response_cost" in usage:
+            self.total_cost += usage["response_cost"]
+
+    async def on_run_continue(
+        self,
+        kwargs: Dict[str, Any],
+        old_items: List[Dict[str, Any]],
+        new_items: List[Dict[str, Any]],
+    ) -> bool:
+        """Check if budget allows continuation."""
+        if self.total_cost >= self.max_budget:
+            if self.raise_error:
+                raise BudgetExceededError(
+                    f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}"
+                )
+            else:
+                print(f"Budget exceeded: ${self.total_cost} >= ${self.max_budget}")
+            return False
+        return True