droidrun 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

droidrun/agent/codeact/codeact_agent.py

@@ -165,16 +165,18 @@ class CodeActAgent(Workflow):
             chat_history = await chat_utils.add_memory_block(self.remembered_info, chat_history)
 
         for context in self.required_context:
-            if model == "DeepSeek":
-                logger.warning(
-                    "[yellow]DeepSeek doesnt support images. Disabling screenshots[/]"
-                )
-            elif self.vision == True and context == "screenshot":
+            if context == "screenshot":
+                # if vision is disabled, screenshot should save to trajectory
                 screenshot = (self.tools.take_screenshot())[1]
                 ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
 
                 await ctx.set("screenshot", screenshot)
-                chat_history = await chat_utils.add_screenshot_image_block(screenshot, chat_history)
+                if model == "DeepSeek":
+                    logger.warning(
+                        "[yellow]DeepSeek doesnt support images. Disabling screenshots[/]"
+                    )
+                elif self.vision == True: # if vision is enabled, add screenshot to chat history
+                    chat_history = await chat_utils.add_screenshot_image_block(screenshot, chat_history)
 
             if context == "ui_state":
                 try:
@@ -248,7 +250,10 @@ class CodeActAgent(Workflow):
         try:
             self.code_exec_counter += 1
             result = await self.executor.execute(ctx, code)
-            logger.info(f"💡 Code execution successful. Result: {result}")
+            logger.info(f"💡 Code execution successful. Result: {result['output']}")
+            screenshots = result['screenshots']
+            for screenshot in screenshots[:-1]: # the last screenshot will be captured by next step
+                ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
 
             if self.tools.finished == True:
                 logger.debug("  - Task completed.")
@@ -260,7 +265,7 @@ class CodeActAgent(Workflow):
             
             self.remembered_info = self.tools.memory
             
-            event = TaskExecutionResultEvent(output=str(result))
+            event = TaskExecutionResultEvent(output=str(result['output']))
             ctx.write_event_to_stream(event)
             return event
 

droidrun/agent/context/task_manager.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Dict
+from typing import List, Dict, Optional
 from dataclasses import dataclass
 import copy
 
@@ -11,6 +11,9 @@ class Task:
     description: str
     status: str
     agent_type: str
+    # Optional fields to carry success/failure context back to the planner
+    message: Optional[str] = None
+    failure_reason: Optional[str] = None
     
 
 class TaskManager:
@@ -40,14 +43,16 @@ class TaskManager:
     def get_task_history(self):
         return self.task_history
 
-    def complete_task(self, task: Task):
+    def complete_task(self, task: Task, message: Optional[str] = None):
         task = copy.deepcopy(task)
         task.status = self.STATUS_COMPLETED
+        task.message = message
         self.task_history.append(task)
 
-    def fail_task(self, task: Task):
+    def fail_task(self, task: Task, failure_reason: Optional[str] = None):
         task = copy.deepcopy(task)
         task.status = self.STATUS_FAILED
+        task.failure_reason = failure_reason
         self.task_history.append(task)
 
     def get_completed_tasks(self) -> list[dict]:

droidrun/agent/droid/droid_agent.py

@@ -68,7 +68,7 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         reflection: bool = False,
         enable_tracing: bool = False,
         debug: bool = False,
-        save_trajectories: bool = False,
+        save_trajectories: str = "none",
         excluded_tools: List[str] = None,
         *args,
         **kwargs
@@ -86,6 +86,10 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             reflection: Whether to reflect on steps the CodeActAgent did to give the PlannerAgent advice
             enable_tracing: Whether to enable Arize Phoenix tracing
             debug: Whether to enable verbose debug logging
+            save_trajectories: Trajectory saving level. Can be:
+                - "none" (no saving)
+                - "step" (save per step)
+                - "action" (save per action)
             **kwargs: Additional keyword arguments to pass to the agents
         """
         self.user_id = kwargs.pop("user_id", None)
@@ -114,7 +118,17 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         self.debug = debug
 
         self.event_counter = 0
-        self.save_trajectories = save_trajectories
+        # Handle backward compatibility: bool -> str mapping
+        if isinstance(save_trajectories, bool):
+            self.save_trajectories = "step" if save_trajectories else "none"
+        else:
+            # Validate string values
+            valid_values = ["none", "step", "action"]
+            if save_trajectories not in valid_values:
+                logger.warning(f"Invalid save_trajectories value: {save_trajectories}. Using 'none' instead.")
+                self.save_trajectories = "none"
+            else:
+                self.save_trajectories = save_trajectories
         
         self.trajectory = Trajectory(goal=goal)
         self.task_manager = TaskManager()
@@ -125,9 +139,12 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         self.current_episodic_memory = None
 
         logger.info("🤖 Initializing DroidAgent...")
+        logger.info(f"💾 Trajectory saving level: {self.save_trajectories}")
         
         self.tool_list = describe_tools(tools, excluded_tools)
         self.tools_instance = tools
+        
+        self.tools_instance.save_trajectories = self.save_trajectories
 
 
         if self.reasoning:
@@ -173,11 +190,11 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
         
         logger.info("✅ DroidAgent initialized successfully.")
 
-    def run(self) -> WorkflowHandler:
+    def run(self, *args, **kwargs) -> WorkflowHandler:
         """
         Run the DroidAgent workflow.
         """
-        return super().run()
+        return super().run(*args, **kwargs)
     
     @step
     async def execute_task(
@@ -237,16 +254,24 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             return CodeActResultEvent(success=False, reason=f"Error: {str(e)}", task=task, steps=[])
     
     @step
-    async def handle_codeact_execute(self, ctx: Context, ev: CodeActResultEvent) -> FinalizeEvent | ReflectionEvent:
+    async def handle_codeact_execute(self, ctx: Context, ev: CodeActResultEvent) -> FinalizeEvent | ReflectionEvent | ReasoningLogicEvent:
         try:
             task = ev.task
             if not self.reasoning:
                 return FinalizeEvent(success=ev.success, reason=ev.reason, output=ev.reason, task=[task], tasks=[task], steps=ev.steps)
             
-            if self.reflection:
+            if self.reflection and ev.success:
                 return ReflectionEvent(task=task)
-            
-            return ReasoningLogicEvent()
+
+            # Reasoning is enabled but reflection is disabled.
+            # Success: mark complete and proceed to next step in reasoning loop.
+            # Failure: mark failed and trigger planner immediately without advancing to the next queued task.
+            if ev.success:
+                self.task_manager.complete_task(task, message=ev.reason)
+                return ReasoningLogicEvent()
+            else:
+                self.task_manager.fail_task(task, failure_reason=ev.reason)
+                return ReasoningLogicEvent(force_planning=True)
 
         except Exception as e:
             logger.error(f"❌ Error during DroidAgent execution: {e}")
@@ -298,7 +323,7 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             if ev.reflection:
                 handler = planner_agent.run(remembered_info=self.tools_instance.memory, reflection=ev.reflection)
             else:
-                if self.task_iter:
+                if not ev.force_planning and self.task_iter:
                     try:
                         task = next(self.task_iter)
                         return CodeActExecuteEvent(task=task, reflection=None)
@@ -387,7 +412,7 @@ A wrapper class that coordinates between PlannerAgent (creates plans) and
             "steps": ev.steps,
         }
 
-        if self.trajectory and self.save_trajectories:
+        if self.trajectory and self.save_trajectories != "none":
             self.trajectory.save_trajectory()
 
         return StopEvent(result)

droidrun/agent/droid/events.py

@@ -13,6 +13,7 @@ class CodeActResultEvent(Event):
 
 class ReasoningLogicEvent(Event):
     reflection: Optional[Reflection] = None
+    force_planning: bool = False
 
 class FinalizeEvent(Event):
     success: bool

droidrun/agent/planner/planner_agent.py

@@ -130,9 +130,10 @@ class PlannerAgent(Workflow):
         self.steps_counter += 1
         logger.info(f"🧠 Thinking about how to plan the goal...")
 
+        # if vision is disabled, screenshot should save to trajectory
+        screenshot = (self.tools_instance.take_screenshot())[1]
+        ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
         if self.vision:
-            screenshot = (self.tools_instance.take_screenshot())[1]
-            ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
             await ctx.set("screenshot", screenshot)
 
         try:
@@ -168,11 +169,15 @@ class PlannerAgent(Workflow):
             try:
                 result = await self.executer.execute(ctx, code)
                 logger.info(f"📝 Planning complete")
-                logger.debug(f"  - Planning code executed. Result: {result}")
+                logger.debug(f"  - Planning code executed. Result: {result['output']}")
+
+                screenshots = result['screenshots']
+                for screenshot in screenshots[:-1]: # the last screenshot will be captured by next step
+                    ctx.write_event_to_stream(ScreenshotEvent(screenshot=screenshot))
 
                 await self.chat_memory.aput(
                     ChatMessage(
-                        role="user", content=f"Execution Result:\n```\n{result}\n```"
+                        role="user", content=f"Execution Result:\n```\n{result['output']}\n```"
                     )
                 )
 
@@ -241,15 +246,15 @@ wrap your code inside this:
             logger.debug(f"  - Sending {len(chat_history)} messages to LLM.")
 
             model = self.llm.class_name()
-            if model == "DeepSeek":
-                logger.warning(
-                    "[yellow]DeepSeek doesnt support images. Disabling screenshots[/]"
-                )
-
-            elif self.vision == True:
-                chat_history = await chat_utils.add_screenshot_image_block(
-                    await ctx.get("screenshot"), chat_history
-                )                   
+            if self.vision == True:
+                if model == "DeepSeek":
+                    logger.warning(
+                        "[yellow]DeepSeek doesnt support images. Disabling screenshots[/]"
+                    )
+                else:
+                    chat_history = await chat_utils.add_screenshot_image_block(
+                        await ctx.get("screenshot"), chat_history
+                    )                   
 
 
 

droidrun/agent/utils/chat_utils.py

@@ -208,7 +208,7 @@ async def add_task_history_block(completed_tasks: list[dict], failed_tasks: list
     all_tasks = completed_tasks + failed_tasks
     
     if all_tasks:
-        task_history += "Task History (chronological order):\n"
+        task_history += "### Task Execution History (chronological):\n"
         for i, task in enumerate(all_tasks, 1):
             if hasattr(task, 'description'):
                 status_indicator = "[success]" if hasattr(task, 'status') and task.status == "completed" else "[failed]"

droidrun/agent/utils/executer.py

@@ -99,10 +99,13 @@ class SimpleCodeExecutor:
         """
         # Update UI elements before execution
         self.globals['ui_state'] = await ctx.get("ui_state", None)
-
+        self.globals['step_screenshots'] = []
+        self.globals['step_ui_states'] = []
+        
         if self.tools_instance and isinstance(self.tools_instance, AdbTools):
             self.tools_instance._set_context(ctx)
 
+        # Capture stdout and stderr
         stdout = io.StringIO()
         stderr = io.StringIO()
 
@@ -137,4 +140,9 @@ class SimpleCodeExecutor:
             output = f"Error: {type(e).__name__}: {str(e)}\n"
             output += traceback.format_exc()
 
-        return output
+        result = {
+            'output': output,
+            'screenshots': self.globals['step_screenshots'],
+            'ui_states': self.globals['step_ui_states']
+        }
+        return result

droidrun/cli/main.py

@@ -22,6 +22,8 @@ from droidrun.portal import (
     enable_portal_accessibility,
     PORTAL_PACKAGE_NAME,
     ping_portal,
+    ping_portal_tcp,
+    ping_portal_content,
 )
 from droidrun.macro.cli import macro_cli
 
@@ -78,9 +80,9 @@ async def run_command(
     reasoning: bool,
     reflection: bool,
     tracing: bool,
-    debug: bool,    
+    debug: bool,
     use_tcp: bool,
-    save_trajectory: bool = False,
+    save_trajectory: str = "none",
     ios: bool = False,
     allow_drag: bool = False,
     **kwargs,
@@ -104,7 +106,7 @@ async def run_command(
             # Device setup
             if device is None and not ios:
                 logger.info("🔍 Finding connected device...")
-                
+
                 devices = adb.list()
                 if not devices:
                     raise ValueError("No connected devices found.")
@@ -117,10 +119,14 @@ async def run_command(
             else:
                 logger.info(f"📱 Using device: {device}")
 
-            tools = AdbTools(serial=device, use_tcp=use_tcp) if not ios else IOSTools(url=device)
+            tools = (
+                AdbTools(serial=device, use_tcp=use_tcp)
+                if not ios
+                else IOSTools(url=device)
+            )
             # Set excluded tools based on CLI flags
             excluded_tools = [] if allow_drag else ["drag"]
-            
+
             # Select personas based on --drag flag
             personas = [BIG_AGENT] if allow_drag else [DEFAULT]
 
@@ -144,7 +150,6 @@ async def run_command(
             if tracing:
                 logger.info("🔍 Tracing enabled")
 
-
             droid_agent = DroidAgent(
                 goal=command,
                 llm=llm,
@@ -254,13 +259,16 @@ class DroidRunCLI(click.Group):
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
 @click.option(
-    "--use-tcp", is_flag=True, help="Use TCP communication for device control", default=False
+    "--use-tcp",
+    is_flag=True,
+    help="Use TCP communication for device control",
+    default=False,
 )
 @click.option(
     "--save-trajectory",
-    is_flag=True,
-    help="Save agent trajectory to file",
-    default=False,
+    type=click.Choice(["none", "step", "action"]),
+    help="Trajectory saving level: none (no saving), step (save per step), action (save per action)",
+    default="none",
 )
 @click.group(cls=DroidRunCLI)
 def cli(
@@ -277,7 +285,7 @@ def cli(
     tracing: bool,
     debug: bool,
     use_tcp: bool,
-    save_trajectory: bool,
+    save_trajectory: str,
 ):
     """DroidRun - Control your Android device through LLM agents."""
     pass
@@ -333,13 +341,16 @@ def cli(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
 @click.option(
-    "--use-tcp", is_flag=True, help="Use TCP communication for device control", default=False
+    "--use-tcp",
+    is_flag=True,
+    help="Use TCP communication for device control",
+    default=False,
 )
 @click.option(
     "--save-trajectory",
-    is_flag=True,
-    help="Save agent trajectory to file",
-    default=False,
+    type=click.Choice(["none", "step", "action"]),
+    help="Trajectory saving level: none (no saving), step (save per step), action (save per action)",
+    default="none",
 )
 @click.option(
     "--drag",
@@ -364,7 +375,7 @@ def run(
     tracing: bool,
     debug: bool,
     use_tcp: bool,
-    save_trajectory: bool,
+    save_trajectory: str,
     allow_drag: bool,
     ios: bool,
 ):
@@ -437,7 +448,11 @@ def disconnect(serial: str):
 
 @cli.command()
 @click.option("--device", "-d", help="Device serial number or IP address", default=None)
-@click.option("--path", help="Path to the Droidrun Portal APK to install on the device. If not provided, the latest portal apk version will be downloaded and installed.", default=None)
+@click.option(
+    "--path",
+    help="Path to the Droidrun Portal APK to install on the device. If not provided, the latest portal apk version will be downloaded and installed.",
+    default=None,
+)
 @click.option(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
@@ -474,7 +489,9 @@ def setup(path: str | None, device: str | None, debug: bool):
 
             console.print(f"[bold blue]Step 1/2: Installing APK:[/] {apk_path}")
             try:
-                device_obj.install(apk_path, uninstall=True, flags=["-g"], silent=not debug)
+                device_obj.install(
+                    apk_path, uninstall=True, flags=["-g"], silent=not debug
+                )
             except Exception as e:
                 console.print(f"[bold red]Installation failed:[/] {e}")
                 return
@@ -499,9 +516,7 @@ def setup(path: str | None, device: str | None, debug: bool):
                     "[yellow]Opening accessibility settings for manual configuration...[/]"
                 )
 
-                device_obj.shell(
-                    "am start -a android.settings.ACCESSIBILITY_SETTINGS"
-                )
+                device_obj.shell("am start -a android.settings.ACCESSIBILITY_SETTINGS")
 
                 console.print(
                     "\n[yellow]Please complete the following steps on your device:[/]"
@@ -530,10 +545,16 @@ def setup(path: str | None, device: str | None, debug: bool):
 
 @cli.command()
 @click.option("--device", "-d", help="Device serial number or IP address", default=None)
+@click.option(
+    "--use-tcp",
+    is_flag=True,
+    help="Use TCP communication for device control",
+    default=False,
+)
 @click.option(
     "--debug", is_flag=True, help="Enable verbose debug logging", default=False
 )
-def ping(device: str | None, debug: bool):
+def ping(device: str | None, use_tcp: bool, debug: bool):
     """Ping a device to check if it is ready and accessible."""
     try:
         device_obj = adb.device(device)
@@ -542,6 +563,12 @@ def ping(device: str | None, debug: bool):
             return
 
         ping_portal(device_obj, debug)
+
+        if use_tcp:
+            ping_portal_tcp(device_obj, debug)
+        else:
+            ping_portal_content(device_obj, debug)
+
         console.print(
             "[bold green]Portal is installed and accessible. You're good to go![/]"
         )

droidrun/macro/cli.py

@@ -215,7 +215,7 @@ def list(directory: str, debug: bool):
         
         # Still use console for table display as it's structured data
         console.print(table)
-        logger.info(f"💡 Use 'droidrun-macro replay {directory}/<folder>' to replay a trajectory")
+        logger.info(f"💡 Use 'droidrun macro replay {directory}/<folder>' to replay a trajectory")
     
     except Exception as e:
         logger.error(f"💥 Error: {e}")

droidrun/portal.py

@@ -3,6 +3,8 @@ import tempfile
 import os
 import contextlib
 from adbutils import adb, AdbDevice
+from droidrun.tools import AdbTools
+from rich.console import Console
 
 REPO = "droidrun/droidrun-portal"
 ASSET_NAME = "droidrun-portal"
@@ -36,8 +38,10 @@ def get_latest_release_assets(debug: bool = False):
 
 @contextlib.contextmanager
 def download_portal_apk(debug: bool = False):
+    console = Console()
     assets = get_latest_release_assets(debug)
 
+    asset_version = None
     asset_url = None
     for asset in assets:
         if (
@@ -46,11 +50,15 @@ def download_portal_apk(debug: bool = False):
             and asset["name"].startswith(ASSET_NAME)
         ):
             asset_url = asset["browser_download_url"]
+            asset_version = asset["name"].split("-")[-1]
+            asset_version = asset_version.removesuffix(".apk")
             break
         elif "downloadUrl" in asset and os.path.basename(
             asset["downloadUrl"]
         ).startswith(ASSET_NAME):
             asset_url = asset["downloadUrl"]
+            asset_version: str = asset["name"].split("-")[-1]
+            asset_version = asset_version.removesuffix(".apk")
             break
         else:
             if debug:
@@ -59,6 +67,10 @@ def download_portal_apk(debug: bool = False):
     if not asset_url:
         raise Exception(f"Asset named '{ASSET_NAME}' not found in the latest release.")
 
+    console.print(f"Found Portal APK [bold]{asset_version}[/bold]")
+    if debug:
+        console.print(f"Asset URL: {asset_url}")
+
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".apk")
     try:
         r = requests.get(asset_url, stream=True)
@@ -81,9 +93,7 @@ def enable_portal_accessibility(device: AdbDevice):
 
 
 def check_portal_accessibility(device: AdbDevice, debug: bool = False) -> bool:
-    a11y_services = device.shell(
-        "settings get secure enabled_accessibility_services"
-    )
+    a11y_services = device.shell("settings get secure enabled_accessibility_services")
     if not A11Y_SERVICE_NAME in a11y_services:
         if debug:
             print(a11y_services)
@@ -118,17 +128,23 @@ def ping_portal(device: AdbDevice, debug: bool = False):
             "Droidrun Portal is not enabled as an accessibility service on the device"
         )
 
+
+def ping_portal_content(device: AdbDevice, debug: bool = False):
     try:
-        state = device.shell(
-            "content query --uri content://com.droidrun.portal/state"
-        )
+        state = device.shell("content query --uri content://com.droidrun.portal/state")
         if not "Row: 0 result=" in state:
             raise Exception("Failed to get state from Droidrun Portal")
-
     except Exception as e:
         raise Exception(f"Droidrun Portal is not reachable: {e}")
 
 
+def ping_portal_tcp(device: AdbDevice, debug: bool = False):
+    try:
+        tools = AdbTools(serial=device.serial, use_tcp=True)
+    except Exception as e:
+        raise Exception(f"Failed to setup TCP forwarding: {e}")
+
+
 def test():
     device = adb.device()
     ping_portal(device, debug=False)

droidrun/telemetry/events.py

@@ -17,7 +17,7 @@ class DroidAgentInitEvent(TelemetryEvent):
     reflection: bool
     enable_tracing: bool
     debug: bool
-    save_trajectories: bool
+    save_trajectories: str
 
 
 class DroidAgentFinalizeEvent(TelemetryEvent):

droidrun/tools/adb.py

@@ -8,7 +8,7 @@ import json
 import time
 import logging
 from llama_index.core.workflow import Context
-from typing_extensions import Optional, Dict, Tuple, List, Any, Type, Self
+from typing import Optional, Dict, Tuple, List, Any
 from droidrun.agent.common.events import (
     InputTextActionEvent,
     KeyPressActionEvent,
@@ -23,13 +23,17 @@ import requests
 import base64
 
 logger = logging.getLogger("droidrun-tools")
+PORTAL_DEFAULT_TCP_PORT = 8080
 
 
 class AdbTools(Tools):
     """Core UI interaction tools for Android device control."""
 
     def __init__(
-        self, serial: str | None = None, use_tcp: bool = False, tcp_port: int = 8080
+        self,
+        serial: str | None = None,
+        use_tcp: bool = False,
+        remote_tcp_port: int = PORTAL_DEFAULT_TCP_PORT,
     ) -> None:
         """Initialize the AdbTools instance.
 
@@ -40,8 +44,7 @@ class AdbTools(Tools):
         """
         self.device = adb.device(serial=serial)
         self.use_tcp = use_tcp
-        self.tcp_port = tcp_port
-        self.tcp_base_url = f"http://localhost:{tcp_port}"
+        self.remote_tcp_port = remote_tcp_port
         self.tcp_forwarded = False
 
         self._ctx = None
@@ -55,6 +58,8 @@ class AdbTools(Tools):
         self.memory: List[str] = []
         # Store all screenshots with timestamps
         self.screenshots: List[Dict[str, Any]] = []
+        # Trajectory saving level
+        self.save_trajectories = "none"
 
         # Set up TCP forwarding if requested
         if self.use_tcp:
@@ -69,18 +74,21 @@ class AdbTools(Tools):
         """
         try:
             logger.debug(
-                f"Setting up TCP port forwarding: tcp:{self.tcp_port} tcp:{self.tcp_port}"
+                f"Setting up TCP port forwarding for port tcp:{self.remote_tcp_port} on device {self.device.serial}"
             )
             # Use adb forward command to set up port forwarding
-            result = self.device.forward(f"tcp:{self.tcp_port}", f"tcp:{self.tcp_port}")
-            self.tcp_forwarded = True
-            logger.debug(f"TCP port forwarding set up successfully: {result}")
+            self.local_tcp_port = self.device.forward_port(self.remote_tcp_port)
+            self.tcp_base_url = f"http://localhost:{self.local_tcp_port}"
+            logger.debug(
+                f"TCP port forwarding set up successfully to {self.tcp_base_url}"
+            )
 
             # Test the connection with a ping
             try:
                 response = requests.get(f"{self.tcp_base_url}/ping", timeout=5)
                 if response.status_code == 200:
                     logger.debug("TCP connection test successful")
+                    self.tcp_forwarded = True
                     return True
                 else:
                     logger.warning(
@@ -105,10 +113,17 @@ class AdbTools(Tools):
         """
         try:
             if self.tcp_forwarded:
-                logger.debug(f"Removing TCP port forwarding for port {self.tcp_port}")
-                result = self.device.forward_remove(f"tcp:{self.tcp_port}")
+                logger.debug(
+                    f"Removing TCP port forwarding for port {self.local_tcp_port}"
+                )
+                # remove forwarding
+                cmd = f"killforward:tcp:{self.local_tcp_port}"
+                logger.debug(f"Removing TCP port forwarding: {cmd}")
+                c = self.device.open_transport(cmd)
+                c.close()
+
                 self.tcp_forwarded = False
-                logger.debug(f"TCP port forwarding removed: {result}")
+                logger.debug(f"TCP port forwarding removed")
                 return True
             return True
         except Exception as e:
@@ -170,6 +185,7 @@ class AdbTools(Tools):
         except json.JSONDecodeError:
             return None
 
+    @Tools.ui_action
     def tap_by_index(self, index: int) -> str:
         """
         Tap on a UI element by its index.
@@ -328,6 +344,7 @@ class AdbTools(Tools):
         """
         return self.tap_by_index(index)
 
+    @Tools.ui_action
     def swipe(
         self,
         start_x: int,
@@ -372,6 +389,7 @@ class AdbTools(Tools):
             print(f"Error: {str(e)}")
             return False
 
+    @Tools.ui_action
     def drag(
         self, start_x: int, start_y: int, end_x: int, end_y: int, duration: float = 3
     ) -> bool:
@@ -413,6 +431,7 @@ class AdbTools(Tools):
             print(f"Error: {str(e)}")
             return False
 
+    @Tools.ui_action
     def input_text(self, text: str) -> str:
         """
         Input text on the device.
@@ -427,59 +446,59 @@ class AdbTools(Tools):
         try:
             logger.debug(f"Inputting text: {text}")
 
-            if self.use_tcp and self.tcp_forwarded:
-                # Use TCP communication
-                encoded_text = base64.b64encode(text.encode()).decode()
-
-                payload = {"base64_text": encoded_text}
-                response = requests.post(
-                    f"{self.tcp_base_url}/keyboard/input",
-                    json=payload,
-                    headers={"Content-Type": "application/json"},
-                    timeout=10,
-                )
-
-                logger.debug(
-                    f"Keyboard input TCP response: {response.status_code}, {response.text}"
-                )
-
-                if response.status_code != 200:
-                    return f"Error: HTTP request failed with status {response.status_code}: {response.text}"
-
-            else:
-                # Fallback to content provider method
-                # Save the current keyboard
-                original_ime = self.device.shell(
-                    "settings get secure default_input_method"
-                )
-                original_ime = original_ime.strip()
+            # if self.use_tcp and self.tcp_forwarded:
+            #     # Use TCP communication
+            #     encoded_text = base64.b64encode(text.encode()).decode()
+
+            #     payload = {"base64_text": encoded_text}
+            #     response = requests.post(
+            #         f"{self.tcp_base_url}/keyboard/input",
+            #         json=payload,
+            #         headers={"Content-Type": "application/json"},
+            #         timeout=10,
+            #     )
+
+            #     logger.debug(
+            #         f"Keyboard input TCP response: {response.status_code}, {response.text}"
+            #     )
+
+            #     if response.status_code != 200:
+            #         return f"Error: HTTP request failed with status {response.status_code}: {response.text}"
+
+            # else:
+            # Fallback to content provider method
+            # Save the current keyboard
+            original_ime = self.device.shell(
+                "settings get secure default_input_method"
+            )
+            original_ime = original_ime.strip()
 
-                # Enable the Droidrun keyboard
-                self.device.shell("ime enable com.droidrun.portal/.DroidrunKeyboardIME")
+            # Enable the Droidrun keyboard
+            self.device.shell("ime enable com.droidrun.portal/.DroidrunKeyboardIME")
 
-                # Set the Droidrun keyboard as the default
-                self.device.shell("ime set com.droidrun.portal/.DroidrunKeyboardIME")
+            # Set the Droidrun keyboard as the default
+            self.device.shell("ime set com.droidrun.portal/.DroidrunKeyboardIME")
 
-                # Wait for keyboard to change
-                time.sleep(1)
+            # Wait for keyboard to change
+            time.sleep(1)
 
-                # Encode the text to Base64
-                encoded_text = base64.b64encode(text.encode()).decode()
+            # Encode the text to Base64
+            encoded_text = base64.b64encode(text.encode()).decode()
 
-                cmd = f'content insert --uri "content://com.droidrun.portal/keyboard/input" --bind base64_text:s:"{encoded_text}"'
-                self.device.shell(cmd)
+            cmd = f'content insert --uri "content://com.droidrun.portal/keyboard/input" --bind base64_text:s:"{encoded_text}"'
+            self.device.shell(cmd)
 
-                # Wait for text input to complete
-                time.sleep(0.5)
+            # Wait for text input to complete
+            time.sleep(0.5)
 
-                # Restore the original keyboard
-                if original_ime and "com.droidrun.portal" not in original_ime:
-                    self.device.shell(f"ime set {original_ime}")
+            # Restore the original keyboard
+            if original_ime and "com.droidrun.portal" not in original_ime:
+                self.device.shell(f"ime set {original_ime}")
 
-                logger.debug(
-                    f"Text input completed: {text[:50]}{'...' if len(text) > 50 else ''}"
-                )
-                return f"Text input completed: {text[:50]}{'...' if len(text) > 50 else ''}"
+            logger.debug(
+                f"Text input completed: {text[:50]}{'...' if len(text) > 50 else ''}"
+            )
+            return f"Text input completed: {text[:50]}{'...' if len(text) > 50 else ''}"
 
             if self._ctx:
                 input_event = InputTextActionEvent(
@@ -501,6 +520,7 @@ class AdbTools(Tools):
         except Exception as e:
             return f"Error sending text input: {str(e)}"
 
+    @Tools.ui_action
     def back(self) -> str:
         """
         Go back on the current view.
@@ -523,6 +543,7 @@ class AdbTools(Tools):
         except ValueError as e:
             return f"Error: {str(e)}"
 
+    @Tools.ui_action
     def press_key(self, keycode: int) -> str:
         """
         Press a key on the Android device.
@@ -561,6 +582,7 @@ class AdbTools(Tools):
         except ValueError as e:
             return f"Error: {str(e)}"
 
+    @Tools.ui_action
     def start_app(self, package: str, activity: str | None = None) -> str:
         """
         Start an app on the device.
@@ -634,48 +656,13 @@ class AdbTools(Tools):
         try:
             logger.debug("Taking screenshot")
 
-            if self.use_tcp and self.tcp_forwarded:
-                # Use TCP communication
-                response = requests.get(f"{self.tcp_base_url}/screenshot", timeout=15)
-
-                if response.status_code == 200:
-                    tcp_response = response.json()
-
-                    # Check if response has the expected format with data field
-                    if isinstance(tcp_response, dict) and "data" in tcp_response:
-                        base64_data = tcp_response["data"]
-                        try:
-                            # Decode base64 to get image bytes
-                            image_bytes = base64.b64decode(base64_data)
-                            img_format = "PNG"  # Assuming PNG format from TCP endpoint
-                            logger.debug("Screenshot taken via TCP")
-                        except Exception as e:
-                            raise ValueError(
-                                f"Failed to decode base64 screenshot data: {str(e)}"
-                            )
-                    else:
-                        # Fallback: assume direct base64 format
-                        try:
-                            image_bytes = base64.b64decode(tcp_response)
-                            img_format = "PNG"
-                            logger.debug("Screenshot taken via TCP (direct base64)")
-                        except Exception as e:
-                            raise ValueError(
-                                f"Failed to decode screenshot response: {str(e)}"
-                            )
-                else:
-                    raise ValueError(
-                        f"HTTP request failed with status {response.status_code}: {response.text}"
-                    )
-
-            else:
-                # Fallback to ADB screenshot method
-                img = self.device.screenshot()
-                img_buf = io.BytesIO()
-                img_format = "PNG"
-                img.save(img_buf, format=img_format)
-                image_bytes = img_buf.getvalue()
-                logger.debug("Screenshot taken via ADB")
+            # Fallback to ADB screenshot method
+            img = self.device.screenshot()
+            img_buf = io.BytesIO()
+            img_format = "PNG"
+            img.save(img_buf, format=img_format)
+            image_bytes = img_buf.getvalue()
+            logger.debug("Screenshot taken via ADB")
 
             # Store screenshot with timestamp
             self.screenshots.append(
@@ -710,6 +697,7 @@ class AdbTools(Tools):
         except ValueError as e:
             raise ValueError(f"Error listing packages: {str(e)}")
 
+    @Tools.ui_action
     def complete(self, success: bool, reason: str = ""):
         """
         Mark the task as finished.

droidrun/tools/tools.py

@@ -2,6 +2,8 @@ from abc import ABC, abstractmethod
 from typing import List, Optional, Dict, Any
 import logging
 from typing import Tuple, Dict, Callable, Any, Optional
+from functools import wraps
+import sys
 
 # Get a logger for this module
 logger = logging.getLogger(__name__)
@@ -13,6 +15,32 @@ class Tools(ABC):
     This class provides a common interface for all tools to implement.
     """
 
+    @staticmethod
+    def ui_action(func):
+        """"
+        Decorator to capture screenshots and UI states for actions that modify the UI.
+        """
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            self = args[0]
+            result = func(*args, **kwargs)
+            
+            # Check if save_trajectories attribute exists and is set to "action"
+            if hasattr(self, 'save_trajectories') and self.save_trajectories == "action":
+                frame = sys._getframe(1)
+                caller_globals = frame.f_globals
+                
+                step_screenshots = caller_globals.get('step_screenshots')
+                step_ui_states = caller_globals.get('step_ui_states')
+                
+                if step_screenshots is not None:
+                    step_screenshots.append(self.take_screenshot()[1])
+                if step_ui_states is not None:
+                    step_ui_states.append(self.get_state())
+            
+            return result
+        return wrapper
+
     @abstractmethod
     def get_state(self) -> Dict[str, Any]:
         """

{droidrun-0.3.3.dist-info → droidrun-0.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: droidrun
-Version: 0.3.3
+Version: 0.3.4
 Summary: A framework for controlling Android devices through LLM agents
 Project-URL: Homepage, https://github.com/droidrun/droidrun
 Project-URL: Bug Tracker, https://github.com/droidrun/droidrun/issues
@@ -46,7 +46,6 @@ Requires-Dist: posthog==6.0.2
 Requires-Dist: pydantic>=2.0.0
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: rich>=13.0.0
-Requires-Dist: typing-extensions
 Provides-Extra: dev
 Requires-Dist: bandit>=1.7.0; extra == 'dev'
 Requires-Dist: black>=23.0.0; extra == 'dev'

{droidrun-0.3.3.dist-info → droidrun-0.3.4.dist-info}/RECORD

@@ -1,9 +1,9 @@
 droidrun/__init__.py,sha256=Cqt4NXZ-753220dlRZXglMkFT2ygcXSErMmqupGsi9A,622
 droidrun/__main__.py,sha256=78o1Wr_Z-NrZy9yLWmEfNfRRhAiJGBr4Xi3lmbgkx3w,105
-droidrun/portal.py,sha256=C0gZK9GHigsM6Vq2fQP_DKaDXqtnQSA2NUrxClVe6hw,3931
+droidrun/portal.py,sha256=Bg3kVoNLJ2bTdk-4VuUr_i77yWJT2ZNLfT5WMumnZV8,4679
 droidrun/agent/__init__.py,sha256=4SqTJeGDvr_wT8rtN9J8hnN6P-pae663mkYr-JmzH4w,208
 droidrun/agent/codeact/__init__.py,sha256=ZLDGT_lTTzyNm7pcBzdyRIGHJ2ZgbInJdhXZRbJLhSQ,278
-droidrun/agent/codeact/codeact_agent.py,sha256=XYyRrDAj2MoNb4WTtCzY8bqlI1WDy9TUIRku4NMUezk,16087
+droidrun/agent/codeact/codeact_agent.py,sha256=6Hltg6J4mdb5pOr5sFQVY-6SMWIcxWp6RZjct5nkiCc,16506
 droidrun/agent/codeact/events.py,sha256=skCfZ-5SR0YhhzZVxx8_VkUjfILk8rCv47k9pHNYhdc,634
 droidrun/agent/codeact/prompts.py,sha256=28HflWMNkC1ky0hGCzAxhJftjU2IIU1ZRUfya3S7M6I,1006
 droidrun/agent/common/default.py,sha256=P07el-PrHsoqQMYsYxSSln6mFl-QY75vzwp1Dll_XmY,259
@@ -13,42 +13,42 @@ droidrun/agent/context/agent_persona.py,sha256=Mxd4HTyirWD-aqNlka1hQBdS-0I-lXJr2
 droidrun/agent/context/context_injection_manager.py,sha256=sA33q2KPtX_4Yap8wM11T6ewlZC_0FIbKPEc400SHrE,2188
 droidrun/agent/context/episodic_memory.py,sha256=1ImeR3jAWOpKwkQt3bMlXVOBiQbIli5fBIlBq2waREQ,394
 droidrun/agent/context/reflection.py,sha256=0hJluOz0hTlHHhReKpIJ9HU5aJbaJsvrjMfraQ84D-M,652
-droidrun/agent/context/task_manager.py,sha256=ESLs4kR6VNYiYQsc4V7WAeoSLwbaZPSWBXpveOfOv8c,4343
+droidrun/agent/context/task_manager.py,sha256=RqIBBLcw_hLotaO0JFsm8HOdJ_UdNttSGRGl_3GqRY4,4648
 droidrun/agent/context/personas/__init__.py,sha256=oSRa8g_xngX7JPIRPu7fLO33m3r7fdEQzIuORuqcw5M,232
 droidrun/agent/context/personas/app_starter.py,sha256=dHeknznxGEPJ7S6VPyEG_MB-HvAvQwUOnRWaShaV8Xo,1585
 droidrun/agent/context/personas/big_agent.py,sha256=Gl_y4ykz3apGc203-KG2UbSOwf7gDUiWh7GOVyiLn-Y,5091
 droidrun/agent/context/personas/default.py,sha256=Xm07YCWoKjvlHAbQRtzE3vn7BVcz6wYcSVeg4FiojJQ,5060
 droidrun/agent/context/personas/ui_expert.py,sha256=j0OKfN1jQSrREHcVeomMTDPCWLsZVX4aeuWN4Y-x3z0,4739
 droidrun/agent/droid/__init__.py,sha256=3BfUVZiUQ8ATAJ_JmqQZQx53WoERRpQ4AyHW5WOgbRI,297
-droidrun/agent/droid/droid_agent.py,sha256=_XtGkkQM27qfxDO8V2MwehWz931Z10CBhtOIoGNsWBk,15252
-droidrun/agent/droid/events.py,sha256=Ks2D6lX5P1rpZ4nIAPXSC83z4AT5OzKt3isP5yk25F4,689
+droidrun/agent/droid/droid_agent.py,sha256=nKr8gV6aV_S5T4llhKokSBqL2uKF9ldxcnE29FOGTL8,16737
+droidrun/agent/droid/events.py,sha256=hjYpWcSffqP83rNv_GyOEc3CNSrdvlVPdUkaRU6QDJc,722
 droidrun/agent/oneflows/reflector.py,sha256=I_tE0PBjvwWbS6SA8Qd41etxJglFgn8oScuKUxc9LEE,11621
 droidrun/agent/planner/__init__.py,sha256=Fu0Ewtd-dIRLgHIL1DB_9EEKvQS_f1vjB8jgO5TbJXg,364
 droidrun/agent/planner/events.py,sha256=oyt2FNrA2uVyUeVT65-N0AC6sWBFxSnwNEqWtnRYoFM,390
-droidrun/agent/planner/planner_agent.py,sha256=fxLeLeN_2kmaWHXLaDLvXOR_uvP_ult_I8SPZbgL_AU,10482
+droidrun/agent/planner/planner_agent.py,sha256=gqXMj1PcTdmKpSGM1b6T_q7ojFptPgx39wAXuW3ZQDQ,10854
 droidrun/agent/planner/prompts.py,sha256=Ci7Oeu3J4TAhx-tKGPZ9l6Wb3a81FSqC8cWW4jW73HI,6046
 droidrun/agent/utils/__init__.py,sha256=JK6ygRjw7gzcQSG0HBEYLoVGH54QQAxJJ7HpIS5mgyc,44
 droidrun/agent/utils/async_utils.py,sha256=IQBcWPwevm89B7R_UdMXk0unWeNCBA232b5kQGqoxNI,336
-droidrun/agent/utils/chat_utils.py,sha256=5oqP2nmKs8sHWP1H_TK82yaxrxWf7FdEbFKASKpR60g,13000
-droidrun/agent/utils/executer.py,sha256=UKVd7asr3RvO-YgRflGdG12FFZxDvU6hFxrdV1sWLcU,4892
+droidrun/agent/utils/chat_utils.py,sha256=SNn-wJSyxiirDWDqAeLehHw4ByNX_6dD5Bx3Inc6Sj0,13008
+droidrun/agent/utils/executer.py,sha256=-iwdZpmpF0w116D7A_eDgeV0ZNXSuotVgBkM5lc0BNI,5202
 droidrun/agent/utils/llm_picker.py,sha256=16tNkNhEM9gD_uivzxLvuaa6s0Tz7Igu-3fxMP2lAtY,5968
 droidrun/agent/utils/trajectory.py,sha256=PU3nI3Zru580_bK0TvqUaf-5kiWvj6hFoedXkDgTqdc,17047
 droidrun/cli/__init__.py,sha256=DuwSRtZ8WILPd-nf-fZ7BaBsRgtofoInOF3JtJ9wag0,167
 droidrun/cli/logs.py,sha256=PsT_VbnOa_sOLXK4KkEJk4AsYCpscqrVoryMmLVwPG0,9714
-droidrun/cli/main.py,sha256=UnTocHNd8za9Fj-ppnvJSD8DT3RQ1aqVPPTvAsWRX5g,17734
+droidrun/cli/main.py,sha256=3RWvY19zY9ewzcasndxpul3BUPpMq4yhrtnuAIBkufM,18311
 droidrun/macro/__init__.py,sha256=333sMt19mA8sfQa4qPKbbCmr8Ej3nvpdxGXUZtVTEqM,344
 droidrun/macro/__main__.py,sha256=-zj42Bj7309oLPZbNsxZeNwIDaEe7Th1I4zF8yAHasw,193
-droidrun/macro/cli.py,sha256=fecato896z9OxZsAzCX7FWgF2TXwM5EnR85EcSGzc4U,9057
+droidrun/macro/cli.py,sha256=GaL1wVWVE_AT0OxHgOJyW-q_kVu4PQ76KZKeYMQdoEk,9057
 droidrun/macro/replay.py,sha256=q_3ZcHVjvsdDfS2xyt_vuuwXGt9_1t38JD1cPsjzIfU,10764
 droidrun/telemetry/__init__.py,sha256=D4Mp02iGJH2Tjpv42Bzyo6_WC3NWj9Qy9hQPWFaCkhA,234
-droidrun/telemetry/events.py,sha256=S6r6_c2bGTZt6F88m_vREDD_MDhw3Pz4I53lLmsI764,518
+droidrun/telemetry/events.py,sha256=FtPMMDAhxvOiCssJLCCr5FX93ulnPCbiFMhkPh8NLw4,517
 droidrun/telemetry/tracker.py,sha256=Ljue6zivX8KnadXI9DivrayuWxAnUwbJKvCvNtY1Y4Y,2717
 droidrun/tools/__init__.py,sha256=9ReauavtSKDQG9ya9_Fr9O0TQnDFixgOPaP5n82_iEk,271
-droidrun/tools/adb.py,sha256=XRa9TnuS_0-vqMJvn6RDjCRzu-MvSbdtBVAuhVHdSzo,42083
+droidrun/tools/adb.py,sha256=zFnVydQi42GZ1vh-4HT79dCQVBeCS6VObywyaIUUCmw,40832
 droidrun/tools/ios.py,sha256=imzojiS6gqz4IKexUEz1ga7-flSOaC5QRpHIJTwcgSQ,21807
-droidrun/tools/tools.py,sha256=eXCaFjb_FNJFPrIfKi9W1WghxfrfnsbfutW3f0XlTzw,3720
-droidrun-0.3.3.dist-info/METADATA,sha256=9aqH7cPVJ24Snsta0Fkyv8tG9b7dMLZ1H3ZhjcvEScM,6718
-droidrun-0.3.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-droidrun-0.3.3.dist-info/entry_points.txt,sha256=o259U66js8TIybQ7zs814Oe_LQ_GpZsp6a9Cr-xm5zE,51
-droidrun-0.3.3.dist-info/licenses/LICENSE,sha256=s-uxn9qChu-kFdRXUp6v_0HhsaJ_5OANmfNOFVm2zdk,1069
-droidrun-0.3.3.dist-info/RECORD,,
+droidrun/tools/tools.py,sha256=rqDe2gRyR45HVM15SwsW1aCTVZo5eleTxlh2hGqCHys,4785
+droidrun-0.3.4.dist-info/METADATA,sha256=C7rh3674hF4xvKES0DfXTQaErn1RckdUBmrNcrito_o,6685
+droidrun-0.3.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+droidrun-0.3.4.dist-info/entry_points.txt,sha256=o259U66js8TIybQ7zs814Oe_LQ_GpZsp6a9Cr-xm5zE,51
+droidrun-0.3.4.dist-info/licenses/LICENSE,sha256=s-uxn9qChu-kFdRXUp6v_0HhsaJ_5OANmfNOFVm2zdk,1069
+droidrun-0.3.4.dist-info/RECORD,,