# 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'." ), }