Files
Artemis-Eapil/mcp_server/tools/inspect_trace.py
T
somew1nd c2c90a6378 feat: add multi-device concurrency, server lifecycle CLI, and unified execution stream
- Implement DevicePool and DeviceExecutionLock concurrency modes (per-device FIFO vs global serialization)
- Extend ArtemisClient SDK and CLI with explicit device targeting and concurrency strategy configuration
- Add server lifecycle management (artemis restart, stop, status) with cross-platform PID discovery
- Update MCP server tools to support multi-device dispatch, trace tagging, and per-device execution
- Enhance Showcase UI with step-by-step visual replay, playback speed controls, and device serial badges
- Add comprehensive unit test coverage and edge-case verification scripts
2026-08-26 15:53:14 -07:00

496 lines
21 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.
"""MCP Tool: mobile_inspect_trace."""
import json
import math
import os
import sqlite3
import sys
from typing import Any
from PIL import Image, ImageDraw
from mcp_server.base import mcp
from mcp_server.utils import env_utils, trace_store
from artemis.utils.task_tree import (
_render_step_detailed,
build_plan_and_history,
)
def _draw_action_overlay(bg_image_path: str, action_obj: dict, output_path: str) -> bool:
"""Draws the action (tap, swipe, etc.) on the image and saves it."""
try:
with Image.open(bg_image_path) as img:
draw_img = img.convert("RGBA")
draw = ImageDraw.Draw(draw_img)
act_type = action_obj.get("action") or action_obj.get("name")
coords = action_obj.get("coordinates") or action_obj.get("target")
if not act_type or not coords:
return False
fill_color = (255, 0, 0, 80)
outline_color = (255, 0, 0, 255)
outline_width = 4
if act_type in (
"tap",
"click",
"long_press",
"focus_and_input_text",
"focus_and_clear_text",
):
if isinstance(coords, list) and len(coords) == 2:
x, y = coords[0], coords[1]
r = 25
draw.ellipse(
[x - r, y - r, x + r, y + r],
fill=fill_color,
outline=outline_color,
width=outline_width,
)
r_center = 4
draw.ellipse(
[
x - r_center,
y - r_center,
x + r_center,
y + r_center,
],
fill=outline_color,
outline=outline_color,
)
elif act_type == "swipe":
if isinstance(coords, list) and len(coords) == 4:
x1, y1, x2, y2 = coords[0], coords[1], coords[2], coords[3]
draw.line([x1, y1, x2, y2], fill=outline_color, width=6)
r_start = 8
draw.ellipse(
[
x1 - r_start,
y1 - r_start,
x1 + r_start,
y1 + r_start,
],
fill=outline_color,
)
dx = x2 - x1
dy = y2 - y1
angle = math.atan2(dy, dx)
arrow_length = 30
arrow_angle = math.pi / 6
p1_x = x2 - arrow_length * math.cos(angle - arrow_angle)
p1_y = y2 - arrow_length * math.sin(angle - arrow_angle)
p2_x = x2 - arrow_length * math.cos(angle + arrow_angle)
p2_y = y2 - arrow_length * math.sin(angle + arrow_angle)
draw.polygon([x2, y2, p1_x, p1_y, p2_x, p2_y], fill=outline_color)
final_img = draw_img.convert("RGB")
final_img.save(output_path, "JPEG", quality=85)
return True
except Exception as e:
print(f"Error drawing action overlay: {e}", file=sys.stderr)
return False
@mcp.tool()
async def mobile_inspect_trace(action: str, trace_id: str, step_number: int | None = None) -> Any:
"""A specialized trace data retrieval tool designed to extract execution details of mobile automation tasks.
Use this tool DURING a running task or AFTER it has finished (or failed) to
monitor progress, verify answers, diagnose agent errors, and debug.
It provides comprehensive runtime information (including the assigned
`device_serial`), raw agent outputs, and visual screenshots for both
**Flash** and **Pro** tasks.
### Available Actions:
- **'view_summary'**: Retrieves a high-level execution summary across all
steps, including the target `device_serial`.
* **Pro Model**: Displays a hierarchical task plan summary, showing the
status, device serial, and descriptions of all completed, active, and pending steps.
* **Flash Model**: Displays an all-in-one execution chain containing the
complete reasoning (thoughts/motivation), device serial, and action taken (`click`,
`input_text`, etc.) for every single step.
- **'view_step_screenshots'**: Retrieves local file paths of screenshots for
a specific step (supports both Flash and Pro), along with `device_serial`. Returns:
1. `before_screenshot`: The original device screenshot observed by the
agent right before taking the action.
2. `after_screenshot`: A post-action screenshot capturing the failure or
mismatch state (if the action failed or was intercepted). Note: This is
conditional/optional and will be `None` (`null`) for normal/successful
steps where no extra post-action screenshot was captured.
3. `action_overlay_screenshot`: The `before_screenshot` with the agent's
action visually marked (e.g., red circle for taps, red arrow for
swipes). CRITICAL for verifying if the agent tapped the correct UI
element.
4. `device_serial`: The serial of the device on which this step was executed.
- **'view_step_details'**: Retrieves deep runtime logs, raw VLM thoughts,
and exact actions for a specific step, along with `device_serial`.
* **Pro Model**: Fully supported. Returns detailed per-step diagnostic
data.
* **Flash Model**: DO NOT USE. (Flash reasoning and actions are already
fully included in 'view_summary', making this action redundant).
Args:
action: STR. **REQUIRED**. The retrieval action to perform
(`"view_summary"`, `"view_step_screenshots"`, or
`"view_step_details"`).
trace_id: STR. **REQUIRED**. The unique session identifier of the task
(returned by `mobile_run_task`).
step_number: INT. **CONDITIONAL**. The 1-indexed step number to query. -
You MUST provide this if `action="view_step_screenshots"` or
`action="view_step_details"`. - You MUST omit or leave this empty if
`action="view_summary"`.
"""
project_root = env_utils.get_project_root()
db_path = os.path.join(trace_store.TRACES_DIR, "data_engine.db")
if not os.path.exists(db_path):
db_path = os.path.join(project_root, "traces", "data_engine.db")
if not os.path.exists(db_path):
return {
"error": "Database not found",
"message": f"Data engine database does not exist at {db_path} yet.",
}
# 1. Handle Action: "view_summary"
if action == "view_summary":
plan_content = None
plan_path = os.path.join(trace_store.get_trace_dir(trace_id), "notes", "task_plan.md")
if not os.path.exists(plan_path):
plan_path = os.path.join(trace_store.TRACES_DIR, "notes", "task_plan.md")
if os.path.exists(plan_path):
try:
with open(plan_path, encoding="utf-8") as f:
plan_content = f.read()
except Exception:
pass
try:
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT start_time FROM sessions WHERE session_id = ?",
(trace_id,),
)
session_row = cursor.fetchone()
session_start_time = session_row["start_time"] if session_row else None
cursor.execute(
"SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC",
(trace_id,),
)
rows = cursor.fetchall()
steps = []
for r in rows:
step_dict = dict(r)
for col in ("action_taken", "last_execution_result", "extra_metadata"):
if step_dict.get(col):
try:
step_dict[col] = json.loads(step_dict[col])
except Exception:
pass
cursor.execute(
"SELECT type, name, status, timestamp, duration, payload FROM traces WHERE step_id = ? ORDER BY timestamp ASC",
(step_dict["step_id"],),
)
trace_rows = cursor.fetchall()
events = []
for tr in trace_rows:
t_dict = dict(tr)
if t_dict.get("payload"):
try:
t_dict["payload"] = json.loads(t_dict["payload"])
except Exception:
pass
events.append(t_dict)
step_dict["interleaved_events"] = events
if session_start_time and step_dict.get("timestamp"):
relative_time_val = step_dict["timestamp"] - session_start_time
step_dict["relative_time"] = f"{relative_time_val:.1f}s"
else:
step_dict["relative_time"] = "N/A"
steps.append(step_dict)
conn.close()
except Exception as e:
return {
"error": "Database query error",
"message": f"Failed to retrieve steps from database: {e}",
}
status_data = trace_store.read_status(trace_id)
is_flash = status_data and status_data.get("model", "").lower() == "flash"
device_serial = status_data.get("device_serial") if status_data else None
if not device_serial and os.path.exists(db_path):
try:
conn_dev = sqlite3.connect(db_path)
conn_dev.row_factory = sqlite3.Row
cur_dev = conn_dev.cursor()
cur_dev.execute(
"SELECT device_info FROM sessions WHERE session_id = ? LIMIT 1",
(trace_id,),
)
row_dev = cur_dev.fetchone()
if row_dev and row_dev["device_info"]:
try:
d_info = json.loads(row_dev["device_info"])
if isinstance(d_info, dict) and d_info.get("device_id"):
device_serial = d_info["device_id"]
except Exception:
pass
conn_dev.close()
except Exception:
pass
device_info_str = f" | **Device Serial:** `{device_serial}`" if device_serial else ""
if is_flash:
summary_lines = [
"# ARTEMIS Flash Execution Summary",
f"**Session ID:** `{trace_id}` | **Model:** Flash{device_info_str}\n",
"---",
"## Step-by-Step Execution Chain\n",
]
if not steps:
summary_lines.append("_No action steps have been recorded yet._")
else:
for s in steps:
s_num = s.get("step_number", "?")
rel_t = s.get("relative_time", "N/A")
dur_val = s.get("duration")
if dur_val is None and s.get("interleaved_events"):
events = s["interleaved_events"]
timestamps = [e.get("timestamp", 0) for e in events if e.get("timestamp")]
if timestamps:
dur_val = max(timestamps) - min(timestamps)
dur_str = f"{dur_val:.1f}s" if dur_val is not None and dur_val > 0 else "< 0.5s"
thought = s.get("operator_raw_thinking")
if not thought or not thought.strip():
for ev in s.get("interleaved_events", []):
if ev.get("type") in ("raw_thinking", "llm_call") and ev.get("payload"):
p = ev["payload"]
if isinstance(p, dict) and p.get("response_text"):
thought = p["response_text"]
break
if not thought or not thought.strip():
thought = "(Direct action without preliminary text reasoning)"
act = s.get("action_taken") or "N/A"
if isinstance(act, dict) and "action" in act:
act_str = f"{act['action']}({act.get('args', act.get('coordinates', ''))})"
else:
act_str = str(act)
res = s.get("last_execution_result") or "N/A"
summary_lines.append(f"### Step {s_num} (at {rel_t}, duration {dur_str})")
summary_lines.append(f"- **Thinking / Motivation:** {thought}")
summary_lines.append(f"- **Action Taken:** `{act_str}` -> `{res}`\n")
return "\n".join(summary_lines)
try:
summary_markdown = build_plan_and_history(
task_plan=plan_content or "No task plan created yet.",
steps=steps,
current_subgoal_hash="default",
min_summaries=len(steps),
last_n_detailed=0,
all_detailed=False,
)
if device_serial and not summary_markdown.startswith("**Session ID:"):
summary_markdown = f"**Session ID:** `{trace_id}` | **Model:** Pro{device_info_str}\n\n" + summary_markdown
return summary_markdown
except Exception as e:
return {
"error": "Formatting error",
"message": f"Failed to format plan and history using Artemis task tree: {e}",
}
# 2. Handle Actions requiring step_number
elif action in ("view_step_screenshots", "view_step_details"):
status_data = trace_store.read_status(trace_id)
is_flash = status_data and status_data.get("model", "").lower() == "flash"
device_serial = status_data.get("device_serial") if status_data else None
if action == "view_step_details" and is_flash:
return {
"error": "Not applicable for Flash model",
"message": (
"Action 'view_step_details' is not needed for Flash model runs. "
"For Flash tasks, 'view_summary' already outputs the complete "
"step-by-step reasoning (thoughts) and executed actions for the entire run."
),
}
if step_number is None:
return {
"error": "Missing parameter",
"message": f"Parameter 'step_number' is required for action '{action}'.",
}
try:
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
cursor.execute(
"SELECT * FROM steps WHERE session_id = ? AND step_number = ?",
(trace_id, step_number),
)
row = cursor.fetchone()
if not row:
conn.close()
return {
"error": "Step not found",
"message": f"Step number {step_number} not found for trace '{trace_id}'.",
}
step_dict = dict(row)
for col in ("action_taken", "last_execution_result", "extra_metadata"):
if step_dict.get(col):
try:
step_dict[col] = json.loads(step_dict[col])
except Exception:
pass
cursor.execute(
"SELECT start_time FROM sessions WHERE session_id = ?",
(trace_id,),
)
session_row = cursor.fetchone()
session_start_time = session_row["start_time"] if session_row else None
cursor.execute(
"SELECT type, name, status, timestamp, duration, payload FROM traces WHERE step_id = ? ORDER BY timestamp ASC",
(step_dict["step_id"],),
)
trace_rows = cursor.fetchall()
events = []
for tr in trace_rows:
t_dict = dict(tr)
if t_dict.get("payload"):
try:
t_dict["payload"] = json.loads(t_dict["payload"])
except Exception:
pass
events.append(t_dict)
step_dict["interleaved_events"] = events
conn.close()
except Exception as e:
return {
"error": "Database error",
"message": f"Database query failed: {e}",
}
# 3. Handle Action: "view_step_screenshots"
if action == "view_step_screenshots":
images_dir = os.path.join(trace_store.TRACES_DIR, "images")
if not os.path.exists(images_dir):
images_dir = os.path.join(project_root, "traces", "images")
pre_image = None
post_image = None
overlay_image = None
if step_dict.get("pre_image_name"):
name = step_dict["pre_image_name"]
if not name.endswith(".jpg"):
name += ".jpg"
path = os.path.join(images_dir, name)
if os.path.exists(path):
pre_image = f"file://{path}"
overlay_dir = trace_store.get_trace_dir(trace_id)
os.makedirs(overlay_dir, exist_ok=True)
overlay_filename = f"step_{step_dict['step_number']}_overlay.jpg"
overlay_abs_path = os.path.join(overlay_dir, overlay_filename)
action_obj = step_dict.get("action_taken")
action_parsed = (
action_obj[0] if isinstance(action_obj, list) and action_obj else action_obj
)
if isinstance(action_parsed, dict) and action_parsed.get("action"):
if not os.path.exists(overlay_abs_path):
success_draw = _draw_action_overlay(
path, action_parsed, overlay_abs_path
)
if success_draw:
overlay_image = f"file://{overlay_abs_path}"
else:
overlay_image = f"file://{overlay_abs_path}"
if step_dict.get("post_image_name"):
name = step_dict["post_image_name"]
if not name.endswith(".jpg"):
name += ".jpg"
path = os.path.join(images_dir, name)
if os.path.exists(path):
post_image = f"file://{path}"
return {
"trace_id": trace_id,
"device_serial": device_serial,
"step_number": step_dict["step_number"],
"before_screenshot": pre_image,
"after_screenshot": post_image,
"action_overlay_screenshot": overlay_image,
}
# 4. Handle Action: "view_step_details"
elif action == "view_step_details":
relative_time_str = "N/A"
if session_start_time and step_dict.get("timestamp"):
relative_time_val = step_dict["timestamp"] - session_start_time
relative_time_str = f"{relative_time_val:.1f}s"
rendered_text = _render_step_detailed(
step=step_dict,
relative_time=relative_time_str,
summary=step_dict.get("summary"),
action=step_dict.get("action_taken"),
result=step_dict.get("last_execution_result"),
is_most_recent=False,
for_failure_analyzer=False,
)
return {
"trace_id": trace_id,
"device_serial": device_serial,
"step_number": step_dict["step_number"],
"details": rendered_text,
}
else:
return {
"error": "Invalid action",
"message": (
f"Action '{action}' is not supported. Supported actions: "
"'view_summary', 'view_step_screenshots', 'view_step_details'."
),
}