From 50a1c8d948fc26d1420291cf8c82ba83ab81ab73 Mon Sep 17 00:00:00 2001
From: johnmalek312 <johnmalek312@users.noreply.github.com>
Date: Wed, 22 Oct 2025 22:19:18 +1100
Subject: [PATCH] macro fix

---
 .gitignore                                |  4 +++-
 droidrun/agent/codeact/codeact_agent.py   |  3 +++
 droidrun/agent/common/events.py           |  6 ++++++
 droidrun/agent/executor/executor_agent.py |  3 +++
 droidrun/agent/utils/tools.py             | 14 +++++++++++++-
 droidrun/macro/cli.py                     |  3 +++
 droidrun/macro/replay.py                  | 10 ++++++++++
 7 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 86df54b..a517368 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,4 +38,6 @@ credentials.yaml
 backend/
 frontend/
 
-reddit_python_posts.json
\ No newline at end of file
+reddit_python_posts.json
+
+test/
\ No newline at end of file
diff --git a/droidrun/agent/codeact/codeact_agent.py b/droidrun/agent/codeact/codeact_agent.py
index ab84a14..b097729 100644
--- a/droidrun/agent/codeact/codeact_agent.py
+++ b/droidrun/agent/codeact/codeact_agent.py
@@ -191,6 +191,9 @@ class CodeActAgent(Workflow):
     @step
     async def prepare_chat(self, ctx: Context, ev: StartEvent) -> TaskInputEvent:
         """Prepare chat history from user input."""
+        # macro tools context
+        self.tools._set_context(ctx)
+
         logger.info("💬 Preparing chat for task execution...")
 
         self.chat_memory: Memory = await ctx.store.get(
diff --git a/droidrun/agent/common/events.py b/droidrun/agent/common/events.py
index 191ad64..91a7e7d 100644
--- a/droidrun/agent/common/events.py
+++ b/droidrun/agent/common/events.py
@@ -64,5 +64,11 @@ class StartAppEvent(MacroEvent):
     activity: str = None
 
 
+class WaitEvent(MacroEvent):
+    """Event for wait/sleep actions"""
+
+    duration: float
+
+
 class RecordUIStateEvent(Event):
     ui_state: list[Dict[str, Any]]
diff --git a/droidrun/agent/executor/executor_agent.py b/droidrun/agent/executor/executor_agent.py
index 0d7528f..fbabd41 100644
--- a/droidrun/agent/executor/executor_agent.py
+++ b/droidrun/agent/executor/executor_agent.py
@@ -94,6 +94,9 @@ class ExecutorAgent(Workflow):
         3. Validates action format (blocks answer actions!)
         4. Returns action event
         """
+        # macro tools context
+        self.tools_instance._set_context(ctx)
+
         subgoal = ev.get("subgoal", "")
         logger.info(f"🧠 Executor thinking about action for: {subgoal}")
 
diff --git a/droidrun/agent/utils/tools.py b/droidrun/agent/utils/tools.py
index bcb08b5..8cf2e4b 100644
--- a/droidrun/agent/utils/tools.py
+++ b/droidrun/agent/utils/tools.py
@@ -260,16 +260,28 @@ async def open_app(text: str, *, tools: "Tools" = None, **kwargs) -> str:
     return result
 
 
-def wait(duration: float, **kwargs) -> str:
+def wait(duration: float = 1.0, *, tools: "Tools" = None, **kwargs) -> str:
     """
     Wait for a specified duration in seconds.
 
     Args:
         duration: Duration to wait in seconds
+        tools: The Tools instance (injected automatically)
 
     Returns:
         Confirmation message
     """
+    # Emit WaitEvent for macro recording if context available
+    if tools is not None and hasattr(tools, "_ctx") and tools._ctx is not None:
+        from droidrun.agent.common.events import WaitEvent
+
+        wait_event = WaitEvent(
+            action_type="wait",
+            description=f"Wait for {duration} seconds",
+            duration=duration,
+        )
+        tools._ctx.write_event_to_stream(wait_event)
+
     time.sleep(duration)
     return f"Waited for {duration} seconds"
 
diff --git a/droidrun/macro/cli.py b/droidrun/macro/cli.py
index fe1bf8c..0d80413 100644
--- a/droidrun/macro/cli.py
+++ b/droidrun/macro/cli.py
@@ -201,6 +201,9 @@ async def _show_dry_run(
         elif action_type == "key_press":
             key_name = action.get("key_name", "UNKNOWN")
             details = f"{key_name}"
+        elif action_type == "wait":
+            duration = action.get("duration", 1.0)
+            details = f"{duration}s"
 
         description = action.get("description", "")
         table.add_row(
diff --git a/droidrun/macro/replay.py b/droidrun/macro/replay.py
index 0f05701..c9f6ea4 100644
--- a/droidrun/macro/replay.py
+++ b/droidrun/macro/replay.py
@@ -7,6 +7,7 @@ that were generated during DroidAgent trajectory recording.
 
 import asyncio
 import logging
+import time
 from typing import Any, Dict, Optional
 
 from droidrun.agent.utils.trajectory import Trajectory
@@ -108,6 +109,8 @@ class MacroPlayer:
                 )
                 result = tools.swipe(start_x, start_y, end_x, end_y, duration_ms)
                 logger.debug(f"   Result: {result}")
+                # Additional wait after swipe for UI to settle
+                time.sleep(2)
                 return True
 
             elif action_type == "drag":
@@ -149,6 +152,13 @@ class MacroPlayer:
                 logger.debug(f"   Result: {result}")
                 return True
 
+            elif action_type == "wait":
+                duration = action.get("duration", 1.0)
+                logger.info(f"⏳ Waiting for {duration} seconds")
+                time.sleep(duration)
+                logger.debug(f"   Waited for {duration} seconds")
+                return True
+
             else:
                 logger.warning(f"⚠️  Unknown action type: {action_type}")
                 return False