From db1d8a4c5dbc4aae5ceb290fe349e6f99fe67133 Mon Sep 17 00:00:00 2001
From: Dillon DuPont <ddupont@mit.edu>
Date: Fri, 8 Aug 2025 19:20:17 -0400
Subject: [PATCH] moved trajectory dir fix

---
 libs/python/agent/agent/integrations/hud/__init__.py | 10 ----------
 libs/python/agent/agent/integrations/hud/agent.py    | 12 +++++++++++-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/libs/python/agent/agent/integrations/hud/__init__.py b/libs/python/agent/agent/integrations/hud/__init__.py
index 81f82925..787613de 100644
--- a/libs/python/agent/agent/integrations/hud/__init__.py
+++ b/libs/python/agent/agent/integrations/hud/__init__.py
@@ -7,7 +7,6 @@ from hud import run_job as hud_run_job
 from .agent import ComputerAgent
 from .adapter import ComputerAgentAdapter
 from .computer_handler import HUDComputerHandler
-from agent.callbacks.trajectory_saver import TrajectorySaverCallback
 
 
 async def run_job(
@@ -49,15 +48,6 @@ async def run_job(
     Returns:
         Job instance from HUD
     """
-    # Handle trajectory_dir by adding TrajectorySaverCallback
-    trajectory_dir = agent_kwargs.pop("trajectory_dir", None)
-    callbacks = agent_kwargs.get("callbacks", [])
-    
-    if trajectory_dir:
-        trajectory_callback = TrajectorySaverCallback(trajectory_dir, reset_on_run=False)
-        callbacks = callbacks + [trajectory_callback]
-        agent_kwargs["callbacks"] = callbacks
-    
     # combine verbose and verbosity kwargs
     if "verbose" in agent_kwargs:
         agent_kwargs["verbosity"] = logging.INFO
diff --git a/libs/python/agent/agent/integrations/hud/agent.py b/libs/python/agent/agent/integrations/hud/agent.py
index 58095417..cb810b36 100644
--- a/libs/python/agent/agent/integrations/hud/agent.py
+++ b/libs/python/agent/agent/integrations/hud/agent.py
@@ -86,6 +86,16 @@ class ComputerAgent(Agent[BaseComputerAgent, dict[str, Any]]):
             dimensions=(self.width, self.height)
         )
 
+        # Handle trajectory_dir by adding TrajectorySaverCallback
+        trajectory_dir = kwargs.pop("trajectory_dir", None)
+        callbacks = kwargs.get("callbacks", [])
+        
+        if trajectory_dir:
+            from agent.callbacks.trajectory_saver import TrajectorySaverCallback
+            trajectory_callback = TrajectorySaverCallback(trajectory_dir, reset_on_run=False)
+            callbacks = callbacks + [trajectory_callback]
+            kwargs["callbacks"] = callbacks
+
         # Initialize ComputerAgent with HUD computer handler
         self.computer_agent = BaseComputerAgent(
             model=model,
@@ -212,7 +222,7 @@ class ComputerAgent(Agent[BaseComputerAgent, dict[str, Any]]):
                 # ComputerAgent.run returns an async generator
                 async for result in self.computer_agent.run(self.conversation_history, stream=False):
                     # if the result has computer_call_output, immediately exit
-                    if result.get("output", [])[-1].get("type") == "computer_call_output":
+                    if result.get("output", []) and result.get("output", [])[-1].get("type") == "computer_call_output":
                         break
                     # otherwise add agent output to conversation history
                     self.conversation_history += result["output"]