feat(idb_client): manage idb_companion lifecycle if no host is provided

Author: charlestheprogrammer

minitap/mobile_use/clients/idb_client.py

@@ -1,4 +1,7 @@
+import asyncio
 import json
+import socket
+import subprocess
 from functools import wraps
 from pathlib import Path
 from typing import Any
@@ -11,6 +14,20 @@
 logger = get_logger(__name__)
 
 
+def _find_available_port(start_port: int = 10882, max_attempts: int = 100) -> int:
+    """Find an available port starting from start_port."""
+    for port in range(start_port, start_port + max_attempts):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.bind(("localhost", port))
+                return port
+            except OSError:
+                continue
+    raise RuntimeError(
+        f"Could not find available port in range {start_port}-{start_port + max_attempts}"
+    )
+
+
 def with_idb_client(func):
     """Decorator that creates a Client.build context and injects it as 'client' parameter."""
 
@@ -23,25 +40,133 @@ async def wrapper(self, *args, **kwargs):
 
 
 class IdbClientWrapper:
-    """Wrapper around fb-idb client for iOS device automation.
+    """Wrapper around fb-idb client for iOS device automation with lifecycle management.
+
+    This wrapper can either manage the idb_companion process lifecycle locally or connect
+    to an external companion server.
+
+    Lifecycle Management:
+    - If host is None (default): Manages companion locally on localhost
+      - Call init_companion() to start the idb_companion process
+      - Call cleanup() to stop the companion process
+      - Or use as async context manager for automatic lifecycle
+    - If host is provided: Connects to external companion server
+      - init_companion() and cleanup() become no-ops
+      - You manage the external companion separately
+
+    Example:
+        # Managed companion (recommended for local development)
+        async with IdbClientWrapper(udid="device-id") as wrapper:
+            await wrapper.tap(100, 200)
 
-    Each method uses the @with_idb_client decorator which:
-    - Creates a fresh Client.build context per operation
-    - Injects the client as the first parameter after self
-    - Ensures proper cleanup after each operation completes
+        # External companion (for production/remote)
+        wrapper = IdbClientWrapper(udid="device-id", host="remote-host", port=10882)
+        await wrapper.tap(100, 200)  # No companion lifecycle management needed
     """
 
-    def __init__(self, udid: str, host: str = "localhost", port: int = 10882):
+    def __init__(self, udid: str, host: str | None = None, port: int | None = None):
+        self.udid = udid
+        self._manage_companion = host is None
+
+        if host is None:
+            actual_port = port if port is not None else _find_available_port()
+            self.address = TCPAddress(host="localhost", port=actual_port)
+            logger.debug(f"Will manage companion for {udid} on port {actual_port}")
+        else:
+            actual_port = port if port is not None else 10882
+            self.address = TCPAddress(host=host, port=actual_port)
+
+        self.companion_process: subprocess.Popen | None = None
+
+    async def init_companion(self, idb_companion_path: str = "idb_companion") -> bool:
         """
-        Initialize IDB Controller
+        Start the idb_companion process for this device.
+        Only starts if managing companion locally (host was None in __init__).
 
         Args:
-            udid: Device UDID
-            host: IDB companion host (default: localhost)
-            port: IDB companion port (default: 10882)
+            idb_companion_path: Path to idb_companion binary (default: "idb_companion" from PATH)
+
+        Returns:
+            True if companion started successfully, False otherwise
         """
-        self.udid = udid
-        self.address = TCPAddress(host=host, port=port)
+        if not self._manage_companion:
+            logger.info(f"Using external idb_companion at {self.address.host}:{self.address.port}")
+            return True
+
+        if self.companion_process is not None:
+            logger.warning(f"idb_companion already running for {self.udid}")
+            return True
+
+        try:
+            cmd = [idb_companion_path, "--udid", self.udid, "--grpc-port", str(self.address.port)]
+
+            logger.info(f"Starting idb_companion: {' '.join(cmd)}")
+            self.companion_process = subprocess.Popen(
+                cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+            )
+
+            await asyncio.sleep(2)
+
+            if self.companion_process.poll() is not None:
+                stdout, stderr = self.companion_process.communicate()
+                logger.error(f"idb_companion failed to start: {stderr}")
+                self.companion_process = None
+                return False
+
+            logger.info(
+                f"idb_companion started successfully for {self.udid} on port {self.address.port}"
+            )
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to start idb_companion: {e}")
+            self.companion_process = None
+            return False
+
+    async def cleanup(self) -> None:
+        if not self._manage_companion:
+            logger.debug(f"Not managing companion for {self.udid}, skipping cleanup")
+            return
+
+        if self.companion_process is None:
+            return
+
+        try:
+            logger.info(f"Stopping idb_companion for {self.udid}")
+
+            self.companion_process.terminate()
+
+            try:
+                await asyncio.wait_for(asyncio.to_thread(self.companion_process.wait), timeout=5.0)
+                logger.info(f"idb_companion stopped gracefully for {self.udid}")
+            except TimeoutError:
+                logger.warning(f"Force killing idb_companion for {self.udid}")
+                self.companion_process.kill()
+                await asyncio.to_thread(self.companion_process.wait)
+
+        except Exception as e:
+            logger.error(f"Error stopping idb_companion: {e}")
+        finally:
+            self.companion_process = None
+
+    def __del__(self):
+        if self.companion_process is not None:
+            try:
+                self.companion_process.terminate()
+                self.companion_process.wait(timeout=2)
+            except Exception:
+                try:
+                    self.companion_process.kill()
+                except Exception:
+                    pass
+
+    async def __aenter__(self):
+        await self.init_companion()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.cleanup()
+        return False
 
     @with_idb_client
     async def tap(self, client: Client, x: int, y: int, duration: float | None = None):