From c1803fcb86fd1a28ba2571c3246dcd9f80ffc779 Mon Sep 17 00:00:00 2001 From: Weilin Liu Date: Thu, 13 Aug 2026 12:11:10 -0700 Subject: [PATCH] feat(console): add Artemis Admin & Trace Console backend and management plane Introduces the comprehensive developer/administrator control plane and trace inspection console backend for the ARTEMIS mobile automation platform. Key Features & Backend Architecture: - Core Server & Modular Routing: * FastAPI modular architecture with dedicated routers for sessions, steps, tasks, media, stream, and replay. * Multi-layer repository pattern with SQLite persistence for sessions, traces, and step data. - Task Control & Queue Management: * FIFO asynchronous task queue service with support for single/batch goal execution and profile selection (Flash / Pro). * Session lifecycle controls (run, pause, resume, stop, clear queue). * Startup orphan cleanup and automatic older replay archiving. - Live Telemetry & IPC Communication: * High-performance IPC server for streaming real-time execution logs, agent thoughts, and tool invocations. * Server-Sent Events (SSE) /api/stream/{session_id} endpoint for live UI updates. - Real-Time Device Streaming: * Live ADB screen capture and MJPEG frame generator (/api/stream/device-live). * Device connectivity and state inspection endpoint (/api/stream/device-state). - Trace Inspection & Step Replay Engine: * Hierarchical call tree inspection, pre/post step screenshots, and action overlay metadata. * Isolated single-step replay manager (/api/replay) for interactive debugging and validation. --- apps/admin_console/README.md | 51 + apps/admin_console/__init__.py | 5 + apps/admin_console/core/__init__.py | 1 + apps/admin_console/core/config.py | 68 + apps/admin_console/core/state.py | 82 + apps/admin_console/database/__init__.py | 1 + apps/admin_console/database/connection.py | 43 + .../database/repositories/__init__.py | 1 + .../repositories/session_repository.py | 202 + .../database/repositories/step_repository.py | 411 ++ .../database/repositories/trace_repository.py | 131 + apps/admin_console/index.html | 3467 +++++++++++++++++ apps/admin_console/replay_manager.py | 2567 ++++++++++++ apps/admin_console/routers/__init__.py | 1 + apps/admin_console/routers/media.py | 116 + apps/admin_console/routers/replay.py | 91 + apps/admin_console/routers/sessions.py | 127 + apps/admin_console/routers/steps.py | 108 + apps/admin_console/routers/stream.py | 56 + apps/admin_console/routers/tasks.py | 189 + apps/admin_console/schemas/__init__.py | 1 + apps/admin_console/schemas/task_schema.py | 28 + apps/admin_console/server.py | 154 + apps/admin_console/services/__init__.py | 1 + .../services/device_stream_service.py | 133 + apps/admin_console/services/ipc_service.py | 204 + apps/admin_console/services/media_service.py | 261 ++ apps/admin_console/services/model_service.py | 133 + .../services/task_queue_service.py | 256 ++ 29 files changed, 8889 insertions(+) create mode 100644 apps/admin_console/README.md create mode 100644 apps/admin_console/__init__.py create mode 100644 apps/admin_console/core/__init__.py create mode 100644 apps/admin_console/core/config.py create mode 100644 apps/admin_console/core/state.py create mode 100644 apps/admin_console/database/__init__.py create mode 100644 apps/admin_console/database/connection.py create mode 100644 apps/admin_console/database/repositories/__init__.py create mode 100644 apps/admin_console/database/repositories/session_repository.py create mode 100644 apps/admin_console/database/repositories/step_repository.py create mode 100644 apps/admin_console/database/repositories/trace_repository.py create mode 100644 apps/admin_console/index.html create mode 100644 apps/admin_console/replay_manager.py create mode 100644 apps/admin_console/routers/__init__.py create mode 100644 apps/admin_console/routers/media.py create mode 100644 apps/admin_console/routers/replay.py create mode 100644 apps/admin_console/routers/sessions.py create mode 100644 apps/admin_console/routers/steps.py create mode 100644 apps/admin_console/routers/stream.py create mode 100644 apps/admin_console/routers/tasks.py create mode 100644 apps/admin_console/schemas/__init__.py create mode 100644 apps/admin_console/schemas/task_schema.py create mode 100644 apps/admin_console/server.py create mode 100644 apps/admin_console/services/__init__.py create mode 100644 apps/admin_console/services/device_stream_service.py create mode 100644 apps/admin_console/services/ipc_service.py create mode 100644 apps/admin_console/services/media_service.py create mode 100644 apps/admin_console/services/model_service.py create mode 100644 apps/admin_console/services/task_queue_service.py diff --git a/apps/admin_console/README.md b/apps/admin_console/README.md new file mode 100644 index 0000000..fd63072 --- /dev/null +++ b/apps/admin_console/README.md @@ -0,0 +1,51 @@ +# 🛠️ Artemis Admin & Trace Console + +A comprehensive developer/administrator control plane and trace inspection console for the Latte/Artemis mobile automation system. + +--- + +## 🎯 Key Features + +1. **Full Trace Inspection (`/api/sessions`, `/api/steps`)**: + - Hierarchical Call Tree exploration (LLM thoughts, tool calls, agent motivations). + - Step-by-step pre/post screenshots and action overlays (tap circles, swipe paths). + - Streaming logs and video recording playback. +2. **Task Control Center (`/api/run`)**: + - Directly launch autonomous tasks with custom goals and model profiles (Flash / Pro). + - Pause, resume, and manage FIFO task execution queues. +3. **Step Replay & Debugging Engine (`/api/replay`)**: + - Replay individual steps with isolated inputs, state inspection, and live device execution. + +--- + +## 🚀 How to Start + +### Option 1: Using `uv` (Recommended) +From the repository root: +```bash +uv run uvicorn apps.admin_console.server:app --host 0.0.0.0 --port 8000 --reload +``` +Or directly using Python module: +```bash +uv run python -m apps.admin_console.server +``` + +### Option 2: From the `apps/admin_console` directory +```bash +cd apps/admin_console +uv run uvicorn server:app --host 0.0.0.0 --port 8000 --reload +``` + +Once started, open your browser and navigate to: +👉 **`http://localhost:8000`** + +--- + +## 🎥 Video Debugging Mode + +By default, `VideoAnalyzer` cleans up local temporary video files after execution. To retain video recordings for full playback in the console: + +```bash +KEEP_VIDEOS=1 artemis run "your task goal" +``` +Or set `ARTEMIS_DEBUG=1` in your environment. diff --git a/apps/admin_console/__init__.py b/apps/admin_console/__init__.py new file mode 100644 index 0000000..495bb41 --- /dev/null +++ b/apps/admin_console/__init__.py @@ -0,0 +1,5 @@ +"""Artemis Admin & Trace Console Package.""" + +from apps.admin_console.server import app + +__all__ = ["app"] diff --git a/apps/admin_console/core/__init__.py b/apps/admin_console/core/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/core/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/core/config.py b/apps/admin_console/core/config.py new file mode 100644 index 0000000..7f80a96 --- /dev/null +++ b/apps/admin_console/core/config.py @@ -0,0 +1,68 @@ +# 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. + +import os +from pathlib import Path +import sys + + +# Automatically write ANTIGRAVITY_LS_ADDRESS to shared synchronization file on startup +# to allow decoupled background processes to communicate with Jetski +def init_ls_address(): + ls_addr = os.environ.get("ANTIGRAVITY_LS_ADDRESS") + if ls_addr: + try: + from artemis.config import write_ls_address + + write_ls_address(ls_addr) + except Exception as e: + sys.stderr.write(f"Failed to write .jetski_ls_address on startup: {e}\n") + + +# Compute real WORKSPACE_ROOT (find repo root containing pyproject.toml or artemis) +_current_p = Path(__file__).resolve().parent +while _current_p != _current_p.parent: + if (_current_p / "pyproject.toml").exists() or (_current_p / "artemis").is_dir(): + WORKSPACE_ROOT = _current_p + break + _current_p = _current_p.parent +else: + WORKSPACE_ROOT = Path(__file__).resolve().parent.parent.parent.parent + +_apps_dir = WORKSPACE_ROOT / "apps" +_admin_console_dir = _apps_dir / "admin_console" +_cloud_service_dir = _apps_dir / "cloud_service" + +for _p in (str(WORKSPACE_ROOT), str(_apps_dir), str(_admin_console_dir), str(_cloud_service_dir)): + if _p not in sys.path: + sys.path.insert(0, _p) + +# Try to import settings from artemis +try: + from artemis.config import settings + + TRACES_PATH = Path(settings.TRACES_PATH) + DB_PATH = Path(settings.DATA_ENGINE_DB_PATH) +except ImportError: + # Fallback to default path if run outside project environment or imports fail + TRACES_PATH = WORKSPACE_ROOT / "traces" + DB_PATH = TRACES_PATH / "data_engine.db" + +# Step Replay Directories (Stored within traces/replay to avoid workspace clutter) +REPLAY_BASE_DIR = TRACES_PATH / "replay" +TEST_DATA_DIR = REPLAY_BASE_DIR / "data" +TEST_OUTPUTS_DIR = REPLAY_BASE_DIR / "outputs" + +IMAGES_DIR = TRACES_PATH / "images" +PAUSE_FILE = WORKSPACE_ROOT / ".artemis_paused" diff --git a/apps/admin_console/core/state.py b/apps/admin_console/core/state.py new file mode 100644 index 0000000..d37dd67 --- /dev/null +++ b/apps/admin_console/core/state.py @@ -0,0 +1,82 @@ +# 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. + +import asyncio +from collections.abc import Callable +from typing import Any + +try: + from admin_console.core.config import PAUSE_FILE +except ImportError: + from apps.admin_console.core.config import PAUSE_FILE + + +class ServerState: + """Encapsulates all runtime states of the debug server.""" + + def __init__(self): + self.ipc_subscribers: list[Callable[[str, Any], None]] = [] + self.ipc_server: asyncio.Server | None = None + self.ipc_port: int | None = None + + self.current_process: asyncio.subprocess.Process | None = None + self.current_goal: str | None = None + self.current_profile: str | None = None + self.active_connections: dict[str, dict[str, Any]] = {} + self.active_session_id: str | None = None + self.was_stopped_manually: bool = False + + # FIFO Task queue + self.task_queue: asyncio.Queue = asyncio.Queue() + self.queue_tasks: list[dict[str, Any]] = [] + self.worker_task: asyncio.Task | None = None + + @property + def queue_goals(self) -> list[str]: + """Backward compatibility helper for queue goals.""" + return [t.get("goal", "") for t in self.queue_tasks if isinstance(t, dict)] + + @queue_goals.setter + def queue_goals(self, val: list[Any]): + # Allow setting if legacy code updates it + pass + + @property + def is_running(self) -> bool: + return self.current_process is not None and self.current_process.returncode is None + + @property + def is_paused(self) -> bool: + return PAUSE_FILE.exists() + + def add_subscriber(self, callback: Callable[[str, Any], None]): + if callback not in self.ipc_subscribers: + self.ipc_subscribers.append(callback) + + def remove_subscriber(self, callback: Callable[[str, Any], None]): + if callback in self.ipc_subscribers: + self.ipc_subscribers.remove(callback) + + def clear_queue(self): + while not self.task_queue.empty(): + try: + self.task_queue.get_nowait() + self.task_queue.task_done() + except (asyncio.QueueEmpty, ValueError): + break + self.queue_tasks.clear() + + +# Global shared instance +state = ServerState() diff --git a/apps/admin_console/database/__init__.py b/apps/admin_console/database/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/database/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/database/connection.py b/apps/admin_console/database/connection.py new file mode 100644 index 0000000..5de31d3 --- /dev/null +++ b/apps/admin_console/database/connection.py @@ -0,0 +1,43 @@ +# 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. + +from contextlib import contextmanager +from pathlib import Path +import sqlite3 +from fastapi import HTTPException + +try: + from admin_console.core.config import DB_PATH +except ImportError: + from apps.admin_console.core.config import DB_PATH + + +def get_db(db_path: Path | None = None) -> sqlite3.Connection: + """Returns a SQLite connection with timeout and Row row_factory.""" + path = db_path or DB_PATH + if not path.exists(): + raise HTTPException(status_code=500, detail=f"Database not found at {path}") + conn = sqlite3.connect(path, timeout=10.0) + conn.row_factory = sqlite3.Row + return conn + + +@contextmanager +def db_session(db_path: Path | None = None): + """Context manager for SQLite connections ensuring clean closure.""" + conn = get_db(db_path) + try: + yield conn + finally: + conn.close() diff --git a/apps/admin_console/database/repositories/__init__.py b/apps/admin_console/database/repositories/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/database/repositories/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/database/repositories/session_repository.py b/apps/admin_console/database/repositories/session_repository.py new file mode 100644 index 0000000..4b0b3cc --- /dev/null +++ b/apps/admin_console/database/repositories/session_repository.py @@ -0,0 +1,202 @@ +# 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. + +import time +from typing import Any + +try: + from admin_console.database.connection import db_session +except ImportError: + from apps.admin_console.database.connection import db_session + + +class SessionRepository: + """Repository handling all database queries and updates for Sessions.""" + + def __init__(self, db_path=None): + self.db_path = db_path + + def get_all_sessions(self) -> list[dict[str, Any]]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT * FROM sessions ORDER BY start_time DESC") + return [dict(r) for r in cursor.fetchall()] + + def get_session_by_id(self, session_id: str) -> dict[str, Any] | None: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT * FROM sessions WHERE session_id = ?", (str(session_id),)) + row = cursor.fetchone() + return dict(row) if row else None + + def get_video_recordings_map(self) -> dict[str, str]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='video_recordings'" + ) + if cursor.fetchone() is None: + return {} + cursor.execute( + "SELECT session_id, local_video_path FROM video_recordings WHERE session_id IS NOT NULL AND local_video_path IS NOT NULL ORDER BY start_time ASC" + ) + return {str(r[0]): str(r[1]) for r in cursor.fetchall() if r[0] and r[1]} + + def get_llm_traces_for_profile(self, session_id: str, limit: int = 3) -> list[str]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT payload FROM traces WHERE session_id = ? AND type = 'llm_call' LIMIT ?", + (session_id, limit), + ) + rows = cursor.fetchall() + return [r[0] for r in rows if r and r[0]] + + def get_agent_trace_names(self, session_id: str) -> list[str]: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT DISTINCT name FROM traces WHERE session_id = ? AND type IN ('agent', 'llm_call')", + (session_id,), + ) + rows = cursor.fetchall() + return [r[0] for r in rows if r and r[0]] + except Exception: + return [] + + def get_session_by_id(self, session_id: str) -> dict[str, Any] | None: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT * FROM sessions WHERE session_id = ?", + (str(session_id),), + ) + row = cursor.fetchone() + return dict(row) if row else None + except Exception: + return None + + def harvest_orphaned_sessions(self, orphaned_ids: list[str]) -> int: + if not orphaned_ids: + return 0 + with db_session(self.db_path) as conn: + cursor = conn.cursor() + now = time.time() + for s_id in orphaned_ids: + cursor.execute( + "UPDATE sessions SET status = ?, end_time = ? " + "WHERE session_id = ? AND status = ?", + ("failed", now, s_id, "running"), + ) + conn.commit() + return len(orphaned_ids) + + def cleanup_orphans_on_startup(self) -> int: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE sessions SET status = ?, end_time = ? WHERE status = ?", + ("failed", time.time(), "running"), + ) + count = cursor.rowcount + conn.commit() + return count + except Exception: + return 0 + + def get_latest_session(self) -> dict[str, Any] | None: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT session_id, status FROM sessions ORDER BY start_time DESC LIMIT 1" + ) + row = cursor.fetchone() + return dict(row) if row else None + except Exception: + return None + + def get_running_session_id(self) -> str | None: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT session_id FROM sessions WHERE status = 'running' " + "ORDER BY start_time DESC LIMIT 1" + ) + row = cursor.fetchone() + return row["session_id"] if row else None + except Exception: + return None + + def get_session_status(self, session_id: str) -> str | None: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT status FROM sessions WHERE session_id = ?", + (str(session_id),), + ) + row = cursor.fetchone() + return row["status"] if row else None + except Exception: + return None + + def update_session_status( + self, session_id: str, status: str, end_time: float | None = None + ) -> bool: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + t = end_time or time.time() + cursor.execute( + "UPDATE sessions SET status = ?, end_time = ? WHERE session_id = ?", + (status, t, str(session_id)), + ) + conn.commit() + return cursor.rowcount > 0 + except Exception: + return False + + def mark_all_running_cancelled(self) -> int: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE sessions SET status = ?, end_time = ? WHERE status = ?", + ("cancelled", time.time(), "running"), + ) + count = cursor.rowcount + conn.commit() + return count + except Exception: + return 0 + + def get_background_tasks(self, session_id: str) -> list[dict[str, Any]]: + try: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT * FROM background_tasks WHERE session_id = ? ORDER BY start_time ASC", + (session_id,), + ) + return [dict(r) for r in cursor.fetchall()] + except Exception: + return [] + + +session_repo = SessionRepository() diff --git a/apps/admin_console/database/repositories/step_repository.py b/apps/admin_console/database/repositories/step_repository.py new file mode 100644 index 0000000..7070999 --- /dev/null +++ b/apps/admin_console/database/repositories/step_repository.py @@ -0,0 +1,411 @@ +# 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. + +import json +from typing import Any + +try: + from admin_console.database.connection import db_session +except ImportError: + from apps.admin_console.database.connection import db_session + + +class StepRepository: + """Repository handling querying and formatting of steps and associated traces.""" + + def __init__(self, db_path=None): + self.db_path = db_path + + def get_step_session_id(self, step_id: str) -> str | None: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT session_id FROM steps WHERE step_id = ?", (step_id,)) + row = cursor.fetchone() + return row["session_id"] if row else None + + @staticmethod + def _clean_value(val: Any) -> Any: + """Recursively cleans values to remove non-serializable and internal + Python runtime representations. + """ + if val is None: + return None + if isinstance(val, (bool, int, float)): + return val + if isinstance(val, str): + if ( + "object at 0x" in val + or val.startswith(" tuple[str, int]: + """Extracts readable text and count of image blocks from LLM content.""" + if not content: + return "", 0 + if isinstance(content, str): + if content.startswith("data:image/") or "base64," in content[:50]: + return "", 1 + return content, 0 + if isinstance(content, dict): + if content.get("type") == "image_url" or "image_url" in content: + return "", 1 + if content.get("type") == "text": + return str(content.get("text", "")), 0 + t = content.get("text") or content.get("thought") or "" + return str(t), 0 + if isinstance(content, list): + total_text = "" + num_images = 0 + for item in content: + txt, imgs = StepRepository._extract_clean_text_and_images(item) + total_text += txt + num_images += imgs + return total_text, num_images + return str(content), 0 + + @staticmethod + def _extract_token_usage_from_trace(payload_raw: Any) -> tuple[int, int, int]: + """Extract prompt, completion, and total tokens from an LLM trace payload with multimodal awareness.""" + if not payload_raw: + return 0, 0, 0 + payload_obj = payload_raw + if isinstance(payload_raw, str): + try: + payload_obj = json.loads(payload_raw) + except Exception: + pass + if not isinstance(payload_obj, dict): + return 0, 0, 0 + + u = payload_obj.get("token_usage") or payload_obj.get("usage_metadata") + if isinstance(u, dict): + pr = u.get("prompt_tokens") or u.get("prompt_token_count") or u.get("input_tokens") or 0 + co = ( + u.get("completion_tokens") + or u.get("candidates_token_count") + or u.get("output_tokens") + or 0 + ) + to = u.get("total_tokens") or u.get("total_token_count") or (pr + co) + if to > 0: + return int(pr), int(co), int(to) + + # Fallback for historical traces: estimate tokens from text and images + msgs = payload_obj.get("messages", []) + resp = payload_obj.get("response", []) + + prompt_text = "" + prompt_images = 0 + if isinstance(msgs, list): + for m in msgs: + if isinstance(m, dict): + t, imgs = StepRepository._extract_clean_text_and_images(m.get("content")) + prompt_text += t + prompt_images += imgs + else: + t, imgs = StepRepository._extract_clean_text_and_images(m) + prompt_text += t + prompt_images += imgs + else: + t, imgs = StepRepository._extract_clean_text_and_images(msgs) + prompt_text += t + prompt_images += imgs + + resp_text = "" + resp_images = 0 + if isinstance(resp, list): + for r in resp: + if isinstance(r, dict): + t, imgs = StepRepository._extract_clean_text_and_images(r.get("content")) + resp_text += t + resp_images += imgs + else: + t, imgs = StepRepository._extract_clean_text_and_images(r) + resp_text += t + resp_images += imgs + + # Standard Gemini multimodal token pricing: 258 tokens per image + pr = (len(prompt_text) // 4) + (prompt_images * 258) + co = (len(resp_text) // 4) + (resp_images * 258) + return int(pr), int(co), int(pr + co) + + def _clean_tool_payload(self, payload_raw: Any, trace_type: str | None = None) -> Any: + if not payload_raw: + return None + payload_obj = payload_raw + if isinstance(payload_raw, str): + try: + payload_obj = json.loads(payload_raw) + except Exception: + pass + + if not isinstance(payload_obj, dict): + return None + + if trace_type == "llm_call": + return ( + {"error": self._clean_value(payload_obj.get("error"))} + if "error" in payload_obj + else None + ) + + args = payload_obj.get("args", payload_obj) + if isinstance(args, dict): + return {"args": self._clean_value(args)} + return {"args": self._clean_value(args)} + + def get_session_steps(self, session_id: str, client: str | None = None) -> list[dict[str, Any]]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC", + (session_id,), + ) + rows = cursor.fetchall() + + steps = [] + for r in rows: + step_dict = dict(r) + + # Strip raw binary bytes for API efficiency (clients fetch images by url/hash) + step_dict["pre_screenshot_bytes"] = None + step_dict["post_screenshot_bytes"] = None + + # Parse JSON fields safely + for json_col in [ + "action_taken", + "last_execution_result", + "extra_metadata", + ]: + if step_dict.get(json_col): + try: + step_dict[json_col] = json.loads(step_dict[json_col]) + except Exception: + pass + + # If post_image_name is identical to pre_image_name, normalize to None + pre_img = step_dict.get("pre_image_name") + if pre_img and step_dict.get("post_image_name") == pre_img: + step_dict["post_image_name"] = None + + # Fetch generic tool traces for this step + step_id = step_dict.get("step_id") + if step_id: + cursor.execute( + "SELECT t1.trace_id, t1.parent_trace_id, t1.step_id," + " t1.type, t1.name, t1.timestamp, t1.duration, t1.status," + " t1.payload, t2.name as agent_name FROM traces t1 LEFT" + " JOIN traces t2 ON t1.parent_trace_id = t2.trace_id WHERE" + " t1.step_id = ? AND (t1.type = 'tool' OR (t1.type =" + " 'llm_call' AND t1.status = 'failed')) ORDER BY" + " t1.timestamp ASC", + (step_id,), + ) + trace_rows = cursor.fetchall() + generic_tools = [] + for tr in trace_rows: + trace_dict = dict(tr) + trace_dict["payload"] = self._clean_tool_payload( + trace_dict.get("payload"), trace_dict.get("type") + ) + generic_tools.append(trace_dict) + step_dict["generic_tools"] = generic_tools + + # Compute token usage for this step from all LLM call traces + cursor.execute( + "SELECT payload FROM traces WHERE step_id = ? AND type = 'llm_call'", + (step_id,), + ) + llm_rows = cursor.fetchall() + step_p, step_c, step_t = 0, 0, 0 + for lr in llm_rows: + p, c, t = self._extract_token_usage_from_trace(lr["payload"]) + step_p += p + step_c += c + step_t += t + + # If step_t is 0, also check extra_metadata for token_usage + if step_t == 0 and step_dict.get("extra_metadata"): + meta = ( + step_dict["extra_metadata"] + if isinstance(step_dict["extra_metadata"], dict) + else {} + ) + u = meta.get("token_usage") + if isinstance(u, dict): + step_p = int(u.get("prompt_tokens") or u.get("input_tokens") or 0) + step_c = int(u.get("completion_tokens") or u.get("output_tokens") or 0) + step_t = int(u.get("total_tokens") or (step_p + step_c)) + + step_dict["token_usage"] = { + "prompt_tokens": step_p, + "completion_tokens": step_c, + "total_tokens": step_t, + } + step_dict["total_tokens"] = step_t + else: + step_dict["generic_tools"] = [] + step_dict["token_usage"] = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + } + step_dict["total_tokens"] = 0 + + try: + from artemis.utils.coordinates import normalize_step_actions + + step_dict = normalize_step_actions(step_dict) + except ImportError: + pass + + steps.append(step_dict) + + # Fetch step-less tool traces (e.g. planner tools) + cursor.execute( + "SELECT t1.trace_id, t1.parent_trace_id, t1.step_id, t1.type," + " t1.name, t1.timestamp, t1.duration, t1.status, t1.payload," + " t2.name as agent_name FROM traces t1 LEFT JOIN traces t2 ON" + " t1.parent_trace_id = t2.trace_id WHERE t1.session_id = ? AND" + " (t1.step_id IS NULL OR t1.step_id = '') AND (t1.type = 'tool' OR" + " (t1.type = 'llm_call' AND t1.status = 'failed')) ORDER BY" + " t1.timestamp ASC", + (session_id,), + ) + stepless_rows = cursor.fetchall() + if stepless_rows: + unassigned_tools = [] + sorted_steps = sorted(steps, key=lambda s: s.get("timestamp") or 0) + + for tr in stepless_rows: + trace_dict = dict(tr) + trace_dict["payload"] = self._clean_tool_payload( + trace_dict.get("payload"), trace_dict.get("type") + ) + + tr_time = trace_dict.get("timestamp") or 0 + target_step = None + for s in sorted_steps: + s_time = s.get("timestamp") or 0 + if s_time <= tr_time: + target_step = s + else: + break + + if target_step: + if "generic_tools" not in target_step or not target_step["generic_tools"]: + target_step["generic_tools"] = [] + target_step["generic_tools"].append(trace_dict) + target_step["generic_tools"].sort(key=lambda t: t.get("timestamp") or 0) + else: + unassigned_tools.append(trace_dict) + + if unassigned_tools: + unassigned_tools.sort(key=lambda t: t.get("timestamp") or 0) + first_timestamp = unassigned_tools[0].get("timestamp") or 0 + + # Calculate token usage ONLY for pre-planning LLM calls (before first step) + first_recorded_step_time = ( + sorted_steps[0].get("timestamp") if sorted_steps else float("inf") + ) + cursor.execute( + "SELECT payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type = 'llm_call'", + (session_id, first_recorded_step_time), + ) + plan_llm_rows = cursor.fetchall() + plan_p, plan_c, plan_t = 0, 0, 0 + for lr in plan_llm_rows: + p, c, t = self._extract_token_usage_from_trace(lr["payload"]) + plan_p += p + plan_c += c + plan_t += t + + # Retrieve planner thinking / thoughts + cursor.execute( + "SELECT type, payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type IN ('raw_thinking', 'thinking') ORDER BY timestamp ASC", + (session_id, first_recorded_step_time), + ) + think_rows = cursor.fetchall() + p_raw_th = [] + p_nat_th = [] + for thr in think_rows: + if thr["payload"]: + try: + th_obj = ( + json.loads(thr["payload"]) + if isinstance(thr["payload"], str) + else thr["payload"] + ) + th_text = th_obj.get("thought") or th_obj.get("text") or "" + if thr["type"] == "raw_thinking" and th_text: + p_raw_th.append(str(th_text).strip()) + elif thr["type"] == "thinking" and th_text: + p_nat_th.append(str(th_text).strip()) + except Exception: + pass + + virtual_step = { + "step_id": "pre-planning", + "step_type": "planning", + "session_id": session_id, + "step_number": 0, + "timestamp": first_timestamp, + "operator_native_thinking": "\n\n".join(p_nat_th) if p_nat_th else "", + "operator_raw_thinking": "\n\n".join(p_raw_th) if p_raw_th else "", + "action_taken": None, + "last_execution_result": None, + "extra_metadata": None, + "generic_tools": unassigned_tools, + "token_usage": { + "prompt_tokens": plan_p, + "completion_tokens": plan_c, + "total_tokens": plan_t, + }, + "total_tokens": plan_t, + } + steps.append(virtual_step) + + steps.sort( + key=lambda x: ( + x.get("timestamp") or 0, + x.get("step_number") or 0, + ) + ) + + return steps + + +step_repo = StepRepository() diff --git a/apps/admin_console/database/repositories/trace_repository.py b/apps/admin_console/database/repositories/trace_repository.py new file mode 100644 index 0000000..95225c5 --- /dev/null +++ b/apps/admin_console/database/repositories/trace_repository.py @@ -0,0 +1,131 @@ +# 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. + +import json +from typing import Any + +try: + from admin_console.database.connection import db_session +except ImportError: + from apps.admin_console.database.connection import db_session + + +class TraceRepository: + """Repository handling querying traces and constructing hierarchical trace trees.""" + + def __init__(self, db_path=None): + self.db_path = db_path + + def get_trace_by_id(self, trace_id: str, db_path=None) -> dict[str, Any] | None: + with db_session(db_path or self.db_path) as conn: + cursor = conn.cursor() + cursor.execute("SELECT * FROM traces WHERE trace_id = ?", (trace_id,)) + row = cursor.fetchone() + return dict(row) if row else None + + def get_trace_tree(self, session_id: str) -> list[dict[str, Any]]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT trace_id, parent_trace_id, type, name, status, timestamp," + " duration, CASE WHEN type IN ('raw_thinking', 'thinking') THEN" + " payload ELSE NULL END AS payload FROM traces WHERE session_id = ?" + " ORDER BY timestamp ASC", + (session_id,), + ) + rows = cursor.fetchall() + + nodes = {} + for r in rows: + row_dict = dict(r) + if row_dict.get("payload"): + try: + row_dict["payload"] = json.loads(row_dict["payload"]) + except Exception: + pass + nodes[row_dict["trace_id"]] = {**row_dict, "children": []} + + root_nodes = [] + for node_id, node in nodes.items(): + parent_id = node["parent_trace_id"] + if parent_id and parent_id in nodes: + nodes[parent_id]["children"].append(node) + else: + root_nodes.append(node) + + def sanitize_tree(nodes_list): + sanitized = [] + for node in nodes_list: + node["children"] = sanitize_tree(node["children"]) + if node["name"] == "node": + sanitized.extend(node["children"]) + else: + sanitized.append(node) + return sanitized + + return sanitize_tree(root_nodes) + + def get_step_traces_tree(self, session_id: str, step_id: str) -> list[dict[str, Any]]: + with db_session(self.db_path) as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT trace_id, parent_trace_id, step_id, type, name, status," + " timestamp, duration, CASE WHEN type IN ('raw_thinking'," + " 'thinking') THEN payload ELSE NULL END AS payload FROM traces" + " WHERE session_id = ? ORDER BY timestamp ASC", + (session_id,), + ) + rows = cursor.fetchall() + + nodes = {} + for r in rows: + row_dict = dict(r) + if row_dict.get("payload"): + try: + row_dict["payload"] = json.loads(row_dict["payload"]) + except Exception: + pass + nodes[row_dict["trace_id"]] = {**row_dict, "children": []} + + step_roots = [] + for n_id, node in nodes.items(): + if node.get("step_id") == step_id: + step_roots.append(node) + + for n_id, node in nodes.items(): + parent_id = node["parent_trace_id"] + if parent_id and parent_id in nodes: + nodes[parent_id]["children"].append(node) + + def sanitize_tree(nodes_list, visited=None): + if visited is None: + visited = set() + sanitized = [] + for node in nodes_list: + node_id = node.get("trace_id") + if node_id and node_id in visited: + continue + if node_id: + visited.add(node_id) + node["children"] = sanitize_tree(node["children"], visited.copy()) + if node["name"] == "node": + sanitized.extend(node["children"]) + else: + sanitized.append(node) + return sanitized + + return sanitize_tree(step_roots) + + +trace_repo = TraceRepository() diff --git a/apps/admin_console/index.html b/apps/admin_console/index.html new file mode 100644 index 0000000..69aad2d --- /dev/null +++ b/apps/admin_console/index.html @@ -0,0 +1,3467 @@ + + + + + + + + Artemis Debug UI + + + + +
+
+
+ +
+

Select a session from the sidebar

+
+
+ +
+ + + + + + + + + + + + +
+ + + + + + + + + + + + + + + + + + + + + + diff --git a/apps/admin_console/replay_manager.py b/apps/admin_console/replay_manager.py new file mode 100644 index 0000000..d12187b --- /dev/null +++ b/apps/admin_console/replay_manager.py @@ -0,0 +1,2567 @@ +# 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. + +import json +import os +from pathlib import Path +import shutil +import sqlite3 +import sys +import traceback +from typing import Any +from uuid import UUID + +WORKSPACE_ROOT = Path(__file__).parent.parent.resolve() + +# Try to import settings from artemis +try: + from artemis.config import settings + + TRACES_PATH = Path(settings.TRACES_PATH) +except ImportError: + # Fallback to default path if run outside project environment or imports fail + TRACES_PATH = WORKSPACE_ROOT / "traces" + +DB_PATH = TRACES_PATH / "data_engine.db" +IMAGES_DIR = TRACES_PATH / "images" + +try: + from artemis.context import ArtemisContext + from artemis.graph.state import State +except ImportError: + ArtemisContext = Any + State = Any + +REPLAY_TOOLS_CONFIG = { + "ask_explorer": { + "display_name": "Ask Explorer", + "description": ( + "Visual parsing agent that executes visual ReAct search loops using multi-modal models." + ), + "module": "artemis.tools.explorer_tool", + "function": "_run_explorer_logic", + "agent_name": "explorer", + "blacklist_args": ["ctx", "state"], + "fallback_mappings": {}, + "llm_span_name": "gemini_explorer_call", + }, + "ask_image_processor": { + "display_name": "Ask Image Processor", + "description": ( + "Visual parsing agent that writes Python code using OpenCV and" + " computer vision to analyze and modify screen images." + ), + "module": "artemis.tools.image_processor_tool", + "function": "_run_image_processor_logic", + "agent_name": "image_processor", + "blacklist_args": ["ctx", "state"], + "fallback_mappings": {}, + "llm_span_name": "gemini_image_processor_call", + }, + # Future tools can easily be registered here using plain dictionaries: + # "ask_validator": { + # "display_name": "Ask Validator", + # "description": "Validation agent that checks if a step or action was completed successfully.", + # "module": "artemis.tools.validator_tool", + # ... + # } +} + + +class ReplayManager: + """Manages sandboxed execution and diagnostics of visual agent steps. + + The ReplayManager is the core orchestrator for the Artemis step replay system. + It reconstructs historical execution states (SQLite database records, trace + histories, and screenshots) and replays a single agent step in an isolated, + deterministic sandbox environment while interacting live with a physical + Android device and Gemini. + + Key Responsibilities: + - Reconstructs and preloads historical SQLite states up to a specific + step. + - Establishes sandboxed `ArtemisContext` and `DataEngine` instances. + - Integrates with physical devices via ADB and UIAutomator. + - Hooks and spies on the Gemini API to capture thoughts and raw outputs. + - Reconciles live and preloaded traces into unified chronological trees. + + Attributes: + workspace_root (Path): The absolute path to the root of the project + workspace. + replay_base_dir (Path): Base directory containing replay data and + outputs. + test_data_dir (Path): Directory containing chunked historical step + traces. + test_outputs_dir (Path): Sandbox directory where replay outputs are + written. + db_path (Path): Path to the master Data Engine SQLite database. + traces_path (Path): Path to the master traces directory. + images_dir (Path): Path to the master images directory. + """ + + def __init__( + self, + workspace_root: Path = WORKSPACE_ROOT, + device_id: str = None, + original_db_path: Path = None, + init_device: bool = False, + ): + """Initializes the ReplayManager and autodetects the connected Android device. + + Args: + workspace_root (Path, optional): The absolute path to the workspace + root. Defaults to WORKSPACE_ROOT. + device_id (str, optional): The serial ID of the target Android + device. If not specified, autodetects the first connected device. + original_db_path (Path, optional): Path to the master SQLite + database. Defaults to DB_PATH. + init_device (bool, optional): Whether to initialize the ADB device + connection on startup. + """ + self.workspace_root = workspace_root + self.traces_path = TRACES_PATH + self.replay_base_dir = self.traces_path / "replay" + self.test_data_dir = self.replay_base_dir / "data" + self.test_outputs_dir = self.replay_base_dir / "outputs" + self.db_path = Path(original_db_path) if original_db_path else DB_PATH + self.images_dir = IMAGES_DIR + + # Device connection and autodetection + self.device_id = None + self.w = None + self.h = None + self.adb = None + self.ui_client = None + + if init_device: + try: + self._init_device(device_id) + except Exception as e: + print( + "Warning: Connected device could not be initialized during" + f" ReplayManager startup: {e}" + ) + + def _init_device(self, device_id: str = None): + """Initializes connection to the target device and queries screen metrics.""" + from adbutils import AdbClient + + try: + from artemis.clients.ui_automator_client import UIAutomatorClient + except ImportError: + raise ImportError( + "Failed to import UIAutomatorClient. Ensure artemis package is installed in path." + ) + + self.adb = AdbClient(host="localhost", port=5037) + try: + devices = self.adb.device_list() + except Exception as adb_err: + raise ConnectionError( + f"Failed to query device list from ADB server: {adb_err}" + ) from adb_err + + if not devices: + raise ConnectionError( + "No active ADB devices connected. Please connect an Android device via ADB." + ) + + if device_id: + matched_device = next((d.serial for d in devices if d.serial == device_id), None) + if not matched_device: + raise ConnectionError( + f"Requested device '{device_id}' is not connected." + f" Connected devices: {[d.serial for d in devices]}" + ) + self.device_id = matched_device + else: + self.device_id = devices[0].serial + + self.ui_client = UIAutomatorClient(device_id=self.device_id) + ui_data = self.ui_client.get_screen_data() + self.w, self.h = ui_data.width, ui_data.height + print(f"Connected to device: {self.device_id} ({self.w}x{self.h})") + + def chunk_session_traces(self, session_id: str, output_dir: Path = None) -> Path: + """Chunks traces, screenshots, and metadata for a specific execution session. + + Queries the master SQLite database and exports step-by-step + subdirectories + (e.g., `step_00/`, `step_01/`) inside the output directory. + + Args: + session_id: The UUID string of the session to extract and chunk. + output_dir: Optional path to the directory where chunked results + should be saved. Defaults to `self.test_data_dir / + f"{session_id}_chunked"`. + """ + import sqlite3 + import json + import shutil + + db_path = self.db_path + traces_dir = self.traces_path + + if not db_path.exists(): + raise FileNotFoundError(f"Master database not found at {db_path}") + + if not output_dir: + output_dir = self.test_data_dir / f"{session_id}_chunked" + else: + output_dir = Path(output_dir) + + if output_dir.exists() or output_dir.is_symlink(): + # Safety check: only delete directories that are named *_chunked + if output_dir.name.endswith("_chunked"): + if output_dir.is_symlink(): + output_dir.unlink() + elif ( + self.test_data_dir in output_dir.parents + or output_dir.parent == self.test_data_dir + or output_dir.parent.name == "inputs" + ): + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Chunking traces for session: {session_id}") + print(f"Using database: {db_path}") + print(f"Output directory: {output_dir}") + + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + + cursor.execute("SELECT * FROM sessions WHERE session_id = ?", (session_id,)) + session_row = cursor.fetchone() + + cursor.execute( + "SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC", + (session_id,), + ) + steps = cursor.fetchall() + print(f"Found {len(steps)} steps for session {session_id}") + + if not session_row and not steps: + conn.close() + raise ValueError(f"Session {session_id} not found in database.") + + for step in steps: + step_number = step["step_number"] + step_id = step["step_id"] + step_dir = output_dir / f"step_{step_number:02d}" + step_dir.mkdir(parents=True, exist_ok=True) + + # 1. Process and save the step's primary metadata (step.json) + step_data = {} + for key in step.keys(): + val = step[key] + if ( + key + in ( + "action_taken", + "last_execution_result", + "extra_metadata", + ) + and val + ): + try: + step_data[key] = json.loads(val) + except Exception: + step_data[key] = val + else: + step_data[key] = val + + step_json_path = step_dir / "step.json" + with open(step_json_path, "w", encoding="utf-8") as f: + json.dump(step_data, f, indent=2, ensure_ascii=False) + + # 2. Retrieve and save trace logs recorded during this step (traces.json) + cursor.execute( + "SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC", + (step_id,), + ) + traces = cursor.fetchall() + + traces_data = [] + for trace_row in traces: + trace_dict = {} + for key in trace_row.keys(): + val = trace_row[key] + if key == "payload" and val: + try: + trace_dict[key] = json.loads(val) + except Exception: + trace_dict[key] = val + else: + trace_dict[key] = val + traces_data.append(trace_dict) + + traces_json_path = step_dir / "traces.json" + with open(traces_json_path, "w", encoding="utf-8") as f: + json.dump(traces_data, f, indent=2, ensure_ascii=False) + + # 3. Copy the step's screenshots (pre-action and post-action screenshots) + pre_image_name = step_data.get("pre_image_name") + post_image_name = step_data.get("post_image_name") + if pre_image_name and post_image_name == pre_image_name: + post_image_name = None + images_dir = traces_dir / "images" + + if pre_image_name: + src_pre = images_dir / f"{pre_image_name}.jpg" + if src_pre.exists(): + shutil.copy2(src_pre, step_dir / "pre.jpg") + else: + print(f"Warning: Pre-action screenshot not found at {src_pre}") + + if post_image_name: + src_post = images_dir / f"{post_image_name}.jpg" + if src_post.exists(): + shutil.copy2(src_post, step_dir / "post.jpg") + else: + print(f"Warning: Post-action screenshot not found at {src_post}") + + # 4. Save metadata of the pre-action screenshot (pre_image_meta.json) + if pre_image_name: + cursor.execute( + "SELECT * FROM images WHERE image_name = ?", + (pre_image_name,), + ) + img_row = cursor.fetchone() + if img_row: + img_meta = {} + for key in img_row.keys(): + val = img_row[key] + if key in ("ocr_result", "ui_tree", "extra_metadata") and val: + try: + img_meta[key] = json.loads(val) + except Exception: + img_meta[key] = val + else: + img_meta[key] = val + + pre_img_meta_path = step_dir / "pre_image_meta.json" + with open(pre_img_meta_path, "w", encoding="utf-8") as f: + json.dump(img_meta, f, indent=2, ensure_ascii=False) + else: + print( + f"Warning: Image metadata row for {pre_image_name} not" + " found in images table." + ) + + # 5. Save metadata of the post-action screenshot (post_image_meta.json) + if post_image_name: + cursor.execute( + "SELECT * FROM images WHERE image_name = ?", + (post_image_name,), + ) + img_row = cursor.fetchone() + if img_row: + img_meta = {} + for key in img_row.keys(): + val = img_row[key] + if key in ("ocr_result", "ui_tree", "extra_metadata") and val: + try: + img_meta[key] = json.loads(val) + except Exception: + img_meta[key] = val + else: + img_meta[key] = val + + post_img_meta_path = step_dir / "post_image_meta.json" + with open(post_img_meta_path, "w", encoding="utf-8") as f: + json.dump(img_meta, f, indent=2, ensure_ascii=False) + else: + print( + f"Warning: Image metadata row for {post_image_name} not" + " found in images table." + ) + + conn.close() + # Write sentinel file to indicate successful chunking + try: + (output_dir / ".chunked").touch() + except Exception as e: + print(f"Warning: Failed to write chunked sentinel file: {e}") + print("Chunking trace completion status: success") + return output_dir + + def list_devices(self) -> list[dict]: + """Dynamically queries the ADB server for connected Android devices.""" + try: + from adbutils import AdbClient + + adb = AdbClient(host="localhost", port=5037) + devices = adb.device_list() + return [{"serial": d.serial, "status": "online"} for d in devices] + except Exception as e: + print(f"Warning: Failed to query device list from ADB: {e}") + return [] + + def load_session_goal( + self, session_id: str, step_dir: Path, original_db_path: str = None + ) -> str: + """Loads the session goal from either session.json or the original session db.""" + session_json_path = step_dir.parent / "session.json" + if session_json_path.exists(): + try: + with open(session_json_path, encoding="utf-8") as f: + data = json.load(f) + if isinstance(data, dict) and "initial_goal" in data and data["initial_goal"]: + return data["initial_goal"] + except Exception as e: + print(f"Warning: Failed to load goal from session.json: {e}") + + if original_db_path and Path(original_db_path).exists(): + try: + conn = sqlite3.connect(original_db_path) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + cursor.execute( + "SELECT initial_goal FROM sessions WHERE session_id = ?", + (str(session_id),), + ) + row = cursor.fetchone() + conn.close() + if row and row["initial_goal"]: + return row["initial_goal"] + except Exception as e: + print(f"Warning: Failed to load goal from original database: {e}") + + raise ValueError( + f"Could not load initial_goal for session {session_id} from" + f" session.json or original database {original_db_path}." + ) + + def extract_replay_data( + self, + step_dir, + original_db_path=None, + agent_name="explorer", + tool_name="ask_explorer", + ): + """Extracts final answer and chronological tool call sequence from traces.json/db.""" + step_json_path = step_dir / "step.json" + if not step_json_path.exists(): + raise FileNotFoundError(f"step.json not found in {step_dir}") + + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + step_id = step_data["step_id"] + + traces_json_path = step_dir / "traces.json" + traces = [] + if traces_json_path.exists(): + with open(traces_json_path, encoding="utf-8") as f: + traces = json.load(f) + + # If traces.json is empty/missing or lacks agent trace, check original DB + agent_trace = next( + (t for t in traces if t.get("type") == "agent" and t.get("name") == agent_name), + None, + ) + if (not agent_trace or not traces) and original_db_path and Path(original_db_path).exists(): + conn = sqlite3.connect(original_db_path) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + cursor.execute( + "SELECT * FROM traces WHERE step_id = ? ORDER BY timestamp ASC", + (step_id,), + ) + rows = cursor.fetchall() + traces = [] + for r in rows: + td = dict(r) + if td.get("payload"): + try: + td["payload"] = json.loads(td["payload"]) + except Exception: + pass + traces.append(td) + conn.close() + agent_trace = next( + (t for t in traces if t.get("type") == "agent" and t.get("name") == agent_name), + None, + ) + + if not agent_trace: + print( + f"Warning: {agent_name} agent trace not found. Replay will" + " fallback to direct submit_answer." + ) + return ( + [], + { + "candidates": [], + "fallback_message": f"No {agent_name} trace found.", + }, + traces, + ) + + payload = agent_trace.get("payload") or {} + result_str = payload.get("result", "{}") + if isinstance(result_str, str): + try: + final_submit_args = json.loads(result_str) + except Exception: + final_submit_args = { + "candidates": [], + "fallback_message": result_str, + } + else: + final_submit_args = result_str + + # Build sequence of intermediate tool calls under the agent trace + agent_trace_id = agent_trace["trace_id"] + + def get_descendants(pid): + child_list = [] + for t in traces: + if t.get("parent_trace_id") == pid: + child_list.append(t) + child_list.extend(get_descendants(t["trace_id"])) + return child_list + + descendants = get_descendants(agent_trace_id) + descendants.sort(key=lambda x: x.get("timestamp", 0)) + + tool_calls_sequence = [] + for t in descendants: + if t.get("type") == "tool" and t.get("name") not in ( + tool_name, + "submit_answer", + ): + t_payload = t.get("payload") or {} + t_args = t_payload.get("args") or {} + tool_calls_sequence.append({"name": t.get("name"), "args": t_args}) + + # Preload traces up to the point when agent is called + traces_to_preload = [] + if agent_trace in traces: + idx = traces.index(agent_trace) + traces_to_preload = traces[:idx] + else: + agent_ts = agent_trace.get("timestamp", 0.0) + traces_to_preload = [t for t in traces if t.get("timestamp", 0.0) < agent_ts] + + return tool_calls_sequence, final_submit_args, traces_to_preload + + def resolve_step_data_path(self, data: dict, path: str): + """Resolves a value from a nested dict/list using a dot-separated path (e.g. + + 'action_taken.0.target_text'). + """ + parts = path.split(".") + curr = data + for p in parts: + if isinstance(curr, dict): + curr = curr.get(p) + elif isinstance(curr, list): + try: + idx = int(p) + if idx < len(curr): + curr = curr[idx] + else: + return None + except ValueError: + return None + else: + return None + return curr + + def setup_temporary_database( + self, + temp_db_path, + temp_traces_dir, + session_id, + step_data, + pre_image_meta, + initial_goal, + traces_to_preload=None, + post_image_meta=None, + replay_id=None, + ): + """Initializes a sandbox SQLite database with preloaded step and screen metadata.""" + if not initial_goal: + raise ValueError("initial_goal must be provided and cannot be empty.") + + if temp_db_path.exists(): + try: + temp_db_path.unlink() + except OSError: + pass + + if temp_traces_dir.exists(): + shutil.rmtree(temp_traces_dir) + temp_traces_dir.mkdir(parents=True, exist_ok=True) + + try: + from artemis.data_engine.storage import StorageManager + from artemis.data_engine.models import ( + SessionMetadata, + ImageRecord, + StepRecord, + TraceRecord, + ) + except ImportError: + raise ImportError( + "Failed to import Artemis Data Engine modules. Ensure you are" + " running in the correct virtual environment." + ) + + storage = StorageManager(temp_db_path, temp_traces_dir) + + # Preload session metadata + session_meta = SessionMetadata( + session_id=UUID(session_id) if isinstance(session_id, str) else session_id, + initial_goal=initial_goal, + status="running", + ) + storage.create_session(session_meta) + + if replay_id: + replay_session_meta = SessionMetadata( + session_id=UUID(replay_id) if isinstance(replay_id, str) else replay_id, + initial_goal=initial_goal, + status="running", + ) + storage.create_session(replay_session_meta) + + # Preload image metadata + if pre_image_meta: + image_record = ImageRecord( + image_name=pre_image_meta["image_name"], + timestamp=pre_image_meta.get("timestamp", 0.0), + ocr_result=pre_image_meta.get("ocr_result"), + ui_tree=pre_image_meta.get("ui_tree"), + extra_metadata=pre_image_meta.get("extra_metadata", {}), + ) + storage.create_image(image_record) + + if post_image_meta: + image_record = ImageRecord( + image_name=post_image_meta["image_name"], + timestamp=post_image_meta.get("timestamp", 0.0), + ocr_result=post_image_meta.get("ocr_result"), + ui_tree=post_image_meta.get("ui_tree"), + extra_metadata=post_image_meta.get("extra_metadata", {}), + ) + storage.create_image(image_record) + + # Preload step metadata + step_record = StepRecord( + step_id=step_data["step_id"], + session_id=step_data["session_id"], + step_number=step_data["step_number"], + timestamp=step_data.get("timestamp", 0.0), + pre_image_name=step_data.get("pre_image_name"), + post_image_name=step_data.get("post_image_name"), + summary=step_data.get("summary"), + action_taken=step_data.get("action_taken"), + operator_raw_thinking=step_data.get("operator_raw_thinking"), + operator_native_thinking=step_data.get("operator_native_thinking"), + last_execution_result=step_data.get("last_execution_result"), + extra_metadata=step_data.get("extra_metadata", {}), + ) + storage.create_step(step_record) + + # Preload trace history + if traces_to_preload: + for t in traces_to_preload: + try: + payload_val = t.get("payload") + if isinstance(payload_val, str): + try: + payload_val = json.loads(payload_val) + except Exception: + payload_val = {} + elif not isinstance(payload_val, dict): + payload_val = {} + + trace_rec = TraceRecord( + trace_id=UUID(t["trace_id"]) + if isinstance(t["trace_id"], str) + else t["trace_id"], + session_id=UUID(t["session_id"]) + if isinstance(t["session_id"], str) + else t["session_id"], + step_id=UUID(t["step_id"]) if t.get("step_id") else None, + parent_trace_id=UUID(t["parent_trace_id"]) + if t.get("parent_trace_id") + else None, + type=t["type"], + name=t["name"], + timestamp=t.get("timestamp", 0.0), + duration=t.get("duration"), + status=t.get("status", "success"), + payload=payload_val, + ) + storage.create_trace(trace_rec) + except Exception as ex: + print(f"Warning: Failed to preload trace {t.get('trace_id')}: {ex}") + + return storage + + def get_preloaded_trace_ids_for_step(self, session_id: str, step_num: int) -> set: + traces_json_path = ( + self.test_data_dir / f"{session_id}_chunked" / f"step_{step_num:02d}" / "traces.json" + ) + if traces_json_path.exists(): + try: + with open(traces_json_path, encoding="utf-8") as f: + traces = json.load(f) + return {str(t.get("trace_id")).lower() for t in traces if t.get("trace_id")} + except Exception: + pass + return set() + + def load_traces_from_db(self, db_path: Path, step_num: int, preloaded_trace_ids: set) -> list: + if not db_path.exists(): + return [] + try: + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + cursor.execute( + "SELECT trace_id, session_id, step_id, parent_trace_id, type," + " name, status, timestamp, duration, payload FROM traces ORDER" + " BY timestamp ASC" + ) + rows = cursor.fetchall() + conn.close() + + preloaded_set = {str(tid).lower() for tid in preloaded_trace_ids if tid} + step_traces = [] + for r in rows: + td = dict(r) + if td.get("payload"): + try: + td["payload"] = json.loads(td["payload"]) + except Exception: + pass + if not isinstance(td.get("payload"), dict): + td["payload"] = {} + + tid = str(r["trace_id"]).lower() + if tid in preloaded_set: + td["is_preloaded"] = True + else: + td["is_preloaded"] = False + + td["step_number"] = step_num + step_traces.append(td) + return step_traces + except Exception as e: + print( + f"Warning: Failed to load traces from database {db_path}: {e}", + file=sys.stderr, + ) + return [] + + def normalize_trace_payload(self, t: dict) -> dict: + if t.get("type") == "tool": + payload = t.get("payload") or {} + args = payload.get("args") or {} + positional_args = payload.get("positional_args") or [] + + if (not args or args == {}) and positional_args: + name = t.get("name") + if name == "search_by_coords" and len(positional_args) >= 2: + try: + args = { + "nx": int(positional_args[0]), + "ny": int(positional_args[1]), + } + except ValueError: + args = { + "nx": positional_args[0], + "ny": positional_args[1], + } + elif name == "search_ui" and len(positional_args) >= 1: + args = {"search_query": positional_args[0]} + elif name == "detect_objects" and len(positional_args) >= 1: + args = {"labels": positional_args[0]} + else: + args = {"positional_args": positional_args} + + payload["args"] = args + t["payload"] = payload + return t + + def load_all_session_traces( + self, + temp_db_path: Path, + preloaded_trace_ids: set, + current_step_number: int, + ) -> list: + """Loads and merges all traces across all steps in the session.""" + session_id = None + try: + session_id = temp_db_path.parents[1].name.split("_")[0] + except Exception: + pass + + if not session_id: + try: + conn = sqlite3.connect(temp_db_path) + cursor = conn.cursor() + cursor.execute("SELECT session_id FROM traces LIMIT 1") + row = cursor.fetchone() + if row: + session_id = str(row[0]) + conn.close() + except Exception: + pass + + if not session_id: + return self.load_traces_from_db(temp_db_path, current_step_number, preloaded_trace_ids) + + session_chunked_dir = self.test_data_dir / f"{session_id}_chunked" + + all_traces = [] + if session_chunked_dir.exists(): + step_subdirs = sorted( + [ + d + for d in session_chunked_dir.iterdir() + if d.is_dir() and d.name.startswith("step_") + ], + key=lambda d: d.name, + ) + + for subdir in step_subdirs: + try: + step_num = int(subdir.name.split("_")[1]) + except (IndexError, ValueError): + continue + + step_outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_num:02d}" + step_db_path = step_outputs_dir / "temp_traces" / "data_engine.db" + + step_preloaded_ids = self.get_preloaded_trace_ids_for_step(session_id, step_num) + + if step_db_path.exists(): + step_traces = self.load_traces_from_db( + step_db_path, step_num, step_preloaded_ids + ) + all_traces.extend(step_traces) + else: + traces_json_path = subdir / "traces.json" + if traces_json_path.exists(): + try: + with open(traces_json_path, encoding="utf-8") as f: + step_traces = json.load(f) + for t in step_traces: + t["is_preloaded"] = True + t["step_number"] = step_num + if not isinstance(t.get("payload"), dict): + t["payload"] = {} + all_traces.append(t) + except Exception as e: + print( + f"Warning: Failed to load traces from {traces_json_path}: {e}", + file=sys.stderr, + ) + else: + all_traces = self.load_traces_from_db( + temp_db_path, current_step_number, preloaded_trace_ids + ) + + # Align live replay traces to the appropriate preloaded step/operator trace + for step_num in range(1, 100): + current_step_id = None + current_step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_num:02d}" + if not current_step_dir.exists(): + continue + + step_json_path = current_step_dir / "step.json" + if step_json_path.exists(): + try: + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + current_step_id = step_data.get("step_id") + except Exception: + pass + + current_step_id_str = str(current_step_id).lower() if current_step_id else None + + # 1. Find live root trace (no parent_trace_id and not preloaded) + live_root_trace = None + for t in all_traces: + if t.get("step_number") != step_num: + continue + if not t.get("is_preloaded") and not t.get("parent_trace_id"): + live_root_trace = t + break + + # 2. Find all preloaded traces for this step + step_preloaded_traces = [ + t + for t in all_traces + if t.get("step_number") == step_num + and t.get("is_preloaded") + and str(t.get("step_id") or "").lower() == current_step_id_str + ] + + # Sort chronologically to identify the final response of the step + step_preloaded_traces.sort(key=lambda x: x.get("timestamp", 0.0)) + final_response_trace = step_preloaded_traces[-1] if step_preloaded_traces else None + + # 3. Branch immediately before the final response of the step + if live_root_trace and final_response_trace: + # Resolve the tool name from the live root trace name + tool_name = "ask_explorer" + agent_name = live_root_trace.get("name") + for t_name, cfg in REPLAY_TOOLS_CONFIG.items(): + if cfg.get("agent_name") == agent_name: + tool_name = t_name + break + + import uuid + + virtual_id = str(uuid.uuid4()) + + # Branch from the same parent as the final response + parent_id = final_response_trace.get("parent_trace_id") + + # Set timestamp to be immediately before the final response + virtual_timestamp = final_response_trace.get("timestamp", 0.0) - 0.001 + + live_query = None + live_cf = None + if live_root_trace.get("payload") and live_root_trace["payload"].get("args"): + args = live_root_trace["payload"]["args"] + live_query = ( + args.get("q") or args.get("query") or args.get("arg2") or args.get("goal") + ) + live_cf = ( + args.get("cf") or args.get("context_feedback") or args.get("arg3") or "" + ) + if not live_query: + live_query = "Replay Execution" + + virtual_tool_trace = { + "trace_id": virtual_id, + "session_id": session_id, + "step_id": current_step_id, + "parent_trace_id": parent_id, + "type": "tool", + "name": tool_name, + "status": live_root_trace.get("status", "success"), + "timestamp": virtual_timestamp, + "duration": live_root_trace.get("duration", 0.0), + "payload": { + "args": { + "query": live_query, + "context_feedback": live_cf, + }, + "result": (live_root_trace.get("payload", {}).get("result")), + }, + "is_preloaded": False, + "step_number": step_num, + } + + # Link the live root trace to our new virtual tool trace + live_root_trace["parent_trace_id"] = virtual_id + all_traces.append(virtual_tool_trace) + + all_traces = [self.normalize_trace_payload(t) for t in all_traces] + all_traces.sort(key=lambda t: t.get("timestamp", 0.0)) + return all_traces + + def enrich_all_session_traces(self, all_traces: list, session_id: str) -> list: + """Enriches all preloaded traces (across all steps) with file URLs to their corresponding screenshots.""" + session_chunked_dir = self.test_data_dir / f"{session_id}_chunked" + + step_dirs = {} + if session_chunked_dir.exists(): + for subdir in session_chunked_dir.iterdir(): + if subdir.is_dir() and subdir.name.startswith("step_"): + step_json_path = subdir / "step.json" + if step_json_path.exists(): + try: + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + sid = step_data.get("step_id") + if sid: + step_dirs[str(sid).lower()] = subdir + except Exception: + pass + + for td in all_traces: + if td.get("is_preloaded"): + step_id = str(td.get("step_id") or "").lower() + if step_id and step_id in step_dirs: + step_dir = step_dirs[step_id] + if td.get("type") == "agent": + pre_jpg = step_dir / "pre.jpg" + if pre_jpg.exists(): + td["payload"]["pre_screenshot_path"] = f"file://{pre_jpg.resolve()}" + elif td.get("type") == "action": + post_jpg = step_dir / "post.jpg" + if post_jpg.exists(): + td["payload"]["post_screenshot_path"] = f"file://{post_jpg.resolve()}" + + return all_traces + + def load_steps_metadata(self, session_id: str) -> dict[int, dict]: + """Loads metadata for all steps in the session from the chunked data.""" + session_chunked_dir = self.test_data_dir / f"{session_id}_chunked" + + steps = {} + if session_chunked_dir.exists(): + for subdir in session_chunked_dir.iterdir(): + if subdir.is_dir() and subdir.name.startswith("step_"): + step_json_path = subdir / "step.json" + if step_json_path.exists(): + try: + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + step_num = step_data.get("step_number") + if step_num is not None: + steps[int(step_num)] = step_data + except Exception as e: + print(f"Warning: Failed to load step data from {step_json_path}: {e}") + return steps + + def make_file_url(self, val): + if isinstance(val, str) and (val.endswith(".jpg") or val.endswith(".png") or "/" in val): + try: + p = Path(val) + if p.exists(): + abs_path = str(p.resolve()) + return ( + f"file://{abs_path}" if abs_path.startswith("/") else f"file:///{abs_path}" + ) + except Exception: + pass + return val + + def clean_payload_paths(self, obj): + if isinstance(obj, dict): + return {k: self.clean_payload_paths(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [self.clean_payload_paths(v) for v in obj] + elif isinstance(obj, str): + return self.make_file_url(obj) + return obj + + def sanitize_nodes(self, nodes_list): + sanitized = [] + for node in nodes_list: + node["children"] = self.sanitize_nodes(node["children"]) + if node["name"] == "node": + sanitized.extend(node["children"]) + else: + sanitized.append(node) + return sanitized + + def sort_tree_nodes_chronologically(self, nodes_list: list[dict]) -> list[dict]: + for node in nodes_list: + if "children" in node and node["children"]: + node["children"] = self.sort_tree_nodes_chronologically(node["children"]) + return sorted(nodes_list, key=lambda x: x.get("timestamp", 0.0) or 0.0) + + def extract_tool_calls_from_db(self, temp_db_path: Path, preloaded_trace_ids: set) -> list: + """Extracts tool calls executed during the current step replay from the sandbox database.""" + preloaded_set = {str(tid).lower() for tid in preloaded_trace_ids if tid} + + if not temp_db_path.exists(): + return [] + + try: + conn = sqlite3.connect(temp_db_path) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + cursor.execute( + "SELECT trace_id, name, status, payload FROM traces WHERE type" + " = 'tool' ORDER BY timestamp ASC" + ) + rows = cursor.fetchall() + conn.close() + except Exception as e: + print( + f"Warning: Failed to query tool calls from database: {e}", + file=sys.stderr, + ) + return [] + + tool_calls = [] + for r in rows: + tid = str(r["trace_id"]).lower() + if tid not in preloaded_set: + payload = {} + if r["payload"]: + try: + payload = json.loads(r["payload"]) + except Exception: + pass + + response_val = payload.get("result") + if response_val is None: + response_val = payload.get("error") + else: + if isinstance(response_val, str) and ( + response_val.startswith("{") or response_val.startswith("[") + ): + try: + response_val = json.loads(response_val) + except Exception: + pass + + args = payload.get("args", {}) + positional_args = payload.get("positional_args", []) + if positional_args: + if r["name"] == "search_by_coords" and len(positional_args) >= 2: + try: + args = { + "nx": int(positional_args[0]), + "ny": int(positional_args[1]), + } + except ValueError: + args = { + "nx": positional_args[0], + "ny": positional_args[1], + } + elif r["name"] == "search_ui" and len(positional_args) >= 1: + args = {"search_query": positional_args[0]} + elif r["name"] == "detect_objects" and len(positional_args) >= 1: + args = {"labels": positional_args[0]} + else: + args = {"positional_args": positional_args} + + tool_calls.append( + { + "name": r["name"], + "args": args, + "status": r["status"], + "response": response_val, + } + ) + return tool_calls + + def extract_state_prepopulation_data(self, step_dir, step_data, pre_image_meta): + """Dynamically extracts state prepopulation fields from chunked test step files.""" + latest_ui_hierarchy = None + if pre_image_meta and "ui_tree" in pre_image_meta: + latest_ui_hierarchy = pre_image_meta["ui_tree"] + + structured_decisions = None + traces_json_path = step_dir / "traces.json" + if traces_json_path.exists(): + try: + with open(traces_json_path, encoding="utf-8") as f: + traces = json.load(f) + operator_trace = next( + (t for t in traces if t.get("type") == "agent" and t.get("name") == "operator"), + None, + ) + if operator_trace: + payload = operator_trace.get("payload") or {} + result_val = payload.get("result") + if isinstance(result_val, str): + try: + res_dict = json.loads(result_val) + structured_decisions = res_dict.get("structured_decisions") + except Exception as e: + print(f"Warning: Failed to parse operator trace result JSON: {e}") + elif isinstance(result_val, dict): + structured_decisions = result_val.get("structured_decisions") + except Exception as e: + print(f"Warning: Failed to read structured_decisions from traces.json: {e}") + pass + + if not structured_decisions and step_data: + action_taken = step_data.get("action_taken") + if action_taken: + try: + structured_decisions = json.dumps(action_taken) + except Exception as e: + print(f"Warning: Failed to serialize action_taken fallback to JSON: {e}") + pass + + focused_app_info = None + if latest_ui_hierarchy: + system_packages = { + "com.android.systemui", + "android", + "com.google.android.apps.nexuslauncher", + } + packages = [] + for node in latest_ui_hierarchy: + if isinstance(node, dict): + pkg = node.get("package") + if pkg and pkg not in system_packages and "launcher" not in pkg.lower(): + packages.append(pkg) + if packages: + from collections import Counter + + focused_app_info = Counter(packages).most_common(1)[0][0] + + device_date = None + timestamp = None + if step_data and "timestamp" in step_data: + timestamp = step_data["timestamp"] + elif pre_image_meta and "timestamp" in pre_image_meta: + timestamp = pre_image_meta["timestamp"] + + if timestamp is not None: + from datetime import datetime, UTC + + try: + dt = datetime.fromtimestamp(timestamp, UTC) + device_date = dt.strftime("%a %b %d %H:%M:%S UTC %Y") + except Exception as e: + print( + f"Warning: Failed to parse timestamp {timestamp} into device date string: {e}" + ) + + return ( + latest_ui_hierarchy, + structured_decisions, + focused_app_info, + device_date, + ) + + def _ensure_session_chunked(self, session_id: str) -> Path: + """Ensures that the session traces are chunked and ready for replay.""" + session_chunked_dir = self.test_data_dir / f"{session_id}_chunked" + + db_steps_count = 0 + if self.db_path.exists(): + try: + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute( + "SELECT COUNT(*) FROM steps WHERE session_id = ?", + (session_id,), + ) + db_steps_count = cursor.fetchone()[0] + conn.close() + except Exception as e: + print(f"Warning: Failed to query step count from DB: {e}") + + chunked_steps_count = 0 + if session_chunked_dir.exists(): + chunked_steps_count = len(list(session_chunked_dir.glob("step_*"))) + + # Check if directory does not exist, does not contain the sentinel file, or count of chunked steps is out of sync + needs_chunking = ( + not session_chunked_dir.exists() + or not (session_chunked_dir / ".chunked").exists() + or db_steps_count != chunked_steps_count + ) + + if needs_chunking: + try: + print( + f"Chunked data needs update for session {session_id} (DB" + f" steps: {db_steps_count}, chunked steps:" + f" {chunked_steps_count}). Chunking on-the-fly..." + ) + self.chunk_session_traces(session_id, session_chunked_dir) + except Exception as e: + print(f"Failed to chunk session {session_id} on-the-fly: {e}") + return session_chunked_dir + + def resolve_master_video_path(self, db_recorded_path_str: str) -> Path | None: + """Resolves the physical master video recording path on disk, searching for renamed folders.""" + db_video_path = Path(db_recorded_path_str) + + # 1. Direct absolute path check + if db_video_path.exists(): + return db_video_path + + # 2. Check if relative to workspace + workspace_rel_path = self.workspace_root / db_video_path + if workspace_rel_path.exists(): + return workspace_rel_path.resolve() + + # 3. Check if under traces_path + traces_rel_path = self.traces_path / db_video_path.name + if traces_rel_path.exists(): + return traces_rel_path.resolve() + + # 4. Resolve folder renaming prefix + parent_dir_name = db_video_path.parent.name + filename = db_video_path.name + for search_dir in [ + self.traces_path, + self.workspace_root / "artemis-traces", + ]: + if parent_dir_name and search_dir.exists(): + try: + for child in search_dir.iterdir(): + if child.is_dir() and ( + child.name.startswith(parent_dir_name) or parent_dir_name in child.name + ): + candidate = child / filename + if candidate.exists(): + return candidate.resolve() + except Exception as e: + print(f"Warning: Error while scanning {search_dir} for renamed folder: {e}") + + # 5. Check if under traces_path or artemis-traces with parent_dir_name + for search_dir in [ + self.traces_path, + self.workspace_root / "artemis-traces", + ]: + fallback_path = search_dir / parent_dir_name / filename + if fallback_path.exists(): + return fallback_path.resolve() + + return None + + def calculate_virtual_start_time(self, session_id: str, step_number: int) -> float | None: + """Calculates the virtual start time based on the database's final response timestamp.""" + if not self.db_path.exists(): + return None + + try: + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + + # Fetch fallback current step timestamp + cursor.execute( + "SELECT timestamp FROM steps WHERE session_id = ? AND step_number = ?", + (session_id, step_number), + ) + row = cursor.fetchone() + step_timestamp = row[0] if (row and row[0] is not None) else None + + if step_number <= 1: + # First step starts at session start + cursor.execute( + "SELECT start_time FROM sessions WHERE session_id = ?", + (session_id,), + ) + row = cursor.fetchone() + db_session_start_time = row[0] if (row and row[0] is not None) else step_timestamp + conn.close() + return db_session_start_time + + # Get step_id of previous step + cursor.execute( + "SELECT step_id FROM steps WHERE session_id = ? AND step_number = ?", + (session_id, step_number - 1), + ) + row = cursor.fetchone() + if not row: + conn.close() + return step_timestamp + + prev_step_id = row[0] + + # Get maximum trace timestamp for that step (final response timestamp) + cursor.execute( + "SELECT MAX(timestamp) FROM traces WHERE step_id = ?", + (prev_step_id,), + ) + row = cursor.fetchone() + conn.close() + + if row and row[0] is not None: + db_final_response_timestamp = row[0] + return db_final_response_timestamp - 0.001 + + return step_timestamp + except Exception as e: + print(f"Warning: Failed to calculate virtual start time from database: {e}") + return None + + def _preemptive_clip_video( + self, + session_id: str, + step_number: int, + step_data: dict, + temp_db_path: Path, + outputs_dir: Path, + replay_id: str = None, + ): + """Locates the master recording, calculates step time bounds, trims the video and updates the temp DB.""" + import subprocess + import uuid + + # 1. Query the master database's video_recordings table + if not self.db_path.exists(): + print( + f"Warning: Master database not found at {self.db_path}." + " Skipping preemptive video clip." + ) + return + + try: + conn_master = sqlite3.connect(self.db_path) + cursor = conn_master.cursor() + cursor.execute( + "SELECT video_id, device_id, start_time, end_time," + " local_video_path FROM video_recordings WHERE session_id = ?", + (session_id,), + ) + row = cursor.fetchone() + conn_master.close() + except Exception as e: + print(f"Warning: Failed to query video_recordings from master database: {e}") + return + + if not row: + print( + f"No video recording found for session {session_id} in master" + " database. Skipping clip." + ) + return + + ( + video_id, + device_id, + video_start_time, + video_end_time, + db_recorded_path, + ) = row + if not db_recorded_path: + print("Warning: Recorded path is empty in video_recordings. Skipping clip.") + return + + # 2. Locate master recording on disk (renaming logic) + master_video_path = self.resolve_master_video_path(db_recorded_path) + if not master_video_path: + print( + f"Warning: Could not locate master recording on disk for path: {db_recorded_path}" + ) + return + + # 3. Calculate step bounds using virtual start time + step_start = self.calculate_virtual_start_time(session_id, step_number) + if step_start is None: + step_start = step_data.get("timestamp") + + if step_start is None: + print("Warning: Step start timestamp could not be resolved. Skipping clip.") + return + + # Look for next step start + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + next_step_dir = step_dir.parent / f"step_{step_number + 1:02d}" + next_step_json_path = next_step_dir / "step.json" + + if next_step_json_path.exists(): + try: + with open(next_step_json_path, encoding="utf-8") as f: + next_step_data = json.load(f) + next_step_start = next_step_data.get("timestamp", step_start + 15.0) + except Exception as e: + print(f"Warning: Failed to read next step json: {e}") + next_step_start = step_start + 15.0 + else: + next_step_start = step_start + 15.0 + + # Calculate offsets + start_offset = step_start - video_start_time + # clamp start_offset to be non-negative + start_offset = max(0.0, start_offset) + end_offset = next_step_start - video_start_time + duration = max(0.1, end_offset - start_offset) + + trimmed_video_path = outputs_dir / "recording_trimmed.mp4" + + # 4. Execute FFmpeg trim + cmd = [ + "ffmpeg", + "-y", + "-i", + str(master_video_path), + "-ss", + f"{start_offset:.3f}", + "-t", + f"{duration:.3f}", + "-c:v", + "libx264", + str(trimmed_video_path), + ] + print( + f"Trimming video segment from {master_video_path}:" + f" offset={start_offset:.3f}s, duration={duration:.3f}s" + ) + try: + res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + if res.returncode != 0: + print( + f"Warning: ffmpeg failed with exit code {res.returncode}. stderr: {res.stderr}" + ) + return + except Exception as e: + print(f"Warning: Failed to execute ffmpeg command: {e}") + return + + # 5. Update Sandbox Database + try: + conn_temp = sqlite3.connect(temp_db_path) + conn_temp.execute(""" + CREATE TABLE IF NOT EXISTS video_recordings ( + video_id TEXT PRIMARY KEY, + session_id TEXT, + device_id TEXT, + start_time REAL, + end_time REAL, + local_video_path TEXT, + FOREIGN KEY(session_id) REFERENCES sessions(session_id) + ) + """) + + # Insert original metadata first if not present + conn_temp.execute( + """ + INSERT OR IGNORE INTO video_recordings (video_id, session_id, device_id, start_time, end_time, local_video_path) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + video_id, + session_id, + device_id, + video_start_time, + video_end_time, + db_recorded_path, + ), + ) + + # Update values + conn_temp.execute( + """ + UPDATE video_recordings + SET local_video_path = ?, start_time = ?, end_time = ? + WHERE session_id = ? + """, + ( + str(trimmed_video_path), + step_start, + next_step_start, + session_id, + ), + ) + + if replay_id: + # Also copy / update for replay_id + conn_temp.execute( + """ + INSERT OR IGNORE INTO video_recordings (video_id, session_id, device_id, start_time, end_time, local_video_path) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + str(uuid.uuid4()), + replay_id, + device_id, + video_start_time, + video_end_time, + db_recorded_path, + ), + ) + + conn_temp.execute( + """ + UPDATE video_recordings + SET local_video_path = ?, start_time = ?, end_time = ? + WHERE session_id = ? + """, + ( + str(trimmed_video_path), + step_start, + next_step_start, + replay_id, + ), + ) + + conn_temp.commit() + conn_temp.close() + print( + f"Successfully updated sandbox database at {temp_db_path} with" + " trimmed video segment." + ) + except Exception as e: + print(f"Warning: Failed to update sandbox SQLite database: {e}") + + def create_ctx( + self, + session_id: str, + step_number: int, + agent_name: str = "explorer", + replay_id: str = None, + ) -> ArtemisContext: + """Creates a sandboxed ArtemisContext and DataEngine for the given session and step.""" + self._ensure_session_chunked(session_id) + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + if not step_dir.exists(): + raise FileNotFoundError(f"Step directory not found at {step_dir}") + + step_json_path = step_dir / "step.json" + pre_image_meta_path = step_dir / "pre_image_meta.json" + post_image_meta_path = step_dir / "post_image_meta.json" + + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + + pre_image_meta = None + if pre_image_meta_path.exists(): + with open(pre_image_meta_path, encoding="utf-8") as f: + pre_image_meta = json.load(f) + + post_image_meta = None + if post_image_meta_path.exists(): + with open(post_image_meta_path, encoding="utf-8") as f: + post_image_meta = json.load(f) + + outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}" + outputs_dir.mkdir(parents=True, exist_ok=True) + + temp_traces_dir = outputs_dir / "temp_traces" + temp_db_path = temp_traces_dir / "data_engine.db" + + _, _, traces_to_preload = self.extract_replay_data( + step_dir, str(self.db_path), agent_name=agent_name + ) + initial_goal = self.load_session_goal(session_id, step_dir, str(self.db_path)) + + self.setup_temporary_database( + temp_db_path, + temp_traces_dir, + session_id, + step_data, + pre_image_meta, + initial_goal, + traces_to_preload, + post_image_meta=post_image_meta, + replay_id=replay_id, + ) + + try: + self._preemptive_clip_video( + session_id=session_id, + step_number=step_number, + step_data=step_data, + temp_db_path=temp_db_path, + outputs_dir=outputs_dir, + replay_id=replay_id, + ) + except Exception as e: + print(f"Warning: Pre-emptive video clipping failed: {e}") + + try: + from artemis.context import ( + ArtemisContext, + DeviceContext, + DevicePlatform, + ExecutionSetup, + ) + from artemis.data_engine.engine import DataEngine + + ArtemisContext.model_rebuild() + except ImportError as import_err: + raise ImportError(f"Failed to import Artemis core modules: {import_err}") + + import sqlite3 + + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute( + "SELECT device_info FROM sessions WHERE session_id = ?", + (session_id,), + ) + row = cursor.fetchone() + conn.close() + + device_info = {} + if row and row[0]: + try: + device_info = json.loads(row[0]) + except Exception: + pass + + sim_device_id = device_info.get("device_id", "replay-device") + sim_w = device_info.get("device_width", 1080) + sim_h = device_info.get("device_height", 2400) + + device_context = DeviceContext( + host_platform="LINUX", + mobile_platform=DevicePlatform.ANDROID, + device_id=sim_device_id, + device_width=sim_w, + device_height=sim_h, + ) + print(f"Replay simulated device: {sim_device_id} ({sim_w}x{sim_h})") + + from artemis.config import get_default_llm_config + + try: + llm_cfg = get_default_llm_config() + except Exception as e: + print(f"Warning: Failed to load default LLM config: {e}") + llm_cfg = None + + ctx = ArtemisContext( + device=device_context, + adb_client=None, + ui_adb_client=None, + execution_setup=ExecutionSetup(traces_path=temp_traces_dir), + llm_config=llm_cfg, + ) + + data_engine = DataEngine(ctx) + data_engine.current_session_id = replay_id or session_id + data_engine.current_step_id = step_data["step_id"] + data_engine.current_step_dir = outputs_dir + ctx.data_engine = data_engine + + settings.DATA_ENGINE_DB_PATH = temp_db_path + settings.TRACES_PATH = temp_traces_dir + os.environ["DATA_ENGINE_DB_PATH"] = str(temp_db_path) + + return ctx + + def instantiate_state( + self, + session_id: str, + step_number: int, + ) -> State: + """Instantiates the State object for the given session and step number.""" + self._ensure_session_chunked(session_id) + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + if not step_dir.exists(): + raise FileNotFoundError(f"Step directory not found at {step_dir}") + + step_json_path = step_dir / "step.json" + pre_image_meta_path = step_dir / "pre_image_meta.json" + + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + + pre_image_meta = None + if pre_image_meta_path.exists(): + with open(pre_image_meta_path, encoding="utf-8") as f: + pre_image_meta = json.load(f) + + initial_goal = self.load_session_goal(session_id, step_dir, str(self.db_path)) + + ui_hier, decisions, app_info, dev_date = self.extract_state_prepopulation_data( + step_dir, step_data, pre_image_meta + ) + + try: + from artemis.graph.state import State + except ImportError as import_err: + raise ImportError(f"Failed to import State module: {import_err}") + + state = State( + messages=[], + initial_goal=initial_goal, + latest_screenshot=str(step_dir / "pre.jpg") + if (step_dir / "pre.jpg").exists() + else str(step_dir / "post.jpg"), + operator_raw_data={"width": self.w, "height": self.h}, + current_step_id=step_data["step_id"], + complete_subgoals_by_ids=[], + validator_messages=[], + subagent_calls=step_data.get("subagent_calls", []), + latest_ui_hierarchy=ui_hier, + focused_app_info=app_info, + device_date=dev_date, + structured_decisions=decisions, + ) + return state + + def generate_replay_id(self) -> str: + """Generates a new unique replay ID.""" + import uuid + + return str(uuid.uuid4()) + + async def replay_step_tool( + self, + session_id: str, + step_number: int, + override_device_id: str = None, + user_submits: dict = None, + tool_name: str = "ask_explorer", + replay_id: str = None, + ) -> list[dict]: + """Replays a specific step tool in a sandboxed execution environment.""" + if not replay_id: + replay_id = self.generate_replay_id() + + replay_result = { + "session_id": session_id, + "step_number": step_number, + "replay_id": replay_id, + "success": False, + } + tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name) + if not tool_cfg: + raise ValueError(f"Unknown tool for replay: {tool_name}") + agent_name = tool_cfg["agent_name"] + + outputs_parent = self.test_outputs_dir + if outputs_parent.exists(): + for item in outputs_parent.iterdir(): + if item.is_dir() and item.name.startswith(f"{session_id}_step_"): + archive_dir = self.replay_base_dir / "data" / "older" + archive_dir.mkdir(parents=True, exist_ok=True) + + from datetime import datetime + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + target_name = f"{item.name}_{timestamp}" + target_path = archive_dir / target_name + + counter = 1 + while target_path.exists(): + target_name = f"{item.name}_{timestamp}_{counter}" + target_path = archive_dir / target_name + counter += 1 + + print( + f"Archiving previous run output directory {item.name} to {target_path}..." + ) + try: + shutil.move(str(item), str(target_path)) + except Exception as archive_err: + print( + f"Warning: Failed to archive {item.name}: {archive_err}", + file=sys.stderr, + ) + + outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}" + temp_db_path = outputs_dir / "temp_traces" / "data_engine.db" + + # Load step data locally for fallback mapping resolutions + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + step_json_path = step_dir / "step.json" + with open(step_json_path, encoding="utf-8") as f: + step_data = json.load(f) + + # Retrieve preloaded trace IDs from chunked trace files to filter new tool calls + _, _, traces_to_preload = self.extract_replay_data( + step_dir, str(self.db_path), agent_name=agent_name + ) + preloaded_trace_ids = {t.get("trace_id") for t in traces_to_preload if t.get("trace_id")} + llm_calls_log = [] + + # Device initialization removed for step replay + + ctx = self.create_ctx( + session_id=session_id, + step_number=step_number, + agent_name=agent_name, + replay_id=replay_id, + ) + + state = self.instantiate_state( + session_id=session_id, + step_number=step_number, + ) + + data_engine = ctx.data_engine + + print("Executing using Live Gemini Client...") + + from google.genai.models import AsyncModels + from unittest.mock import patch + + original_generate_content = AsyncModels.generate_content + + async def wrapped_generate_content(self, *, model, contents, config=None, **kwargs): + system_instruction = None + if config and config.system_instruction: + if isinstance(config.system_instruction, str): + system_instruction = config.system_instruction + elif hasattr(config.system_instruction, "parts"): + parts = [] + for p in config.system_instruction.parts: + if hasattr(p, "text") and p.text: + parts.append(p.text) + system_instruction = " ".join(parts) + else: + system_instruction = str(config.system_instruction) + + prompt_data = { + "model": model, + "system_instruction": system_instruction, + "contents": [], + } + + contents_list = contents if isinstance(contents, list) else [contents] + for content in contents_list: + parts_desc = [] + if hasattr(content, "parts") and content.parts: + for part in content.parts: + if hasattr(part, "text") and part.text: + parts_desc.append({"type": "text", "text": part.text}) + elif hasattr(part, "inline_data") and part.inline_data: + parts_desc.append( + { + "type": "image", + "mime_type": part.inline_data.mime_type, + "length": len(part.inline_data.data), + } + ) + elif hasattr(part, "function_call") and part.function_call: + parts_desc.append( + { + "type": "function_call", + "name": part.function_call.name, + "args": ( + dict(part.function_call.args) + if part.function_call.args + else {} + ), + } + ) + elif hasattr(part, "function_call") and part.function_response: + parts_desc.append( + { + "type": "function_response", + "name": part.function_response.name, + "response": ( + dict(part.function_response.response) + if part.function_response.response + else {} + ), + } + ) + else: + parts_desc.append({"type": "raw", "value": str(content)}) + + prompt_data["contents"].append( + { + "role": getattr(content, "role", "unknown"), + "parts": parts_desc, + } + ) + + record = {"prompt": prompt_data, "response": None, "error": None} + llm_calls_log.append(record) + + try: + res = await original_generate_content( + self, + model=model, + contents=contents, + config=config, + **kwargs, + ) + + thinking_parts = [] + if hasattr(res, "candidates") and res.candidates and res.candidates[0].content: + for part in res.candidates[0].content.parts: + if getattr(part, "thought", False): + thinking_parts.append(part.text) + + resp_data = { + "text": res.text if hasattr(res, "text") else None, + "function_calls": [], + "thought": ("\n".join(thinking_parts) if thinking_parts else None), + } + if hasattr(res, "function_calls") and res.function_calls: + for fc in res.function_calls: + resp_data["function_calls"].append( + { + "name": fc.name, + "args": dict(fc.args) if fc.args else {}, + } + ) + record["response"] = resp_data + return res + except Exception as err: + err_str = str(err) + record["error"] = ( + f"{err.__class__.__name__}: {err_str}" if err_str else err.__class__.__name__ + ) + raise err + + try: + try: + import inspect + import importlib + + mod = importlib.import_module(tool_cfg["module"]) + tool_fn = getattr(mod, tool_cfg["function"]) + + user_submits = user_submits or {} + sig = inspect.signature(tool_fn) + tool_args = {} + blacklist = tool_cfg.get("blacklist_args", []) + fallback_mappings = tool_cfg.get("fallback_mappings", {}) + + user_submits = user_submits or {} + sig = inspect.signature(tool_fn) + tool_args = {} + for name, param in sig.parameters.items(): + if name in blacklist: + if name == "ctx": + tool_args["ctx"] = ctx + elif name == "state": + tool_args["state"] = state + continue + + if name in user_submits: + val = user_submits[name] + if param.annotation == bool and isinstance(val, str): + val = val.lower() == "true" + tool_args[name] = val + else: + val = "" + if param.default is not inspect.Parameter.empty: + val = param.default + if name in fallback_mappings: + fallback_path = fallback_mappings[name] + resolved_val = self.resolve_step_data_path(step_data, fallback_path) + if resolved_val is not None: + val = resolved_val + tool_args[name] = val + + with patch.object(AsyncModels, "generate_content", wrapped_generate_content): + outcome = await tool_fn(**tool_args) + + replay_result["success"] = True + replay_result["result"] = outcome + print(f"Replay Execution complete. Outcome: {outcome}") + + except Exception as replay_err: + print(f"Error during step replay: {replay_err}", file=sys.stderr) + traceback.print_exc() + + replay_result["success"] = False + replay_result["error"] = ( + f"{replay_err.__class__.__name__}: {str(replay_err)}" + if str(replay_err) + else replay_err.__class__.__name__ + ) + raise replay_err + finally: + await data_engine.shutdown() + + tool_calls = [] + if temp_db_path.exists(): + try: + tool_calls = self.extract_tool_calls_from_db(temp_db_path, preloaded_trace_ids) + except Exception as ex: + print( + f"Warning: Failed to extract tool calls: {ex}", + file=sys.stderr, + ) + + replay_result["tool_calls"] = tool_calls + replay_result["llm_calls"] = llm_calls_log + + outcome_path = outputs_dir / "result.json" + try: + with open(outcome_path, "w", encoding="utf-8") as f: + json.dump(replay_result, f, indent=2, ensure_ascii=False) + except Exception as write_err: + print( + f"Warning: Failed to write result.json: {write_err}", + file=sys.stderr, + ) + + return llm_calls_log + + def get_replay_config(self, tool_name: str = "ask_explorer") -> list[dict]: + """Returns the parameter list and types of the selected tool logic.""" + tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name) + if not tool_cfg: + raise ValueError(f"Unknown tool for replay: {tool_name}") + + try: + import inspect + import importlib + + mod = importlib.import_module(tool_cfg["module"]) + tool_fn = getattr(mod, tool_cfg["function"]) + sig = inspect.signature(tool_fn) + + blacklist = tool_cfg.get("blacklist_args", []) + explorer_params = [] + for name, param in sig.parameters.items(): + if name in blacklist: + continue + default_val = "" + if param.default is not inspect.Parameter.empty: + default_val = param.default + + input_type = "text" + if any( + keyword in name.lower() + for keyword in ( + "feedback", + "context", + "description", + "note", + ) + ): + input_type = "textarea" + elif param.annotation == bool: + input_type = "checkbox" + elif param.annotation in (int, float): + input_type = "number" + + explorer_params.append( + { + "name": name, + "default": default_val, + "input_type": input_type, + } + ) + return explorer_params + except Exception as e: + raise Exception(f"Failed to load signature for {tool_name}: {str(e)}") + + def get_replay_tools(self) -> list[dict]: + """Returns the list of registered tools available for replay.""" + return [ + { + "name": name, + "display_name": cfg["display_name"], + "description": cfg["description"], + } + for name, cfg in REPLAY_TOOLS_CONFIG.items() + ] + + def get_replay_steps(self, session_id: str) -> list[dict]: + """Loads and formats metadata for all chunked steps in the session.""" + session_chunked_dir = self._ensure_session_chunked(session_id) + if not session_chunked_dir.exists(): + raise FileNotFoundError("Step replay chunked data not found for this session.") + + try: + steps_metadata = self.load_steps_metadata(session_id) + step_numbers = sorted(list(steps_metadata.keys())) + + import inspect + import importlib + + # Pre-load signatures and configs for all registered tools + tool_signatures = {} + for t_name, t_cfg in REPLAY_TOOLS_CONFIG.items(): + try: + mod = importlib.import_module(t_cfg["module"]) + fn = getattr(mod, t_cfg["function"]) + tool_signatures[t_name] = inspect.signature(fn) + except Exception as e: + print( + f"Warning: Failed to load signature for tool {t_name}: {e}", + file=sys.stderr, + ) + + steps_list = [] + for num in step_numbers: + s_data = steps_metadata.get(num, {}) + sid = s_data.get("step_id") + + pre_url = None + post_url = None + step_dir = session_chunked_dir / f"step_{num:02d}" + if step_dir.exists(): + pre_jpg = step_dir / "pre.jpg" + if pre_jpg.exists(): + pre_url = f"file://{pre_jpg.resolve()}" + post_jpg = step_dir / "post.jpg" + if post_jpg.exists(): + post_url = f"file://{post_jpg.resolve()}" + if pre_url and post_url and pre_url == post_url: + post_url = None + + # Generate param_defaults for all registered tools + all_tool_defaults = {} + for t_name, sig in tool_signatures.items(): + t_cfg = REPLAY_TOOLS_CONFIG[t_name] + blacklist = t_cfg.get("blacklist_args", []) + fallback_mappings = t_cfg.get("fallback_mappings", {}) + + defaults = {} + for name, param in sig.parameters.items(): + if name in blacklist: + continue + + val = "" + if param.default is not inspect.Parameter.empty: + val = param.default + + # Resolve fallback using mappings + if name in fallback_mappings: + mapping_path = fallback_mappings[name] + parts = mapping_path.split(".") + curr = s_data + for part in parts: + if isinstance(curr, dict): + curr = curr.get(part) + elif isinstance(curr, list): + try: + idx = int(part) + if idx < len(curr): + curr = curr[idx] + else: + curr = None + except ValueError: + curr = None + else: + curr = None + + if curr is not None: + val = curr + + defaults[name] = val + all_tool_defaults[t_name] = defaults + + # Check if this step has already been replayed by scanning outputs folder + step_outputs_dir = self.test_outputs_dir / f"{session_id}_step_{num:02d}" + step_db_path = step_outputs_dir / "temp_traces" / "data_engine.db" + is_replayed = step_db_path.exists() + + steps_list.append( + { + "step_id": str(sid).lower() if sid else None, + "step_number": num, + "timestamp": s_data.get("timestamp"), + "summary": s_data.get("summary", "No summary available."), + "action_taken": s_data.get("action_taken", []), + "query": "", + "param_defaults": all_tool_defaults, + "pre_screenshot": pre_url, + "post_screenshot": post_url, + "operator_raw_thinking": s_data.get("operator_raw_thinking"), + "operator_native_thinking": s_data.get("operator_native_thinking"), + "is_replayed": is_replayed, + } + ) + + return steps_list + except Exception as e: + traceback.print_exc() + raise e + + async def run_step_replay( + self, + session_id: str, + step_number: int, + device_id: str, + user_submits: dict, + tool_name: str = "ask_explorer", + replay_id: str = None, + ) -> dict: + """Triggers sandbox execution of step replay and returns the resulting traces trees.""" + tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name) + if not tool_cfg: + raise ValueError(f"Unknown tool for replay: {tool_name}") + agent_name = tool_cfg["agent_name"] + + # 1. Run live step replay + llm_calls_log = await self.replay_step_tool( + session_id=session_id, + step_number=step_number, + override_device_id=device_id, + user_submits=user_submits, + tool_name=tool_name, + replay_id=replay_id, + ) + + # 2. Extract freshly generated traces + outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}" + temp_db_path = outputs_dir / "temp_traces" / "data_engine.db" + + self._ensure_session_chunked(session_id) + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + _, _, traces_to_preload = self.extract_replay_data( + step_dir, str(self.db_path), agent_name=agent_name + ) + preloaded_trace_ids = {t.get("trace_id") for t in traces_to_preload if t.get("trace_id")} + + all_traces = self.load_all_session_traces(temp_db_path, preloaded_trace_ids, step_number) + + # Align LLM calls and extract final tool result from result.json if available + replay_outcome = None + replay_id = None + result_json_path = outputs_dir / "result.json" + if result_json_path.exists(): + try: + with open(result_json_path, encoding="utf-8") as f: + res_data = json.load(f) + replay_outcome = res_data.get("result") + replay_id = res_data.get("replay_id") + except Exception: + pass + + # Update virtual tool trace result with the actual formatted outcome if available + if replay_outcome is not None: + for node in all_traces: + if ( + not node.get("is_preloaded") + and node["type"] == "tool" + and node["name"] == tool_name + ): + node["payload"]["result"] = replay_outcome + + # Collect all valid LLM span names from REPLAY_TOOLS_CONFIG to align subagent calls too + all_llm_span_names = { + cfg.get("llm_span_name") + for cfg in REPLAY_TOOLS_CONFIG.values() + if cfg.get("llm_span_name") + } + + llm_call_index = 0 + for node in all_traces: + if ( + not node.get("is_preloaded") + and node["type"] == "span" + and node["name"] in all_llm_span_names + and llm_calls_log + and llm_call_index < len(llm_calls_log) + ): + node["type"] = "llm_call" + node["llm_call"] = llm_calls_log[llm_call_index] + llm_call_index += 1 + + # Clean payload paths + for node in all_traces: + if node.get("payload"): + node["payload"] = self.clean_payload_paths(node["payload"]) + if node.get("llm_call"): + node["llm_call"] = self.clean_payload_paths(node["llm_call"]) + + # Filter all_traces for ONLY this step's traces + step_traces = [t for t in all_traces if t.get("step_number") == step_number] + + preloaded_traces = [t for t in step_traces if t.get("is_preloaded")] + live_traces = [t for t in step_traces if not t.get("is_preloaded")] + + # Build preloaded hierarchy + preloaded_nodes = { + str(t["trace_id"]).lower(): {**t, "children": []} for t in preloaded_traces + } + preloaded_root_nodes = [] + for node_id, node in preloaded_nodes.items(): + parent_id = node.get("parent_trace_id") + if parent_id: + parent_id_str = str(parent_id).lower() + if parent_id_str in preloaded_nodes: + preloaded_nodes[parent_id_str]["children"].append(node) + else: + preloaded_root_nodes.append(node) + else: + preloaded_root_nodes.append(node) + preloaded_root_nodes = self.sanitize_nodes(preloaded_root_nodes) + preloaded_root_nodes = self.sort_tree_nodes_chronologically(preloaded_root_nodes) + + # Build live hierarchy + live_nodes = {str(t["trace_id"]).lower(): {**t, "children": []} for t in live_traces} + + # Add a virtual "Response" child node at the end of the virtual tool trace to show the final outcome at the bottom of the tree + if replay_outcome is not None: + virtual_tool_node_id = None + for node_id, node in live_nodes.items(): + if node.get("type") == "tool" and node.get("name") == tool_name: + virtual_tool_node_id = node_id + break + + if virtual_tool_node_id: + import uuid + + outcome_node_id = str(uuid.uuid4()) + max_ts = ( + max([t.get("timestamp", 0.0) for t in live_traces]) + if live_traces + else live_nodes[virtual_tool_node_id].get("timestamp", 0.0) + ) + + outcome_node = { + "trace_id": outcome_node_id, + "session_id": session_id, + "step_id": live_nodes[virtual_tool_node_id].get("step_id"), + "parent_trace_id": live_nodes[virtual_tool_node_id]["trace_id"], + "type": "tool", + "name": f"{tool_name} Response", + "status": (live_nodes[virtual_tool_node_id].get("status", "success")), + "timestamp": max_ts + 0.001, + "duration": 0.0, + "payload": {"result": replay_outcome}, + "is_preloaded": False, + "step_number": step_number, + "children": [], + } + live_nodes[outcome_node_id] = outcome_node + + live_root_nodes = [] + for node_id, node in live_nodes.items(): + parent_id = node.get("parent_trace_id") + if parent_id: + parent_id_str = str(parent_id).lower() + if parent_id_str in live_nodes: + live_nodes[parent_id_str]["children"].append(node) + else: + live_root_nodes.append(node) + else: + live_root_nodes.append(node) + live_root_nodes = self.sanitize_nodes(live_root_nodes) + live_root_nodes = self.sort_tree_nodes_chronologically(live_root_nodes) + + return { + "success": True, + "preloaded": preloaded_root_nodes, + "live": live_root_nodes, + "replay_id": replay_id, + } + + def get_step_replay_traces( + self, session_id: str, step_number: int, tool_name: str = "ask_explorer" + ) -> dict: + """Retrieves previously generated step replay traces if they exist.""" + tool_cfg = REPLAY_TOOLS_CONFIG.get(tool_name) + if not tool_cfg: + raise ValueError(f"Unknown tool for replay: {tool_name}") + agent_name = tool_cfg["agent_name"] + + outputs_dir = self.test_outputs_dir / f"{session_id}_step_{step_number:02d}" + temp_db_path = outputs_dir / "temp_traces" / "data_engine.db" + + if not temp_db_path.exists(): + return {"success": True, "live": []} + + self._ensure_session_chunked(session_id) + step_dir = self.test_data_dir / f"{session_id}_chunked" / f"step_{step_number:02d}" + preloaded_trace_ids = set() + if step_dir.exists(): + _, _, traces_to_preload = self.extract_replay_data( + step_dir, + str(self.db_path), + agent_name=agent_name, + tool_name=tool_name, + ) + preloaded_trace_ids = { + t.get("trace_id") for t in traces_to_preload if t.get("trace_id") + } + + all_traces = self.load_all_session_traces(temp_db_path, preloaded_trace_ids, step_number) + + llm_span_name = tool_cfg.get("llm_span_name") + if not llm_span_name: + raise ValueError(f"Tool {tool_name} config is missing 'llm_span_name'") + + # Align LLM calls and extract final tool result from result.json if available + llm_calls_log = [] + replay_outcome = None + replay_id = None + result_json_path = outputs_dir / "result.json" + if result_json_path.exists(): + try: + with open(result_json_path, encoding="utf-8") as f: + res_data = json.load(f) + llm_calls_log = res_data.get("llm_calls") or [] + replay_outcome = res_data.get("result") + replay_id = res_data.get("replay_id") + except Exception: + pass + + # Update virtual tool trace result with the actual formatted outcome if available + if replay_outcome is not None: + for node in all_traces: + if ( + not node.get("is_preloaded") + and node["type"] == "tool" + and node["name"] == tool_name + ): + node["payload"]["result"] = replay_outcome + + # Collect all valid LLM span names from REPLAY_TOOLS_CONFIG to align subagent calls too + all_llm_span_names = { + cfg.get("llm_span_name") + for cfg in REPLAY_TOOLS_CONFIG.values() + if cfg.get("llm_span_name") + } + + llm_call_index = 0 + for node in all_traces: + if ( + not node.get("is_preloaded") + and node["type"] == "span" + and node["name"] in all_llm_span_names + and llm_calls_log + and llm_call_index < len(llm_calls_log) + ): + node["type"] = "llm_call" + node["llm_call"] = llm_calls_log[llm_call_index] + llm_call_index += 1 + + # Clean payload paths + for node in all_traces: + if node.get("payload"): + node["payload"] = self.clean_payload_paths(node["payload"]) + if node.get("llm_call"): + node["llm_call"] = self.clean_payload_paths(node["llm_call"]) + + # Filter all_traces for ONLY this step's traces + step_traces = [t for t in all_traces if t.get("step_number") == step_number] + live_traces = [t for t in step_traces if not t.get("is_preloaded")] + + # Build live hierarchy + live_nodes = {str(t["trace_id"]).lower(): {**t, "children": []} for t in live_traces} + + # Add a virtual "Response" child node at the end of the virtual tool trace to show the final outcome at the bottom of the tree + if replay_outcome is not None: + virtual_tool_node_id = None + for node_id, node in live_nodes.items(): + if node.get("type") == "tool" and node.get("name") == tool_name: + virtual_tool_node_id = node_id + break + + if virtual_tool_node_id: + import uuid + + outcome_node_id = str(uuid.uuid4()) + max_ts = ( + max([t.get("timestamp", 0.0) for t in live_traces]) + if live_traces + else live_nodes[virtual_tool_node_id].get("timestamp", 0.0) + ) + + outcome_node = { + "trace_id": outcome_node_id, + "session_id": session_id, + "step_id": live_nodes[virtual_tool_node_id].get("step_id"), + "parent_trace_id": live_nodes[virtual_tool_node_id]["trace_id"], + "type": "tool", + "name": f"{tool_name} Response", + "status": (live_nodes[virtual_tool_node_id].get("status", "success")), + "timestamp": max_ts + 0.001, + "duration": 0.0, + "payload": {"result": replay_outcome}, + "is_preloaded": False, + "step_number": step_number, + "children": [], + } + live_nodes[outcome_node_id] = outcome_node + + live_root_nodes = [] + for node_id, node in live_nodes.items(): + parent_id = node.get("parent_trace_id") + if parent_id: + parent_id_str = str(parent_id).lower() + if parent_id_str in live_nodes: + live_nodes[parent_id_str]["children"].append(node) + else: + live_root_nodes.append(node) + else: + live_root_nodes.append(node) + live_root_nodes = self.sanitize_nodes(live_root_nodes) + live_root_nodes = self.sort_tree_nodes_chronologically(live_root_nodes) + + return { + "success": True, + "live": live_root_nodes, + "replay_id": replay_id, + } diff --git a/apps/admin_console/routers/__init__.py b/apps/admin_console/routers/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/routers/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/routers/media.py b/apps/admin_console/routers/media.py new file mode 100644 index 0000000..62806b8 --- /dev/null +++ b/apps/admin_console/routers/media.py @@ -0,0 +1,116 @@ +# 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. + +from pathlib import Path +from fastapi import APIRouter, HTTPException +from fastapi.responses import FileResponse, HTMLResponse + +try: + from admin_console.core.config import IMAGES_DIR, WORKSPACE_ROOT + from admin_console.services.media_service import media_service +except ImportError: + from apps.admin_console.core.config import IMAGES_DIR, WORKSPACE_ROOT + from apps.admin_console.services.media_service import media_service + +router = APIRouter(tags=["media"]) + + +@router.get("/", response_class=HTMLResponse) +@router.get("/admin", response_class=HTMLResponse) +@router.get("/debug", response_class=HTMLResponse) +async def get_index(): + admin_index = Path(__file__).resolve().parent.parent / "index.html" + if admin_index.exists(): + return admin_index.read_text(encoding="utf-8") + + admin_index_alt = WORKSPACE_ROOT / "apps" / "admin_console" / "index.html" + if admin_index_alt.exists(): + return admin_index_alt.read_text(encoding="utf-8") + + return "

Admin Console index.html not found

" + + +@router.get("/images/{image_name}") +@router.get("/api/images/{image_name}") +async def get_image(image_name: str): + if not image_name.endswith(".jpg"): + image_name += ".jpg" + image_path = IMAGES_DIR / image_name + if not image_path.exists(): + raise HTTPException(status_code=404, detail="Image not found") + + return FileResponse(image_path, media_type="image/jpeg") + + +@router.get("/videos/{video_path:path}") +async def get_video(video_path: str): + path = Path(video_path) + if not path.exists(): + if not path.is_absolute(): + candidate = WORKSPACE_ROOT / path + if candidate.exists(): + path = candidate + else: + candidate_root = Path("/" + video_path.lstrip("/")) + if candidate_root.exists(): + path = candidate_root + else: + path = candidate + + if not path.exists(): + raise HTTPException(status_code=404, detail=f"Video not found at {path}") + + suffix = path.suffix.lower() + media_type = "video/mp4" + if suffix == ".webm": + media_type = "video/webm" + elif suffix == ".mkv": + media_type = "video/x-matroska" + + return FileResponse(path, media_type=media_type) + + +@router.get("/api/sessions/{session_id}/video") +async def get_session_video(session_id: str): + try: + from admin_console.database.repositories.session_repository import session_repo + except ImportError: + from apps.admin_console.database.repositories.session_repository import session_repo + + video_rec_map = session_repo.get_video_recordings_map() + video_idx = media_service.build_video_index() + row = session_repo.get_session_by_id(session_id) + row_dict = dict(row) if row else {"session_id": session_id} + v_url = media_service.resolve_video_url(row_dict, video_rec_map, video_idx) + return { + "session_id": session_id, + "has_video": bool(v_url), + "video_url": v_url, + } + + +@router.get("/local_file") +async def get_local_file(path: str): + p, media_type = media_service.get_safe_local_file(path) + return FileResponse(p, media_type=media_type) + + +@router.get("/api/sessions/{session_id}/plan") +async def get_task_plan(session_id: str): + return {"plan": media_service.get_task_plan_content(session_id)} + + +@router.get("/api/sessions/{session_id}/notes") +async def get_all_notes(session_id: str): + return {"notes": media_service.get_session_notes_content(session_id)} diff --git a/apps/admin_console/routers/replay.py b/apps/admin_console/routers/replay.py new file mode 100644 index 0000000..cdf5112 --- /dev/null +++ b/apps/admin_console/routers/replay.py @@ -0,0 +1,91 @@ +# 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. + +import traceback +from fastapi import APIRouter, HTTPException + +try: + from admin_console.core.config import WORKSPACE_ROOT + from admin_console.schemas.task_schema import ReplayRequest + from admin_console.replay_manager import ReplayManager +except ImportError: + from apps.admin_console.core.config import WORKSPACE_ROOT + from apps.admin_console.schemas.task_schema import ReplayRequest + from apps.admin_console.replay_manager import ReplayManager + +replay_manager = ReplayManager(WORKSPACE_ROOT) + + +router = APIRouter(tags=["replay"]) + + +@router.get("/api/devices") +async def list_devices(): + """Dynamically queries the ADB server for connected Android devices.""" + return replay_manager.list_devices() + + +@router.get("/api/replay/tools") +async def get_replay_tools(): + """Returns the list of registered tools available for replay.""" + return replay_manager.get_replay_tools() + + +@router.get("/api/replay/config") +async def get_replay_config(tool_name: str = "ask_explorer"): + """Returns the parameter list and types of the selected tool logic.""" + try: + return replay_manager.get_replay_config(tool_name) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/sessions/{session_id}/replay_steps") +async def get_replay_steps(session_id: str): + """Loads and formats metadata for all chunked steps in the session.""" + try: + return replay_manager.get_replay_steps(session_id) + except FileNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post("/api/sessions/{session_id}/steps/{step_number}/replay") +async def trigger_step_replay_endpoint(session_id: str, step_number: int, req: ReplayRequest): + """Triggers sandbox execution of step replay and returns the resulting traces trees.""" + try: + return await replay_manager.run_step_replay( + session_id=session_id, + step_number=step_number, + device_id=req.device_id, + user_submits=req.user_submits, + tool_name=req.tool_name, + replay_id=req.replay_id, + ) + except Exception as e: + traceback.print_exc() + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/sessions/{session_id}/steps/{step_number}/replay_traces") +async def get_step_replay_traces_endpoint( + session_id: str, step_number: int, tool_name: str = "ask_explorer" +): + """Retrieves previously generated step replay traces if they exist.""" + try: + return replay_manager.get_step_replay_traces(session_id, step_number, tool_name) + except Exception as e: + traceback.print_exc() + raise HTTPException(status_code=500, detail=str(e)) diff --git a/apps/admin_console/routers/sessions.py b/apps/admin_console/routers/sessions.py new file mode 100644 index 0000000..69258e9 --- /dev/null +++ b/apps/admin_console/routers/sessions.py @@ -0,0 +1,127 @@ +# 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. + +import time +from uuid import UUID +from fastapi import APIRouter, HTTPException + +try: + from admin_console.core.config import DB_PATH, TRACES_PATH + from admin_console.core.state import state + from admin_console.database.repositories.session_repository import session_repo + from admin_console.database.repositories.trace_repository import trace_repo + from admin_console.services.media_service import media_service + from admin_console.services.model_service import model_service +except ImportError: + from apps.admin_console.core.config import DB_PATH, TRACES_PATH + from apps.admin_console.core.state import state + from apps.admin_console.database.repositories.session_repository import session_repo + from apps.admin_console.database.repositories.trace_repository import trace_repo + from apps.admin_console.services.media_service import media_service + from apps.admin_console.services.model_service import model_service + + +router = APIRouter(tags=["sessions"]) + + +@router.get("/api/sessions") +async def list_sessions(): + try: + rows = session_repo.get_all_sessions() + video_rec_map = session_repo.get_video_recordings_map() + video_idx = media_service.build_video_index() + + orphaned_ids = [] + result = [] + + for row_dict in rows: + s_id = str(row_dict.get("session_id")) + if row_dict.get("status") == "running": + is_active = s_id in state.active_connections or ( + state.is_running and s_id == str(state.active_session_id) + ) + if not is_active: + row_dict["status"] = "failed" + row_dict["end_time"] = time.time() + orphaned_ids.append(s_id) + + row_dict["video_url"] = media_service.resolve_video_url( + row_dict, video_rec_map, video_idx + ) + + llm_traces = session_repo.get_llm_traces_for_profile(s_id) + agent_names = session_repo.get_agent_trace_names(s_id) + sess_profile = model_service.resolve_session_profile( + row_dict, llm_traces, state.current_profile, agent_names=agent_names + ) + + row_dict["model_info"] = model_service.get_active_model_info(sess_profile) + result.append(row_dict) + + if orphaned_ids: + try: + harvested = session_repo.harvest_orphaned_sessions(orphaned_ids) + print(f"[list_sessions] Auto-harvested {harvested} orphaned running session(s).") + except Exception as e: + print(f"[list_sessions] Failed to update orphaned sessions: {e}") + + return result + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/sessions/{session_id}/tree") +async def get_tree(session_id: str): + try: + return trace_repo.get_trace_tree(session_id) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/sessions/{session_id}/background_tasks") +async def get_session_background_tasks(session_id: str): + try: + return session_repo.get_background_tasks(session_id) + except Exception: + return [] + + +@router.post("/api/cleanup") +async def cleanup_history_endpoint(): + try: + from artemis.data_engine.storage import StorageManager + + storage = StorageManager(DB_PATH, TRACES_PATH) + storage.clear_all_data() + return { + "status": "success", + "message": "History cleaned up successfully", + } + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post("/api/sessions/{session_id}/delete") +async def delete_session_endpoint(session_id: str): + try: + from artemis.data_engine.storage import StorageManager + + storage = StorageManager(DB_PATH, TRACES_PATH) + storage.delete_session(UUID(session_id)) + return { + "status": "success", + "message": f"Session {session_id} deleted successfully", + } + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/apps/admin_console/routers/steps.py b/apps/admin_console/routers/steps.py new file mode 100644 index 0000000..331b86a --- /dev/null +++ b/apps/admin_console/routers/steps.py @@ -0,0 +1,108 @@ +# 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. + +import json +import traceback +from fastapi import APIRouter, HTTPException +from fastapi.responses import PlainTextResponse + +try: + from admin_console.core.config import TEST_OUTPUTS_DIR + from admin_console.database.repositories.step_repository import step_repo + from admin_console.database.repositories.trace_repository import trace_repo + from admin_console.services.media_service import media_service +except ImportError: + from apps.admin_console.core.config import TEST_OUTPUTS_DIR + from apps.admin_console.database.repositories.step_repository import step_repo + from apps.admin_console.database.repositories.trace_repository import trace_repo + from apps.admin_console.services.media_service import media_service + + +router = APIRouter(tags=["steps"]) + + +@router.get("/api/sessions/{session_id}/steps") +async def get_session_steps(session_id: str, client: str | None = None): + try: + return step_repo.get_session_steps(session_id, client=client) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/steps/{step_id}/traces") +async def get_step_traces_endpoint(step_id: str): + try: + session_id = step_repo.get_step_session_id(step_id) + if not session_id: + raise HTTPException(status_code=404, detail="Step not found") + return trace_repo.get_step_traces_tree(session_id, step_id) + except HTTPException as e: + raise e + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +def _resolve_trace_db_path(session_id: str | None, step_number: int | None): + if session_id and step_number is not None: + sandbox_db_path = ( + TEST_OUTPUTS_DIR + / f"{session_id}_step_{step_number:02d}" + / "temp_traces" + / "data_engine.db" + ) + if sandbox_db_path.exists(): + return sandbox_db_path + return None + + +@router.get("/api/traces/{trace_id}") +async def get_trace(trace_id: str, session_id: str = None, step_number: int = None): + try: + db_path = _resolve_trace_db_path(session_id, step_number) + trace_dict = trace_repo.get_trace_by_id(trace_id, db_path=db_path) + if not trace_dict: + raise HTTPException(status_code=404, detail="Trace not found") + + if trace_dict.get("payload"): + try: + payload_obj = json.loads(trace_dict["payload"]) + trace_dict["payload"] = media_service.unwrap_payload(payload_obj) + except Exception: + pass + + return trace_dict + except HTTPException as e: + raise e + except Exception as e: + traceback.print_exc() + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get("/api/traces/{trace_id}/download") +async def download_trace(trace_id: str, session_id: str = None, step_number: int = None): + try: + db_path = _resolve_trace_db_path(session_id, step_number) + trace_dict = trace_repo.get_trace_by_id(trace_id, db_path=db_path) + if not trace_dict or not trace_dict.get("payload"): + raise HTTPException(status_code=404, detail="Payload not found") + + return PlainTextResponse( + trace_dict["payload"], + headers={"Content-Disposition": f'attachment; filename="trace_{trace_id}.json"'}, + ) + except HTTPException as e: + raise e + except Exception as e: + traceback.print_exc() + raise HTTPException(status_code=500, detail=str(e)) diff --git a/apps/admin_console/routers/stream.py b/apps/admin_console/routers/stream.py new file mode 100644 index 0000000..5bc3c68 --- /dev/null +++ b/apps/admin_console/routers/stream.py @@ -0,0 +1,56 @@ +# 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. + +"""Device Screen Live Streaming Router. + +Exposes real-time screen streaming endpoints for the Web UI. +""" + +from fastapi import APIRouter +from fastapi.responses import JSONResponse, StreamingResponse + +try: + from admin_console.services.device_stream_service import device_stream_service +except ImportError: + from apps.admin_console.services.device_stream_service import device_stream_service + +router = APIRouter(tags=["stream"]) + + +@router.get("/api/stream/device-live") +async def stream_device_live(): + """Streams live device screen frames as multipart MJPEG.""" + return StreamingResponse( + device_stream_service.mjpeg_frame_generator(), + media_type="multipart/x-mixed-replace; boundary=frame", + headers={ + "Cache-Control": "no-cache, no-store, must-revalidate, max-age=0", + "Pragma": "no-cache", + "Expires": "0", + "Connection": "keep-alive", + }, + ) + + +@router.get("/api/stream/device-state") +async def get_device_stream_state(): + """Returns whether an ADB device is connected and live streaming is available.""" + serial = await device_stream_service.get_device_serial() + return JSONResponse( + { + "connected": serial is not None, + "serial": serial, + "live_stream_url": "/api/stream/device-live" if serial else None, + } + ) diff --git a/apps/admin_console/routers/tasks.py b/apps/admin_console/routers/tasks.py new file mode 100644 index 0000000..a989ace --- /dev/null +++ b/apps/admin_console/routers/tasks.py @@ -0,0 +1,189 @@ +# 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. + +import asyncio +import json +from fastapi import APIRouter, HTTPException +from fastapi.responses import StreamingResponse + +try: + from admin_console.core.state import state + from admin_console.database.repositories.session_repository import session_repo + from admin_console.schemas.task_schema import RunRequest + from admin_console.services.ipc_service import ipc_service + from admin_console.services.model_service import model_service + from admin_console.services.task_queue_service import task_queue_service +except ImportError: + from apps.admin_console.core.state import state + from apps.admin_console.database.repositories.session_repository import session_repo + from apps.admin_console.schemas.task_schema import RunRequest + from apps.admin_console.services.ipc_service import ipc_service + from apps.admin_console.services.model_service import model_service + from apps.admin_console.services.task_queue_service import task_queue_service + + +router = APIRouter(tags=["tasks"]) + + +@router.post("/api/run") +async def run_task(request: RunRequest): + incoming_goals = [] + if request.goals: + incoming_goals = request.goals + elif request.goal: + incoming_goals = [request.goal] + + if not incoming_goals: + raise HTTPException( + status_code=400, + detail="Either 'goal' or 'goals' list must be provided.", + ) + + return await task_queue_service.enqueue_tasks( + incoming_goals, profile=request.profile or "flash" + ) + + +@router.post("/api/stop") +async def stop_task(all: bool = False): + stopped = task_queue_service.stop_tasks(clear_all=all) + if stopped: + return {"status": "stopped"} + return {"status": "no_running_task"} + + +@router.post("/api/resume") +async def resume_task(): + resumed = task_queue_service.resume_task() + if resumed: + return {"status": "resumed"} + return {"status": "not_paused"} + + +@router.get("/api/status") +async def get_status(): + latest_session = session_repo.get_latest_session() + latest_session_id = latest_session.get("session_id") if latest_session else None + bg_tasks = session_repo.get_background_tasks(latest_session_id) if latest_session_id else [] + + is_running = state.is_running + running_sid = state.active_session_id or ( + session_repo.get_running_session_id() if is_running else None + ) + active_profile = state.current_profile + if not active_profile and (running_sid or latest_session_id): + check_sid = running_sid or latest_session_id + sess_row = session_repo.get_session_by_id(check_sid) + if sess_row: + llm_traces = session_repo.get_llm_traces_for_profile(check_sid) + agent_names = session_repo.get_agent_trace_names(check_sid) + active_profile = model_service.resolve_session_profile( + sess_row, llm_traces, agent_names=agent_names + ) + + model_info = model_service.get_active_model_info(active_profile) + queue_data = state.queue_tasks + + if state.current_process: + if state.current_process.returncode is None: + is_paused = state.is_paused + return { + "status": "paused" if is_paused else "running", + "goal": state.current_goal, + "pid": state.current_process.pid, + "session_id": running_sid, + "background_tasks": bg_tasks, + "queue": queue_data, + "model_info": model_info, + } + else: + completed_status = { + "status": "completed", + "goal": state.current_goal, + "returncode": state.current_process.returncode, + "session_id": latest_session_id, + "background_tasks": bg_tasks, + "queue": queue_data, + "model_info": model_info, + } + state.current_process = None + return completed_status + + if latest_session_id and str(latest_session_id) in state.active_connections: + conn_info = state.active_connections[str(latest_session_id)] + is_paused = state.is_paused + conn_profile = conn_info.get("profile") or active_profile + conn_model_info = model_service.get_active_model_info(conn_profile) + return { + "status": "paused" if is_paused else "running", + "goal": conn_info.get("goal"), + "pid": conn_info.get("pid"), + "session_id": latest_session_id, + "background_tasks": bg_tasks, + "queue": queue_data, + "model_info": conn_model_info, + } + + return { + "status": "idle", + "session_id": latest_session_id, + "background_tasks": bg_tasks, + "queue": queue_data, + "model_info": model_info, + } + + +@router.get("/api/stream/{session_id}") +async def stream_events(session_id: str, client: str | None = None): + async def event_generator(): + queue = asyncio.Queue() + + def callback(event_type, data): + try: + # Filter events by session_id when subscribed to a specific session + if session_id and session_id not in ("all", "active"): + evt_session_id = None + if isinstance(data, dict): + evt_session_id = data.get("session_id") + + if evt_session_id and str(evt_session_id) != str(session_id): + return + + if ( + not evt_session_id + and state.active_session_id + and str(state.active_session_id) != str(session_id) + ): + return + + sanitized_data = ipc_service.sanitize_event_data(event_type, data) + loop = asyncio.get_running_loop() + loop.call_soon_threadsafe(queue.put_nowait, (event_type, sanitized_data)) + except RuntimeError: + pass + + state.add_subscriber(callback) + yield f'event: info\ndata: {{"message": "Subscribed to session {session_id}"}}\n\n' + + try: + while True: + try: + event_type, data = await asyncio.wait_for(queue.get(), timeout=5.0) + yield f"event: {event_type}\ndata: {json.dumps(data, default=str)}\n\n" + except TimeoutError: + yield "event: keep-alive\ndata: {}\n\n" + except asyncio.CancelledError: + state.remove_subscriber(callback) + + return StreamingResponse(event_generator(), media_type="text/event-stream") diff --git a/apps/admin_console/schemas/__init__.py b/apps/admin_console/schemas/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/schemas/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/schemas/task_schema.py b/apps/admin_console/schemas/task_schema.py new file mode 100644 index 0000000..9b6d491 --- /dev/null +++ b/apps/admin_console/schemas/task_schema.py @@ -0,0 +1,28 @@ +# 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. + +from pydantic import BaseModel + + +class RunRequest(BaseModel): + goal: str | None = None + goals: list[str] | None = None + profile: str | None = "flash" + + +class ReplayRequest(BaseModel): + device_id: str + user_submits: dict + tool_name: str = "ask_explorer" + replay_id: str | None = None diff --git a/apps/admin_console/server.py b/apps/admin_console/server.py new file mode 100644 index 0000000..be39134 --- /dev/null +++ b/apps/admin_console/server.py @@ -0,0 +1,154 @@ +# 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. + +"""Artemis Admin & Trace Console Server + +Modular entrypoint for full trace inspection, step replay, and task execution management. +""" + +import asyncio +import os +from pathlib import Path +import sys + +# Bootstrap sys.path to allow running from any CWD +_current_p = Path(__file__).resolve().parent +while _current_p != _current_p.parent: + if (_current_p / "pyproject.toml").exists() or (_current_p / "artemis").is_dir(): + _workspace_root = _current_p + break + _current_p = _current_p.parent +else: + _workspace_root = Path(__file__).resolve().parent.parent.parent + +_apps_dir = _workspace_root / "apps" +_admin_console_dir = _apps_dir / "admin_console" +_cloud_service_dir = _apps_dir / "cloud_service" + +for _p in (str(_workspace_root), str(_apps_dir), str(_admin_console_dir), str(_cloud_service_dir)): + if _p not in sys.path: + sys.path.insert(0, _p) + +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware + +try: + from admin_console.core.config import ( + DB_PATH, + IMAGES_DIR, + REPLAY_BASE_DIR, + TEST_DATA_DIR, + TEST_OUTPUTS_DIR, + TRACES_PATH, + WORKSPACE_ROOT, + init_ls_address, + ) + from admin_console.core.state import ServerState, state + from admin_console.database.connection import db_session, get_db + from admin_console.database.repositories.session_repository import session_repo + from admin_console.database.repositories.step_repository import step_repo + from admin_console.database.repositories.trace_repository import trace_repo + from admin_console.routers import media, replay, sessions, steps, stream, tasks + from admin_console.routers.replay import replay_manager + from admin_console.services.ipc_service import ipc_service + from admin_console.services.media_service import media_service + from admin_console.services.model_service import model_service + from admin_console.services.task_queue_service import task_queue_service +except ImportError: + from apps.admin_console.core.config import ( + init_ls_address, + ) + from apps.admin_console.core.state import state + from apps.admin_console.database.repositories.session_repository import session_repo + from apps.admin_console.routers import media, replay, sessions, steps, stream, tasks + from apps.admin_console.routers.replay import replay_manager + from apps.admin_console.services.ipc_service import ipc_service + from apps.admin_console.services.model_service import model_service + from apps.admin_console.services.task_queue_service import task_queue_service + +# Initialize language server synchronization address +init_ls_address() + +# Initialize FastAPI application +app = FastAPI(title="Artemis Admin & Trace Console") + +# Configure CORS +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.on_event("startup") +async def on_startup(): + """Startup lifecycle hooks.""" + asyncio.create_task(asyncio.to_thread(task_queue_service.archive_older_replays_on_launch)) + asyncio.create_task( + asyncio.to_thread(task_queue_service.verify_chunks_exist_on_launch, replay_manager) + ) + + cleaned = session_repo.cleanup_orphans_on_startup() + if cleaned > 0: + print(f"[ServerStartup] Marked {cleaned} orphaned running session(s) as failed.") + + await ipc_service.start_server() + state.worker_task = asyncio.create_task(task_queue_service.queue_worker()) + + +# Mount modular routers +app.include_router(stream.router) +app.include_router(media.router) +app.include_router(sessions.router) +app.include_router(steps.router) +app.include_router(tasks.router) +app.include_router(replay.router) + +# Mount cloud gateway router for Frappe / Cloud integration if present +try: + from cloud_service.gateway import cloud_router + + app.include_router(cloud_router) +except ImportError: + try: + from apps.cloud_service.gateway import cloud_router + + app.include_router(cloud_router) + except ImportError: + pass + + +# ------------------------------------------------------------------------------ +# Backward compatibility exports +# ------------------------------------------------------------------------------ +get_active_model_info = model_service.get_active_model_info +start_ipc_server = ipc_service.start_server +queue_worker = task_queue_service.queue_worker +active_connections = state.active_connections +ipc_subscribers = state.ipc_subscribers +task_queue = state.task_queue +queue_goals = state.queue_goals + + +def main(): + import uvicorn + + port = int(os.environ.get("ANTIGRAVITY_SIDECAR_WEB_PORT", 8000)) + uvicorn.run(app, host="0.0.0.0", port=port) + + +if __name__ == "__main__": + main() diff --git a/apps/admin_console/services/__init__.py b/apps/admin_console/services/__init__.py new file mode 100644 index 0000000..aee27a2 --- /dev/null +++ b/apps/admin_console/services/__init__.py @@ -0,0 +1 @@ +# Copyright 2026 Google LLC diff --git a/apps/admin_console/services/device_stream_service.py b/apps/admin_console/services/device_stream_service.py new file mode 100644 index 0000000..ad88041 --- /dev/null +++ b/apps/admin_console/services/device_stream_service.py @@ -0,0 +1,133 @@ +# 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. + +"""Device Live Screen Streaming Service. + +Provides real-time, low-latency device screen frames over HTTP MJPEG and WebSocket. +Operates concurrently with ADB agent actions with zero interference. +""" + +import asyncio +import logging +import subprocess +import time +from collections.abc import AsyncGenerator + +logger = logging.getLogger("artemis.stream_service") + + +class DeviceStreamService: + """Manages real-time screen capture and distribution to web clients.""" + + def __init__(self): + self._active_listeners = 0 + self._lock = asyncio.Lock() + self._latest_frame: bytes | None = None + self._last_frame_time: float = 0.0 + self._is_capturing = False + self._capture_task: asyncio.Task | None = None + + async def get_device_serial(self) -> str | None: + """Find the currently connected active ADB device serial.""" + try: + proc = await asyncio.create_subprocess_exec( + "adb", + "devices", + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + lines = stdout.decode().strip().splitlines() + for line in lines[1:]: + parts = line.strip().split() + if len(parts) >= 2 and parts[1] == "device": + return parts[0] + except Exception as e: + logger.warning(f"Error checking adb devices: {e}") + return None + + async def _capture_loop(self): + """Background frame capture loop that runs while listeners > 0.""" + logger.info("[StreamService] Starting live screen capture loop...") + while self._active_listeners > 0: + try: + start_t = time.time() + serial = await self.get_device_serial() + cmd = ( + ["adb", "-s", serial, "exec-out", "screencap", "-p"] + if serial + else ["adb", "exec-out", "screencap", "-p"] + ) + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + if proc.returncode == 0 and len(stdout) > 1000: + self._latest_frame = stdout + self._last_frame_time = time.time() + + elapsed = time.time() - start_t + delay = max(0.03, 0.08 - elapsed) + await asyncio.sleep(delay) + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"[StreamService] Capture error: {e}") + await asyncio.sleep(0.5) + + logger.info("[StreamService] Stopping live screen capture loop (0 listeners).") + self._is_capturing = False + + async def start_capturing(self): + """Register a new listener and start background capture if needed.""" + async with self._lock: + self._active_listeners += 1 + if not self._is_capturing or self._capture_task is None or self._capture_task.done(): + self._is_capturing = True + self._capture_task = asyncio.create_task(self._capture_loop()) + + async def stop_capturing(self): + """Deregister a listener and stop capture loop when count reaches 0.""" + async with self._lock: + self._active_listeners = max(0, self._active_listeners - 1) + if self._active_listeners == 0 and self._capture_task and not self._capture_task.done(): + self._capture_task.cancel() + self._is_capturing = False + + async def mjpeg_frame_generator(self) -> AsyncGenerator[bytes, None]: + """Async generator streaming MJPEG multipart bytes to HTTP response.""" + await self.start_capturing() + try: + last_sent_time = 0.0 + while True: + if self._latest_frame and self._last_frame_time > last_sent_time: + last_sent_time = self._last_frame_time + frame_bytes = self._latest_frame + yield ( + b"--frame\r\n" + b"Content-Type: image/png\r\n" + b"Content-Length: " + + str(len(frame_bytes)).encode() + + b"\r\n\r\n" + + frame_bytes + + b"\r\n" + ) + await asyncio.sleep(0.04) + finally: + await self.stop_capturing() + + +device_stream_service = DeviceStreamService() diff --git a/apps/admin_console/services/ipc_service.py b/apps/admin_console/services/ipc_service.py new file mode 100644 index 0000000..24618cd --- /dev/null +++ b/apps/admin_console/services/ipc_service.py @@ -0,0 +1,204 @@ +# 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. + +import asyncio +import json +from typing import Any + +try: + from admin_console.core.state import state +except ImportError: + from apps.admin_console.core.state import state + + +class IPCService: + """Service managing internal IPC socket server and event distribution.""" + + @staticmethod + def _clean_value(val: Any) -> Any: + """Recursively cleans values to remove non-serializable and internal + Python runtime representations. + """ + if val is None: + return None + if isinstance(val, (bool, int, float)): + return val + if isinstance(val, str): + if ( + "object at 0x" in val + or val.startswith(" Any: + """Sanitizes stream event payloads for all clients, stripping bulky + raw bytes and internal memory references. + """ + if not isinstance(data, dict): + return data + + if event_type == "trace_recorded": + tr_type = data.get("type") + payload = data.get("payload") + if isinstance(payload, str): + try: + payload = json.loads(payload) + except Exception: + pass + + if tr_type == "tool" and isinstance(payload, dict): + args = payload.get("args", payload) + data["payload"] = {"args": cls._clean_value(args)} + elif tr_type == "llm_call" and isinstance(payload, dict): + data["payload"] = {"error": cls._clean_value(payload.get("error"))} + elif payload is not None: + data["payload"] = cls._clean_value(payload) + + elif event_type in ("step_recorded", "step_updated"): + # Strip heavy raw image bytes from SSE broadcast (clients fetch via image URLs) + data["pre_screenshot_bytes"] = None + data["post_screenshot_bytes"] = None + + if "generic_tools" in data and isinstance(data["generic_tools"], list): + sanitized_tools = [] + for tr in data["generic_tools"]: + if not isinstance(tr, dict): + continue + tr_dict = dict(tr) + pl = tr_dict.get("payload") + if isinstance(pl, str): + try: + pl = json.loads(pl) + except Exception: + pass + + if tr_dict.get("type") == "llm_call": + tr_dict["payload"] = ( + {"error": cls._clean_value(pl.get("error"))} + if isinstance(pl, dict) + else None + ) + elif isinstance(pl, dict): + args = pl.get("args", pl) + tr_dict["payload"] = {"args": cls._clean_value(args)} + else: + tr_dict["payload"] = cls._clean_value(pl) + sanitized_tools.append(tr_dict) + data["generic_tools"] = sanitized_tools + + try: + from artemis.utils.coordinates import normalize_step_actions + + data = normalize_step_actions(data) + except ImportError: + pass + + return data + + # Backward compatibility alias + filter_event_for_angular = sanitize_event_data + + @classmethod + async def start_server(cls): + async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter): + current_session_id = None + while True: + line = await reader.readline() + if not line: + break + try: + payload = json.loads(line.decode("utf-8")) + event_type = payload.get("event_type") + data = payload.get("data") + + if event_type == "session_started" and data: + state.active_session_id = data.get("session_id") + current_session_id = data.get("session_id") + pid = data.get("pid") + goal = data.get("initial_goal") + device_info = data.get("device_info") + if isinstance(device_info, dict) and device_info.get("profile"): + state.current_profile = device_info.get("profile") + if goal: + state.current_goal = goal + print( + f"[IPC] Session started: {state.active_session_id} (profile: {state.current_profile})" + ) + if current_session_id: + state.active_connections[str(current_session_id)] = { + "writer": writer, + "pid": pid, + "goal": goal, + "profile": state.current_profile, + } + elif event_type == "session_ended" and data: + ended_sid = data.get("session_id") + if state.active_session_id and str(state.active_session_id) == str( + ended_sid + ): + state.active_session_id = None + state.current_goal = None + state.current_profile = None + + for cb in list(state.ipc_subscribers): + try: + cb(event_type, data) + except Exception: + pass + except Exception as e: + print(f"IPC parse error: {e}") + + if current_session_id and str(current_session_id) in state.active_connections: + try: + del state.active_connections[str(current_session_id)] + except KeyError: + pass + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(handle_client, "127.0.0.1", 0, limit=1024 * 1024 * 100) + state.ipc_port = server.sockets[0].getsockname()[1] + state.ipc_server = server + asyncio.create_task(server.serve_forever()) + print(f"Internal IPC server started on port {state.ipc_port}") + try: + from artemis.config import write_ipc_port + + write_ipc_port(state.ipc_port) + except Exception as e: + print(f"Failed to write IPC port file: {e}") + + +ipc_service = IPCService() diff --git a/apps/admin_console/services/media_service.py b/apps/admin_console/services/media_service.py new file mode 100644 index 0000000..2e566de --- /dev/null +++ b/apps/admin_console/services/media_service.py @@ -0,0 +1,261 @@ +# 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. + +import base64 +import hashlib +import json +import os +from pathlib import Path +from typing import Any +import urllib.parse +from fastapi import HTTPException + +try: + from admin_console.core.config import IMAGES_DIR, TRACES_PATH, WORKSPACE_ROOT +except ImportError: + from apps.admin_console.core.config import IMAGES_DIR, TRACES_PATH, WORKSPACE_ROOT + + +class MediaService: + """Service handling media indexing, local file security, payload unwrapping and notes/plans.""" + + @staticmethod + def path_to_video_url(p: Path) -> str: + resolved = p.resolve() + try: + rel = resolved.relative_to(WORKSPACE_ROOT) + return f"/videos/{rel}" + except ValueError: + return f"/videos/{str(resolved).lstrip('/')}" + + @classmethod + def build_video_index(cls) -> dict[str, str]: + idx = {} + dirs = [ + TRACES_PATH, + WORKSPACE_ROOT / "artemis-traces", + WORKSPACE_ROOT / ".benchmarks" / "diagnoser" / "outputs" / "artifacts", + ] + for d in dirs: + if not d.exists(): + continue + try: + for item in d.iterdir(): + if item.is_dir(): + for ext in [".mp4", ".mkv", ".webm"]: + vfile = item / f"recording{ext}" + if vfile.exists(): + url = cls.path_to_video_url(vfile) + idx[item.name] = url + prefix = item.name.split("_PASS_")[0].split("_FAIL_")[0] + idx[prefix] = url + break + elif item.suffix in [".mp4", ".mkv", ".webm"]: + url = cls.path_to_video_url(item) + idx[item.stem] = url + idx[item.name] = url + except Exception: + pass + return idx + + @classmethod + def resolve_video_url( + cls, + row_dict: dict[str, Any], + video_rec_map: dict[str, str], + video_idx: dict[str, str], + ) -> str | None: + raw_sid = row_dict.get("session_id") + s_id = ( + str(raw_sid).strip() + if raw_sid is not None and str(raw_sid).strip().lower() != "none" + else "" + ) + v_url = None + v_fp = row_dict.get("video_filepath") + if v_fp and os.path.exists(v_fp): + v_url = cls.path_to_video_url(Path(v_fp)) + + orig_rec = video_rec_map.get(s_id) + if not v_url and orig_rec: + p = Path(orig_rec) + if p.exists(): + v_url = cls.path_to_video_url(p) + else: + folder_name = p.parent.name + v_url = video_idx.get(folder_name) or video_idx.get(folder_name.split("_")[0]) + + if not v_url: + v_url = video_idx.get(s_id) + + if not v_url and s_id: + sess_dir = TRACES_PATH / s_id + if sess_dir.is_dir(): + for ext in [".mp4", ".webm", ".mkv"]: + vfile = sess_dir / f"recording{ext}" + if vfile.exists(): + v_url = cls.path_to_video_url(vfile) + break + if not v_url: + try: + for item in sess_dir.iterdir(): + if item.is_file() and item.suffix.lower() in [".mp4", ".webm", ".mkv"]: + v_url = cls.path_to_video_url(item) + break + except Exception: + pass + + if not v_url and row_dict.get("initial_goal"): + goal_str = row_dict["initial_goal"] + task_hint = ( + goal_str.split("Task: ")[1].split()[0].strip() + if "Task: " in goal_str + else goal_str.split("\n")[0].strip() + ) + v_url = video_idx.get(task_hint) + + return v_url + + @staticmethod + def unwrap_payload(obj: Any) -> Any: + """Unwraps trace payloads, storing base64 images as disk files and + truncating large dumps. + """ + if isinstance(obj, str): + try: + if (obj.startswith("{") and obj.endswith("}")) or ( + obj.startswith("[") and obj.endswith("]") + ): + return MediaService.unwrap_payload(json.loads(obj)) + except Exception: + pass + + # Check if the string is a Base64 image + is_base64_img = False + base64_data = None + if obj.startswith("data:image/") and ";base64," in obj: + try: + _, base64_data = obj.split(";base64,", 1) + is_base64_img = True + except Exception: + pass + elif obj.startswith("iVBORw0KGgo") or obj.startswith("/9j/"): + is_base64_img = True + base64_data = obj + + if is_base64_img and base64_data: + try: + image_bytes = base64.b64decode(base64_data) + image_hash = hashlib.sha256(image_bytes).hexdigest() + + IMAGES_DIR.mkdir(parents=True, exist_ok=True) + image_path = IMAGES_DIR / f"{image_hash}.jpg" + if not image_path.exists(): + image_path.write_bytes(image_bytes) + + return f"image://{image_hash}" + except Exception: + pass + + if len(obj) > 2000: + preview_start = obj[:100].replace("\n", " ") + preview_end = obj[-100:].replace("\n", " ") + return ( + f"" + ) + return obj + elif isinstance(obj, dict): + return {k: MediaService.unwrap_payload(v) for k, v in obj.items()} + elif isinstance(obj, list): + if len(obj) > 20 and any( + isinstance(x, dict) and ("bounds" in x or "resource-id" in x) for x in obj + ): + return [f""] + return [MediaService.unwrap_payload(x) for x in obj] + return obj + + @staticmethod + def get_safe_local_file(path_str: str) -> tuple[Path, str]: + decoded_path = urllib.parse.unquote(path_str) + if decoded_path.startswith("file://"): + decoded_path = decoded_path[len("file://") :] + p = Path(decoded_path).resolve() + + if not p.is_file(): + raise HTTPException(status_code=404, detail="File not found") + + if not str(p).startswith(str(WORKSPACE_ROOT)): + raise HTTPException( + status_code=403, + detail="Access denied. Path is outside workspace root.", + ) + + ext = p.suffix.lower() + media_type = "application/octet-stream" + if ext in (".jpg", ".jpeg"): + media_type = "image/jpeg" + elif ext == ".png": + media_type = "image/png" + elif ext == ".mp4": + media_type = "video/mp4" + + return p, media_type + + @staticmethod + def get_task_plan_content(session_id: str) -> str: + plan_path = TRACES_PATH / session_id / "notes" / "task_plan.md" + if not plan_path.exists(): + plan_path = TRACES_PATH / "notes" / "task_plan.md" + + if plan_path.exists(): + try: + return plan_path.read_text(encoding="utf-8") + except Exception as e: + return f"Error reading task plan: {e}" + return "No task plan created yet." + + @staticmethod + def get_session_notes_content(session_id: str) -> dict[str, str]: + notes_dir = TRACES_PATH / session_id / "notes" + notes_content = {} + + if notes_dir.exists() and notes_dir.is_dir(): + for file in notes_dir.iterdir(): + if file.is_file() and file.suffix in [".md", ".txt", ".json", ".yaml"]: + try: + notes_content[file.name] = file.read_text(encoding="utf-8") + except Exception as e: + notes_content[file.name] = f"Error reading file: {e}" + + if not notes_content: + global_notes_dir = TRACES_PATH / "notes" + if global_notes_dir.exists() and global_notes_dir.is_dir(): + for file in global_notes_dir.iterdir(): + if file.is_file() and file.suffix in [ + ".md", + ".txt", + ".json", + ".yaml", + ]: + try: + notes_content[file.name] = file.read_text(encoding="utf-8") + except Exception as e: + notes_content[file.name] = f"Error reading file: {e}" + + return notes_content + + +media_service = MediaService() diff --git a/apps/admin_console/services/model_service.py b/apps/admin_console/services/model_service.py new file mode 100644 index 0000000..405f6b9 --- /dev/null +++ b/apps/admin_console/services/model_service.py @@ -0,0 +1,133 @@ +# 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. + +import json +from typing import Any + + +class ModelService: + """Service handling model profile detection and metadata formatting.""" + + @staticmethod + def get_active_model_info(profile: str | None = None) -> dict[str, str]: + if profile: + p_lower = str(profile).lower() + if "pro" in p_lower: + return { + "name": "Pro", + "id": "gemini-3.6-pro", + "provider": "google", + } + elif "flash" in p_lower: + return { + "name": "Flash", + "id": "gemini-3.6-flash", + "provider": "google", + } + + model_name = "Flash" + model_id = "gemini-3.6-flash" + provider = "google" + try: + from artemis.config import parse_llm_config + + llm_cfg = parse_llm_config() + if llm_cfg and llm_cfg.operator: + op_model = str(llm_cfg.operator.model).lower() + provider = str(llm_cfg.operator.provider) + model_id = llm_cfg.operator.model + if "pro" in op_model: + model_name = "Pro" + elif "flash" in op_model: + model_name = "Flash" + else: + model_name = llm_cfg.operator.model.capitalize() + except Exception: + pass + return {"name": model_name, "id": model_id, "provider": provider} + + @staticmethod + def resolve_session_profile( + row_dict: dict[str, Any], + llm_trace_payloads: list[str] | None = None, + running_profile: str | None = None, + agent_names: list[str] | None = None, + ) -> str | None: + sess_profile = None + d_info_raw = row_dict.get("device_info") + if d_info_raw: + try: + d_info = json.loads(d_info_raw) if isinstance(d_info_raw, str) else d_info_raw + if isinstance(d_info, dict): + p = d_info.get("profile") + if p: + p_str = str(p).lower() + if "pro" in p_str: + return "pro" + elif "flash" in p_str: + return "flash" + except Exception: + pass + + if (row_dict.get("status") == "running") and running_profile: + p_str = str(running_profile).lower() + if "pro" in p_str: + return "pro" + elif "flash" in p_str: + return "flash" + + # Check Agent/Trace names - the most definitive indicator of architecture + if agent_names: + agent_names_lower = [str(a).lower() for a in agent_names] + if any( + name in ("planner", "validator", "summarizer", "operator", "checker", "diagnoser") + for name in agent_names_lower + ): + return "pro" + if any("flashrunner" in name for name in agent_names_lower): + return "flash" + + # Check LLM traces for agent name or model overrides + if llm_trace_payloads: + for tr_payload in llm_trace_payloads: + if tr_payload: + try: + p_obj = ( + json.loads(tr_payload) if isinstance(tr_payload, str) else tr_payload + ) + if isinstance(p_obj, dict): + agent_name = str(p_obj.get("agent") or p_obj.get("name") or "").lower() + if "flashrunner" in agent_name: + return "flash" + if any( + x in agent_name + for x in ("planner", "validator", "operator", "checker") + ): + return "pro" + + m_str = str(p_obj.get("model") or p_obj.get("model_name") or "").lower() + if ( + "gemini-3.6-pro" in m_str + or "gemini-1.5-pro" in m_str + or "gpt-4" in m_str + or "claude" in m_str + ): + return "pro" + except Exception: + pass + + return None + + +model_service = ModelService() diff --git a/apps/admin_console/services/task_queue_service.py b/apps/admin_console/services/task_queue_service.py new file mode 100644 index 0000000..12c64e1 --- /dev/null +++ b/apps/admin_console/services/task_queue_service.py @@ -0,0 +1,256 @@ +# 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. + +import asyncio +from datetime import datetime +import os +import shutil +import signal +import sys +import time +from typing import Any +import uuid + +try: + from admin_console.core.config import ( + PAUSE_FILE, + TEST_DATA_DIR, + TEST_OUTPUTS_DIR, + WORKSPACE_ROOT, + ) + from admin_console.core.state import state + from admin_console.database.repositories.session_repository import session_repo +except ImportError: + from apps.admin_console.core.config import ( + PAUSE_FILE, + TEST_DATA_DIR, + TEST_OUTPUTS_DIR, + WORKSPACE_ROOT, + ) + from apps.admin_console.core.state import state + from apps.admin_console.database.repositories.session_repository import session_repo + + +class TaskQueueService: + """Service managing FIFO task execution, background worker, subprocess lifecycle, + and startup tasks. + """ + + @classmethod + async def queue_worker(cls): + while True: + goal = None + sess_id = None + try: + item = await state.task_queue.get() + if isinstance(item, dict): + goal = item.get("goal") + profile = item.get("profile", "flash") + sess_id = item.get("session_id") + else: + goal = item + profile = "flash" + sess_id = None + + state.current_goal = goal + state.current_profile = profile + if sess_id: + state.active_session_id = sess_id + + # Remove from state.queue_tasks + state.queue_tasks = [ + t + for t in state.queue_tasks + if not ( + isinstance(t, dict) + and (t.get("session_id") == sess_id or t.get("goal") == goal) + ) + ] + + test_name = f"web_{int(time.time())}" + env = os.environ.copy() + pythonpath_parts = [ + str(WORKSPACE_ROOT), + str(WORKSPACE_ROOT / "apps/admin_console"), + str(WORKSPACE_ROOT / "apps/cloud_service"), + env.get("PYTHONPATH", ""), + ] + if state.ipc_port is not None: + env["ARTEMIS_IPC_PORT"] = str(state.ipc_port) + if sess_id: + env["ARTEMIS_SESSION_ID"] = str(sess_id) + + print(f"[QueueWorker] Starting task: {goal} with profile: {profile}") + state.current_process = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + "artemis.main", + goal, + "--profile", + profile, + "--test-name", + test_name, + cwd=str(WORKSPACE_ROOT), + env=env, + ) + await state.current_process.wait() + print( + f"[QueueWorker] Task finished with code " + f"{state.current_process.returncode}: {goal}" + ) + + # Fallback status update + if state.active_session_id: + curr_status = session_repo.get_session_status(state.active_session_id) + if curr_status == "running": + if state.was_stopped_manually: + new_status = "cancelled" + else: + new_status = ( + "completed" if state.current_process.returncode == 0 else "failed" + ) + session_repo.update_session_status( + state.active_session_id, new_status, time.time() + ) + print( + f"[QueueWorker] Fallback: Updated session " + f"{state.active_session_id} status to {new_status}" + ) + state.active_session_id = None + + state.current_profile = None + state.was_stopped_manually = False + + except asyncio.CancelledError: + print("[QueueWorker] Background worker task was cancelled.") + break + except Exception as e: + print(f"[QueueWorker] Error executing task: {e}") + finally: + if goal is not None: + state.task_queue.task_done() + + @classmethod + async def enqueue_tasks(cls, goals: list[str], profile: str = "flash") -> dict[str, Any]: + enqueued_tasks = [] + now = time.time() + for i, goal in enumerate(goals): + sess_id = str(uuid.uuid4()) + task_item = { + "session_id": sess_id, + "goal": goal, + "profile": profile or "flash", + "status": "pending", + "created_at": now + i * 0.001, + "start_time": now + i * 0.001, + } + await state.task_queue.put(task_item) + state.queue_tasks.append(task_item) + enqueued_tasks.append(task_item) + + return { + "status": "queued" if state.is_running else "started", + "tasks": enqueued_tasks, + "enqueued_count": len(goals), + "total_queued": len(state.queue_tasks), + } + + @classmethod + def stop_tasks(cls, clear_all: bool = False) -> bool: + stopped = False + if clear_all: + state.clear_queue() + + state.was_stopped_manually = True + if state.is_running and state.current_process: + try: + state.current_process.kill() + stopped = True + except Exception: + pass + + for _, conn_info in list(state.active_connections.items()): + pid = conn_info.get("pid") + if pid: + try: + os.kill(pid, signal.SIGINT) + stopped = True + except Exception as e: + print(f"Failed to kill external process {pid}: {e}") + + session_repo.mark_all_running_cancelled() + + if PAUSE_FILE.exists(): + try: + PAUSE_FILE.unlink() + except Exception: + pass + + return stopped + + @classmethod + def resume_task(cls) -> bool: + if PAUSE_FILE.exists(): + PAUSE_FILE.unlink() + return True + return False + + @staticmethod + def archive_older_replays_on_launch(): + """Archives all existing step replay output folders on server launch.""" + if TEST_OUTPUTS_DIR.exists(): + archive_dir = TEST_DATA_DIR / "older" + for item in TEST_OUTPUTS_DIR.iterdir(): + if item.is_dir(): + archive_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + target_name = f"{item.name}_{timestamp}" + target_path = archive_dir / target_name + + counter = 1 + while target_path.exists(): + target_name = f"{item.name}_{timestamp}_{counter}" + target_path = archive_dir / target_name + counter += 1 + + print( + f"Archiving older replay record on server launch: " + f"{item.name} -> {target_path}" + ) + try: + shutil.move(str(item), str(target_path)) + except Exception as e: + print( + f"Warning: Failed to archive {item.name} during server launch: {e}", + file=sys.stderr, + ) + + @staticmethod + def verify_chunks_exist_on_launch(replay_manager): + """Verifies that chunked directories exist for all sessions in the database.""" + print("Verifying session chunks on launch...") + try: + sessions = session_repo.get_all_sessions() + session_ids = [s.get("session_id") for s in sessions if s.get("session_id")] + print(f"Found {len(session_ids)} sessions in database to verify.") + for session_id in session_ids: + replay_manager._ensure_session_chunked(str(session_id)) + except Exception as e: + print( + f"Warning: Failed to verify chunks during server launch: {e}", + file=sys.stderr, + ) + + +task_queue_service = TaskQueueService()