Revert "refactor: clean up legacy code, rename hopper to entity_extractor, and add third-party notices"

This reverts commit 54fcef9887.

Please do not submit unrelated docs and codes in the same commit. I am reverting this commit so that the related files can be submitted in better-organized future commits.
This commit is contained in:
yaoyao-open-source committed 2026-09-11 19:32:22 -07:00
1 parent 54fcef9887
commit 1d82d59ec8
83 files changed
+3677 -2013

No files matched your search

+327
View File
@@ -0,0 +1,327 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""ADB WebSocket tunnel implementation.
Bridges local TCP connections to a remote ADB server via WebSocket.
This replaces the need for the `lim` CLI for ADB tunneling.
The tunnel runs in a dedicated thread with its own event loop to ensure
it cannot be blocked by synchronous operations in the main thread.
Based on the Go SDK implementation which:
1. Creates a TCP listener
2. Accepts TCP connections from ADB
3. Dials the WebSocket
4. Forwards traffic bidirectionally
"""
import asyncio
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import socket
import threading
from artemis.utils.logger import get_logger
import websockets
from websockets.asyncio.client import ClientConnection
logger = get_logger(__name__)
BUFFER_SIZE = 32 * 1024 # 32KB, same as Go SDK
PING_INTERVAL = 30 # seconds
class AdbTunnel:
"""ADB tunnel that bridges local TCP connections to a remote WebSocket endpoint.
Runs in a dedicated thread with its own event loop to ensure the tunnel
cannot be blocked by synchronous operations in the caller's thread.
Usage:
tunnel = AdbTunnel(remote_url, token)
addr = await tunnel.start() # Returns "127.0.0.1:PORT"
# Use addr with ADB: adb connect 127.0.0.1:PORT
await tunnel.stop()
"""
def __init__(self, remote_url: str, token: str):
"""Initialize the ADB tunnel.
Args:
remote_url: WebSocket URL for the remote ADB endpoint
token: Bearer token for authentication
"""
self.remote_url = remote_url
self.token = token
self._local_port: int | None = None
self._thread: threading.Thread | None = None
self._loop: asyncio.AbstractEventLoop | None = None
self._started = threading.Event()
self._stop_event: asyncio.Event | None = None
@property
def local_addr(self) -> str:
"""Get the local address to connect ADB to."""
if self._local_port is None:
raise RuntimeError("Tunnel not started")
return f"127.0.0.1:{self._local_port}"
@property
def local_port(self) -> int:
"""Get the local port number."""
if self._local_port is None:
raise RuntimeError("Tunnel not started")
return self._local_port
async def start(self) -> str:
"""Start the ADB tunnel in a dedicated thread.
Returns:
The local address (host:port) to connect ADB to.
"""
if self._thread is not None and self._thread.is_alive():
return self.local_addr
# Create listener in main thread to get the port immediately
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(5)
listener.setblocking(False)
self._local_port = listener.getsockname()[1]
logger.info(f"ADB tunnel listening on {self.local_addr}")
# Start tunnel thread
self._thread = threading.Thread(
target=self._run_in_thread,
args=(listener,),
name=f"adb-tunnel-{self._local_port}",
daemon=True,
)
self._thread.start()
# Wait for thread to initialize its event loop
self._started.wait(timeout=5.0)
if not self._started.is_set():
await self.stop()
raise RuntimeError("Tunnel thread failed to start")
return self.local_addr
async def stop(self) -> None:
"""Stop the ADB tunnel and cleanup resources."""
if self._loop and self._stop_event:
# Signal the tunnel loop to stop
self._loop.call_soon_threadsafe(self._stop_event.set)
if self._thread and self._thread.is_alive():
await asyncio.to_thread(self._thread.join, 5.0)
if self._thread.is_alive():
logger.warning("ADB tunnel thread did not stop within 5s")
self._thread = None
self._loop = None
self._stop_event = None
logger.info("ADB tunnel stopped")
def _run_in_thread(self, listener: socket.socket) -> None:
"""Run the tunnel event loop in a dedicated thread."""
try:
self._loop = asyncio.new_event_loop()
asyncio.set_event_loop(self._loop)
self._stop_event = asyncio.Event()
self._started.set()
self._loop.run_until_complete(self._run_tunnel(listener))
except Exception as e:
logger.error(f"Tunnel thread error: {e}")
finally:
listener.close()
if self._loop:
self._loop.close()
async def _run_tunnel(self, listener: socket.socket) -> None:
"""Run the tunnel - accept TCP connections and bridge to WebSocket."""
loop = asyncio.get_event_loop()
connection_tasks: list[asyncio.Task] = []
stop_event = self._stop_event
assert stop_event is not None
logger.debug("Tunnel: ready to accept connections")
while not stop_event.is_set():
try:
# Use wait_for to periodically check stop_event
try:
tcp_conn, addr = await asyncio.wait_for(
loop.sock_accept(listener),
timeout=1.0,
)
except TimeoutError:
# Check if we should stop
continue
logger.debug(f"Tunnel: new TCP connection from {addr}")
task = asyncio.create_task(
self._handle_connection(tcp_conn, loop),
name=f"conn-{addr[1]}",
)
task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
connection_tasks.append(task)
connection_tasks = [t for t in connection_tasks if not t.done()]
except Exception as e:
if not stop_event.is_set():
logger.error(f"Tunnel accept error: {e}")
break
# Cleanup pending connections
for task in connection_tasks:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def _handle_connection(
self,
tcp_conn: socket.socket,
loop: asyncio.AbstractEventLoop,
) -> None:
"""Handle a single TCP connection."""
tcp_conn.setblocking(False)
ws: ClientConnection | None = None
try:
logger.debug(f"Connecting to WebSocket: {self.remote_url[:80]}...")
ws = await websockets.connect(
self.remote_url,
additional_headers={"Authorization": f"Bearer {self.token}"},
ping_interval=PING_INTERVAL,
ping_timeout=10,
)
logger.debug("WebSocket connected to remote ADB")
await self._bridge(tcp_conn, ws, loop)
except Exception as e:
logger.error(f"Connection error: {e}")
finally:
if ws is not None:
await ws.close()
tcp_conn.close()
logger.debug("Connection closed")
async def _bridge(
self,
tcp_conn: socket.socket,
ws: ClientConnection,
loop: asyncio.AbstractEventLoop,
) -> None:
"""Bridge TCP and WebSocket bidirectionally."""
stop_event = self._stop_event
assert stop_event is not None
async def tcp_to_ws() -> None:
"""Forward data from TCP to WebSocket."""
msg_count = 0
try:
while not stop_event.is_set():
try:
data = await asyncio.wait_for(
loop.sock_recv(tcp_conn, BUFFER_SIZE),
timeout=1.0,
)
except TimeoutError:
continue
except (OSError, ConnectionError) as e:
logger.debug(f"tcp->ws: recv error after {msg_count} msgs: {e}")
break
if not data:
logger.debug(f"tcp->ws: closed by client after {msg_count} msgs")
break
msg_count += 1
logger.debug(f"tcp->ws: [{msg_count}] sending {len(data)} bytes")
await ws.send(data)
logger.debug(f"tcp->ws: loop ended after {msg_count} messages")
except asyncio.CancelledError:
logger.debug(f"tcp->ws: cancelled after {msg_count} messages")
raise
except Exception as e:
logger.error(f"tcp->ws: error after {msg_count} messages: {e}")
async def ws_to_tcp() -> None:
"""Forward data from WebSocket to TCP."""
msg_count = 0
try:
logger.debug("ws->tcp: starting message loop")
async for message in ws:
if stop_event.is_set():
break
if isinstance(message, bytes) and message:
msg_count += 1
logger.debug(f"ws->tcp: [{msg_count}] received {len(message)} bytes")
await loop.sock_sendall(tcp_conn, message)
elif isinstance(message, str):
logger.debug(f"ws->tcp: received string: {message[:100]}")
logger.debug(f"ws->tcp: loop ended after {msg_count} messages")
except asyncio.CancelledError:
logger.debug(f"ws->tcp: cancelled after {msg_count} messages")
raise
except Exception as e:
logger.error(f"ws->tcp: error after {msg_count} messages: {e}")
tcp_task = asyncio.create_task(tcp_to_ws())
ws_task = asyncio.create_task(ws_to_tcp())
done, pending = await asyncio.wait(
[tcp_task, ws_task],
return_when=asyncio.FIRST_COMPLETED,
)
for task in pending:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
@asynccontextmanager
async def adb_tunnel(remote_url: str, token: str) -> AsyncGenerator[AdbTunnel, None]:
"""Context manager for creating and managing an ADB tunnel.
Args:
remote_url: WebSocket URL for the remote ADB endpoint
token: Bearer token for authentication
Yields:
AdbTunnel instance with the tunnel running
Example:
async with adb_tunnel(ws_url, token) as tunnel:
# Connect ADB to tunnel.local_addr
adb_client = AdbClient()
adb_client.connect(tunnel.local_addr)
"""
tunnel = AdbTunnel(remote_url, token)
try:
await tunnel.start()
yield tunnel
finally:
await tunnel.stop()
+109 -7
View File
@@ -11,13 +11,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Portions of this file are derived from mobile-use (https://github.com/minitap-ai/mobile-use)
# Copyright 2025-2026 Minitap, Inc. Licensed under the Apache License 2.0.
"""UIAutomator2 client for Android device screen data retrieval.
Provides direct device access for hierarchy and screenshot retrieval.
Provides an alternative to the Maestro-based screen API with direct device
access.
Handles Maestro blocker detection and removal before connecting.
"""
import base64
@@ -42,9 +41,32 @@ if TYPE_CHECKING:
logger = get_logger(__name__)
MAESTRO_PACKAGE = "dev.mobile.maestro"
class _CyFunctionDetectorMeta(type):
def __instancecheck__(self, instance):
name = type(instance).__name__
return (
name
in (
"cyfunction",
"cython_function_or_method",
"builtin_function_or_method",
)
or "cyfunction" in name.lower()
)
class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta):
pass
class UIAutomatorScreenData(BaseModel):
"""Screen data response from UIAutomator2."""
model_config = {"ignored_types": (CyFunctionDetector,)}
base64: str
hierarchy_xml: str
elements: list[dict]
@@ -148,11 +170,88 @@ def _pil_to_base64(img: Image.Image, format: str = "JPEG", quality: int = 80) ->
return base64.b64encode(buffer.getvalue()).decode("utf-8")
def _is_package_installed(device_id: str, pkg: str) -> bool:
"""Check if a package is installed on the device."""
try:
result = subprocess.run(
adb_command(["-s", device_id, "shell", "pm", "list", "packages"]),
stdin=subprocess.DEVNULL,
capture_output=True,
text=True,
timeout=10,
)
if result.returncode != 0:
logger.warning(f"Failed to list packages: {result.stderr}")
return False
lines = result.stdout.splitlines()
target = f"package:{pkg}"
return target in lines
except subprocess.TimeoutExpired:
logger.warning("Timeout checking installed packages")
return False
except (OSError, subprocess.SubprocessError) as e:
logger.warning(f"Error checking installed packages: {e}")
return False
def _uninstall_package(device_id: str, pkg: str) -> bool:
"""Uninstall a package from the device for the current user."""
try:
result = subprocess.run(
adb_command(
[
"-s",
device_id,
"shell",
"pm",
"uninstall",
"--user",
"0",
pkg,
]
),
stdin=subprocess.DEVNULL,
capture_output=True,
text=True,
timeout=30,
)
if result.returncode == 0:
logger.info(f"Successfully uninstalled {pkg}")
return True
else:
logger.warning(f"Failed to uninstall {pkg}: {result.stderr}")
return False
except subprocess.TimeoutExpired:
logger.warning(f"Timeout uninstalling {pkg}")
return False
except Exception as e:
logger.warning(f"Error uninstalling {pkg}: {e}")
return False
def _ensure_maestro_not_installed(device_id: str) -> None:
"""Check if Maestro is installed and uninstall it.
Maestro conflicts with uiautomator2 - if Maestro's package is installed,
u2.connect() will fail. This function ensures Maestro is removed before
attempting to connect.
"""
if _is_package_installed(device_id, MAESTRO_PACKAGE):
logger.warning(
f"Maestro ({MAESTRO_PACKAGE}) detected - uninstalling to enable UIAutomator2..."
)
_uninstall_package(device_id, MAESTRO_PACKAGE)
class UIAutomatorClient:
"""UIAutomator2 client for Android screen data retrieval.
This client uses the uiautomator2 library for direct device communication,
providing high-performance hierarchy and screenshot retrieval.
This client uses uiautomator2 library for direct device communication,
providing faster hierarchy and screenshot retrieval compared to Maestro.
Important: Maestro must not be installed on the device as it conflicts
with uiautomator2. This client automatically handles Maestro removal.
"""
def __init__(self, device_id: str):
@@ -166,7 +265,7 @@ class UIAutomatorClient:
self._awake_strategy: str | None = None
def _ensure_connected(self) -> "Device":
"""Ensure connection to the device.
"""Ensure connection to the device, handling Maestro blocker.
Returns:
Connected uiautomator2 Device instance
@@ -180,6 +279,9 @@ class UIAutomatorClient:
logger.warning("UIAutomator2 connection lost, reconnecting...")
self._device = None
# Ensure Maestro is not blocking us
_ensure_maestro_not_installed(self._device_id)
# Enroll the device in the containing Artemis service lifetime before
# UIAutomator2 reads the screen. Client disconnects do not release it.
if self._awake_strategy is None: