From 9ed95abf83897d6ed5ce105c3c98e7e3a13e9541 Mon Sep 17 00:00:00 2001
From: Engel Nyst <enyst@users.noreply.github.com>
Date: Sat, 27 Jul 2024 19:35:40 +0200
Subject: [PATCH] Fix max budget per task error in headless mode (#3147)

* set agent in ERROR instead of PAUSED when in headless mode

* fallback to config value for budget
---
 opendevin/controller/agent_controller.py | 16 ++++++++++++----
 opendevin/core/main.py                   |  7 +++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/opendevin/controller/agent_controller.py b/opendevin/controller/agent_controller.py
index ad9bdc4ffb..350ae49d5c 100644
--- a/opendevin/controller/agent_controller.py
+++ b/opendevin/controller/agent_controller.py
@@ -380,10 +380,18 @@ class AgentController:
                     self.state.traffic_control_state = TrafficControlState.NORMAL
                 else:
                     self.state.traffic_control_state = TrafficControlState.THROTTLING
-                    await self.report_error(
-                        f'Task budget exceeded. Current cost: {current_cost:.2f}, Max budget: {self.max_budget_per_task:.2f}, task paused. {TRAFFIC_CONTROL_REMINDER}'
-                    )
-                    await self.set_agent_state_to(AgentState.PAUSED)
+                    if self.headless_mode:
+                        # set to ERROR state if running in headless mode
+                        # there is no way to resume
+                        await self.report_error(
+                            f'Task budget exceeded. Current cost: {current_cost:.2f}, max budget: {self.max_budget_per_task:.2f}, task stopped.'
+                        )
+                        await self.set_agent_state_to(AgentState.ERROR)
+                    else:
+                        await self.report_error(
+                            f'Task budget exceeded. Current cost: {current_cost:.2f}, Max budget: {self.max_budget_per_task:.2f}, task paused. {TRAFFIC_CONTROL_REMINDER}'
+                        )
+                        await self.set_agent_state_to(AgentState.PAUSED)
                     return
 
         self.update_state_before_step()
diff --git a/opendevin/core/main.py b/opendevin/core/main.py
index 34f5808297..f51e0d5178 100644
--- a/opendevin/core/main.py
+++ b/opendevin/core/main.py
@@ -177,6 +177,13 @@ if __name__ == '__main__':
     AgentCls: Type[Agent] = Agent.get_cls(args.agent_cls)
     agent = AgentCls(llm=llm)
 
+    # if max budget per task is not sent on the command line, use the config value
+    max_budget_per_task = (
+        args.max_budget_per_task
+        if args.max_budget_per_task
+        else config.max_budget_per_task
+    )
+
     asyncio.run(
         run_agent_controller(
             agent=agent,