565 lines
23 KiB
Python
565 lines
23 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
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 interrupted.
|
|
|
|
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 = ?",
|
|
("interrupted", 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"]),
|
|
"interrupted",
|
|
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)
|