somew1nd
12d7902a82
Fix queue cancellation, concurrency limits and worker status updates. Use direct Explorer imports, split the Diagnoser loop into helpers and remove unused tool registration and prompt-builder code. Correct model caching, endpoint selection and retry handling.
1210 lines
44 KiB
Python
1210 lines
44 KiB
Python
# 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.
|
|
|
|
"""Universal ADB command and background task management tools for ARTEMIS.
|
|
|
|
State model
|
|
-----------
|
|
Every automation task owns one :class:`AdbTaskRegistry` (stored on its
|
|
:class:`~artemis.context.ArtemisContext`). The registry holds the running
|
|
background processes, the recently finished task logs, and the persistent
|
|
terminal environments. Keeping this per context matters because the MCP
|
|
daemon runs concurrent tasks for different devices inside one process: a
|
|
module-level registry would leak device A's tasks and completion notices into
|
|
device B's Operator prompt. A module-level default registry still exists for
|
|
callers that have no context (tests, ad-hoc scripts).
|
|
|
|
Hang protection
|
|
---------------
|
|
Commands are passed to ``adb shell`` as an argument, never through stdin, so
|
|
the remote shell exits as soon as the script does. ``stdin`` is closed
|
|
(``DEVNULL``) unless the caller asks for ``Interactive=True``; a command that
|
|
reads stdin therefore sees EOF instead of blocking forever. Anything that
|
|
outlives ``WaitMsBeforeAsync`` is moved to a background task with a bounded
|
|
log buffer, and :func:`shutdown_adb_background_tasks` kills whatever is still
|
|
running when the automation task ends.
|
|
"""
|
|
|
|
import asyncio
|
|
from asyncio.subprocess import Process as AsyncProcess
|
|
from collections import deque
|
|
import os
|
|
import shlex
|
|
from typing import Any, Literal
|
|
import uuid
|
|
|
|
from langchain_core.tools import BaseTool
|
|
from pydantic import BaseModel, Field
|
|
|
|
from artemis.agents.log_analyzer.output_analyzer import TaskOutputAnalyzerNode
|
|
from artemis.context import ArtemisContext
|
|
from artemis.controllers.platform_specific_commands_controller import (
|
|
get_adb_device,
|
|
)
|
|
from artemis.data_engine.trace import trace_langchain_tool
|
|
from artemis.drivers.base import BaseDeviceDriver
|
|
from artemis.tools.base import ArtemisTool
|
|
from artemis.tools.tool_wrapper import ToolWrapper
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
#: Maximum number of finished task records kept per registry.
|
|
_FINISHED_TASKS_CAP = 50
|
|
#: Maximum number of output lines retained for a background task.
|
|
_BACKGROUND_LOG_MAX_LINES = 20000
|
|
#: Seconds to wait for a terminated process before killing it.
|
|
_TERMINATE_GRACE_SECONDS = 3.0
|
|
|
|
_EXIT_CODE_MARKER = "===EXIT_CODE==="
|
|
_ENV_START_MARKER = "===ENV_START==="
|
|
|
|
#: Variables that every Android shell already defines. They are never
|
|
#: re-exported into a persistent terminal: re-exporting PATH/HOME etc. is
|
|
#: pointless and some of them (BOOTCLASSPATH, LD_*) are long and fragile.
|
|
_PERSISTENT_ENV_EXCLUDE = frozenset(
|
|
{
|
|
"PATH",
|
|
"HOME",
|
|
"HOSTNAME",
|
|
"USER",
|
|
"LOGNAME",
|
|
"SHELL",
|
|
"TERM",
|
|
"PWD",
|
|
"OLDPWD",
|
|
"SHLVL",
|
|
"_",
|
|
"TMPDIR",
|
|
"HISTFILE",
|
|
"LD_LIBRARY_PATH",
|
|
"LD_PRELOAD",
|
|
"EXTERNAL_STORAGE",
|
|
"ASEC_MOUNTPOINT",
|
|
"DOWNLOAD_CACHE",
|
|
}
|
|
)
|
|
_PERSISTENT_ENV_EXCLUDE_PREFIXES = (
|
|
"ANDROID_",
|
|
"BOOTCLASSPATH",
|
|
"DEX2OAT",
|
|
"SYSTEMSERVER",
|
|
"STANDALONE_",
|
|
)
|
|
|
|
#: Guidance shared by the tool descriptions: the commands that hang an
|
|
#: ``adb shell`` and how to bound them.
|
|
HANG_GUIDANCE = (
|
|
"Streaming or open-ended commands (logcat, top, getevent, screenrecord,"
|
|
" tail -f, an interactive shell) never return on their own: bound them"
|
|
" (logcat -d, top -n 1, --time-limit) or expect them to be moved to a"
|
|
" background task. Detach daemons explicitly with"
|
|
" `nohup CMD >/dev/null 2>&1 &`, otherwise the shell waits for them."
|
|
)
|
|
|
|
|
|
def _adb_binary() -> str:
|
|
"""Resolves the adb executable through the shared toolchain resolver."""
|
|
try:
|
|
from artemis.toolchain import toolchain
|
|
|
|
return toolchain.find_adb()
|
|
except Exception: # pylint: disable=broad-exception-caught
|
|
return "adb"
|
|
|
|
|
|
def _is_output_long(output: str) -> bool:
|
|
"""Checks if command output exceeds the max threshold for inline display."""
|
|
return len(output.splitlines()) > 500 or len(output) > 50000
|
|
|
|
|
|
def _format_long_output_response(_task_id: str, output: str, intro: str) -> str:
|
|
"""Formats truncated output with guidance to use the analyzer tool."""
|
|
lines = output.splitlines()
|
|
last_200 = "\n".join(lines[-200:])
|
|
return (
|
|
f"{intro}\nWarning: Output is too long ({len(lines)} lines,"
|
|
f" {len(output)} chars) and has been truncated.\n--- Last 200 lines of"
|
|
f" output ---\n{last_200}\n------------------------\nYou can use the"
|
|
" 'analyze_task_output' tool to get more information."
|
|
)
|
|
|
|
|
|
def _filter_persistent_env(env: dict[str, str]) -> dict[str, str]:
|
|
"""Drops the platform-provided variables from a captured shell environment."""
|
|
kept: dict[str, str] = {}
|
|
for key, value in env.items():
|
|
if not key or key in _PERSISTENT_ENV_EXCLUDE:
|
|
continue
|
|
if key.startswith(_PERSISTENT_ENV_EXCLUDE_PREFIXES):
|
|
continue
|
|
if not key.replace("_", "a").isalnum():
|
|
continue
|
|
kept[key] = value
|
|
return kept
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AdbTaskRegistry:
|
|
"""Per-task bookkeeping for background processes, finished logs and envs."""
|
|
|
|
def __init__(self) -> None:
|
|
self.background: dict[str, BackgroundTask] = {}
|
|
self.finished: dict[str, dict[str, Any]] = {}
|
|
self.persistent_envs: dict[str, dict[str, str]] = {}
|
|
|
|
# -- finished task cache -------------------------------------------------
|
|
|
|
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
|
def register_finished(
|
|
self,
|
|
task_id: str,
|
|
command: str,
|
|
cwd: str | None,
|
|
terminal_id: str | None,
|
|
status: str,
|
|
exit_code: int | None,
|
|
output: str,
|
|
) -> None:
|
|
"""Caches execution details for a finished task."""
|
|
self.finished[task_id] = {
|
|
"task_id": task_id,
|
|
"command": command,
|
|
"cwd": cwd,
|
|
"terminal_id": terminal_id,
|
|
"status": status,
|
|
"exit_code": exit_code,
|
|
"output": output,
|
|
"notified": False,
|
|
}
|
|
while len(self.finished) > _FINISHED_TASKS_CAP:
|
|
oldest_key = next(iter(self.finished))
|
|
del self.finished[oldest_key]
|
|
|
|
def get_task_info(self, task_id: str) -> dict[str, Any] | None:
|
|
"""Returns metadata and logs for a running or finished task, if known."""
|
|
if task_id in self.background:
|
|
t = self.background[task_id]
|
|
return {
|
|
"command": t.command,
|
|
"cwd": t.cwd,
|
|
"terminal_id": t.terminal_id,
|
|
"status": t.status,
|
|
"output": "".join(t.stdout_log),
|
|
}
|
|
return self.finished.get(task_id)
|
|
|
|
# -- prompt helpers ------------------------------------------------------
|
|
|
|
def active_task_summaries(self) -> list[dict[str, Any]]:
|
|
"""Compact view of running tasks for the Operator prompt."""
|
|
return [
|
|
{
|
|
"task_id": tid,
|
|
"command": t.command,
|
|
"cwd": t.cwd,
|
|
"terminal_id": t.terminal_id,
|
|
"output_line_count": len(t.stdout_log),
|
|
}
|
|
for tid, t in self.background.items()
|
|
]
|
|
|
|
def pop_unnotified_finished(self) -> list[dict[str, Any]]:
|
|
"""Returns finished tasks not yet shown to the model and marks them shown."""
|
|
newly_finished: list[dict[str, Any]] = []
|
|
for tid, tinfo in self.finished.items():
|
|
if tinfo.get("notified", False):
|
|
continue
|
|
newly_finished.append(
|
|
{
|
|
"task_id": tid,
|
|
"command": tinfo.get("command", ""),
|
|
"status": tinfo.get("status", "completed"),
|
|
"output_text": tinfo.get("output", ""),
|
|
}
|
|
)
|
|
tinfo["notified"] = True
|
|
return newly_finished
|
|
|
|
# -- lifecycle -----------------------------------------------------------
|
|
|
|
async def shutdown(self, ctx: ArtemisContext | None = None) -> int:
|
|
"""Terminates every running background task. Returns the number killed."""
|
|
tasks = list(self.background.values())
|
|
for task in tasks:
|
|
try:
|
|
await task.stop(ctx, status="killed")
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.warning(f"Failed to stop background ADB task {task.task_id}: {e}")
|
|
return len(tasks)
|
|
|
|
|
|
#: Registry used when no context is available.
|
|
_DEFAULT_REGISTRY = AdbTaskRegistry()
|
|
|
|
# Backward-compatibility aliases onto the default registry's dictionaries.
|
|
_BACKGROUND_TASKS = _DEFAULT_REGISTRY.background
|
|
_FINISHED_TASKS_LOGS = _DEFAULT_REGISTRY.finished
|
|
_PERSISTENT_ENVIRONMENTS = _DEFAULT_REGISTRY.persistent_envs
|
|
|
|
|
|
def get_adb_task_registry(ctx: ArtemisContext | None) -> AdbTaskRegistry:
|
|
"""Returns the registry bound to ``ctx``, creating it on first use."""
|
|
if ctx is None:
|
|
return _DEFAULT_REGISTRY
|
|
existing = getattr(ctx, "adb_task_registry", None)
|
|
if isinstance(existing, AdbTaskRegistry):
|
|
return existing
|
|
registry = AdbTaskRegistry()
|
|
try:
|
|
setattr(ctx, "adb_task_registry", registry)
|
|
except (AttributeError, TypeError):
|
|
return _DEFAULT_REGISTRY
|
|
return registry
|
|
|
|
|
|
async def shutdown_adb_background_tasks(ctx: ArtemisContext | None) -> int:
|
|
"""Kills the background ADB tasks of ``ctx``. Safe to call when none exist."""
|
|
existing = getattr(ctx, "adb_task_registry", None) if ctx is not None else None
|
|
if not isinstance(existing, AdbTaskRegistry):
|
|
return 0
|
|
killed = await existing.shutdown(ctx)
|
|
if killed:
|
|
logger.info(f"Terminated {killed} background ADB task(s) at task end.")
|
|
return killed
|
|
|
|
|
|
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
|
def _register_finished_task(
|
|
task_id: str,
|
|
command: str,
|
|
cwd: str | None,
|
|
terminal_id: str | None,
|
|
status: str,
|
|
exit_code: int | None,
|
|
output: str,
|
|
ctx: ArtemisContext | None = None,
|
|
):
|
|
"""Caches execution details for finished tasks (registry of ``ctx``)."""
|
|
get_adb_task_registry(ctx).register_finished(
|
|
task_id, command, cwd, terminal_id, status, exit_code, output
|
|
)
|
|
|
|
|
|
def _get_task_info(task_id: str, ctx: ArtemisContext | None) -> dict[str, Any] | None:
|
|
"""Retrieves metadata and logs for a running or finished task.
|
|
|
|
Looks in the context registry, then the default registry, then the
|
|
persisted background task table.
|
|
"""
|
|
registry = get_adb_task_registry(ctx)
|
|
info = registry.get_task_info(task_id)
|
|
if info is None and registry is not _DEFAULT_REGISTRY:
|
|
info = _DEFAULT_REGISTRY.get_task_info(task_id)
|
|
if info is not None:
|
|
return info
|
|
if ctx and ctx.data_engine:
|
|
db_tasks = ctx.data_engine.get_all_background_tasks()
|
|
for db_t in db_tasks:
|
|
if db_t.get("task_id") == task_id:
|
|
return {
|
|
"command": db_t.get("summary") or "",
|
|
"cwd": None,
|
|
"terminal_id": None,
|
|
"status": db_t.get("status") or "completed",
|
|
"output": db_t.get("logs") or "",
|
|
}
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Background task
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# pylint: disable=too-many-instance-attributes
|
|
class BackgroundTask:
|
|
"""Represents an active or completing background ADB shell task."""
|
|
|
|
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
|
def __init__(
|
|
self,
|
|
task_id: str,
|
|
command: str,
|
|
process: AsyncProcess,
|
|
terminal_id: str | None = None,
|
|
cwd: str | None = None,
|
|
interactive: bool = False,
|
|
registry: AdbTaskRegistry | None = None,
|
|
):
|
|
self.task_id = task_id
|
|
self.command = command
|
|
self.process = process
|
|
self.terminal_id = terminal_id
|
|
self.cwd = cwd
|
|
self.interactive = interactive
|
|
self.registry = registry
|
|
self.stdout_log: deque[str] = deque(maxlen=_BACKGROUND_LOG_MAX_LINES)
|
|
self.status = "running"
|
|
self.exit_code: int | None = None
|
|
self.trace_id: uuid.UUID | None = None
|
|
self.listener_task: asyncio.Task | None = None
|
|
self._stop_requested = False
|
|
self._finalized = False
|
|
|
|
async def start(self, ctx: ArtemisContext):
|
|
"""Starts the background output listener task."""
|
|
self.listener_task = asyncio.create_task(self._listen(ctx))
|
|
|
|
async def stop(self, ctx: ArtemisContext | None, status: str = "killed") -> None:
|
|
"""Terminates the process (kill after a grace period) and records the result."""
|
|
self._stop_requested = True
|
|
self.status = status
|
|
proc = self.process
|
|
if getattr(proc, "returncode", None) is None:
|
|
try:
|
|
proc.terminate()
|
|
except (ProcessLookupError, OSError):
|
|
pass
|
|
try:
|
|
await asyncio.wait_for(proc.wait(), timeout=_TERMINATE_GRACE_SECONDS)
|
|
except TimeoutError:
|
|
try:
|
|
proc.kill()
|
|
except (ProcessLookupError, OSError):
|
|
pass
|
|
try:
|
|
await asyncio.wait_for(proc.wait(), timeout=_TERMINATE_GRACE_SECONDS)
|
|
except TimeoutError:
|
|
pass
|
|
except Exception as exc: # pylint: disable=broad-exception-caught
|
|
# Shutdown must never raise; the exit code is read back below.
|
|
logger.debug(
|
|
f"Waiting for background ADB task {self.task_id} to exit failed: {exc}",
|
|
exc_info=True,
|
|
)
|
|
if self.listener_task is not None and not self.listener_task.done():
|
|
self.listener_task.cancel()
|
|
try:
|
|
await self.listener_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self.exit_code = getattr(proc, "returncode", None)
|
|
self._finalize(ctx)
|
|
|
|
def _finalize(self, ctx: ArtemisContext | None) -> None:
|
|
"""Moves the task from the active map to the finished cache exactly once."""
|
|
if self._finalized:
|
|
return
|
|
self._finalized = True
|
|
registry = self.registry or get_adb_task_registry(ctx)
|
|
|
|
output_text = "".join(self.stdout_log)
|
|
_, parsed_env, _ = self.parse_output(output_text)
|
|
if self.terminal_id and parsed_env is not None:
|
|
registry.persistent_envs[self.terminal_id] = _filter_persistent_env(parsed_env)
|
|
logger.info(
|
|
"Updated persistent env for Android terminal"
|
|
f" '{self.terminal_id}' from background task."
|
|
)
|
|
|
|
if ctx and getattr(ctx, "data_engine", None):
|
|
try:
|
|
ctx.data_engine.unregister_background_task(self.task_id, self.status)
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.debug(f"Failed to unregister background task {self.task_id}: {e}")
|
|
|
|
registry.register_finished(
|
|
task_id=self.task_id,
|
|
command=self.command,
|
|
cwd=self.cwd,
|
|
terminal_id=self.terminal_id,
|
|
status=self.status,
|
|
exit_code=self.exit_code,
|
|
output=output_text,
|
|
)
|
|
registry.background.pop(self.task_id, None)
|
|
|
|
async def _listen(self, ctx: ArtemisContext):
|
|
"""Reads process output streams and updates lifecycle status."""
|
|
try:
|
|
while True:
|
|
line = await self.process.stdout.readline()
|
|
if not line:
|
|
break
|
|
decoded = line.decode(errors="replace")
|
|
self.stdout_log.append(decoded)
|
|
if ctx and getattr(ctx, "data_engine", None) and self.trace_id:
|
|
ctx.data_engine.stream_output(self.trace_id, decoded)
|
|
|
|
self.exit_code = await self.process.wait()
|
|
if self._stop_requested:
|
|
return
|
|
_, _, script_exit = self.parse_output("".join(self.stdout_log))
|
|
if self.terminal_id and script_exit is not None:
|
|
self.exit_code = script_exit
|
|
self.status = "completed" if self.exit_code == 0 else "failed"
|
|
logger.info(
|
|
f"Background ADB task {self.task_id} completed with exit code {self.exit_code}"
|
|
)
|
|
except asyncio.CancelledError:
|
|
return
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.error(f"Error in background task {self.task_id} listener: {e}")
|
|
self.status = "error"
|
|
finally:
|
|
if not self._stop_requested:
|
|
self._finalize(ctx)
|
|
|
|
@staticmethod
|
|
def parse_output(output: str) -> tuple[str, dict[str, str] | None, int | None]:
|
|
"""Separates stdout lines from the exit-code marker and exported env.
|
|
|
|
Returns ``(clean_output, env_or_None, script_exit_code_or_None)``.
|
|
"""
|
|
clean_lines: list[str] = []
|
|
env_lines: list[str] = []
|
|
exit_code: int | None = None
|
|
in_env = False
|
|
|
|
for line in output.splitlines():
|
|
if _EXIT_CODE_MARKER in line:
|
|
raw = line.split(_EXIT_CODE_MARKER, 1)[1].strip()
|
|
try:
|
|
exit_code = int(raw)
|
|
except ValueError:
|
|
exit_code = None
|
|
continue
|
|
if _ENV_START_MARKER in line:
|
|
in_env = True
|
|
continue
|
|
if in_env:
|
|
env_lines.append(line)
|
|
else:
|
|
clean_lines.append(line)
|
|
|
|
clean_output = "\n".join(clean_lines)
|
|
|
|
parsed_env: dict[str, str] | None = None
|
|
if in_env:
|
|
parsed_env = {}
|
|
for el in env_lines:
|
|
if "=" in el:
|
|
k, v = el.split("=", 1)
|
|
parsed_env[k] = v
|
|
|
|
return clean_output, parsed_env, exit_code
|
|
|
|
# Backwards compatibility alias
|
|
_parse_output = parse_output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# run_adb_command
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RunAdbCommandArgs(BaseModel):
|
|
"""Arguments schema for executing ADB shell commands."""
|
|
|
|
CommandLine: str = Field(
|
|
...,
|
|
description=(
|
|
"The exact shell command line string to execute inside the Android device shell."
|
|
),
|
|
)
|
|
Cwd: str = Field(
|
|
default="/data/local/tmp",
|
|
description="The directory on the Android device to run the command in.",
|
|
)
|
|
RunPersistent: bool = Field(
|
|
default=False,
|
|
description=(
|
|
"Set to true to run this command in a persistent terminal session"
|
|
" that preserves environment variables between invocations on the"
|
|
" phone."
|
|
),
|
|
)
|
|
RequestedTerminalID: str | None = Field(
|
|
default=None,
|
|
description=(
|
|
"Optional ID of a persistent terminal to reuse. Reuses environment"
|
|
" variables if provided."
|
|
),
|
|
)
|
|
WaitMsBeforeAsync: int = Field(
|
|
default=5000,
|
|
description=(
|
|
"Number of milliseconds to wait for the command to finish"
|
|
" synchronously before sending it to the background."
|
|
),
|
|
)
|
|
Interactive: bool = Field(
|
|
default=False,
|
|
description=(
|
|
"Keep stdin open so 'manage_task send_input' can feed the command"
|
|
" later. Default false: stdin is closed and commands that read it"
|
|
" see EOF instead of hanging."
|
|
),
|
|
)
|
|
|
|
|
|
def _build_phone_script(
|
|
cmd_line: str,
|
|
cwd: str,
|
|
env_vars: dict[str, str],
|
|
run_persistent: bool,
|
|
) -> str:
|
|
"""Builds the script passed as the single ``adb shell`` argument."""
|
|
script = f"cd {shlex.quote(cwd)}\n"
|
|
for k, v in env_vars.items():
|
|
script += f"export {k}={shlex.quote(v)}\n"
|
|
script += f"{cmd_line}\n"
|
|
if run_persistent:
|
|
script += (
|
|
"_artemis_ec=$?\n"
|
|
f'echo "{_EXIT_CODE_MARKER}$_artemis_ec"\n'
|
|
f'echo "{_ENV_START_MARKER}"\n'
|
|
"env\n"
|
|
)
|
|
return script
|
|
|
|
|
|
class RunAdbCommandTool(ArtemisTool):
|
|
"""Universal tool for executing ADB shell commands on an Android device."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="run_adb_command",
|
|
description=(
|
|
"[SHELL] Executes a shell command directly on the Android mobile "
|
|
"device via ADB shell. Runs synchronously and transitions to a "
|
|
"background task if execution takes longer than WaitMsBeforeAsync. "
|
|
"Supports persistent environments (environment variables) on the "
|
|
f"phone across invocations. {HANG_GUIDANCE}"
|
|
),
|
|
args_schema=RunAdbCommandArgs,
|
|
category="system",
|
|
)
|
|
|
|
# pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,too-many-return-statements,too-many-positional-arguments,too-many-boolean-expressions
|
|
async def execute(
|
|
self,
|
|
driver: BaseDeviceDriver | None = None,
|
|
ctx: ArtemisContext | None = None,
|
|
CommandLine: str | None = None,
|
|
Cwd: str = "/data/local/tmp",
|
|
RunPersistent: bool = False,
|
|
RequestedTerminalID: str | None = None,
|
|
WaitMsBeforeAsync: int = 5000,
|
|
Interactive: bool = False,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
# Extract command line supporting alternative key formats
|
|
cmd_line = (
|
|
CommandLine
|
|
if CommandLine is not None
|
|
else (kwargs.get("command_line") or kwargs.get("command") or kwargs.get("cmd") or "")
|
|
)
|
|
cwd = Cwd or kwargs.get("cwd") or "/data/local/tmp"
|
|
run_persistent = (
|
|
RunPersistent if RunPersistent is not False else kwargs.get("run_persistent", False)
|
|
)
|
|
requested_terminal_id = (
|
|
RequestedTerminalID
|
|
if RequestedTerminalID is not None
|
|
else kwargs.get("requested_terminal_id")
|
|
)
|
|
wait_ms = (
|
|
WaitMsBeforeAsync
|
|
if WaitMsBeforeAsync != 5000
|
|
else kwargs.get("wait_ms_before_async", WaitMsBeforeAsync)
|
|
)
|
|
interactive = bool(Interactive or kwargs.get("interactive", False))
|
|
wait_seconds = max(float(wait_ms) / 1000.0, 0.05)
|
|
|
|
# Fallback to driver's execute_shell if no full ArtemisContext is available
|
|
if (
|
|
(ctx is None or not hasattr(ctx, "device") or not ctx.device)
|
|
and driver is not None
|
|
and hasattr(driver, "execute_shell")
|
|
and not run_persistent
|
|
):
|
|
return await driver.execute_shell(cmd_line, timeout_seconds=wait_seconds)
|
|
|
|
registry = get_adb_task_registry(ctx)
|
|
|
|
device_id = None
|
|
if ctx and hasattr(ctx, "device") and ctx.device:
|
|
device_id = getattr(ctx.device, "device_id", None)
|
|
elif driver and hasattr(driver, "device_id"):
|
|
device_id = getattr(driver, "device_id", None)
|
|
if not device_id:
|
|
device_id = "default_device"
|
|
|
|
# Prepare environment
|
|
terminal_id = None
|
|
run_env = os.environ.copy() # Local host env to run adb command
|
|
android_env_vars: dict[str, str] = {}
|
|
|
|
if run_persistent:
|
|
terminal_id = requested_terminal_id or f"term_{uuid.uuid4().hex[:8]}"
|
|
android_env_vars = registry.persistent_envs.setdefault(terminal_id, {})
|
|
|
|
phone_script = _build_phone_script(cmd_line, cwd, android_env_vars, run_persistent)
|
|
|
|
# In Cloud Mode or when adb client is virtualized: no host subprocess,
|
|
# therefore no background handoff. The wait window is a hard timeout.
|
|
adb_client = getattr(ctx, "adb_client", None) if ctx else None
|
|
if os.environ.get("ARTEMIS_CLOUD_MODE") == "1" or (
|
|
adb_client is not None and hasattr(adb_client, "_bridge")
|
|
):
|
|
try:
|
|
device = get_adb_device(ctx)
|
|
output = await asyncio.wait_for(
|
|
asyncio.to_thread(device.shell, phone_script),
|
|
timeout=wait_seconds,
|
|
)
|
|
except TimeoutError:
|
|
return (
|
|
f"Error: Command did not finish within {wait_seconds:.1f}s."
|
|
" Background tasks are not available on this device"
|
|
" connection; bound the command (logcat -d, top -n 1,"
|
|
" --time-limit) and retry."
|
|
)
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.error(f"Failed to execute virtual adb shell for '{cmd_line}': {e}")
|
|
return f"Failed to execute adb command: {e}"
|
|
clean_output, new_env, script_exit = BackgroundTask.parse_output(str(output))
|
|
if run_persistent and new_env is not None:
|
|
registry.persistent_envs[terminal_id] = _filter_persistent_env(new_env)
|
|
intro = "ADB command completed"
|
|
if script_exit is not None:
|
|
intro += f" with exit code {script_exit}"
|
|
intro += "."
|
|
if run_persistent:
|
|
intro += f" TerminalID: {terminal_id}."
|
|
body = clean_output if clean_output.strip() else "(empty output)"
|
|
return f"{intro}\nOutput:\n{body}"
|
|
|
|
# We spawn the host subprocess: adb -s <serial> shell <script>.
|
|
# The script travels as an argument so the remote shell exits with it;
|
|
# stdin is closed unless the caller wants to feed input later.
|
|
try:
|
|
process = await asyncio.create_subprocess_exec(
|
|
_adb_binary(),
|
|
"-s",
|
|
device_id,
|
|
"shell",
|
|
phone_script,
|
|
cwd=os.getcwd(),
|
|
env=run_env,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.STDOUT,
|
|
stdin=asyncio.subprocess.PIPE if interactive else asyncio.subprocess.DEVNULL,
|
|
)
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.error(f"Failed to spawn adb shell process for command '{cmd_line}': {e}")
|
|
return f"Failed to spawn adb process: {e}"
|
|
|
|
stdout_data: list[str] = []
|
|
try:
|
|
|
|
async def read_limit():
|
|
while True:
|
|
line = await process.stdout.readline()
|
|
if not line:
|
|
break
|
|
stdout_data.append(line.decode(errors="replace"))
|
|
await process.wait()
|
|
|
|
await asyncio.wait_for(read_limit(), timeout=wait_seconds)
|
|
except TimeoutError:
|
|
return await self._hand_off_to_background(
|
|
ctx,
|
|
registry,
|
|
process,
|
|
cmd_line,
|
|
cwd,
|
|
terminal_id,
|
|
interactive,
|
|
stdout_data,
|
|
)
|
|
|
|
# The process completed within the wait window.
|
|
exit_code = process.returncode
|
|
output_text = "".join(stdout_data)
|
|
|
|
clean_output = output_text
|
|
if run_persistent:
|
|
clean_output, parsed_env, script_exit = BackgroundTask.parse_output(output_text)
|
|
if parsed_env is not None:
|
|
registry.persistent_envs[terminal_id] = _filter_persistent_env(parsed_env)
|
|
if script_exit is not None:
|
|
exit_code = script_exit
|
|
|
|
intro = f"ADB command completed with exit code {exit_code}."
|
|
if run_persistent:
|
|
intro += f" TerminalID: {terminal_id}."
|
|
|
|
if _is_output_long(clean_output):
|
|
task_id = f"task_sync_{uuid.uuid4().hex[:8]}"
|
|
intro += f" TaskId: {task_id}."
|
|
registry.register_finished(
|
|
task_id=task_id,
|
|
command=cmd_line,
|
|
cwd=cwd,
|
|
terminal_id=terminal_id,
|
|
status="completed" if exit_code == 0 else "failed",
|
|
exit_code=exit_code,
|
|
output=clean_output,
|
|
)
|
|
# Already shown inline; do not re-announce it as newly finished.
|
|
registry.finished[task_id]["notified"] = True
|
|
return _format_long_output_response(task_id, clean_output, intro)
|
|
|
|
return f"{intro}\nOutput:\n{clean_output}"
|
|
|
|
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
|
async def _hand_off_to_background(
|
|
self,
|
|
ctx: ArtemisContext | None,
|
|
registry: AdbTaskRegistry,
|
|
process: AsyncProcess,
|
|
cmd_line: str,
|
|
cwd: str,
|
|
terminal_id: str | None,
|
|
interactive: bool,
|
|
stdout_so_far: list[str],
|
|
) -> str:
|
|
"""Registers a still-running process as a background task."""
|
|
task_id = f"task_{process.pid}_{uuid.uuid4().hex[:6]}"
|
|
trace_id = uuid.uuid4()
|
|
|
|
bg_task = BackgroundTask(
|
|
task_id=task_id,
|
|
command=cmd_line,
|
|
process=process,
|
|
terminal_id=terminal_id,
|
|
cwd=cwd,
|
|
interactive=interactive,
|
|
registry=registry,
|
|
)
|
|
bg_task.stdout_log.extend(stdout_so_far)
|
|
bg_task.trace_id = trace_id
|
|
registry.background[task_id] = bg_task
|
|
|
|
if ctx and getattr(ctx, "data_engine", None):
|
|
ctx.data_engine.register_background_task(
|
|
task_id=task_id,
|
|
summary=f"ADB Command: {cmd_line[:50]}...",
|
|
trace_id=trace_id,
|
|
)
|
|
if stdout_so_far:
|
|
ctx.data_engine.stream_output(trace_id, "".join(stdout_so_far))
|
|
|
|
await bg_task.start(ctx)
|
|
|
|
response_msg = f"ADB command was sent to the background as a task.\nTaskId: {task_id}\n"
|
|
if terminal_id:
|
|
response_msg += f"TerminalID: {terminal_id}\n"
|
|
response_msg += "Use the 'manage_task' tool to check its status or terminate it."
|
|
if not interactive:
|
|
response_msg += " (stdin is closed; relaunch with Interactive=true to send input.)"
|
|
return response_msg
|
|
|
|
|
|
# Universal tool instance & aliases
|
|
run_adb_command = RunAdbCommandTool()
|
|
RunAdbCommand = RunAdbCommandTool
|
|
|
|
|
|
def get_run_adb_command_tool(ctx: ArtemisContext) -> BaseTool:
|
|
"""Exports run_adb_command as a LangChain BaseTool."""
|
|
return trace_langchain_tool(run_adb_command.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# manage_task
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class ManageTaskArgs(BaseModel):
|
|
"""Arguments schema for managing background ADB tasks."""
|
|
|
|
Action: Literal["list", "list_finished", "kill", "status", "send_input"] = Field(
|
|
...,
|
|
description=(
|
|
"The action to perform: 'list' (list running tasks),"
|
|
" 'list_finished' (list completed tasks), 'kill', 'status',"
|
|
" 'send_input'."
|
|
),
|
|
)
|
|
TaskId: str | None = Field(
|
|
default=None,
|
|
description=("The task ID to manage (required for kill, status, and send_input)."),
|
|
)
|
|
Input: str | None = Field(
|
|
default=None,
|
|
description=("The input to send to the task (required when Action is 'send_input')."),
|
|
)
|
|
|
|
|
|
class ManageTaskTool(ArtemisTool):
|
|
"""Universal tool for managing background ADB shell tasks."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="manage_task",
|
|
description=("[SHELL] Manage background ADB shell tasks launched via run_adb_command."),
|
|
args_schema=ManageTaskArgs,
|
|
category="system",
|
|
)
|
|
|
|
# pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements,too-many-return-statements,too-many-positional-arguments
|
|
async def execute(
|
|
self,
|
|
driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument
|
|
ctx: ArtemisContext | None = None,
|
|
Action: Literal["list", "list_finished", "kill", "status", "send_input"] = "list",
|
|
TaskId: str | None = None,
|
|
Input: str | None = None,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
action = Action or kwargs.get("action", "list")
|
|
task_id = TaskId if TaskId is not None else kwargs.get("task_id")
|
|
input_str = Input if Input is not None else kwargs.get("input")
|
|
registry = get_adb_task_registry(ctx)
|
|
|
|
if action == "list":
|
|
active_tasks = []
|
|
for tid, t in registry.background.items():
|
|
active_tasks.append(
|
|
f"- {tid}: Command='{t.command[:40]}...', Phone"
|
|
f" Cwd='{t.cwd}', TerminalID='{t.terminal_id}',"
|
|
f" Interactive={t.interactive}"
|
|
)
|
|
if not active_tasks:
|
|
return "No active background tasks running."
|
|
return "Active background tasks:\n" + "\n".join(active_tasks)
|
|
|
|
if action == "list_finished":
|
|
finished_tasks = []
|
|
for tid, t in registry.finished.items():
|
|
finished_tasks.append(
|
|
f"- {tid}: Command='{t['command'][:40]}...', Status='{t['status']}'"
|
|
)
|
|
if not finished_tasks:
|
|
return "No recently finished tasks."
|
|
return f"Recently finished tasks (up to {_FINISHED_TASKS_CAP}):\n" + "\n".join(
|
|
finished_tasks
|
|
)
|
|
|
|
if not task_id:
|
|
return "Error: TaskId is required for status, kill, or send_input actions."
|
|
|
|
task = registry.background.get(task_id)
|
|
|
|
if action == "status":
|
|
task_info = _get_task_info(task_id, ctx)
|
|
if not task_info:
|
|
return f"Task {task_id} not found."
|
|
|
|
output_text = task_info.get("output", "")
|
|
if _ENV_START_MARKER in output_text:
|
|
output_text = output_text.split(_ENV_START_MARKER)[0]
|
|
|
|
intro = (
|
|
f"Task: {task_id}\n"
|
|
f"Status: {task_info.get('status')}\n"
|
|
f"Command: {task_info.get('command')}\n"
|
|
f"Cwd: {task_info.get('cwd')}\n"
|
|
f"TerminalID: {task_info.get('terminal_id')}"
|
|
)
|
|
exit_code = task_info.get("exit_code")
|
|
if exit_code is not None:
|
|
intro += f"\nExitCode: {exit_code}"
|
|
|
|
if _is_output_long(output_text):
|
|
return _format_long_output_response(task_id, output_text, intro)
|
|
|
|
return f"{intro}\nAccumulated output:\n{output_text}"
|
|
|
|
if action == "kill":
|
|
if not task:
|
|
return f"Task {task_id} is not active or already finished."
|
|
try:
|
|
await task.stop(ctx, status="killed")
|
|
return f"Task {task_id} successfully terminated."
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
return f"Failed to terminate task {task_id}: {e}"
|
|
|
|
if action == "send_input":
|
|
if not task:
|
|
return f"Task {task_id} is not active."
|
|
if not input_str:
|
|
return "Error: Input is required for send_input action."
|
|
if not task.interactive or task.process.stdin is None:
|
|
return (
|
|
f"Error: Task {task_id} was launched without Interactive=true, so its"
|
|
" stdin is closed. Relaunch the command with Interactive=true."
|
|
)
|
|
try:
|
|
task.process.stdin.write(input_str.encode())
|
|
await task.process.stdin.drain()
|
|
return f"Input successfully sent to task {task_id}."
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
return f"Failed to send input: {e}"
|
|
|
|
return "Unsupported action."
|
|
|
|
|
|
# Universal tool instance & aliases
|
|
manage_task = ManageTaskTool()
|
|
ManageTask = ManageTaskTool
|
|
|
|
|
|
def get_manage_task_tool(ctx: ArtemisContext) -> BaseTool:
|
|
"""Exports manage_task as a LangChain BaseTool."""
|
|
return trace_langchain_tool(manage_task.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# run_short_adb_command (Diagnoser)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_SHORT_COMMAND_TIMEOUT_SECONDS = 3.0
|
|
|
|
|
|
class RunShortAdbCommandArgs(BaseModel):
|
|
"""Arguments schema for short ADB shell commands."""
|
|
|
|
CommandLine: str = Field(
|
|
...,
|
|
description=(
|
|
"The exact shell command line string to execute inside the Android device shell."
|
|
),
|
|
)
|
|
|
|
|
|
class RunShortAdbCommandTool(ArtemisTool):
|
|
"""Universal tool for executing short, synchronous ADB shell commands."""
|
|
|
|
def __init__(self, name: str = "run_short_adb_command"):
|
|
super().__init__(
|
|
name=name,
|
|
description=(
|
|
"[SHELL] Executes a short, synchronous shell command directly on the Android "
|
|
"mobile device via ADB shell. Returns the command output. Times out after 3 "
|
|
"seconds, so bound streaming commands (logcat -d, top -n 1). Do not use this "
|
|
"tool to execute actions that change the phone state such as clicking, "
|
|
"swiping, or navigating back."
|
|
),
|
|
args_schema=RunShortAdbCommandArgs,
|
|
category="system",
|
|
)
|
|
|
|
# pylint: disable=too-many-locals,too-many-return-statements
|
|
async def execute(
|
|
self,
|
|
driver: BaseDeviceDriver | None = None,
|
|
ctx: ArtemisContext | None = None,
|
|
CommandLine: str | None = None,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
cmd_line = (
|
|
CommandLine
|
|
if CommandLine is not None
|
|
else (kwargs.get("command_line") or kwargs.get("command") or kwargs.get("cmd") or "")
|
|
)
|
|
timeout = _SHORT_COMMAND_TIMEOUT_SECONDS
|
|
|
|
# Fallback to driver's execute_shell if no full ArtemisContext is available
|
|
if (
|
|
(ctx is None or not hasattr(ctx, "device") or not ctx.device)
|
|
and driver is not None
|
|
and hasattr(driver, "execute_shell")
|
|
):
|
|
return await driver.execute_shell(cmd_line, timeout_seconds=timeout)
|
|
|
|
device_id = None
|
|
if ctx and hasattr(ctx, "device") and ctx.device:
|
|
device_id = getattr(ctx.device, "device_id", None)
|
|
elif driver and hasattr(driver, "device_id"):
|
|
device_id = getattr(driver, "device_id", None)
|
|
if not device_id:
|
|
device_id = "default_device"
|
|
|
|
run_env = os.environ.copy()
|
|
phone_script = f"cd /data/local/tmp\n{cmd_line}\n"
|
|
|
|
adb_client = getattr(ctx, "adb_client", None) if ctx else None
|
|
if os.environ.get("ARTEMIS_CLOUD_MODE") == "1" or (
|
|
adb_client is not None and hasattr(adb_client, "_bridge")
|
|
):
|
|
try:
|
|
device = get_adb_device(ctx)
|
|
output = await asyncio.wait_for(
|
|
asyncio.to_thread(device.shell, phone_script),
|
|
timeout=timeout,
|
|
)
|
|
return f"ADB command completed.\nOutput:\n{output}"
|
|
except TimeoutError:
|
|
return f"Error: Command timed out after {timeout:.0f} seconds."
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
return f"Error running command: {e}"
|
|
|
|
try:
|
|
process = await asyncio.create_subprocess_exec(
|
|
_adb_binary(),
|
|
"-s",
|
|
device_id,
|
|
"shell",
|
|
phone_script,
|
|
cwd=os.getcwd(),
|
|
env=run_env,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.STDOUT,
|
|
stdin=asyncio.subprocess.DEVNULL,
|
|
)
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
logger.error(f"Failed to spawn adb shell process for command '{cmd_line}': {e}")
|
|
return f"Failed to spawn adb process: {e}"
|
|
|
|
try:
|
|
stdout_data, _ = await asyncio.wait_for(process.communicate(), timeout=timeout)
|
|
exit_code = process.returncode
|
|
output_text = stdout_data.decode(errors="replace")
|
|
|
|
return f"ADB command completed with exit code {exit_code}.\nOutput:\n{output_text}"
|
|
except TimeoutError:
|
|
try:
|
|
process.terminate()
|
|
await asyncio.wait_for(process.wait(), timeout=_TERMINATE_GRACE_SECONDS)
|
|
except Exception: # pylint: disable=broad-exception-caught
|
|
try:
|
|
process.kill()
|
|
except (ProcessLookupError, OSError):
|
|
pass
|
|
return f"Error: Command timed out after {timeout:.0f} seconds."
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
return f"Error running command: {e}"
|
|
|
|
|
|
# Universal tool instance & aliases
|
|
run_short_adb_command = RunShortAdbCommandTool()
|
|
RunShortAdbCommand = RunShortAdbCommandTool
|
|
|
|
|
|
def get_run_short_adb_command_tool(ctx: ArtemisContext) -> BaseTool:
|
|
"""Exports run_short_adb_command as a LangChain BaseTool."""
|
|
return trace_langchain_tool(
|
|
run_short_adb_command.to_langchain_tool(ctx, name="run_adb_command"),
|
|
ctx,
|
|
)
|
|
|
|
|
|
# Expose Wrappers
|
|
run_adb_command_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_run_adb_command_tool,
|
|
on_success_fn=lambda output: output,
|
|
on_failure_fn=lambda err: f"ADB command failed: {err}",
|
|
)
|
|
|
|
manage_task_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_manage_task_tool,
|
|
on_success_fn=lambda output: output,
|
|
on_failure_fn=lambda err: f"Task management failed: {err}",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# analyze_task_output
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AnalyzeTaskOutputArgs(BaseModel):
|
|
"""Arguments schema for output analysis requests."""
|
|
|
|
TaskId: str = Field(
|
|
...,
|
|
description="The task ID whose output you want to analyze.",
|
|
)
|
|
Query: str = Field(
|
|
...,
|
|
description=(
|
|
"The question or query to ask the output analyzer about the task's full output log."
|
|
),
|
|
)
|
|
|
|
|
|
class AnalyzeTaskOutputTool(ArtemisTool):
|
|
"""Universal tool for analyzing the full, long output of an ADB task."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="analyze_task_output",
|
|
description=(
|
|
"[ANALYSIS] Analyzes the full, long output of a completed or running ADB task. "
|
|
"Use this tool when you receive a message indicating that a task's output "
|
|
"is too long and was truncated."
|
|
),
|
|
args_schema=AnalyzeTaskOutputArgs,
|
|
category="system",
|
|
)
|
|
|
|
async def execute(
|
|
self,
|
|
driver: BaseDeviceDriver | None = None, # pylint: disable=unused-argument
|
|
ctx: ArtemisContext | None = None,
|
|
TaskId: str | None = None,
|
|
Query: str | None = None,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
task_id = TaskId or kwargs.get("task_id") or ""
|
|
query = Query or kwargs.get("query") or ""
|
|
|
|
task_info = _get_task_info(task_id, ctx)
|
|
if not task_info:
|
|
return f"Error: Task {task_id} not found."
|
|
|
|
analyzer = TaskOutputAnalyzerNode(ctx)
|
|
result = await analyzer.run(
|
|
command=task_info.get("command", ""),
|
|
output_text=task_info.get("output", ""),
|
|
query=query,
|
|
)
|
|
return f"Analysis result for Task {task_id}:\n{result}"
|
|
|
|
|
|
# Universal tool instance & aliases
|
|
analyze_task_output = AnalyzeTaskOutputTool()
|
|
AnalyzeTaskOutput = AnalyzeTaskOutputTool
|
|
|
|
|
|
def get_analyze_task_output_tool(ctx: ArtemisContext) -> BaseTool:
|
|
"""Exports analyze_task_output as a LangChain BaseTool."""
|
|
return trace_langchain_tool(analyze_task_output.to_langchain_tool(ctx), ctx)
|
|
|
|
|
|
analyze_task_output_wrapper = ToolWrapper(
|
|
tool_fn_getter=get_analyze_task_output_tool,
|
|
on_success_fn=lambda output: output,
|
|
on_failure_fn=lambda err: f"Task output analysis failed: {err}",
|
|
)
|