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

droidrun/tools/adb.py

@@ -3,31 +3,50 @@ UI Actions - Core UI interaction tools for Android device control.
 """
 
 import os
+import io
 import json
 import time
-import asyncio
 import logging
-from typing import Optional, Dict, Tuple, List, Any, Type, Self
-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-adb-tools")
+logger = logging.getLogger("droidrun-tools")
 
 
 class AdbTools(Tools):
     """Core UI interaction tools for Android device control."""
 
-    def __init__(self, serial: str) -> 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_manager = DeviceManager()
+        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.last_screenshot = None
         self.reason = None
         self.success = None
@@ -37,46 +56,72 @@ class AdbTools(Tools):
         # Store all screenshots with timestamps
         self.screenshots: List[Dict[str, Any]] = []
 
-    @classmethod
-    async def create(cls: Type[Self], serial: str = None) -> Self:
-        """Create an AdbTools instance.
+        # Set up TCP forwarding if requested
+        if self.use_tcp:
+            self.setup_tcp_forward()
 
-        Args:
-            serial: Optional device serial number. If not provided, the first device found will be used.
+    def setup_tcp_forward(self) -> bool:
+        """
+        Set up ADB TCP port forwarding for communication with the portal app.
 
         Returns:
-            AdbTools instance
+            bool: True if forwarding was set up successfully, False otherwise
         """
-        if not serial:
-            dvm = DeviceManager()
-            devices = await dvm.list_devices()
-            if not devices or len(devices) < 1:
-                raise ValueError("No devices found")
-            serial = devices[0].serial
+        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 AdbTools(serial)
+            # 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 _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
+        except Exception as e:
+            logger.error(f"Failed to set up TCP port forwarding: {e}")
+            self.tcp_forwarded = False
+            return False
 
-    async def _get_device(self) -> Optional[Device]:
-        """Get the device instance using the instance's serial or from environment variable.
+    def teardown_tcp_forward(self) -> bool:
+        """
+        Remove ADB TCP port forwarding.
 
         Returns:
-            Device instance or None if not found
+            bool: True if forwarding was removed successfully, False otherwise
         """
-        serial = self._get_device_serial()
-        if not serial:
-            raise ValueError("No device serial specified - set device_serial parameter")
+        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
 
-        device = await self.device_manager.get_device(serial)
-        if not device:
-            raise ValueError(f"Device {serial} not found")
+    def __del__(self):
+        """Cleanup when the object is destroyed."""
+        if hasattr(self, "tcp_forwarded") and self.tcp_forwarded:
+            self.teardown_tcp_forward()
 
-        return device
+    def _set_context(self, ctx: Context):
+        self._ctx = ctx
 
     def _parse_content_provider_output(
         self, raw_output: str
@@ -125,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.
 
@@ -197,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})")
+
+            # Emit coordinate action event for trajectory recording
 
-            await device.tap(x, y)
+            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 = []
@@ -233,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.
 
@@ -245,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.
 
@@ -273,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.
@@ -286,29 +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
 
-    async def input_text(self, text: str, serial: Optional[str] = None) -> str:
+    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
+
+    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.
@@ -320,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()
+            logger.debug(f"Inputting text: {text}")
 
-            # Save the current keyboard
-            original_ime = await device._adb.shell(
-                device._serial, "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()
 
-            # 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"
-            )
-
-            # Wait for keyboard to change
-            await asyncio.sleep(1)
+                payload = {"base64_text": encoded_text}
+                response = requests.post(
+                    f"{self.tcp_base_url}/keyboard/input",
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                    timeout=10,
+                )
 
-            # Encode the text to Base64
-            import base64
+                logger.debug(
+                    f"Keyboard input TCP response: {response.status_code}, {response.text}"
+                )
 
-            encoded_text = base64.b64encode(text.encode()).decode()
-
-            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.
 
@@ -399,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",
@@ -414,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.
 
@@ -428,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:
         """
@@ -452,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.
 
@@ -506,14 +705,8 @@ 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()
-
-            return await device.list_packages(include_system_apps)
+            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)}")
 
@@ -536,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.
 
@@ -572,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.
 
@@ -584,40 +777,61 @@ class AdbTools(Tools):
         """
 
         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)
+                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",
-                }
-
-            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:
@@ -655,17 +869,271 @@ class AdbTools(Tools):
                 "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)}",
             }
 
+    def get_a11y_tree(self) -> Dict[str, Any]:
+        """
+        Get just the accessibility tree using the /a11y_tree endpoint.
 
-if __name__ == "__main__":
+        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", [])}
 
-    async def main():
-        tools = await AdbTools.create()
-        print(tools.serial)
+        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.
 
-    asyncio.run(main())
+        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.
+
+        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()