android-midscene-automation 0.2.3 → 1.1.0

package/artemis-main/apps/admin_console/database/repositories/session_repository.py

@@ -0,0 +1,564 @@
+# 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 logging
+import os
+import sqlite3
+import time
+from typing import Any
+
+from artemis.runtime import trace_store
+
+logger = logging.getLogger(__name__)
+
+try:
+    from admin_console.database.connection import db_session
+except ImportError:
+    from apps.admin_console.database.connection import db_session
+
+
+def _canonicalize_status(row: dict[str, Any]) -> dict[str, Any]:
+    """Map the legacy "success" terminal status to the canonical "completed".
+
+    Rows written before the vocabulary was unified may still carry "success";
+    API consumers must only ever see "completed". get_session_status is left
+    raw on purpose so the queue reconcile can detect and persist the rewrite.
+    """
+    if row.get("status") == "success":
+        row["status"] = "completed"
+    return row
+
+
+class SessionRepository:
+    """Repository handling all database queries and updates for Sessions."""
+
+    def __init__(self, db_path=None):
+        self.db_path = db_path
+
+    @staticmethod
+    def process_is_alive(pid: Any) -> bool:
+        """Return whether a session's worker PID is still alive.
+
+        Delegates to the repository-wide probe: indeterminate liveness
+        (e.g. AccessDenied) counts as alive so a running session is never
+        reaped by mistake.
+        """
+        try:
+            pid_int = int(pid)
+        except (TypeError, ValueError):
+            return False
+        from artemis.runtime.process_probe import pid_is_alive
+
+        return pid_is_alive(pid_int)
+
+    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 [_canonicalize_status(dict(r)) for r in cursor.fetchall()]
+
+    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 {}
+            columns = {str(row[1]) for row in cursor.execute("PRAGMA table_info(video_recordings)")}
+            # Failed recordings may still be recoverable from disk.
+            # Callers exclude in-progress sessions before resolving paths.
+            ready_clause = (
+                "status IN ('ready', 'failed')" if "status" in columns else "end_time IS NOT NULL"
+            )
+            cursor.execute(
+                "SELECT session_id, local_video_path FROM video_recordings "
+                "WHERE session_id IS NOT NULL AND local_video_path IS NOT NULL "
+                f"AND {ready_clause} 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_video_recording_for_session(self, session_id: str) -> dict[str, Any] | None:
+        try:
+            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 None
+                cursor.execute(
+                    "SELECT * FROM video_recordings WHERE session_id = ? "
+                    "ORDER BY start_time DESC LIMIT 1",
+                    (str(session_id),),
+                )
+                row = cursor.fetchone()
+                if not row:
+                    return None
+                result = dict(row)
+                if "status" not in result:
+                    result["status"] = (
+                        "ready" if result.get("end_time") is not None else "recording"
+                    )
+                result.setdefault("error", None)
+                return result
+        except Exception:
+            return None
+
+    def get_latest_video_recordings_map(self) -> dict[str, dict[str, Any]]:
+        """Return the newest recording row for every session in one query.
+
+        The task-list endpoint renders every session at once. Calling
+        ``get_video_recording_for_session`` in that loop used to open a new
+        SQLite connection for every row, which made a refresh progressively
+        slower as history grew.
+        """
+        try:
+            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 {}
+
+                rows = cursor.execute(
+                    "SELECT * FROM video_recordings WHERE session_id IS NOT NULL "
+                    "ORDER BY start_time DESC"
+                ).fetchall()
+                latest: dict[str, dict[str, Any]] = {}
+                for row in rows:
+                    recording = dict(row)
+                    session_id = str(recording.get("session_id") or "")
+                    if not session_id or session_id in latest:
+                        continue
+                    if "status" not in recording:
+                        recording["status"] = (
+                            "ready" if recording.get("end_time") is not None else "recording"
+                        )
+                    recording.setdefault("error", None)
+                    latest[session_id] = recording
+                return latest
+        except Exception:
+            return {}
+
+    def get_unfinalized_video_recordings(self) -> list[dict[str, Any]]:
+        """Recordings whose worker never finalized them, oldest first.
+
+        Rows of sessions that are still ``running`` are skipped: their worker
+        (possibly owned by another daemon) may still be writing the file.
+        """
+        try:
+            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 []
+                columns = {
+                    str(row[1]) for row in cursor.execute("PRAGMA table_info(video_recordings)")
+                }
+                pending_clause = (
+                    "v.status IN ('recording', 'finalizing', 'failed')"
+                    if "status" in columns
+                    else "v.end_time IS NULL"
+                )
+                cursor.execute(
+                    "SELECT v.session_id, v.local_video_path, v.start_time, "
+                    "s.status AS session_status, s.start_time AS session_start_time "
+                    "FROM video_recordings v LEFT JOIN sessions s ON s.session_id = v.session_id "
+                    f"WHERE v.session_id IS NOT NULL AND v.local_video_path IS NOT NULL "
+                    f"AND {pending_clause} ORDER BY v.start_time ASC"
+                )
+                rows = [dict(r) for r in cursor.fetchall()]
+        except (OSError, sqlite3.Error):
+            return []
+        return [row for row in rows if str(row.get("session_status") or "").lower() != "running"]
+
+    def mark_recording_failed_if_pending(self, session_id: str, error: str) -> bool:
+        """Close a recording lifecycle when its worker exits before finalization."""
+        try:
+            with db_session(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "UPDATE video_recordings SET status = 'failed', error = ?, "
+                    "end_time = COALESCE(end_time, ?) "
+                    "WHERE session_id = ? AND status IN ('recording', 'finalizing')",
+                    (error, time.time(), str(session_id)),
+                )
+                conn.commit()
+                return cursor.rowcount > 0
+        except Exception:
+            return False
+
+    def mark_recording_ready(self, session_id: str, local_video_path: str) -> bool:
+        """Mark a recording as ready with its finalized video path."""
+        try:
+            with db_session(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "UPDATE video_recordings SET status = 'ready', error = NULL, "
+                    "local_video_path = ?, end_time = COALESCE(end_time, ?) "
+                    "WHERE session_id = ?",
+                    (str(local_video_path), time.time(), str(session_id)),
+                )
+                updated = cursor.rowcount > 0
+                if os.path.exists(str(local_video_path)):
+                    # Mirror the worker's own finalization so the session row
+                    # and trace tooling see the recovered file too.
+                    cursor.execute(
+                        "UPDATE sessions SET video_filepath = ? WHERE session_id = ?",
+                        (str(local_video_path), str(session_id)),
+                    )
+                conn.commit()
+                return updated
+        except Exception:
+            return False
+
+    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_llm_traces_for_profiles_map(
+        self, session_ids: list[str] | None = None, limit: int = 3
+    ) -> dict[str, list[str]]:
+        """Return bounded LLM payloads for only the sessions needing fallback."""
+        if session_ids is not None and not session_ids:
+            return {}
+
+        session_filter = ""
+        params: list[Any] = []
+        if session_ids is not None:
+            placeholders = ",".join("?" for _ in session_ids)
+            session_filter = f" AND session_id IN ({placeholders})"
+            params.extend(session_ids)
+        params.append(limit)
+
+        with db_session(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT session_id, payload FROM ("
+                " SELECT session_id, payload, ROW_NUMBER() OVER ("
+                "  PARTITION BY session_id ORDER BY timestamp ASC"
+                " ) AS row_number FROM traces WHERE type = 'llm_call'"
+                f"{session_filter}"
+                ") WHERE row_number <= ?",
+                params,
+            )
+            result: dict[str, list[str]] = {}
+            for row in cursor.fetchall():
+                session_id = str(row[0] or "")
+                payload = row[1]
+                if session_id and payload:
+                    result.setdefault(session_id, []).append(payload)
+            return result
+
+    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_agent_trace_names_map(self) -> dict[str, list[str]]:
+        """Return distinct agent/LLM trace names grouped by session."""
+        with db_session(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT session_id, name FROM traces "
+                "WHERE type IN ('agent', 'llm_call') AND session_id IS NOT NULL "
+                "GROUP BY session_id, name"
+            )
+            result: dict[str, list[str]] = {}
+            for row in cursor.fetchall():
+                session_id = str(row[0] or "")
+                name = row[1]
+                if session_id and name:
+                    result.setdefault(session_id, []).append(name)
+            return result
+
+    def get_session_usage(self, session_id: str) -> dict[str, Any]:
+        """Aggregate measured LLM usage for one session.
+
+        Sums every ``llm_usage`` trace the token meter recorded (all agents,
+        Flash and Pro alike) and reports the prompt size of the latest
+        Operator/FlashRunner call as the live executor context. Also echoes
+        the profile and Pro tuning the run was started with, read from the
+        session's stored ``device_info``.
+        """
+        from artemis.services.token_meter import (
+            OPERATOR_CONTEXT_WINDOW_TOKENS,
+            OPERATOR_NODE_NAMES,
+        )
+
+        def _int(value: Any) -> int:
+            try:
+                return int(value or 0)
+            except (TypeError, ValueError):
+                return 0
+
+        llm_calls = prompt = completion = total = cached = 0
+        operator_context: int | None = None
+        operator_context_at: float | None = None
+        with db_session(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT payload, timestamp FROM traces WHERE session_id = ?"
+                " AND type = 'llm_call' AND name = 'llm_usage' ORDER BY timestamp ASC",
+                (session_id,),
+            )
+            for row in cursor.fetchall():
+                raw = row["payload"]
+                try:
+                    payload = json.loads(raw) if isinstance(raw, str) else raw
+                except (ValueError, TypeError):
+                    continue
+                if not isinstance(payload, dict):
+                    continue
+                p = _int(payload.get("prompt_tokens"))
+                c = _int(payload.get("completion_tokens"))
+                t = _int(payload.get("total_tokens")) or (p + c)
+                if t <= 0:
+                    continue
+                llm_calls += 1
+                prompt += p
+                completion += c
+                total += t
+                cached += _int(payload.get("cached_tokens"))
+                node = str(payload.get("node") or "").strip().lower()
+                if node in OPERATOR_NODE_NAMES:
+                    operator_context = _int(payload.get("context_base_tokens")) or p
+                    operator_context_at = row["timestamp"]
+
+            cursor.execute("SELECT device_info FROM sessions WHERE session_id = ?", (session_id,))
+            session_row = cursor.fetchone()
+
+        profile: str | None = None
+        run_tuning: dict[str, Any] | None = None
+        if session_row and session_row["device_info"]:
+            try:
+                d_info = json.loads(session_row["device_info"])
+            except (ValueError, TypeError):
+                d_info = None
+            if isinstance(d_info, dict):
+                profile = d_info.get("profile") or None
+                tuning = d_info.get("run_tuning")
+                if isinstance(tuning, dict) and tuning:
+                    run_tuning = dict(tuning)
+
+        return {
+            "session_id": session_id,
+            "llm_calls": llm_calls,
+            "prompt_tokens": prompt,
+            "completion_tokens": completion,
+            "total_tokens": total,
+            "cached_tokens": cached,
+            "operator_context_tokens": operator_context,
+            "operator_context_window_tokens": OPERATOR_CONTEXT_WINDOW_TOKENS,
+            "operator_context_updated_at": operator_context_at,
+            "profile": profile,
+            "run_tuning": run_tuning,
+        }
+
+    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 _canonicalize_status(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()
+            count = 0
+            for s_id in orphaned_ids:
+                row = cursor.execute(
+                    "SELECT pid FROM sessions WHERE session_id = ? AND status = ?",
+                    (s_id, "running"),
+                ).fetchone()
+                if row is None or self.process_is_alive(row["pid"]):
+                    continue
+                cursor.execute(
+                    "UPDATE sessions SET status = ?, end_time = ? "
+                    "WHERE session_id = ? AND status = ?",
+                    ("failed", now, s_id, "running"),
+                )
+                count += cursor.rowcount
+            conn.commit()
+            return count
+
+    def reconcile_orphaned_sessions(self) -> int:
+        """Mark running sessions with no live worker process as failed.
+
+        This is safe to call both during startup and immediately after a
+        forced server stop: live workers, including workers owned by another
+        server instance, are left untouched.
+        """
+        try:
+            with db_session(self.db_path) as conn:
+                cursor = conn.cursor()
+                rows = cursor.execute(
+                    "SELECT session_id, pid FROM sessions WHERE status = ?", ("running",)
+                ).fetchall()
+                count = 0
+                now = time.time()
+                for row in rows:
+                    if self.process_is_alive(row["pid"]):
+                        continue
+                    cursor.execute(
+                        "UPDATE sessions SET status = ?, end_time = ? "
+                        "WHERE session_id = ? AND status = ?",
+                        ("failed", now, row["session_id"], "running"),
+                    )
+                    count += cursor.rowcount
+                    try:
+                        if trace_store.read_status(str(row["session_id"])):
+                            trace_store.update_trace_status(
+                                str(row["session_id"]),
+                                "failed",
+                                error="Process terminated prior to server startup.",
+                            )
+                    except OSError as exc:
+                        # read_status itself never raises; this guards the
+                        # lock/write side of update_trace_status.
+                        logger.warning(
+                            "Could not mark trace %s failed during orphan reconciliation: %s",
+                            row["session_id"],
+                            exc,
+                        )
+                conn.commit()
+                return count
+        except Exception:
+            # Reconciliation is best-effort at startup, but a silent abort
+            # would leave every orphaned "running" row untouched -- log it.
+            logger.warning("Orphan-session reconciliation aborted", exc_info=True)
+            return 0
+
+    def cleanup_orphans_on_startup(self) -> int:
+        """Backward-compatible startup entrypoint for orphan reconciliation."""
+        return self.reconcile_orphaned_sessions()
+
+    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 _canonicalize_status(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()
+
+# Give the base runtime package access to orphan-session reconciliation without
+# letting it import this application-layer module (inverted dependency): any
+# process that loads the admin console DB layer -- the server itself, or an
+# entry point stopping one -- thereby arms stop_server()'s post-kill cleanup.
+from artemis.runtime.server_lifecycle import register_session_reconciler  # noqa: E402
+
+register_session_reconciler(session_repo.reconcile_orphaned_sessions)