1616 lines
63 KiB
Python
1616 lines
63 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 asyncio
|
|
from collections.abc import Callable
|
|
import functools
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import socket
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
from uuid import UUID, uuid4
|
|
|
|
from artemis.config import PAUSE_FILE, get_ipc_port_file, read_ipc_port, settings
|
|
from artemis.context import ArtemisContext
|
|
from artemis.data_engine.models import (
|
|
BackgroundTaskRecord,
|
|
FailedOutputRecord,
|
|
HistoryChunkRecord,
|
|
ImageRecord,
|
|
SessionMetadata,
|
|
StepRecord,
|
|
TraceRecord,
|
|
VideoRecordingRecord,
|
|
)
|
|
from artemis.data_engine.storage import StorageManager
|
|
from artemis.data_engine.trace import CURRENT_TRACE_ID
|
|
from artemis.utils.coordinates import (
|
|
normalize_any_structure,
|
|
normalize_step_actions,
|
|
)
|
|
from artemis.utils.logger import get_logger
|
|
from artemis.utils.text import safe_extract_text
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_CURRENT_DATA_ENGINE = None
|
|
|
|
# Overlay packages that never identify the app the user is working in.
|
|
_FOREGROUND_APP_IGNORED = {"com.android.systemui"}
|
|
|
|
|
|
def _derive_foreground_app(ui_tree: Any | None) -> str | None:
|
|
"""Best-effort foreground package from a perception UI tree (M5).
|
|
|
|
Counts ``package``/``packageName`` attributes across the node list
|
|
(recursing into ``children``), ignores system-overlay packages, and
|
|
returns the most frequent value. Pure string work — never raises, never
|
|
touches the device.
|
|
"""
|
|
counts: dict[str, int] = {}
|
|
|
|
def _visit(node: Any, depth: int = 0) -> None:
|
|
if depth > 50:
|
|
return
|
|
if isinstance(node, dict):
|
|
pkg = node.get("package") or node.get("packageName")
|
|
if isinstance(pkg, str) and pkg and pkg not in _FOREGROUND_APP_IGNORED:
|
|
counts[pkg] = counts.get(pkg, 0) + 1
|
|
children = node.get("children")
|
|
if isinstance(children, (list, tuple)):
|
|
for child in children:
|
|
_visit(child, depth + 1)
|
|
elif isinstance(node, (list, tuple)):
|
|
for child in node:
|
|
_visit(child, depth + 1)
|
|
|
|
try:
|
|
_visit(ui_tree)
|
|
except Exception:
|
|
return None
|
|
if not counts:
|
|
return None
|
|
return max(counts.items(), key=lambda kv: kv[1])[0]
|
|
|
|
|
|
#: ``smart_serialize`` stores any image payload inside a trace as this
|
|
#: reference; the hash is the SHA-256 of the decoded image bytes — the same
|
|
#: scheme ``record_step`` uses for ``image_name`` — so a referenced image can
|
|
#: be resolved back to the step screenshot it came from.
|
|
_IMAGE_REF_RE = re.compile(r"<ImageRef: sha256=([0-9a-fA-F]{64})")
|
|
|
|
#: Text stand-in for an image whose origin is unknown (e.g. an annotated
|
|
#: Explorer overlay that was never stored as a step screenshot).
|
|
GENERIC_IMAGE_LABEL = "[image attached]"
|
|
|
|
|
|
def image_ref_hash(value: Any) -> str | None:
|
|
"""The SHA-256 named by an ``<ImageRef: …>`` placeholder, if any."""
|
|
if not isinstance(value, str):
|
|
return None
|
|
match = _IMAGE_REF_RE.search(value)
|
|
return match.group(1).lower() if match else None
|
|
|
|
|
|
def build_image_describer(steps: list[Any]) -> Callable[[str | None], dict[str, Any]]:
|
|
"""Describer that resolves an image hash to the step screenshot it is.
|
|
|
|
Returns a callable producing a text block: ``[screenshot: pre-action of
|
|
Step 3]`` (with the ``image_name`` attached so search can pull that
|
|
screenshot's OCR/XML as the image's description) or the generic label
|
|
when the hash matches no recorded step screenshot.
|
|
"""
|
|
index: dict[str, tuple[int, str]] = {}
|
|
for step in steps:
|
|
number = getattr(step, "step_number", None)
|
|
if number is None and isinstance(step, dict):
|
|
number = step.get("step_number")
|
|
for which in ("pre", "post"):
|
|
name = getattr(step, f"{which}_image_name", None)
|
|
if name is None and isinstance(step, dict):
|
|
name = step.get(f"{which}_image_name")
|
|
if name and name not in index:
|
|
index[str(name).lower()] = (number, which)
|
|
|
|
def describe(image_name: str | None) -> dict[str, Any]:
|
|
if image_name and image_name.lower() in index:
|
|
number, which = index[image_name.lower()]
|
|
return {
|
|
"type": "text",
|
|
"text": f"[screenshot: {which}-action of Step {number}]",
|
|
"image_name": image_name.lower(),
|
|
}
|
|
return {"type": "text", "text": GENERIC_IMAGE_LABEL}
|
|
|
|
return describe
|
|
|
|
|
|
def describe_result_images(result: Any, describer: Callable[[str | None], dict[str, Any]]) -> Any:
|
|
"""Rewrites image blocks inside a stored tool result into text blocks.
|
|
|
|
Images are never dropped silently from the text views: a block becomes
|
|
the description the describer gives it (which step screenshot it is, or
|
|
a generic label), so replay and search both see that an image was there
|
|
and, when known, what it showed.
|
|
"""
|
|
if isinstance(result, list):
|
|
out = []
|
|
for block in result:
|
|
if isinstance(block, dict) and block.get("type") in ("image_url", "image"):
|
|
url = block.get("image_url")
|
|
if isinstance(url, dict):
|
|
url = url.get("url")
|
|
if url is None:
|
|
url = block.get("source") or block.get("data")
|
|
out.append(describer(image_ref_hash(url)))
|
|
else:
|
|
out.append(describe_result_images(block, describer))
|
|
return out
|
|
if isinstance(result, dict):
|
|
return {k: describe_result_images(v, describer) for k, v in result.items()}
|
|
return result
|
|
|
|
|
|
def build_interleaved_events(
|
|
traces: list[Any],
|
|
image_describer: Callable[[str | None], dict[str, Any]] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Turns one step's raw trace rows into the agent-facing event stream.
|
|
|
|
Operator LLM calls become ``thought`` / ``native_thought`` events; top-level
|
|
tool and sub-agent traces become ``tool_call`` events carrying ``name``,
|
|
``args`` and ``result``. Nested calls made *inside* a kept tool/agent trace
|
|
(e.g. the Explorer's own perception calls) are dropped so only the
|
|
conclusion the Operator actually saw is replayed. Shared by the live
|
|
DataEngine and the offline MCP trace inspector so both replay the same
|
|
events.
|
|
|
|
Image blocks inside tool results are rewritten into text descriptions via
|
|
``image_describer`` (see ``build_image_describer``); without one they
|
|
become the generic label.
|
|
"""
|
|
describer = image_describer or build_image_describer([])
|
|
interleaved_events: list[dict[str, Any]] = []
|
|
denylist = (
|
|
"operator",
|
|
"perception",
|
|
"planner",
|
|
"validator",
|
|
"summarizer",
|
|
"checker",
|
|
)
|
|
|
|
relevant_traces = []
|
|
for t in traces:
|
|
if t.status not in ("success", "failed"):
|
|
continue
|
|
|
|
is_relevant_llm = False
|
|
if t.type == "llm_call":
|
|
curr_parent = t.parent_trace_id
|
|
visited = set()
|
|
while curr_parent:
|
|
if curr_parent in visited:
|
|
break
|
|
visited.add(curr_parent)
|
|
parent_trace = next(
|
|
(x for x in traces if x.trace_id == curr_parent),
|
|
None,
|
|
)
|
|
if parent_trace:
|
|
if parent_trace.name == "operator":
|
|
is_relevant_llm = True
|
|
t.name = parent_trace.name
|
|
break
|
|
curr_parent = parent_trace.parent_trace_id
|
|
else:
|
|
break
|
|
|
|
if is_relevant_llm:
|
|
relevant_traces.append(t)
|
|
elif t.type in ("tool", "agent") and t.name not in denylist:
|
|
relevant_traces.append(t)
|
|
|
|
# Sort by timestamp to ensure exact chronological order
|
|
relevant_traces.sort(key=lambda x: x.timestamp)
|
|
|
|
# Separate set of candidate tool/agent IDs to do top-level filtering
|
|
# (only keep top-level tool calls)
|
|
candidate_ids = {t.trace_id for t in relevant_traces if t.type in ("tool", "agent")}
|
|
|
|
for t in relevant_traces:
|
|
if t.type == "llm_call":
|
|
payload = t.payload or {}
|
|
response_list = payload.get("response") or []
|
|
thought_text = ""
|
|
has_structured_blocks = False
|
|
if response_list:
|
|
first_gen = response_list[0]
|
|
if isinstance(first_gen, dict):
|
|
content_val = first_gen.get("content")
|
|
if isinstance(content_val, list):
|
|
has_structured_blocks = True
|
|
for block in content_val:
|
|
if isinstance(block, dict):
|
|
if (
|
|
block.get("type") == "thinking"
|
|
and block.get("thinking", "").strip()
|
|
):
|
|
interleaved_events.append(
|
|
{
|
|
"type": "native_thought",
|
|
"content": (block["thinking"].strip()),
|
|
}
|
|
)
|
|
elif block.get("type") == "text" and block.get("text", "").strip():
|
|
interleaved_events.append(
|
|
{
|
|
"type": "thought",
|
|
"content": (block["text"].strip()),
|
|
}
|
|
)
|
|
if not has_structured_blocks:
|
|
thought_text = first_gen.get("content") or first_gen.get("text") or ""
|
|
else:
|
|
thought_text = str(first_gen)
|
|
|
|
if not has_structured_blocks:
|
|
if isinstance(thought_text, list):
|
|
thought_text = safe_extract_text(thought_text)
|
|
|
|
if thought_text and thought_text.strip():
|
|
interleaved_events.append(
|
|
{
|
|
"type": "thought",
|
|
"content": thought_text.strip(),
|
|
}
|
|
)
|
|
elif t.type in ("tool", "agent"):
|
|
# Keep only top-level candidate traces (no ancestor in candidates)
|
|
is_sub_call = False
|
|
curr_parent = t.parent_trace_id
|
|
visited = set()
|
|
while curr_parent:
|
|
if curr_parent in visited:
|
|
break
|
|
visited.add(curr_parent)
|
|
if curr_parent in candidate_ids:
|
|
is_sub_call = True
|
|
break
|
|
parent_trace = next(
|
|
(x for x in traces if x.trace_id == curr_parent),
|
|
None,
|
|
)
|
|
if parent_trace:
|
|
curr_parent = parent_trace.parent_trace_id
|
|
else:
|
|
break
|
|
if is_sub_call:
|
|
continue
|
|
|
|
payload = t.payload or {}
|
|
tc_args = payload.get("args") or {}
|
|
|
|
if isinstance(tc_args, dict):
|
|
filtered_args = {
|
|
k: v for k, v in tc_args.items() if k not in ("state", "tool_call_id")
|
|
}
|
|
else:
|
|
filtered_args = tc_args
|
|
|
|
tc_result = payload.get("result") or payload.get("error") or "No result"
|
|
tc_result = describe_result_images(tc_result, describer)
|
|
|
|
interleaved_events.append(
|
|
{
|
|
"type": "tool_call",
|
|
"name": t.name,
|
|
"args": filtered_args,
|
|
"result": tc_result,
|
|
}
|
|
)
|
|
|
|
return interleaved_events
|
|
|
|
|
|
def relative_time_label(timestamp: float | None, session_start_time: float | None) -> str:
|
|
"""Session-relative ``"12.3s"`` label of a step timestamp (``"0.0s"`` when
|
|
the session clock is unknown, matching ``DataEngine.get_relative_time``)."""
|
|
if session_start_time is None or timestamp is None:
|
|
return "0.0s"
|
|
return f"{timestamp - session_start_time:.1f}s"
|
|
|
|
|
|
def friendly_step(
|
|
step: StepRecord,
|
|
traces: list[TraceRecord],
|
|
image_describer: Callable[[str | None], dict[str, Any]] | None = None,
|
|
*,
|
|
session_start_time: float | None = None,
|
|
default_width: int = 1080,
|
|
default_height: int = 2400,
|
|
) -> dict[str, Any]:
|
|
"""Agent-friendly record of one stored step.
|
|
|
|
JSON columns are already parsed (``StepRecord``), the step's traces are
|
|
expanded into ``interleaved_events`` (see ``build_interleaved_events``),
|
|
``tool_calls`` lists the tool events, ``relative_time`` is the
|
|
session-relative label and physical coordinates are normalized using the
|
|
dimensions stamped on the step (falling back to the defaults). Shared by
|
|
the live ``DataEngine`` and the offline ``OfflineHistoryReader`` so both
|
|
replay the same record.
|
|
"""
|
|
step_dict = step.model_dump()
|
|
step_dict["relative_time"] = relative_time_label(step.timestamp, session_start_time)
|
|
|
|
try:
|
|
step_dict["interleaved_events"] = build_interleaved_events(traces, image_describer)
|
|
except Exception as e:
|
|
logger.error(f"Failed to retrieve traces for step {step.step_id}: {e}")
|
|
step_dict["interleaved_events"] = []
|
|
step_dict["tool_calls"] = [
|
|
e for e in step_dict["interleaved_events"] if e["type"] == "tool_call"
|
|
]
|
|
|
|
# Normalize physical coordinates to normalized ones for agent consumption
|
|
try:
|
|
extra = step_dict.get("extra_metadata") or {}
|
|
width = extra.get("width") or default_width
|
|
height = extra.get("height") or default_height
|
|
|
|
if step_dict.get("action_taken"):
|
|
step_dict["action_taken"] = normalize_any_structure(
|
|
step_dict["action_taken"], width, height
|
|
)
|
|
if step_dict.get("last_execution_result"):
|
|
step_dict["last_execution_result"] = normalize_any_structure(
|
|
step_dict["last_execution_result"], width, height
|
|
)
|
|
except Exception as norm_err:
|
|
logger.error(f"Failed to normalize coordinates in friendly steps history: {norm_err}")
|
|
|
|
return step_dict
|
|
|
|
|
|
class DataEngine:
|
|
"""Core engine for managing runtime data, storage, and streaming."""
|
|
|
|
def __init__(self, ctx: ArtemisContext):
|
|
self.ctx = ctx
|
|
|
|
# Determine paths from context
|
|
if ctx.execution_setup and ctx.execution_setup.traces_path:
|
|
self.global_base_dir = Path(ctx.execution_setup.traces_path)
|
|
self.db_path = self.global_base_dir / "data_engine.db"
|
|
else:
|
|
# Fallback to default unified settings
|
|
self.global_base_dir = settings.TRACES_PATH
|
|
self.db_path = settings.DATA_ENGINE_DB_PATH
|
|
|
|
self.storage = StorageManager(self.db_path, self.global_base_dir)
|
|
|
|
self.current_session_id: UUID | None = None
|
|
self.session_start_time: float | None = None
|
|
self.current_step_id: UUID | None = None
|
|
self.last_recorded_step_id: UUID | None = None
|
|
self.current_step_dir: Path | None = None
|
|
self.subscribers: list[Callable[[str, Any], None]] = []
|
|
self.lock = threading.Lock()
|
|
self._trace_counter = 0
|
|
self._pending_tasks = set()
|
|
self._pending_threads = []
|
|
self._accumulated_logs = {}
|
|
self._bg_task_to_trace_id = {}
|
|
self._trace_name_cache = {}
|
|
self._step_number_cache: dict[str, int] = {}
|
|
# Background tasks are written directly to SQLite storage
|
|
|
|
self.ipc_socket = None
|
|
self._ipc_socket_lock = threading.Lock()
|
|
self._ipc_shutdown = False
|
|
self._ipc_retry_after = 0.0
|
|
self._ipc_failure_count = 0
|
|
with self._ipc_socket_lock:
|
|
self._connect_ipc_locked()
|
|
|
|
@staticmethod
|
|
def _ipc_port_candidates() -> list[int]:
|
|
"""Return inherited and freshly published UI ports, newest file included.
|
|
|
|
A worker inherits ``ARTEMIS_IPC_PORT`` when it starts. If the Windows UI
|
|
server restarts, that environment value is permanently stale while the
|
|
shared port file is updated by the new server.
|
|
"""
|
|
candidates: list[int] = []
|
|
inherited_or_file_port = read_ipc_port()
|
|
if inherited_or_file_port:
|
|
candidates.append(inherited_or_file_port)
|
|
|
|
port_file = get_ipc_port_file()
|
|
if port_file.exists():
|
|
try:
|
|
file_value = port_file.read_text(encoding="utf-8").strip()
|
|
if file_value.isdigit() and int(file_value) not in candidates:
|
|
candidates.append(int(file_value))
|
|
except OSError as exc:
|
|
logger.debug(f"Could not read refreshed IPC port from {port_file}: {exc}")
|
|
return candidates
|
|
|
|
def _connect_ipc_locked(self, *, force: bool = False) -> bool:
|
|
"""Connect to the UI event bridge while holding ``_ipc_socket_lock``."""
|
|
if self._ipc_shutdown:
|
|
return False
|
|
if self.ipc_socket is not None:
|
|
return True
|
|
if not force and time.monotonic() < self._ipc_retry_after:
|
|
return False
|
|
|
|
for ipc_port in self._ipc_port_candidates():
|
|
try:
|
|
# A localhost listener either accepts immediately or is absent.
|
|
# A short connect cap prevents telemetry from stalling the
|
|
# automation/data worker when the desktop event bridge is stale.
|
|
sock = socket.create_connection(("127.0.0.1", ipc_port), timeout=0.2)
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
sock.settimeout(None)
|
|
self.ipc_socket = sock
|
|
self._ipc_failure_count = 0
|
|
self._ipc_retry_after = 0.0
|
|
logger.info(f"Connected to IPC server on port {ipc_port} with TCP_NODELAY")
|
|
return True
|
|
except OSError as exc:
|
|
logger.warning(f"Failed to connect to IPC server on port {ipc_port}: {exc}")
|
|
|
|
self.ipc_socket = None
|
|
self._ipc_failure_count += 1
|
|
retry_delay = min(30.0, 0.5 * (2 ** min(self._ipc_failure_count - 1, 6)))
|
|
self._ipc_retry_after = time.monotonic() + retry_delay
|
|
return False
|
|
|
|
def _send_ipc_event(self, event_type: str, data: Any) -> None:
|
|
"""Send one complete event, reconnecting once if a stale socket is found.
|
|
|
|
Publishing can originate from LLM callbacks and background worker threads.
|
|
Serializing socket access keeps newline-delimited JSON frames from being
|
|
interleaved on Windows and prevents the event that detects a dead socket
|
|
from being silently lost.
|
|
"""
|
|
payload = json.dumps({"event_type": event_type, "data": data}, default=str) + "\n"
|
|
encoded_payload = payload.encode("utf-8")
|
|
|
|
with self._ipc_socket_lock:
|
|
for attempt in range(2):
|
|
if not self._connect_ipc_locked(force=attempt > 0):
|
|
return
|
|
try:
|
|
self.ipc_socket.sendall(encoded_payload)
|
|
return
|
|
except OSError as exc:
|
|
logger.warning(f"IPC send failed; reconnecting: {exc}")
|
|
try:
|
|
self.ipc_socket.close()
|
|
except OSError:
|
|
pass
|
|
self.ipc_socket = None
|
|
self._ipc_retry_after = 0.0
|
|
|
|
@property
|
|
def base_dir(self) -> Path:
|
|
if self.current_session_id:
|
|
return self.global_base_dir / str(self.current_session_id)
|
|
return self.global_base_dir
|
|
|
|
def subscribe(self, callback: Callable[[str, Any], None]):
|
|
"""Subscribe to real-time events (e.g., for SSE)."""
|
|
self.subscribers.append(callback)
|
|
|
|
def unsubscribe(self, callback: Callable[[str, Any], None]):
|
|
"""Unsubscribe from real-time events."""
|
|
if callback in self.subscribers:
|
|
self.subscribers.remove(callback)
|
|
|
|
def _publish(self, event_type: str, data: Any):
|
|
"""Publish event to all subscribers."""
|
|
if isinstance(data, dict) and "session_id" not in data and self.current_session_id:
|
|
data["session_id"] = str(self.current_session_id)
|
|
|
|
for callback in self.subscribers:
|
|
try:
|
|
callback(event_type, data)
|
|
except Exception as e:
|
|
logger.error(f"Error in subscriber callback: {e}")
|
|
|
|
# Bridge to Cloud Gateway SSE stream in Cloud Mode
|
|
if os.environ.get("ARTEMIS_CLOUD_MODE") == "1":
|
|
try:
|
|
import asyncio
|
|
from debug_pub.cloud_gateway import cloud_manager
|
|
|
|
sid = os.environ.get("ARTEMIS_CLOUD_SESSION_ID") or (
|
|
str(self.current_session_id) if self.current_session_id else None
|
|
)
|
|
if sid and sid in cloud_manager.sessions:
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
if loop.is_running():
|
|
loop.create_task(
|
|
cloud_manager.sessions[sid].emit_event(
|
|
event_type,
|
|
data if isinstance(data, dict) else {"payload": data},
|
|
)
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
except Exception as exc:
|
|
logger.debug(f"Cloud gateway event bridge skipped: {exc}", exc_info=True)
|
|
|
|
self._send_ipc_event(event_type, data)
|
|
|
|
def start_session(
|
|
self,
|
|
goal: str,
|
|
device_info: dict[str, Any] | None = None,
|
|
session_id: UUID | str | None = None,
|
|
) -> UUID:
|
|
"""Start a new session."""
|
|
global _CURRENT_DATA_ENGINE
|
|
_CURRENT_DATA_ENGINE = self
|
|
|
|
if session_id is not None:
|
|
if isinstance(session_id, str):
|
|
try:
|
|
session_id = UUID(session_id)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
env_session_id = os.getenv("ARTEMIS_CLOUD_SESSION_ID") or os.getenv(
|
|
"ARTEMIS_SESSION_ID"
|
|
)
|
|
if env_session_id:
|
|
try:
|
|
session_id = UUID(env_session_id)
|
|
except ValueError:
|
|
session_id = env_session_id
|
|
else:
|
|
session_id = uuid4()
|
|
self.current_session_id = session_id
|
|
self.session_start_time = time.time()
|
|
|
|
# Clear old pause file if it exists
|
|
pause_file = PAUSE_FILE
|
|
if pause_file.exists():
|
|
try:
|
|
pause_file.unlink()
|
|
logger.info("Removed old pause file on session start.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete old pause file: {e}")
|
|
|
|
session = SessionMetadata(
|
|
session_id=session_id,
|
|
initial_goal=goal,
|
|
start_time=self.session_start_time,
|
|
device_info=device_info or {},
|
|
pid=os.getpid(),
|
|
)
|
|
self.storage.create_session(session)
|
|
try:
|
|
from artemis.runtime import DeviceExecutionLock
|
|
|
|
DeviceExecutionLock.annotate_active_owner(
|
|
session_id=str(session_id),
|
|
ingress=os.getenv("ARTEMIS_TASK_INGRESS") or "sdk",
|
|
)
|
|
except Exception as exc:
|
|
logger.debug(f"Could not annotate active device owner: {exc}")
|
|
# Initialize step counter from existing steps if resuming an existing session
|
|
existing_steps = []
|
|
if self.storage and hasattr(self.storage, "get_steps"):
|
|
try:
|
|
existing_steps = self.storage.get_steps(session_id) or []
|
|
except Exception:
|
|
existing_steps = []
|
|
if existing_steps:
|
|
self._step_number_cache.clear()
|
|
for s in existing_steps:
|
|
if s.step_id and s.step_number is not None:
|
|
self._step_number_cache[str(s.step_id)] = s.step_number
|
|
self.current_step_number = max([s.step_number for s in existing_steps], default=0)
|
|
else:
|
|
self._step_number_cache.clear()
|
|
self.current_step_number = 0
|
|
logger.info(
|
|
f"Session started: {session_id} (current step counter: {self.current_step_number})"
|
|
)
|
|
self._publish("session_started", session.model_dump())
|
|
return session_id
|
|
|
|
def end_session(self, status: str = "completed"):
|
|
"""End the current session, updating its status and end time."""
|
|
if not self.current_session_id:
|
|
return
|
|
|
|
# Session-level terminal statuses are canonically "completed" /
|
|
# "failed" / "cancelled"; "success" is a legacy alias some callers
|
|
# still pass and must never reach the sessions table.
|
|
if status == "success":
|
|
status = "completed"
|
|
|
|
# Clear pause file if it exists
|
|
pause_file = PAUSE_FILE
|
|
if pause_file.exists():
|
|
try:
|
|
pause_file.unlink()
|
|
logger.info("Removed pause file on session end.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete pause file on session end: {e}")
|
|
|
|
session_id = self.current_session_id
|
|
end_time = time.time()
|
|
session = self.storage.get_session(session_id)
|
|
if session and session.end_time is not None and session.status not in ("running", "paused"):
|
|
logger.debug(f"Session end already published for {session_id}; skipping duplicate")
|
|
return
|
|
# Session-level LLM usage line (cache-hit ratios per source), best-effort.
|
|
try:
|
|
from artemis.services.token_meter import log_session_summary
|
|
|
|
log_session_summary(session_id)
|
|
except Exception as e:
|
|
logger.debug(f"Session usage summary skipped: {e}")
|
|
if session:
|
|
session.end_time = end_time
|
|
session.status = status
|
|
else:
|
|
session = SessionMetadata(
|
|
session_id=session_id,
|
|
initial_goal="",
|
|
start_time=self.session_start_time or end_time,
|
|
end_time=end_time,
|
|
status=status,
|
|
device_info=self.ctx.device.model_dump()
|
|
if getattr(self, "ctx", None) and self.ctx.device
|
|
else {},
|
|
)
|
|
|
|
try:
|
|
self.storage.update_session(session)
|
|
logger.info(f"Session ended: {session_id} with status: {status}")
|
|
self._publish("session_ended", session.model_dump())
|
|
except Exception as e:
|
|
logger.error(f"Failed to end session in DataEngine: {e}")
|
|
|
|
def record_video_start(
|
|
self,
|
|
video_id: UUID,
|
|
device_id: str,
|
|
local_video_path: str | Path,
|
|
start_time: float | None = None,
|
|
):
|
|
"""Record the start of a video recording."""
|
|
if not self.storage:
|
|
return
|
|
try:
|
|
record = VideoRecordingRecord(
|
|
video_id=video_id,
|
|
session_id=self.current_session_id,
|
|
device_id=device_id,
|
|
start_time=start_time or time.time(),
|
|
local_video_path=str(local_video_path),
|
|
)
|
|
self.storage.create_video_recording(record)
|
|
except Exception as e:
|
|
logger.error(f"Failed to record video start in DataEngine: {e}")
|
|
|
|
def record_video_stop(
|
|
self,
|
|
video_id: UUID,
|
|
device_id: str,
|
|
local_video_path: str | Path,
|
|
start_time: float,
|
|
end_time: float | None = None,
|
|
):
|
|
"""Record the completion of a video recording and sync to session metadata."""
|
|
if not self.storage:
|
|
return
|
|
try:
|
|
record = VideoRecordingRecord(
|
|
video_id=video_id,
|
|
session_id=self.current_session_id,
|
|
device_id=device_id,
|
|
start_time=start_time,
|
|
end_time=end_time or time.time(),
|
|
local_video_path=str(local_video_path),
|
|
status="ready",
|
|
)
|
|
self.storage.update_video_recording(record)
|
|
self._publish(
|
|
"recording_ready",
|
|
{
|
|
"session_id": str(self.current_session_id),
|
|
"video_id": str(video_id),
|
|
"local_video_path": str(local_video_path),
|
|
"end_time": record.end_time,
|
|
},
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to record video stop in DataEngine: {e}")
|
|
|
|
def record_video_failure(
|
|
self,
|
|
video_id: UUID,
|
|
device_id: str,
|
|
local_video_path: str | Path | None,
|
|
start_time: float,
|
|
error: str,
|
|
):
|
|
"""Persist and publish a terminal recording failure."""
|
|
if not self.storage:
|
|
return
|
|
try:
|
|
record = VideoRecordingRecord(
|
|
video_id=video_id,
|
|
session_id=self.current_session_id,
|
|
device_id=device_id,
|
|
start_time=start_time,
|
|
end_time=time.time(),
|
|
local_video_path=str(local_video_path) if local_video_path else None,
|
|
status="failed",
|
|
error=error,
|
|
)
|
|
self.storage.update_video_recording(record)
|
|
self._publish(
|
|
"recording_failed",
|
|
{
|
|
"session_id": str(self.current_session_id),
|
|
"video_id": str(video_id),
|
|
"error": error,
|
|
},
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to record video failure in DataEngine: {e}")
|
|
|
|
def update_session_device_info(self, **fields: Any) -> None:
|
|
"""Merge ``fields`` into the running session's ``device_info`` record."""
|
|
if not self.storage or not self.current_session_id:
|
|
return
|
|
try:
|
|
session = self.storage.get_session(self.current_session_id)
|
|
if session is None:
|
|
return
|
|
session.device_info = {**(session.device_info or {}), **fields}
|
|
self.storage.update_session(session)
|
|
except (OSError, ValueError, sqlite3.Error) as e:
|
|
logger.warning(f"Failed to update session device_info in DataEngine: {e}")
|
|
|
|
def update_video_path(self, local_video_path: str | Path):
|
|
"""Update the video path across tables when traces or videos are moved."""
|
|
if not self.storage or not self.current_session_id:
|
|
return
|
|
try:
|
|
self.storage.update_session_video_path(self.current_session_id, str(local_video_path))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to update video path in DataEngine: {e}")
|
|
|
|
def get_or_create_image(
|
|
self,
|
|
image_bytes: bytes,
|
|
ui_tree: Any | None = None,
|
|
ocr_result: Any | None = None,
|
|
) -> str:
|
|
"""Get image name by hash, or create new record if not exists."""
|
|
|
|
hasher = hashlib.sha256()
|
|
hasher.update(image_bytes)
|
|
image_name = hasher.hexdigest()
|
|
|
|
image_record = self.storage.get_image(image_name)
|
|
if image_record:
|
|
# If the image exists but is missing OCR or UI tree, and we now have them, update the record.
|
|
if (ocr_result is not None and image_record.ocr_result is None) or (
|
|
ui_tree is not None and image_record.ui_tree is None
|
|
):
|
|
self.storage.update_image_data(image_name, ocr_result, ui_tree)
|
|
return image_name
|
|
|
|
images_dir = self.global_base_dir / "images"
|
|
images_dir.mkdir(parents=True, exist_ok=True)
|
|
file_path = images_dir / f"{image_name}.jpg"
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(image_bytes)
|
|
|
|
new_record = ImageRecord(
|
|
image_name=image_name,
|
|
ui_tree=ui_tree,
|
|
ocr_result=ocr_result,
|
|
extra_metadata={},
|
|
)
|
|
self.storage.create_image(new_record)
|
|
|
|
return image_name
|
|
|
|
def get_image_path(self, image_name: str) -> Path:
|
|
"""Get the absolute file path of an image by its name/hash."""
|
|
return self.global_base_dir / "images" / f"{image_name}.jpg"
|
|
|
|
def record_step(
|
|
self,
|
|
pre_screenshot_bytes: bytes | None = None,
|
|
post_screenshot_bytes: bytes | None = None,
|
|
ui_tree: Any | None = None,
|
|
ocr_result: Any | None = None,
|
|
foreground_app: str | None = None,
|
|
action_taken: dict[str, Any] | None = None,
|
|
summary: str | None = None,
|
|
operator_raw_thinking: str | None = None,
|
|
operator_native_thinking: str | None = None,
|
|
last_execution_result: Any | None = None,
|
|
extra_metadata: dict[str, Any] | None = None,
|
|
) -> UUID:
|
|
"""Record a step, saving screenshots and metadata."""
|
|
if not self.current_session_id:
|
|
raise ValueError("No active session. Call start_session first.")
|
|
|
|
with self.lock:
|
|
existing_steps = []
|
|
if self.storage and hasattr(self.storage, "get_steps") and self.current_session_id:
|
|
try:
|
|
existing_steps = self.storage.get_steps(self.current_session_id) or []
|
|
except Exception:
|
|
existing_steps = []
|
|
max_existing = (
|
|
max([s.step_number for s in existing_steps], default=0) if existing_steps else 0
|
|
)
|
|
self.current_step_number = max(self.current_step_number, max_existing) + 1
|
|
step_number = self.current_step_number
|
|
|
|
step_id = self.current_step_id or uuid4()
|
|
self.last_recorded_step_id = step_id
|
|
self._step_number_cache[str(step_id)] = step_number
|
|
# Reset current_step_id so subsequent steps will allocate a new one
|
|
self.current_step_id = None
|
|
|
|
# Process images and step creation in background thread
|
|
|
|
pre_image_name = (
|
|
hashlib.sha256(pre_screenshot_bytes).hexdigest() if pre_screenshot_bytes else None
|
|
)
|
|
post_image_name = (
|
|
hashlib.sha256(post_screenshot_bytes).hexdigest() if post_screenshot_bytes else None
|
|
)
|
|
if pre_image_name and post_image_name and pre_image_name == post_image_name:
|
|
post_image_name = None
|
|
|
|
# Persist the foreground app for the recall search surface (the
|
|
# parameter was historically accepted but never stored). The explicit
|
|
# parameter wins; otherwise it is derived best-effort from the UI
|
|
# tree's package attributes — a pure string scan, no device call.
|
|
extra_metadata = dict(extra_metadata or {})
|
|
try:
|
|
app = foreground_app or _derive_foreground_app(ui_tree)
|
|
if app:
|
|
extra_metadata.setdefault("foreground_app", app)
|
|
except (AttributeError, TypeError, ValueError):
|
|
pass
|
|
|
|
# Stamp perceptual hashes for the local screen-similarity hint
|
|
# (pure PIL dHash; best-effort, computed synchronously to keep the
|
|
# record immutable once handed to the background writer).
|
|
try:
|
|
from artemis.utils.image_hash import dhash_hex
|
|
|
|
pre_dhash = dhash_hex(pre_screenshot_bytes)
|
|
if pre_dhash:
|
|
extra_metadata.setdefault("pre_image_dhash", pre_dhash)
|
|
post_dhash = dhash_hex(post_screenshot_bytes) if post_image_name else pre_dhash
|
|
if post_dhash:
|
|
extra_metadata.setdefault("post_image_dhash", post_dhash)
|
|
except ImportError:
|
|
pass
|
|
|
|
step = StepRecord(
|
|
step_id=step_id,
|
|
session_id=self.current_session_id,
|
|
step_number=step_number,
|
|
timestamp=time.time(),
|
|
pre_image_name=pre_image_name,
|
|
post_image_name=post_image_name,
|
|
summary=summary,
|
|
action_taken=action_taken,
|
|
operator_raw_thinking=operator_raw_thinking,
|
|
operator_native_thinking=operator_native_thinking,
|
|
last_execution_result=last_execution_result,
|
|
extra_metadata=extra_metadata,
|
|
)
|
|
|
|
def _run_background_storage():
|
|
if pre_screenshot_bytes:
|
|
self.get_or_create_image(pre_screenshot_bytes, ui_tree, ocr_result)
|
|
if post_screenshot_bytes:
|
|
self.get_or_create_image(post_screenshot_bytes)
|
|
self.storage.create_step(step)
|
|
|
|
# Run storage in the shared background scheduler. It chooses an
|
|
# asyncio worker when a loop is active and a tracked thread otherwise.
|
|
self._run_in_background(_run_background_storage)
|
|
|
|
logger.info(f"Recorded step {step_number} for session {self.current_session_id}")
|
|
|
|
# Publish event with relative time
|
|
step_dict = step.model_dump()
|
|
step_dict["relative_time"] = self.get_relative_time(step.timestamp)
|
|
|
|
# Include generic tools in the real-time event
|
|
try:
|
|
traces = self.storage.get_step_traces(step_id)
|
|
generic_tools = []
|
|
for t in traces:
|
|
if t.type in ("tool", "action") or (t.type == "llm_call" and t.status == "failed"):
|
|
t_dict = t.model_dump()
|
|
if t.parent_trace_id and "agent_name" not in t_dict:
|
|
parent_name = getattr(self, "_trace_name_cache", {}).get(
|
|
t.parent_trace_id
|
|
) or getattr(self, "_trace_name_cache", {}).get(str(t.parent_trace_id))
|
|
if parent_name:
|
|
t_dict["agent_name"] = parent_name
|
|
generic_tools.append(t_dict)
|
|
step_dict["generic_tools"] = generic_tools
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch step traces for SSE: {e}")
|
|
step_dict["generic_tools"] = []
|
|
|
|
# Extract and attach token usage for real-time SSE stream
|
|
step_p, step_c, step_t = 0, 0, 0
|
|
if extra_metadata and isinstance(extra_metadata, dict) and "token_usage" in extra_metadata:
|
|
u = extra_metadata["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))
|
|
|
|
if step_t == 0:
|
|
token_info = self._get_step_token_usage_for_sse(step_id)
|
|
if token_info:
|
|
step_p = token_info["prompt_tokens"]
|
|
step_c = token_info["completion_tokens"]
|
|
step_t = token_info["total_tokens"]
|
|
|
|
if step_t > 0:
|
|
step_dict["token_usage"] = {
|
|
"prompt_tokens": step_p,
|
|
"completion_tokens": step_c,
|
|
"total_tokens": step_t,
|
|
}
|
|
step_dict["total_tokens"] = step_t
|
|
|
|
step_dict = normalize_step_actions(step_dict)
|
|
|
|
self._publish("step_recorded", step_dict)
|
|
|
|
return step_id
|
|
|
|
def record_trace(
|
|
self,
|
|
type: str,
|
|
name: str,
|
|
payload: dict[str, Any],
|
|
step_id: UUID | None = None,
|
|
parent_trace_id: UUID | None = None,
|
|
status: str = "success",
|
|
duration: float | None = None,
|
|
trace_id: UUID | None = None,
|
|
) -> UUID:
|
|
"""Record a trace (agent, tool, or log), non-blocking."""
|
|
if not self.current_session_id:
|
|
raise ValueError("No active session. Call start_session first.")
|
|
|
|
trace_id = trace_id or uuid4()
|
|
|
|
with self.lock:
|
|
self._trace_counter = getattr(self, "_trace_counter", 0) + 1
|
|
trace_ts = time.time() + (self._trace_counter * 1e-7)
|
|
if not hasattr(self, "_trace_name_cache"):
|
|
self._trace_name_cache = {}
|
|
self._trace_name_cache[trace_id] = name
|
|
self._trace_name_cache[str(trace_id)] = name
|
|
|
|
trace = TraceRecord(
|
|
trace_id=trace_id,
|
|
session_id=self.current_session_id,
|
|
step_id=step_id,
|
|
parent_trace_id=parent_trace_id,
|
|
type=type,
|
|
name=name,
|
|
timestamp=trace_ts,
|
|
duration=duration,
|
|
status=status,
|
|
payload=payload,
|
|
)
|
|
|
|
self._run_in_background(self.storage.create_trace, trace)
|
|
|
|
# Publish event
|
|
trace_dict = trace.model_dump()
|
|
if parent_trace_id:
|
|
parent_name = getattr(self, "_trace_name_cache", {}).get(parent_trace_id) or getattr(
|
|
self, "_trace_name_cache", {}
|
|
).get(str(parent_trace_id))
|
|
if parent_name:
|
|
trace_dict["agent_name"] = parent_name
|
|
self._publish("trace_recorded", trace_dict)
|
|
|
|
return trace_id
|
|
|
|
def record_failed_output(
|
|
self,
|
|
trace_id: UUID,
|
|
model_name: str,
|
|
prompt: str,
|
|
raw_output: str,
|
|
error_message: str,
|
|
):
|
|
"""Record a malformed output (crime scene) for later training."""
|
|
if not self.current_session_id:
|
|
return
|
|
|
|
record = FailedOutputRecord(
|
|
session_id=self.current_session_id,
|
|
trace_id=trace_id,
|
|
model_name=model_name,
|
|
prompt=prompt,
|
|
raw_output=raw_output,
|
|
error_message=error_message,
|
|
)
|
|
|
|
self._run_in_background(self.storage.create_failed_output, record)
|
|
|
|
def has_pending_operations(self) -> bool:
|
|
"""Check if there are any pending background tasks or threads."""
|
|
return (
|
|
len(self._pending_tasks) > 0
|
|
or len([t for t in self._pending_threads if t.is_alive()]) > 0
|
|
)
|
|
|
|
async def shutdown(self):
|
|
"""Wait for all pending background tasks and threads to complete."""
|
|
if self._pending_tasks:
|
|
logger.info(
|
|
f"Waiting for {len(self._pending_tasks)} pending async tasks to complete..."
|
|
)
|
|
await asyncio.gather(*self._pending_tasks, return_exceptions=True)
|
|
|
|
if self._pending_threads:
|
|
logger.info(f"Waiting for {len(self._pending_threads)} pending threads to complete...")
|
|
threads_to_join = list(self._pending_threads)
|
|
for thread in threads_to_join:
|
|
if thread.is_alive():
|
|
await asyncio.to_thread(thread.join)
|
|
|
|
with self._ipc_socket_lock:
|
|
self._ipc_shutdown = True
|
|
if self.ipc_socket is not None:
|
|
try:
|
|
self.ipc_socket.close()
|
|
except OSError:
|
|
pass
|
|
self.ipc_socket = None
|
|
|
|
logger.info("DataEngine shutdown complete. All data persisted.")
|
|
|
|
def stream_output(
|
|
self,
|
|
execution_id: UUID,
|
|
chunk: str,
|
|
stream_type: str | None = None,
|
|
is_thinking: bool | None = None,
|
|
):
|
|
"""Stream LLM output chunks."""
|
|
if stream_type is None:
|
|
if is_thinking is True:
|
|
stream_type = "thinking"
|
|
else:
|
|
stream_type = "text"
|
|
|
|
parent_id = CURRENT_TRACE_ID.get()
|
|
session_id_str = str(self.current_session_id) if self.current_session_id else None
|
|
|
|
self._publish(
|
|
"llm_stream",
|
|
{
|
|
"execution_id": str(execution_id),
|
|
"session_id": session_id_str,
|
|
"parent_trace_id": str(parent_id) if parent_id else None,
|
|
"step_id": (str(self.current_step_id) if self.current_step_id else None),
|
|
"chunk": chunk,
|
|
"stream_type": stream_type,
|
|
"is_thinking": (
|
|
is_thinking if is_thinking is not None else (stream_type == "thinking")
|
|
),
|
|
},
|
|
)
|
|
|
|
exec_str = str(execution_id)
|
|
if not hasattr(self, "_accumulated_logs"):
|
|
self._accumulated_logs = {}
|
|
if exec_str not in self._accumulated_logs:
|
|
self._accumulated_logs[exec_str] = ""
|
|
self._accumulated_logs[exec_str] += chunk
|
|
|
|
if session_id_str:
|
|
if session_id_str not in self._accumulated_logs:
|
|
self._accumulated_logs[session_id_str] = ""
|
|
self._accumulated_logs[session_id_str] += chunk
|
|
|
|
def _get_generic_tools_for_sse(self, step_id: UUID) -> list[dict[str, Any]]:
|
|
try:
|
|
traces = self.storage.get_step_traces(step_id)
|
|
return [t.model_dump() for t in traces if t.type == "tool"]
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch step traces for SSE: {e}")
|
|
return []
|
|
|
|
def _get_step_token_usage_for_sse(self, step_id: UUID) -> dict[str, Any] | None:
|
|
try:
|
|
traces = self.storage.get_step_traces(step_id)
|
|
step_p, step_c, step_t = 0, 0, 0
|
|
for t in traces:
|
|
if t.type == "llm_call" and isinstance(t.payload, dict):
|
|
u = t.payload.get("token_usage") or t.payload.get("usage_metadata")
|
|
if isinstance(u, dict):
|
|
p = int(
|
|
u.get("prompt_tokens")
|
|
or u.get("prompt_token_count")
|
|
or u.get("input_tokens")
|
|
or 0
|
|
)
|
|
c = int(
|
|
u.get("completion_tokens")
|
|
or u.get("candidates_token_count")
|
|
or u.get("output_tokens")
|
|
or 0
|
|
)
|
|
tot = int(u.get("total_tokens") or u.get("total_token_count") or (p + c))
|
|
step_p += p
|
|
step_c += c
|
|
step_t += tot
|
|
if step_t > 0:
|
|
return {
|
|
"prompt_tokens": step_p,
|
|
"completion_tokens": step_c,
|
|
"total_tokens": step_t,
|
|
}
|
|
except Exception as exc:
|
|
logger.debug(
|
|
f"Step token usage aggregation for SSE skipped for {step_id}: {exc}",
|
|
exc_info=True,
|
|
)
|
|
return None
|
|
|
|
def _on_background_task_done(self, task: asyncio.Task):
|
|
self._pending_tasks.discard(task)
|
|
if not task.cancelled() and task.exception():
|
|
exc = task.exception()
|
|
logger.error(f"Background task failed with exception: {exc}", exc_info=exc)
|
|
|
|
def _run_in_background(self, fn, *args):
|
|
"""Safely offload blocking storage/disk operation to background without stalling event loop."""
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
|
|
def run_and_cleanup():
|
|
try:
|
|
fn(*args)
|
|
except Exception as exc:
|
|
logger.error(f"Background storage operation failed: {exc}", exc_info=exc)
|
|
finally:
|
|
curr_thread = threading.current_thread()
|
|
if curr_thread in self._pending_threads:
|
|
self._pending_threads.remove(curr_thread)
|
|
|
|
thread = threading.Thread(target=run_and_cleanup, name="artemis-storage", daemon=True)
|
|
self._pending_threads.append(thread)
|
|
thread.start()
|
|
else:
|
|
task = loop.create_task(asyncio.to_thread(fn, *args))
|
|
self._pending_tasks.add(task)
|
|
task.add_done_callback(self._on_background_task_done)
|
|
|
|
def update_step_action(
|
|
self,
|
|
action_taken: dict[str, Any],
|
|
post_screenshot_bytes: bytes | None = None,
|
|
):
|
|
"""Update the current step with action taken and post-action screenshot in background."""
|
|
with self.lock:
|
|
step_id = self.current_step_id or self.last_recorded_step_id
|
|
if not step_id:
|
|
raise ValueError("No active step to update.")
|
|
if self.base_dir:
|
|
step_dir = self.base_dir / "steps" / str(step_id)
|
|
self.current_step_dir = step_dir
|
|
step_dir.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
step_dir = None
|
|
|
|
def _update_and_write():
|
|
self.storage.update_step_action(step_id, action_taken)
|
|
if post_screenshot_bytes and step_dir:
|
|
step_dir.mkdir(parents=True, exist_ok=True)
|
|
with open(step_dir / "post.jpg", "wb") as f:
|
|
f.write(post_screenshot_bytes)
|
|
|
|
self._run_in_background(_update_and_write)
|
|
|
|
step_num = self.get_step_number(step_id)
|
|
update_payload = {
|
|
"step_id": str(step_id),
|
|
"action_taken": action_taken,
|
|
"generic_tools": self._get_generic_tools_for_sse(step_id),
|
|
}
|
|
if step_num is not None:
|
|
update_payload["step_number"] = step_num
|
|
|
|
tokens = self._get_step_token_usage_for_sse(step_id)
|
|
if tokens:
|
|
update_payload["token_usage"] = tokens
|
|
update_payload["total_tokens"] = tokens["total_tokens"]
|
|
|
|
self._publish("step_updated", update_payload)
|
|
|
|
def get_step_number(self, step_id: UUID | int | str | None) -> int | None:
|
|
"""Resolve the 1-based sequential step number for a step ID."""
|
|
if step_id is None:
|
|
return None
|
|
if isinstance(step_id, int):
|
|
return step_id
|
|
sid = str(step_id)
|
|
cached = self._step_number_cache.get(sid)
|
|
if cached is not None:
|
|
return cached
|
|
if self.storage and hasattr(self.storage, "get_step"):
|
|
try:
|
|
target_uuid = UUID(sid) if isinstance(step_id, str) else step_id
|
|
step_record = self.storage.get_step(target_uuid)
|
|
if step_record and getattr(step_record, "step_number", None) is not None:
|
|
self._step_number_cache[sid] = int(step_record.step_number)
|
|
return int(step_record.step_number)
|
|
except Exception as exc:
|
|
logger.debug(
|
|
f"Step number lookup from storage skipped for {sid}: {exc}", exc_info=True
|
|
)
|
|
return None
|
|
|
|
def get_step_record(self, step_number: int):
|
|
"""Read-only lookup of a step record by its 1-based step number."""
|
|
if not self.current_session_id:
|
|
return None
|
|
try:
|
|
steps = self.storage.get_steps(self.current_session_id) or []
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch steps for step lookup: {e}")
|
|
return None
|
|
for s in steps:
|
|
if s.step_number == step_number:
|
|
return s
|
|
return None
|
|
|
|
def get_step_image_path(self, step_number: int, which: str = "pre") -> Path | None:
|
|
"""Read-only path of a step's pre/post screenshot, if recorded."""
|
|
step = self.get_step_record(step_number)
|
|
if step is None:
|
|
return None
|
|
image_name = step.pre_image_name if which == "pre" else step.post_image_name
|
|
if not image_name:
|
|
return None
|
|
path = self.get_image_path(image_name)
|
|
return path if path.exists() else None
|
|
|
|
def update_step_summary(
|
|
self,
|
|
step_id: UUID | int | str,
|
|
summary: str | None,
|
|
*,
|
|
source: str | None = None,
|
|
version: int | None = None,
|
|
model: str | None = None,
|
|
status: str = "ready",
|
|
):
|
|
"""Update the step summary in background, with status/version metadata.
|
|
|
|
Writes ``summary_status`` / ``summary_source`` / ``summary_version`` /
|
|
``summary_model`` into the step's ``extra_metadata`` alongside the
|
|
summary text; concurrent same-step writes are ordered by version in the
|
|
storage layer. ``status`` may be ``pending`` / ``ready`` / ``failed`` /
|
|
``stale``; only ``ready`` writes carry summary text and publish SSE.
|
|
"""
|
|
target_uuid = None
|
|
if isinstance(step_id, int):
|
|
if self.current_session_id:
|
|
steps = self.storage.get_steps(self.current_session_id)
|
|
for s in steps:
|
|
if s.step_number == step_id:
|
|
target_uuid = s.step_id
|
|
break
|
|
elif isinstance(step_id, str):
|
|
try:
|
|
target_uuid = UUID(step_id)
|
|
except ValueError:
|
|
target_uuid = step_id
|
|
else:
|
|
target_uuid = step_id
|
|
|
|
if not target_uuid:
|
|
logger.debug(f"Could not resolve step_id for {step_id} to update summary.")
|
|
return
|
|
|
|
self._run_in_background(
|
|
functools.partial(
|
|
self.storage.update_step_summary,
|
|
target_uuid,
|
|
summary,
|
|
source=source,
|
|
version=version,
|
|
model=model,
|
|
status=status,
|
|
)
|
|
)
|
|
if status != "ready" or summary is None:
|
|
return
|
|
step_num = self.get_step_number(target_uuid if not isinstance(step_id, int) else step_id)
|
|
payload = {
|
|
"step_id": str(target_uuid),
|
|
"summary": summary,
|
|
"generic_tools": self._get_generic_tools_for_sse(target_uuid),
|
|
}
|
|
if step_num is not None:
|
|
payload["step_number"] = step_num
|
|
self._publish("step_updated", payload)
|
|
|
|
def record_history_chunk(
|
|
self,
|
|
*,
|
|
start_step_number: int,
|
|
end_step_number: int,
|
|
version: int,
|
|
status: str,
|
|
start_step_id: str | None = None,
|
|
end_step_id: str | None = None,
|
|
source_step_ids: list[str] | None = None,
|
|
subgoal_hash: str | None = None,
|
|
band1: dict[str, Any] | None = None,
|
|
band2: str | None = None,
|
|
band3: str | None = None,
|
|
rendered_text: str | None = None,
|
|
) -> UUID | None:
|
|
"""Persist one history-chunk version row in the background (append-only).
|
|
|
|
Chunk identity is the step range; callers own version numbering (a
|
|
newer version for the same range supersedes older ones at read time).
|
|
"""
|
|
if not self.current_session_id:
|
|
return None
|
|
chunk = HistoryChunkRecord(
|
|
chunk_id=uuid4(),
|
|
session_id=self.current_session_id,
|
|
start_step_id=start_step_id,
|
|
end_step_id=end_step_id,
|
|
start_step_number=start_step_number,
|
|
end_step_number=end_step_number,
|
|
source_step_ids=source_step_ids or [],
|
|
subgoal_hash=subgoal_hash,
|
|
version=version,
|
|
status=status,
|
|
band1=band1 or {},
|
|
band2=band2,
|
|
band3=band3,
|
|
rendered_text=rendered_text,
|
|
)
|
|
self._run_in_background(self.storage.create_history_chunk, chunk)
|
|
return chunk.chunk_id
|
|
|
|
def get_history_chunks(self, *, all_versions: bool = False) -> list[HistoryChunkRecord]:
|
|
"""History chunks of the active session (newest version per step range)."""
|
|
if not self.current_session_id:
|
|
return []
|
|
try:
|
|
return self.storage.get_history_chunks(
|
|
self.current_session_id, all_versions=all_versions
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch history chunks: {e}")
|
|
return []
|
|
|
|
def update_step_thinking(self, step_id: UUID, operator_raw_thinking: str):
|
|
"""Update step's thinking in SQLite in background."""
|
|
self._run_in_background(self.storage.update_step_thinking, step_id, operator_raw_thinking)
|
|
step_num = self.get_step_number(step_id)
|
|
payload = {
|
|
"step_id": str(step_id),
|
|
"operator_raw_thinking": operator_raw_thinking,
|
|
"generic_tools": self._get_generic_tools_for_sse(step_id),
|
|
}
|
|
if step_num is not None:
|
|
payload["step_number"] = step_num
|
|
self._publish("step_updated", payload)
|
|
|
|
def update_step_native_thinking(self, step_id: UUID, operator_native_thinking: str):
|
|
"""Update step's native thinking in SQLite in background."""
|
|
self._run_in_background(
|
|
self.storage.update_step_native_thinking,
|
|
step_id,
|
|
operator_native_thinking,
|
|
)
|
|
step_num = self.get_step_number(step_id)
|
|
payload = {
|
|
"step_id": str(step_id),
|
|
"operator_native_thinking": operator_native_thinking,
|
|
"generic_tools": self._get_generic_tools_for_sse(step_id),
|
|
}
|
|
if step_num is not None:
|
|
payload["step_number"] = step_num
|
|
self._publish("step_updated", payload)
|
|
|
|
def update_step_execution_result(
|
|
self,
|
|
step_id: UUID,
|
|
last_execution_result: dict,
|
|
post_image_name: str | None = None,
|
|
):
|
|
"""Update step's execution result and post_image_name in SQLite in background."""
|
|
self._run_in_background(
|
|
self.storage.update_step_execution_result,
|
|
step_id,
|
|
last_execution_result,
|
|
post_image_name,
|
|
)
|
|
step_num = self.get_step_number(step_id)
|
|
payload = {
|
|
"step_id": str(step_id),
|
|
"last_execution_result": last_execution_result,
|
|
"post_image_name": post_image_name,
|
|
"generic_tools": self._get_generic_tools_for_sse(step_id),
|
|
}
|
|
if step_num is not None:
|
|
payload["step_number"] = step_num
|
|
self._publish("step_updated", payload)
|
|
|
|
def get_relative_time(self, timestamp: float) -> str:
|
|
"""Compute relative time since session start."""
|
|
if self.session_start_time is None:
|
|
return "0.0s"
|
|
diff = timestamp - self.session_start_time
|
|
return f"{diff:.1f}s"
|
|
|
|
def allocate_step_id(self) -> UUID:
|
|
"""Pre-allocate a step ID for the upcoming step to synchronize traces."""
|
|
self.current_step_id = uuid4()
|
|
if self.base_dir:
|
|
self.current_step_dir = self.base_dir / "steps" / str(self.current_step_id)
|
|
self.current_step_dir.mkdir(parents=True, exist_ok=True)
|
|
return self.current_step_id
|
|
|
|
def get_agent_friendly_steps(self) -> list[dict[str, Any]]:
|
|
"""Retrieve steps formatted for agent consumption (relative time)."""
|
|
if not self.current_session_id:
|
|
return []
|
|
steps_with_traces = self.storage.get_steps_with_traces(self.current_session_id)
|
|
describer = build_image_describer([step for step, _ in steps_with_traces])
|
|
return [self._friendly_step(step, traces, describer) for step, traces in steps_with_traces]
|
|
|
|
def get_agent_friendly_steps_in_range(
|
|
self, start_step: int, end_step: int | None = None
|
|
) -> list[dict[str, Any]]:
|
|
"""Agent-friendly records of steps ``start_step``..``end_step`` (inclusive,
|
|
1-based) without loading the traces of the whole session. ``end_step``
|
|
defaults to ``start_step`` (a single step)."""
|
|
if not self.current_session_id:
|
|
return []
|
|
end_step = start_step if end_step is None else end_step
|
|
if start_step > end_step:
|
|
start_step, end_step = end_step, start_step
|
|
try:
|
|
steps = self.storage.get_steps(self.current_session_id) or []
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch steps for range lookup: {e}")
|
|
return []
|
|
describer = build_image_describer(steps)
|
|
result = []
|
|
for step in steps:
|
|
if start_step <= step.step_number <= end_step:
|
|
traces = self.storage.get_traces_for_step(step.step_id)
|
|
result.append(self._friendly_step(step, traces, describer))
|
|
return result
|
|
|
|
def get_agent_friendly_step(self, step_number: int) -> dict[str, Any] | None:
|
|
"""Agent-friendly record of one step by its 1-based number, or None."""
|
|
rows = self.get_agent_friendly_steps_in_range(step_number, step_number)
|
|
return rows[0] if rows else None
|
|
|
|
def _friendly_step(
|
|
self,
|
|
step: StepRecord,
|
|
traces: list[TraceRecord],
|
|
image_describer: Callable[[str | None], dict[str, Any]] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Live twin of :func:`friendly_step`: the session clock and the device
|
|
dimensions come from this engine's context."""
|
|
device = self.ctx.device if self.ctx else None
|
|
return friendly_step(
|
|
step,
|
|
traces,
|
|
image_describer,
|
|
session_start_time=self.session_start_time,
|
|
default_width=getattr(device, "device_width", 1080) if device else 1080,
|
|
default_height=getattr(device, "device_height", 2400) if device else 2400,
|
|
)
|
|
|
|
def register_background_task(self, task_id: str, summary: str, trace_id: UUID | None = None):
|
|
if not self.current_session_id:
|
|
return
|
|
|
|
record = BackgroundTaskRecord(
|
|
task_id=task_id,
|
|
session_id=self.current_session_id,
|
|
summary=summary,
|
|
status="running",
|
|
start_time=time.time(),
|
|
trace_id=str(trace_id) if trace_id else None,
|
|
)
|
|
|
|
if trace_id:
|
|
if not hasattr(self, "_bg_task_to_trace_id"):
|
|
self._bg_task_to_trace_id = {}
|
|
self._bg_task_to_trace_id[task_id] = str(trace_id)
|
|
|
|
self._run_in_background(self.storage.create_background_task, record)
|
|
|
|
self._publish("background_tasks_updated", self.get_all_background_tasks())
|
|
|
|
def unregister_background_task(self, task_id: str, status: str = "completed"):
|
|
end_time = time.time()
|
|
|
|
trace_id_str = None
|
|
if hasattr(self, "_bg_task_to_trace_id"):
|
|
trace_id_str = self._bg_task_to_trace_id.pop(task_id, None)
|
|
|
|
logs = ""
|
|
if trace_id_str and hasattr(self, "_accumulated_logs"):
|
|
logs = self._accumulated_logs.pop(trace_id_str, "")
|
|
|
|
self._run_in_background(
|
|
self.storage.update_background_task_status_and_logs,
|
|
task_id,
|
|
status,
|
|
end_time,
|
|
logs,
|
|
)
|
|
|
|
self._publish("background_tasks_updated", self.get_all_background_tasks())
|
|
|
|
def get_all_background_tasks(self) -> list[dict]:
|
|
if not self.current_session_id:
|
|
return []
|
|
records = self.storage.get_background_tasks(self.current_session_id)
|
|
return [r.model_dump() for r in records]
|
|
|
|
def get_active_background_tasks(self) -> list[dict]:
|
|
return self.get_all_background_tasks()
|