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

droidrun/tools/adb.py

@@ -3,29 +3,50 @@ UI Actions - Core UI interaction tools for Android device control.
 """
 
 import os
-import re
+import io
 import json
 import time
-import tempfile
-import asyncio
-import aiofiles
-import contextlib
 import logging
-from typing import Optional, Dict, Tuple, List, Any
-from droidrun.adb.device import Device
-from droidrun.adb.manager import DeviceManager
+from llama_index.core.workflow import Context
+from typing_extensions import Optional, Dict, Tuple, List, Any, Type, Self
+from droidrun.agent.common.events import (
+    InputTextActionEvent,
+    KeyPressActionEvent,
+    StartAppEvent,
+    SwipeActionEvent,
+    TapActionEvent,
+    DragActionEvent,
+)
 from droidrun.tools.tools import Tools
+from adbutils import adb
+import requests
+import base64
+
+logger = logging.getLogger("droidrun-tools")
 
-logger = logging.getLogger("droidrun-adb-tools")
 
 class AdbTools(Tools):
     """Core UI interaction tools for Android device control."""
 
-    def __init__(self, serial: str = "emulator-5554") -> None:
+    def __init__(
+        self, serial: str | None = None, use_tcp: bool = False, tcp_port: int = 8080
+    ) -> None:
+        """Initialize the AdbTools instance.
+
+        Args:
+            serial: Device serial number
+            use_tcp: Whether to use TCP communication (default: False)
+            tcp_port: TCP port for communication (default: 8080)
+        """
+        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.tcp_forwarded = False
+
+        self._ctx = None
         # Instance‐level cache for clickable elements (index-based tapping)
         self.clickable_elements_cache: List[Dict[str, Any]] = []
-        self.serial = serial
-        self.device_manager = DeviceManager()
         self.last_screenshot = None
         self.reason = None
         self.success = None
@@ -35,85 +56,113 @@ class AdbTools(Tools):
         # Store all screenshots with timestamps
         self.screenshots: List[Dict[str, Any]] = []
 
-    def get_device_serial(self) -> str:
-        """Get the device serial from the instance or environment variable."""
-        # First try using the instance's serial
-        if self.serial:
-            return self.serial
+        # Set up TCP forwarding if requested
+        if self.use_tcp:
+            self.setup_tcp_forward()
 
-    async def get_device(self) -> Optional[Device]:
-        """Get the device instance using the instance's serial or from environment variable.
+    def setup_tcp_forward(self) -> bool:
+        """
+        Set up ADB TCP port forwarding for communication with the portal app.
 
         Returns:
-            Device instance or None if not found
+            bool: True if forwarding was set up successfully, False otherwise
         """
-        serial = self.get_device_serial()
-        if not serial:
-            raise ValueError("No device serial specified - set device_serial parameter")
-
-        device = await self.device_manager.get_device(serial)
-        if not device:
-            raise ValueError(f"Device {serial} not found")
+        try:
+            logger.debug(
+                f"Setting up TCP port forwarding: tcp:{self.tcp_port} tcp:{self.tcp_port}"
+            )
+            # 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}")
 
-        return device
+            # 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")
+                    return True
+                else:
+                    logger.warning(
+                        f"TCP connection test failed with status: {response.status_code}"
+                    )
+                    return False
+            except requests.exceptions.RequestException as e:
+                logger.warning(f"TCP connection test failed: {e}")
+                return False
 
-    def parse_package_list(self, output: str) -> List[Dict[str, str]]:
-        """Parse the output of 'pm list packages -f' command.
+        except Exception as e:
+            logger.error(f"Failed to set up TCP port forwarding: {e}")
+            self.tcp_forwarded = False
+            return False
 
-        Args:
-            output: Raw command output from 'pm list packages -f'
+    def teardown_tcp_forward(self) -> bool:
+        """
+        Remove ADB TCP port forwarding.
 
         Returns:
-            List of dictionaries containing package info with 'package' and 'path' keys
+            bool: True if forwarding was removed successfully, False otherwise
         """
-        apps = []
-        for line in output.splitlines():
-            if line.startswith("package:"):
-                # Format is: "package:/path/to/base.apk=com.package.name"
-                path_and_pkg = line[8:]  # Strip "package:"
-                if "=" in path_and_pkg:
-                    path, package = path_and_pkg.rsplit("=", 1)
-                    apps.append({"package": package.strip(), "path": path.strip()})
-        return apps
+        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}")
+                self.tcp_forwarded = False
+                logger.debug(f"TCP port forwarding removed: {result}")
+                return True
+            return True
+        except Exception as e:
+            logger.error(f"Failed to remove TCP port forwarding: {e}")
+            return False
+
+    def __del__(self):
+        """Cleanup when the object is destroyed."""
+        if hasattr(self, "tcp_forwarded") and self.tcp_forwarded:
+            self.teardown_tcp_forward()
 
-    def _parse_content_provider_output(self, raw_output: str) -> Optional[Dict[str, Any]]:
+    def _set_context(self, ctx: Context):
+        self._ctx = ctx
+
+    def _parse_content_provider_output(
+        self, raw_output: str
+    ) -> Optional[Dict[str, Any]]:
         """
         Parse the raw ADB content provider output and extract JSON data.
-        
+
         Args:
             raw_output (str): Raw output from ADB content query command
-            
+
         Returns:
             dict: Parsed JSON data or None if parsing failed
         """
         # The ADB content query output format is: "Row: 0 result={json_data}"
         # We need to extract the JSON part after "result="
-        lines = raw_output.strip().split('\n')
-        
+        lines = raw_output.strip().split("\n")
+
         for line in lines:
             line = line.strip()
-            
+
             # Look for lines that contain "result=" pattern
             if "result=" in line:
                 # Extract everything after "result="
                 result_start = line.find("result=") + 7
                 json_str = line[result_start:]
-                
+
                 try:
                     # Parse the JSON string
                     json_data = json.loads(json_str)
                     return json_data
                 except json.JSONDecodeError:
                     continue
-            
+
             # Fallback: try to parse lines that start with { or [
-            elif line.startswith('{') or line.startswith('['):
+            elif line.startswith("{") or line.startswith("["):
                 try:
                     json_data = json.loads(line)
                     return json_data
                 except json.JSONDecodeError:
                     continue
-        
+
         # If no valid JSON found in individual lines, try the entire output
         try:
             json_data = json.loads(raw_output.strip())
@@ -121,7 +170,7 @@ class AdbTools(Tools):
         except json.JSONDecodeError:
             return None
 
-    async def tap_by_index(self, index: int, serial: Optional[str] = None) -> str:
+    def tap_by_index(self, index: int) -> str:
         """
         Tap on a UI element by its index.
 
@@ -193,18 +242,32 @@ class AdbTools(Tools):
             x = (left + right) // 2
             y = (top + bottom) // 2
 
+            logger.debug(
+                f"Tapping element with index {index} at coordinates ({x}, {y})"
+            )
             # Get the device and tap at the coordinates
-            if serial:
-                device = await self.device_manager.get_device(serial)
-                if not device:
-                    return f"Error: Device {serial} not found"
-            else:
-                device = await self.get_device()
+            self.device.click(x, y)
+            logger.debug(f"Tapped element with index {index} at coordinates ({x}, {y})")
 
-            await device.tap(x, y)
+            # Emit coordinate action event for trajectory recording
+
+            if self._ctx:
+                element_text = element.get("text", "No text")
+                element_class = element.get("className", "Unknown class")
+
+                tap_event = TapActionEvent(
+                    action_type="tap",
+                    description=f"Tap element at index {index}: '{element_text}' ({element_class}) at coordinates ({x}, {y})",
+                    x=x,
+                    y=y,
+                    element_index=index,
+                    element_text=element_text,
+                    element_bounds=bounds_str,
+                )
+                self._ctx.write_event_to_stream(tap_event)
 
             # Add a small delay to allow UI to update
-            await asyncio.sleep(0.5)
+            time.sleep(0.5)
 
             # Create a descriptive response
             response_parts = []
@@ -229,7 +292,7 @@ class AdbTools(Tools):
             return f"Error: {str(e)}"
 
     # Rename the old tap function to tap_by_coordinates for backward compatibility
-    async def tap_by_coordinates(self, x: int, y: int) -> bool:
+    def tap_by_coordinates(self, x: int, y: int) -> bool:
         """
         Tap on the device screen at specific coordinates.
 
@@ -241,22 +304,16 @@ class AdbTools(Tools):
             Bool indicating success or failure
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
-
-            await device.tap(x, y)
-            print(f"Tapped at coordinates ({x}, {y})")
+            logger.debug(f"Tapping at coordinates ({x}, {y})")
+            self.device.click(x, y)
+            logger.debug(f"Tapped at coordinates ({x}, {y})")
             return True
         except ValueError as e:
-            print(f"Error: {str(e)}")
+            logger.debug(f"Error: {str(e)}")
             return False
 
     # Replace the old tap function with the new one
-    async def tap(self, index: int) -> str:
+    def tap(self, index: int) -> str:
         """
         Tap on a UI element by its index.
 
@@ -269,10 +326,15 @@ class AdbTools(Tools):
         Returns:
             Result message
         """
-        return await self.tap_by_index(index)
-
-    async def swipe(
-        self, start_x: int, start_y: int, end_x: int, end_y: int, duration_ms: int = 300
+        return self.tap_by_index(index)
+
+    def swipe(
+        self,
+        start_x: int,
+        start_y: int,
+        end_x: int,
+        end_y: int,
+        duration_ms: float = 300,
     ) -> bool:
         """
         Performs a straight-line swipe gesture on the device screen.
@@ -282,27 +344,76 @@ class AdbTools(Tools):
             start_y: Starting Y coordinate
             end_x: Ending X coordinate
             end_y: Ending Y coordinate
-            duration_ms: Duration of swipe in milliseconds
+            duration: Duration of swipe in seconds
         Returns:
             Bool indicating success or failure
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
 
-            await device.swipe(start_x, start_y, end_x, end_y, duration_ms)
-            await asyncio.sleep(1)
-            print(f"Swiped from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration_ms}ms")
+            if self._ctx:
+                swipe_event = SwipeActionEvent(
+                    action_type="swipe",
+                    description=f"Swipe from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration_ms} milliseconds",
+                    start_x=start_x,
+                    start_y=start_y,
+                    end_x=end_x,
+                    end_y=end_y,
+                    duration_ms=duration_ms,
+                )
+                self._ctx.write_event_to_stream(swipe_event)
+
+            self.device.swipe(start_x, start_y, end_x, end_y, float(duration_ms / 1000))
+            time.sleep(duration_ms / 1000)
+            logger.debug(
+                f"Swiped from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration_ms} milliseconds"
+            )
+            return True
+        except ValueError as e:
+            print(f"Error: {str(e)}")
+            return False
+
+    def drag(
+        self, start_x: int, start_y: int, end_x: int, end_y: int, duration: float = 3
+    ) -> bool:
+        """
+        Performs a straight-line drag and drop gesture on the device screen.
+        Args:
+            start_x: Starting X coordinate
+            start_y: Starting Y coordinate
+            end_x: Ending X coordinate
+            end_y: Ending Y coordinate
+            duration: Duration of swipe in seconds
+        Returns:
+            Bool indicating success or failure
+        """
+        try:
+            logger.debug(
+                f"Dragging from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration} seconds"
+            )
+            self.device.drag(start_x, start_y, end_x, end_y, duration)
+
+            if self._ctx:
+                drag_event = DragActionEvent(
+                    action_type="drag",
+                    description=f"Drag from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration} seconds",
+                    start_x=start_x,
+                    start_y=start_y,
+                    end_x=end_x,
+                    end_y=end_y,
+                    duration=duration,
+                )
+                self._ctx.write_event_to_stream(drag_event)
+
+            time.sleep(duration)
+            logger.debug(
+                f"Dragged from ({start_x}, {start_y}) to ({end_x}, {end_y}) in {duration} seconds"
+            )
             return True
         except ValueError as e:
             print(f"Error: {str(e)}")
             return False
 
-    async def input_text(self, text: str, serial: Optional[str] = None) -> str:
+    def input_text(self, text: str) -> str:
         """
         Input text on the device.
         Always make sure that the Focused Element is not None before inputting text.
@@ -314,72 +425,105 @@ class AdbTools(Tools):
             Result message
         """
         try:
-            if serial:
-                device = await self.device_manager.get_device(serial)
-                if not device:
-                    return f"Error: Device {serial} not found"
-            else:
-                device = await self.get_device()
-
-            # Save the current keyboard
-            original_ime = await device._adb.shell(
-                device._serial, "settings get secure default_input_method"
-            )
-            original_ime = original_ime.strip()
-
-            # Enable the Droidrun keyboard
-            await device._adb.shell(
-                device._serial, "ime enable com.droidrun.portal/.DroidrunKeyboardIME"
-            )
-
-            # Set the Droidrun keyboard as the default
-            await device._adb.shell(
-                device._serial, "ime set com.droidrun.portal/.DroidrunKeyboardIME"
-            )
+            logger.debug(f"Inputting text: {text}")
 
-            # Wait for keyboard to change
-            await asyncio.sleep(1)
+            if self.use_tcp and self.tcp_forwarded:
+                # Use TCP communication
+                encoded_text = base64.b64encode(text.encode()).decode()
 
-            # Encode the text to Base64
-            import base64
+                payload = {"base64_text": encoded_text}
+                response = requests.post(
+                    f"{self.tcp_base_url}/keyboard/input",
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                    timeout=10,
+                )
 
-            encoded_text = base64.b64encode(text.encode()).decode()
+                logger.debug(
+                    f"Keyboard input TCP response: {response.status_code}, {response.text}"
+                )
 
-            cmd = f'content insert --uri "content://com.droidrun.portal/keyboard/input" --bind base64_text:s:"{encoded_text}"'
-            await device._adb.shell(device._serial, cmd)
-
-            # Wait for text input to complete
-            await asyncio.sleep(0.5)
-
-            # Restore the original keyboard
-            if original_ime and "com.droidrun.portal" not in original_ime:
-                await device._adb.shell(device._serial, f"ime set {original_ime}")
+                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")
+
+                # Set the Droidrun keyboard as the default
+                self.device.shell("ime set com.droidrun.portal/.DroidrunKeyboardIME")
+
+                # Wait for keyboard to change
+                time.sleep(1)
+
+                # 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)
+
+                # 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}")
+
+                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(
+                    action_type="input_text",
+                    description=f"Input text: '{text[:50]}{'...' if len(text) > 50 else ''}'",
+                    text=text,
+                )
+                self._ctx.write_event_to_stream(input_event)
+
+            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 ''}"
+
+        except requests.exceptions.RequestException as e:
+            return f"Error: TCP request failed: {str(e)}"
         except ValueError as e:
             return f"Error: {str(e)}"
         except Exception as e:
             return f"Error sending text input: {str(e)}"
 
-    async def back(self) -> str:
+    def back(self) -> str:
         """
         Go back on the current view.
         This presses the Android back button.
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
+            logger.debug("Pressing key BACK")
+            self.device.keyevent(3)
+
+            if self._ctx:
+                key_event = KeyPressActionEvent(
+                    action_type="key_press",
+                    description=f"Pressed key BACK",
+                    keycode=3,
+                    key_name="BACK",
+                )
+                self._ctx.write_event_to_stream(key_event)
 
-            await device.press_key(3)
             return f"Pressed key BACK"
         except ValueError as e:
             return f"Error: {str(e)}"
 
-    async def press_key(self, keycode: int) -> str:
+    def press_key(self, keycode: int) -> str:
         """
         Press a key on the Android device.
 
@@ -393,13 +537,6 @@ class AdbTools(Tools):
             keycode: Android keycode to press
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
-
             key_names = {
                 66: "ENTER",
                 4: "BACK",
@@ -408,12 +545,23 @@ class AdbTools(Tools):
             }
             key_name = key_names.get(keycode, str(keycode))
 
-            await device.press_key(keycode)
+            if self._ctx:
+                key_event = KeyPressActionEvent(
+                    action_type="key_press",
+                    description=f"Pressed key {key_name}",
+                    keycode=keycode,
+                    key_name=key_name,
+                )
+                self._ctx.write_event_to_stream(key_event)
+
+            logger.debug(f"Pressing key {key_name}")
+            self.device.keyevent(keycode)
+            logger.debug(f"Pressed key {key_name}")
             return f"Pressed key {key_name}"
         except ValueError as e:
             return f"Error: {str(e)}"
 
-    async def start_app(self, package: str, activity: str = "") -> str:
+    def start_app(self, package: str, activity: str | None = None) -> str:
         """
         Start an app on the device.
 
@@ -422,19 +570,32 @@ class AdbTools(Tools):
             activity: Optional activity name
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
 
-            result = await device.start_app(package, activity)
-            return result
-        except ValueError as e:
+            logger.debug(f"Starting app {package} with activity {activity}")
+            if not activity:
+                dumpsys_output = self.device.shell(
+                    f"cmd package resolve-activity --brief {package}"
+                )
+                activity = dumpsys_output.splitlines()[1].split("/")[1]
+
+            if self._ctx:
+                start_app_event = StartAppEvent(
+                    action_type="start_app",
+                    description=f"Start app {package}",
+                    package=package,
+                    activity=activity,
+                )
+                self._ctx.write_event_to_stream(start_app_event)
+
+            print(f"Activity: {activity}")
+
+            self.device.app_start(package, activity)
+            logger.debug(f"App started: {package} with activity {activity}")
+            return f"App started: {package} with activity {activity}"
+        except Exception as e:
             return f"Error: {str(e)}"
 
-    async def install_app(
+    def install_app(
         self, apk_path: str, reinstall: bool = False, grant_permissions: bool = True
     ) -> str:
         """
@@ -446,50 +607,94 @@ class AdbTools(Tools):
             grant_permissions: Whether to grant all permissions
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    return f"Error: Device {self.serial} not found"
-            else:
-                device = await self.get_device()
-
             if not os.path.exists(apk_path):
                 return f"Error: APK file not found at {apk_path}"
 
-            result = await device.install_app(apk_path, reinstall, grant_permissions)
+            logger.debug(
+                f"Installing app: {apk_path} with reinstall: {reinstall} and grant_permissions: {grant_permissions}"
+            )
+            result = self.device.install(
+                apk_path,
+                nolaunch=True,
+                uninstall=reinstall,
+                flags=["-g"] if grant_permissions else [],
+                silent=True,
+            )
+            logger.debug(f"Installed app: {apk_path} with result: {result}")
             return result
         except ValueError as e:
             return f"Error: {str(e)}"
 
-    async def take_screenshot(self) -> Tuple[str, bytes]:
+    def take_screenshot(self) -> Tuple[str, bytes]:
         """
         Take a screenshot of the device.
         This function captures the current screen and adds the screenshot to context in the next message.
         Also stores the screenshot in the screenshots list with timestamp for later GIF creation.
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    raise ValueError(f"Device {self.serial} not found")
+            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:
-                device = await self.get_device()
-            screen_tuple = await device.take_screenshot()
-            self.last_screenshot = screen_tuple[1]
+                # 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(
                 {
                     "timestamp": time.time(),
-                    "image_data": screen_tuple[1],
-                    "format": screen_tuple[0],  # Usually 'PNG'
+                    "image_data": image_bytes,
+                    "format": img_format,
                 }
             )
-            return screen_tuple
+            return img_format, image_bytes
+
+        except requests.exceptions.RequestException as e:
+            raise ValueError(f"Error taking screenshot via TCP: {str(e)}")
         except ValueError as e:
             raise ValueError(f"Error taking screenshot: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Unexpected error taking screenshot: {str(e)}")
 
-    async def list_packages(self, include_system_apps: bool = False) -> List[str]:
+    def list_packages(self, include_system_apps: bool = False) -> List[str]:
         """
         List installed packages on the device.
 
@@ -500,163 +705,11 @@ class AdbTools(Tools):
             List of package names
         """
         try:
-            if self.serial:
-                device = await self.device_manager.get_device(self.serial)
-                if not device:
-                    raise ValueError(f"Device {self.serial} not found")
-            else:
-                device = await self.get_device()
-
-            # Use the direct ADB command to get packages with paths
-            cmd = ["pm", "list", "packages", "-f"]
-            if not include_system_apps:
-                cmd.append("-3")
-
-            output = await device._adb.shell(device._serial, " ".join(cmd))
-
-            # Parse the package list using the function
-            packages = self.parse_package_list(output)
-            # Format package list for better readability
-            package_list = [pack["package"] for pack in packages]
-            for package in package_list:
-                print(package)
-            return package_list
+            logger.debug("Listing packages")
+            return self.device.list_packages(["-3"] if not include_system_apps else [])
         except ValueError as e:
             raise ValueError(f"Error listing packages: {str(e)}")
 
-    async def extract(self, filename: Optional[str] = None) -> str:
-        """Extract and save the current UI state to a JSON file.
-
-        This function captures the current UI state including all UI elements
-        and saves it to a JSON file for later analysis or reference.
-
-        Args:
-            filename: Optional filename to save the UI state (defaults to ui_state_TIMESTAMP.json)
-
-        Returns:
-            Path to the saved JSON file
-        """
-        try:
-            # Generate default filename if not provided
-            if not filename:
-                timestamp = int(time.time())
-                filename = f"ui_state_{timestamp}.json"
-
-            # Ensure the filename ends with .json
-            if not filename.endswith(".json"):
-                filename += ".json"
-
-            # Get the UI elements
-            ui_elements = await self.get_all_elements(self.serial)
-
-            # Save to file
-            save_path = os.path.abspath(filename)
-            async with aiofiles.open(save_path, "w", encoding="utf-8") as f:
-                await f.write(json.dumps(ui_elements, indent=2))
-
-            return f"UI state extracted and saved to {save_path}"
-
-        except Exception as e:
-            return f"Error extracting UI state: {e}"
-
-    async def get_all_elements(self) -> Dict[str, Any]:
-        """
-        Get all UI elements from the device, including non-interactive elements.
-
-        This function interacts with the TopViewService app installed on the device
-        to capture all UI elements, even those that are not interactive. This provides
-        a complete view of the UI hierarchy for analysis or debugging purposes.
-
-        Returns:
-            Dictionary containing all UI elements extracted from the device screen
-        """
-        try:
-            # Get the device
-            device = await self.device_manager.get_device(self.serial)
-            if not device:
-                raise ValueError(f"Device {self.serial} not found")
-
-            # Create a temporary file for the JSON
-            with tempfile.NamedTemporaryFile(suffix=".json") as temp:
-                local_path = temp.name
-
-                try:
-                    # Clear logcat to make it easier to find our output
-                    await device._adb.shell(device._serial, "logcat -c")
-
-                    # Trigger the custom service via broadcast to get ALL elements
-                    await device._adb.shell(
-                        device._serial,
-                        "am broadcast -a com.droidrun.portal.GET_ALL_ELEMENTS",
-                    )
-
-                    # Poll for the JSON file path
-                    start_time = asyncio.get_event_loop().time()
-                    max_wait_time = 10  # Maximum wait time in seconds
-                    poll_interval = 0.2  # Check every 200ms
-
-                    device_path = None
-                    while asyncio.get_event_loop().time() - start_time < max_wait_time:
-                        # Check logcat for the file path
-                        logcat_output = await device._adb.shell(
-                            device._serial,
-                            'logcat -d | grep "DROIDRUN_FILE" | grep "JSON data written to" | tail -1',
-                        )
-
-                        # Parse the file path if present
-                        match = re.search(r"JSON data written to: (.*)", logcat_output)
-                        if match:
-                            device_path = match.group(1).strip()
-                            break
-
-                        # Wait before polling again
-                        await asyncio.sleep(poll_interval)
-
-                    # Check if we found the file path
-                    if not device_path:
-                        raise ValueError(
-                            f"Failed to find the JSON file path in logcat after {max_wait_time} seconds"
-                        )
-
-                    logger.debug(f"Pulling file from {device_path} to {local_path}")
-                    # Pull the JSON file from the device
-                    await device._adb.pull_file(device._serial, device_path, local_path)
-
-                    # Read the JSON file
-                    async with aiofiles.open(local_path, "r", encoding="utf-8") as f:
-                        json_content = await f.read()
-
-                    # Clean up the temporary file
-                    with contextlib.suppress(OSError):
-                        os.unlink(local_path)
-
-                    # Try to parse the JSON
-                    import json
-
-                    try:
-                        ui_data = json.loads(json_content)
-
-                        return {
-                            "all_elements": ui_data,
-                            "count": (
-                                len(ui_data)
-                                if isinstance(ui_data, list)
-                                else sum(1 for _ in ui_data.get("elements", []))
-                            ),
-                            "message": "Retrieved all UI elements from the device screen",
-                        }
-                    except json.JSONDecodeError:
-                        raise ValueError("Failed to parse UI elements JSON data")
-
-                except Exception as e:
-                    # Clean up in case of error
-                    with contextlib.suppress(OSError):
-                        os.unlink(local_path)
-                    raise ValueError(f"Error retrieving all UI elements: {e}")
-
-        except Exception as e:
-            raise ValueError(f"Error getting all UI elements: {e}")
-
     def complete(self, success: bool, reason: str = ""):
         """
         Mark the task as finished.
@@ -676,7 +729,7 @@ class AdbTools(Tools):
             self.reason = reason
             self.finished = True
 
-    async def remember(self, information: str) -> str:
+    def remember(self, information: str) -> str:
         """
         Store important information to remember for future context.
 
@@ -712,7 +765,7 @@ class AdbTools(Tools):
         """
         return self.memory.copy()
 
-    async def get_state(self, serial: Optional[str] = None) -> Dict[str, Any]:
+    def get_state(self, serial: Optional[str] = None) -> Dict[str, Any]:
         """
         Get both the a11y tree and phone state in a single call using the combined /state endpoint.
 
@@ -722,54 +775,75 @@ class AdbTools(Tools):
         Returns:
             Dictionary containing both 'a11y_tree' and 'phone_state' data
         """
-        
+
         try:
-            if serial:
-                device = await self.device_manager.get_device(serial)
-                if not device:
-                    raise ValueError(f"Device {serial} not found")
+            logger.debug("Getting state")
+
+            if self.use_tcp and self.tcp_forwarded:
+                # Use TCP communication
+                response = requests.get(f"{self.tcp_base_url}/state", timeout=10)
+
+                if response.status_code == 200:
+                    tcp_response = response.json()
+
+                    # Check if response has the expected format
+                    if isinstance(tcp_response, dict) and "data" in tcp_response:
+                        data_str = tcp_response["data"]
+                        try:
+                            combined_data = json.loads(data_str)
+                        except json.JSONDecodeError:
+                            return {
+                                "error": "Parse Error",
+                                "message": "Failed to parse JSON data from TCP response data field",
+                            }
+                    else:
+                        # Fallback: assume direct JSON format
+                        combined_data = tcp_response
+                else:
+                    return {
+                        "error": "HTTP Error",
+                        "message": f"HTTP request failed with status {response.status_code}",
+                    }
             else:
-                device = await self.get_device()
-
-            adb_output = await device._adb.shell(
-                device._serial,
-                'content query --uri content://com.droidrun.portal/state'
-            )
+                # Fallback to content provider method
+                adb_output = self.device.shell(
+                    "content query --uri content://com.droidrun.portal/state",
+                )
 
-            state_data = self._parse_content_provider_output(adb_output)
-            
-            if state_data is None:
-                return {
-                    "error": "Parse Error",
-                    "message": "Failed to parse state data from ContentProvider response"
-                }
+                state_data = self._parse_content_provider_output(adb_output)
 
-            if isinstance(state_data, dict) and "data" in state_data:
-                data_str = state_data["data"]
-                try:
-                    combined_data = json.loads(data_str)
-                except json.JSONDecodeError:
+                if state_data is None:
                     return {
                         "error": "Parse Error",
-                        "message": "Failed to parse JSON data from ContentProvider data field"
+                        "message": "Failed to parse state data from ContentProvider response",
+                    }
+
+                if isinstance(state_data, dict) and "data" in state_data:
+                    data_str = state_data["data"]
+                    try:
+                        combined_data = json.loads(data_str)
+                    except json.JSONDecodeError:
+                        return {
+                            "error": "Parse Error",
+                            "message": "Failed to parse JSON data from ContentProvider data field",
+                        }
+                else:
+                    return {
+                        "error": "Format Error",
+                        "message": f"Unexpected state data format: {type(state_data)}",
                     }
-            else:
-                return {
-                    "error": "Format Error",
-                    "message": f"Unexpected state data format: {type(state_data)}"
-                }
 
             # Validate that both a11y_tree and phone_state are present
             if "a11y_tree" not in combined_data:
                 return {
                     "error": "Missing Data",
-                    "message": "a11y_tree not found in combined state data"
+                    "message": "a11y_tree not found in combined state data",
                 }
-            
+
             if "phone_state" not in combined_data:
                 return {
-                    "error": "Missing Data", 
-                    "message": "phone_state not found in combined state data"
+                    "error": "Missing Data",
+                    "message": "phone_state not found in combined state data",
                 }
 
             # Filter out the "type" attribute from all a11y_tree elements
@@ -777,9 +851,7 @@ class AdbTools(Tools):
             filtered_elements = []
             for element in elements:
                 # Create a copy of the element without the "type" attribute
-                filtered_element = {
-                    k: v for k, v in element.items() if k != "type"
-                }
+                filtered_element = {k: v for k, v in element.items() if k != "type"}
 
                 # Also filter children if present
                 if "children" in filtered_element:
@@ -789,19 +861,279 @@ class AdbTools(Tools):
                     ]
 
                 filtered_elements.append(filtered_element)
-            
+
             self.clickable_elements_cache = filtered_elements
-            
+
             return {
                 "a11y_tree": filtered_elements,
-                "phone_state": combined_data["phone_state"]
+                "phone_state": combined_data["phone_state"],
             }
 
+        except requests.exceptions.RequestException as e:
+            return {
+                "error": "TCP Error",
+                "message": f"TCP request failed: {str(e)}",
+            }
         except Exception as e:
-            return {"error": str(e), "message": f"Error getting combined state: {str(e)}"}
+            return {
+                "error": str(e),
+                "message": f"Error getting combined state: {str(e)}",
+            }
 
-if __name__ == "__main__":
-    async def main():
-        tools = AdbTools()
+    def get_a11y_tree(self) -> Dict[str, Any]:
+        """
+        Get just the accessibility tree using the /a11y_tree endpoint.
+
+        Returns:
+            Dictionary containing accessibility tree data
+        """
+        try:
+            if self.use_tcp and self.tcp_forwarded:
+                response = requests.get(f"{self.tcp_base_url}/a11y_tree", timeout=10)
+
+                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:
+                        data_str = tcp_response["data"]
+                        try:
+                            return json.loads(data_str)
+                        except json.JSONDecodeError:
+                            return {
+                                "error": "Parse Error",
+                                "message": "Failed to parse JSON data from TCP response data field",
+                            }
+                    else:
+                        # Fallback: assume direct JSON format
+                        return tcp_response
+                else:
+                    return {
+                        "error": "HTTP Error",
+                        "message": f"HTTP request failed with status {response.status_code}",
+                    }
+            else:
+                # Fallback: use get_state and extract a11y_tree
+                state = self.get_state()
+                if "error" in state:
+                    return state
+                return {"a11y_tree": state.get("a11y_tree", [])}
+
+        except requests.exceptions.RequestException as e:
+            return {
+                "error": "TCP Error",
+                "message": f"TCP request failed: {str(e)}",
+            }
+        except Exception as e:
+            return {
+                "error": str(e),
+                "message": f"Error getting a11y tree: {str(e)}",
+            }
+
+    def get_phone_state(self) -> Dict[str, Any]:
+        """
+        Get just the phone state using the /phone_state endpoint.
+
+        Returns:
+            Dictionary containing phone state data
+        """
+        try:
+            if self.use_tcp and self.tcp_forwarded:
+                response = requests.get(f"{self.tcp_base_url}/phone_state", timeout=10)
+
+                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:
+                        data_str = tcp_response["data"]
+                        try:
+                            return json.loads(data_str)
+                        except json.JSONDecodeError:
+                            return {
+                                "error": "Parse Error",
+                                "message": "Failed to parse JSON data from TCP response data field",
+                            }
+                    else:
+                        # Fallback: assume direct JSON format
+                        return tcp_response
+                else:
+                    return {
+                        "error": "HTTP Error",
+                        "message": f"HTTP request failed with status {response.status_code}",
+                    }
+            else:
+                # Fallback: use get_state and extract phone_state
+                state = self.get_state()
+                if "error" in state:
+                    return state
+                return {"phone_state": state.get("phone_state", {})}
+
+        except requests.exceptions.RequestException as e:
+            return {
+                "error": "TCP Error",
+                "message": f"TCP request failed: {str(e)}",
+            }
+        except Exception as e:
+            return {
+                "error": str(e),
+                "message": f"Error getting phone state: {str(e)}",
+            }
+
+    def ping(self) -> Dict[str, Any]:
+        """
+        Test the TCP connection using the /ping endpoint.
 
-    asyncio.run(main())
+        Returns:
+            Dictionary with ping result
+        """
+        try:
+            if self.use_tcp and self.tcp_forwarded:
+                response = requests.get(f"{self.tcp_base_url}/ping", timeout=5)
+
+                if response.status_code == 200:
+                    try:
+                        tcp_response = response.json() if response.content else {}
+                        logger.debug(f"Ping TCP response: {tcp_response}")
+                        return {
+                            "status": "success",
+                            "message": "Ping successful",
+                            "response": tcp_response,
+                        }
+                    except json.JSONDecodeError:
+                        return {
+                            "status": "success",
+                            "message": "Ping successful (non-JSON response)",
+                            "response": response.text,
+                        }
+                else:
+                    return {
+                        "status": "error",
+                        "message": f"Ping failed with status {response.status_code}: {response.text}",
+                    }
+            else:
+                return {
+                    "status": "error",
+                    "message": "TCP communication is not enabled",
+                }
+
+        except requests.exceptions.RequestException as e:
+            return {
+                "status": "error",
+                "message": f"Ping failed: {str(e)}",
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"Error during ping: {str(e)}",
+            }
+
+
+def _shell_test_cli(serial: str, command: str) -> tuple[str, float]:
+    """
+    Run an adb shell command using the adb CLI and measure execution time.
+    Args:
+        serial: Device serial number
+        command: Shell command to run
+    Returns:
+        Tuple of (output, elapsed_time)
+    """
+    import time
+    import subprocess
+
+    adb_cmd = ["adb", "-s", serial, "shell", command]
+    start = time.perf_counter()
+    result = subprocess.run(adb_cmd, capture_output=True, text=True)
+    elapsed = time.perf_counter() - start
+    output = result.stdout.strip() if result.returncode == 0 else result.stderr.strip()
+    return output, elapsed
+
+
+def _shell_test():
+    device = adb.device("emulator-5554")
+    # Native Python adb client
+    start = time.time()
+    res = device.shell("echo 'Hello, World!'")
+    end = time.time()
+    print(f"[Native] Shell execution took {end - start:.3f} seconds: {res}")
+
+    start = time.time()
+    res = device.shell("content query --uri content://com.droidrun.portal/state")
+    end = time.time()
+    print(f"[Native] Shell execution took {end - start:.3f} seconds: phone_state")
+
+    # CLI version
+    output, elapsed = _shell_test_cli("emulator-5554", "echo 'Hello, World!'")
+    print(f"[CLI] Shell execution took {elapsed:.3f} seconds: {output}")
+
+    output, elapsed = _shell_test_cli(
+        "emulator-5554", "content query --uri content://com.droidrun.portal/state"
+    )
+    print(f"[CLI] Shell execution took {elapsed:.3f} seconds: phone_state")
+
+
+def _list_packages():
+    tools = AdbTools()
+    print(tools.list_packages())
+
+
+def _start_app():
+    tools = AdbTools()
+    tools.start_app("com.android.settings", ".Settings")
+
+
+def _shell_test_cli(serial: str, command: str) -> tuple[str, float]:
+    """
+    Run an adb shell command using the adb CLI and measure execution time.
+    Args:
+        serial: Device serial number
+        command: Shell command to run
+    Returns:
+        Tuple of (output, elapsed_time)
+    """
+    import time
+    import subprocess
+
+    adb_cmd = ["adb", "-s", serial, "shell", command]
+    start = time.perf_counter()
+    result = subprocess.run(adb_cmd, capture_output=True, text=True)
+    elapsed = time.perf_counter() - start
+    output = result.stdout.strip() if result.returncode == 0 else result.stderr.strip()
+    return output, elapsed
+
+
+def _shell_test():
+    device = adb.device("emulator-5554")
+    # Native Python adb client
+    start = time.time()
+    res = device.shell("echo 'Hello, World!'")
+    end = time.time()
+    print(f"[Native] Shell execution took {end - start:.3f} seconds: {res}")
+
+    start = time.time()
+    res = device.shell("content query --uri content://com.droidrun.portal/state")
+    end = time.time()
+    print(f"[Native] Shell execution took {end - start:.3f} seconds: phone_state")
+
+    # CLI version
+    output, elapsed = _shell_test_cli("emulator-5554", "echo 'Hello, World!'")
+    print(f"[CLI] Shell execution took {elapsed:.3f} seconds: {output}")
+
+    output, elapsed = _shell_test_cli(
+        "emulator-5554", "content query --uri content://com.droidrun.portal/state"
+    )
+    print(f"[CLI] Shell execution took {elapsed:.3f} seconds: phone_state")
+
+
+def _list_packages():
+    tools = AdbTools()
+    print(tools.list_packages())
+
+
+def _start_app():
+    tools = AdbTools()
+    tools.start_app("com.android.settings", ".Settings")
+
+
+if __name__ == "__main__":
+    _start_app()