From 8dc4cdd64fea19e070f8d48dd7ff2802a7cec093 Mon Sep 17 00:00:00 2001
From: lidia9 <lidia@aliasrobotics.com>
Date: Mon, 19 May 2025 10:42:11 +0200
Subject: [PATCH] FIX CAI_PRICE_LIMIT

---
 src/cai/repl/ui/toolbar.py       |  2 ++
 src/cai/sdk/agents/exceptions.py |  6 ++++++
 src/cai/sdk/agents/run.py        |  9 +++++++++
 src/cai/util.py                  | 20 +++++++++++++-------
 4 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py
index eb44d9a0..6f3edde6 100644
--- a/src/cai/repl/ui/toolbar.py
+++ b/src/cai/repl/ui/toolbar.py
@@ -123,6 +123,8 @@ def update_toolbar_in_background():
                 os.getenv('CAI_MODEL', 'default')}</ansigreen> | "
             f"<ansicyan><b>Max Turns:</b></ansicyan> <ansiblue>{
                 os.getenv('CAI_MAX_TURNS', 'inf')}</ansiblue> | "
+            f"<ansiyellow><b>Price Limit:</b></ansiyellow> <ansiblue>{
+                os.getenv('CAI_PRICE_LIMIT', 'inf')}</ansiblue> | "
             f"<ansigray>{current_time_with_tz}</ansigray>"
         )
         toolbar_cache['last_update'] = datetime.datetime.now()
diff --git a/src/cai/sdk/agents/exceptions.py b/src/cai/sdk/agents/exceptions.py
index 820390db..170bdd76 100644
--- a/src/cai/sdk/agents/exceptions.py
+++ b/src/cai/sdk/agents/exceptions.py
@@ -61,3 +61,9 @@ class OutputGuardrailTripwireTriggered(AgentsException):
         super().__init__(
             f"Guardrail {guardrail_result.guardrail.__class__.__name__} triggered tripwire"
         )
+
+
+class PriceLimitExceeded(AgentsException):
+    """Raised when the maximum price limit is exceeded."""
+    def __init__(self, current_cost: float, price_limit: float):
+        super().__init__(f"Maximum price limit (${price_limit:.4f}) exceeded. Current cost: ${current_cost:.4f}")
diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py
index 6089d331..15d6eb70 100644
--- a/src/cai/sdk/agents/run.py
+++ b/src/cai/sdk/agents/run.py
@@ -58,6 +58,15 @@ if max_turns_env is not None:
 else:
     DEFAULT_MAX_TURNS = float("inf")
 
+price_limit_env = os.getenv("CAI_PRICE_LIMIT")
+if price_limit_env is not None:
+    try:
+        DEFAULT_PRICE_LIMIT = float(price_limit_env)
+    except ValueError:
+        DEFAULT_PRICE_LIMIT = float("inf")
+else:
+    DEFAULT_PRICE_LIMIT = float("inf")
+
 
 @dataclass
 class RunConfig:
diff --git a/src/cai/util.py b/src/cai/util.py
index a33cb05d..cf94f2ab 100644
--- a/src/cai/util.py
+++ b/src/cai/util.py
@@ -228,16 +228,22 @@ class CostTracker:
     # Track the last calculation to debug inconsistencies
     last_interaction_cost: float = 0.0
     last_total_cost: float = 0.0
-    
-    def reset_interaction_stats(self):
-        """Reset stats for a new interaction"""
-        self.interaction_input_tokens = 0
-        self.interaction_output_tokens = 0
-        self.interaction_reasoning_tokens = 0
-        self.interaction_cost = 0.0
+
+    def check_price_limit(self, new_cost: float) -> None:
+        """Check if adding the new cost would exceed the price limit."""
+        from cai.sdk.agents.exceptions import PriceLimitExceeded
+        from cai.sdk.agents.run import DEFAULT_PRICE_LIMIT
+        
+        if DEFAULT_PRICE_LIMIT != float("inf"):
+            total_cost = self.session_total_cost + new_cost
+            if total_cost > DEFAULT_PRICE_LIMIT:
+                raise PriceLimitExceeded(total_cost, DEFAULT_PRICE_LIMIT)
     
     def update_session_cost(self, new_cost: float) -> None:
         """Add cost to session total and log the update"""
+        # Check price limit before updating
+        self.check_price_limit(new_cost)
+        
         old_total = self.session_total_cost
         self.session_total_cost += new_cost