Fix: Bump max_iterations when resuming due to throttling (#3410)

* Fix: Reset iteration count when resuming due to throttling

* Fix inadvertent additions

* WIP

* Changing max_iterations instead of iteration count

* Now adjusting max_iterations or max_budget_per_task as appropriate

* Fix check on iterations

* Fix linter issues

* AgentController: remember initial max_iterations and use it to extend state's iterations

* increase task budget by initial value (not doubling it)

---------

Co-authored-by: Tim O'Farrell <[email protected]>
Co-authored-by: tobitege <[email protected]>
Co-authored-by: mamoodi <[email protected]>

Author: 4 people

openhands/controller/agent_controller.py

@@ -108,6 +108,8 @@ def __init__(
         self.max_budget_per_task = max_budget_per_task
         self.agent_to_llm_config = agent_to_llm_config if agent_to_llm_config else {}
         self.agent_configs = agent_configs if agent_configs else {}
+        self._initial_max_iterations = max_iterations
+        self._initial_max_budget_per_task = max_budget_per_task
 
         # stuck helper
         self._stuck_detector = StuckDetector(self.state)
@@ -245,6 +247,21 @@ async def set_agent_state_to(self, new_state: AgentState):
         ):
             # user intends to interrupt traffic control and let the task resume temporarily
             self.state.traffic_control_state = TrafficControlState.PAUSED
+            # User has chosen to deliberately continue - lets double the max iterations
+            if (
+                self.state.iteration is not None
+                and self.state.max_iterations is not None
+                and self._initial_max_iterations is not None
+            ):
+                if self.state.iteration >= self.state.max_iterations:
+                    self.state.max_iterations += self._initial_max_iterations
+            if (
+                self.state.metrics.accumulated_cost is not None
+                and self.max_budget_per_task is not None
+                and self._initial_max_budget_per_task is not None
+            ):
+                if self.state.metrics.accumulated_cost >= self.max_budget_per_task:
+                    self.max_budget_per_task += self._initial_max_budget_per_task
 
         self.state.agent_state = new_state
         if new_state == AgentState.STOPPED or new_state == AgentState.ERROR: