somew1nd
0f5fb1d32c
Persist and publish recoverable provider retries so temporary 503 and throttling events remain visible in live and historical task streams. Aggregate repeated retries into an expandable action card, preserve terminal pause recovery, harden Windows worker isolation and output forwarding, and retain reliable Android hierarchy fallbacks with regression coverage.
486 lines
20 KiB
Python
486 lines
20 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 re
|
|
from typing import Any
|
|
|
|
try:
|
|
from admin_console.database.connection import db_session
|
|
except ImportError:
|
|
from apps.admin_console.database.connection import db_session
|
|
|
|
|
|
class StepRepository:
|
|
"""Repository handling querying and formatting of steps and associated traces."""
|
|
|
|
def __init__(self, db_path=None):
|
|
self.db_path = db_path
|
|
|
|
def get_step_session_id(self, step_id: str) -> str | None:
|
|
with db_session(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT session_id FROM steps WHERE step_id = ?", (step_id,))
|
|
row = cursor.fetchone()
|
|
return row["session_id"] if row else None
|
|
|
|
@staticmethod
|
|
def _clean_value(val: Any) -> Any:
|
|
"""Recursively cleans values to remove non-serializable and internal
|
|
Python runtime representations.
|
|
"""
|
|
if val is None:
|
|
return None
|
|
if isinstance(val, (bool, int, float)):
|
|
return val
|
|
if isinstance(val, str):
|
|
if (
|
|
"object at 0x" in val
|
|
or val.startswith("<artemis.")
|
|
or val.startswith("<controller")
|
|
):
|
|
return None
|
|
return val
|
|
if isinstance(val, dict):
|
|
cleaned = {}
|
|
for k, v in val.items():
|
|
cv = StepRepository._clean_value(v)
|
|
if cv is not None:
|
|
cleaned[k] = cv
|
|
return cleaned
|
|
if isinstance(val, (list, tuple)):
|
|
cleaned_list = []
|
|
for item in val:
|
|
cv = StepRepository._clean_value(item)
|
|
if cv is not None:
|
|
cleaned_list.append(cv)
|
|
return cleaned_list
|
|
val_str = str(val)
|
|
if "object at 0x" in val_str or val_str.startswith("<"):
|
|
return None
|
|
return val_str
|
|
|
|
@staticmethod
|
|
def _extract_clean_text_and_images(content: Any) -> tuple[str, int]:
|
|
"""Extracts readable text and count of image blocks from LLM content."""
|
|
if not content:
|
|
return "", 0
|
|
if isinstance(content, str):
|
|
if content.startswith("data:image/") or "base64," in content[:50]:
|
|
return "", 1
|
|
return content, 0
|
|
if isinstance(content, dict):
|
|
if content.get("type") == "image_url" or "image_url" in content:
|
|
return "", 1
|
|
if content.get("type") == "text":
|
|
return str(content.get("text", "")), 0
|
|
t = content.get("text") or content.get("thought") or ""
|
|
return str(t), 0
|
|
if isinstance(content, list):
|
|
total_text = ""
|
|
num_images = 0
|
|
for item in content:
|
|
txt, imgs = StepRepository._extract_clean_text_and_images(item)
|
|
total_text += txt
|
|
num_images += imgs
|
|
return total_text, num_images
|
|
return str(content), 0
|
|
|
|
@staticmethod
|
|
def _extract_token_usage_from_trace(payload_raw: Any) -> tuple[int, int, int]:
|
|
"""Extract prompt, completion, and total tokens from an LLM trace payload with multimodal awareness."""
|
|
if not payload_raw:
|
|
return 0, 0, 0
|
|
payload_obj = payload_raw
|
|
if isinstance(payload_raw, str):
|
|
try:
|
|
payload_obj = json.loads(payload_raw)
|
|
except Exception:
|
|
pass
|
|
if not isinstance(payload_obj, dict):
|
|
return 0, 0, 0
|
|
|
|
u = payload_obj.get("token_usage") or payload_obj.get("usage_metadata")
|
|
if isinstance(u, dict):
|
|
pr = u.get("prompt_tokens") or u.get("prompt_token_count") or u.get("input_tokens") or 0
|
|
co = (
|
|
u.get("completion_tokens")
|
|
or u.get("candidates_token_count")
|
|
or u.get("output_tokens")
|
|
or 0
|
|
)
|
|
to = u.get("total_tokens") or u.get("total_token_count") or (pr + co)
|
|
if to > 0:
|
|
return int(pr), int(co), int(to)
|
|
|
|
# Fallback for historical traces: estimate tokens from text and images
|
|
msgs = payload_obj.get("messages", [])
|
|
resp = payload_obj.get("response", [])
|
|
|
|
prompt_text = ""
|
|
prompt_images = 0
|
|
if isinstance(msgs, list):
|
|
for m in msgs:
|
|
if isinstance(m, dict):
|
|
t, imgs = StepRepository._extract_clean_text_and_images(m.get("content"))
|
|
prompt_text += t
|
|
prompt_images += imgs
|
|
else:
|
|
t, imgs = StepRepository._extract_clean_text_and_images(m)
|
|
prompt_text += t
|
|
prompt_images += imgs
|
|
else:
|
|
t, imgs = StepRepository._extract_clean_text_and_images(msgs)
|
|
prompt_text += t
|
|
prompt_images += imgs
|
|
|
|
resp_text = ""
|
|
resp_images = 0
|
|
if isinstance(resp, list):
|
|
for r in resp:
|
|
if isinstance(r, dict):
|
|
t, imgs = StepRepository._extract_clean_text_and_images(r.get("content"))
|
|
resp_text += t
|
|
resp_images += imgs
|
|
else:
|
|
t, imgs = StepRepository._extract_clean_text_and_images(r)
|
|
resp_text += t
|
|
resp_images += imgs
|
|
|
|
# Standard Gemini multimodal token pricing: 258 tokens per image
|
|
pr = (len(prompt_text) // 4) + (prompt_images * 258)
|
|
co = (len(resp_text) // 4) + (resp_images * 258)
|
|
return int(pr), int(co), int(pr + co)
|
|
|
|
def _clean_tool_payload(self, payload_raw: Any, trace_type: str | None = None) -> Any:
|
|
if not payload_raw:
|
|
return None
|
|
payload_obj = payload_raw
|
|
if isinstance(payload_raw, str):
|
|
try:
|
|
payload_obj = json.loads(payload_raw)
|
|
except Exception:
|
|
pass
|
|
|
|
if not isinstance(payload_obj, dict):
|
|
return None
|
|
|
|
if trace_type == "llm_call":
|
|
allowed_keys = (
|
|
"error",
|
|
"delay",
|
|
"attempt",
|
|
"max_retries",
|
|
"provider",
|
|
"source",
|
|
"recoverable",
|
|
"pause",
|
|
)
|
|
cleaned = {
|
|
key: self._clean_value(payload_obj.get(key))
|
|
for key in allowed_keys
|
|
if key in payload_obj
|
|
}
|
|
return cleaned or None
|
|
|
|
args = payload_obj.get("args", payload_obj)
|
|
cleaned_args = (
|
|
self._clean_value(args) if isinstance(args, dict) else self._clean_value(args)
|
|
)
|
|
result_payload: dict[str, Any] = {"args": cleaned_args}
|
|
|
|
# Extract structured action execution results
|
|
res = payload_obj.get("result")
|
|
if isinstance(res, dict):
|
|
if res.get("post_image_name"):
|
|
result_payload["post_image_name"] = res["post_image_name"]
|
|
if res.get("pre_image_name"):
|
|
result_payload["pre_image_name"] = res["pre_image_name"]
|
|
if res.get("outcome"):
|
|
result_payload["outcome"] = res["outcome"]
|
|
elif res:
|
|
res_str = str(res)
|
|
match = re.search(r"([a-f0-9]{64})", res_str)
|
|
if match:
|
|
result_payload["post_image_name"] = match.group(1)
|
|
|
|
return result_payload
|
|
|
|
@staticmethod
|
|
def _extract_legacy_pause_error(payload_raw: Any) -> str | None:
|
|
"""Extract an LLM error from the warning trace used by older runners."""
|
|
payload_obj = payload_raw
|
|
if isinstance(payload_raw, str):
|
|
try:
|
|
payload_obj = json.loads(payload_raw)
|
|
except Exception:
|
|
return None
|
|
if not isinstance(payload_obj, dict):
|
|
return None
|
|
|
|
message = payload_obj.get("message")
|
|
if not isinstance(message, str) or "LLM Error:" not in message:
|
|
return None
|
|
pause_marker = ". Pausing execution"
|
|
if pause_marker not in message:
|
|
return None
|
|
|
|
error = message.split("LLM Error:", 1)[1].split(pause_marker, 1)[0].strip()
|
|
return error or None
|
|
|
|
def _normalize_display_trace(self, trace_dict: dict[str, Any]) -> dict[str, Any]:
|
|
"""Normalize traces consumed by the task stream's tool/error cards."""
|
|
if trace_dict.get("type") == "log":
|
|
error = self._extract_legacy_pause_error(trace_dict.get("payload"))
|
|
if error:
|
|
trace_dict.update(
|
|
{
|
|
"type": "llm_call",
|
|
"name": "llm_pause",
|
|
"status": "failed",
|
|
"payload": {"error": error, "pause": True},
|
|
}
|
|
)
|
|
return trace_dict
|
|
|
|
trace_dict["payload"] = self._clean_tool_payload(
|
|
trace_dict.get("payload"), trace_dict.get("type")
|
|
)
|
|
return trace_dict
|
|
|
|
def get_session_steps(self, session_id: str, client: str | None = None) -> list[dict[str, Any]]:
|
|
with db_session(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT * FROM steps WHERE session_id = ? ORDER BY step_number ASC",
|
|
(session_id,),
|
|
)
|
|
rows = cursor.fetchall()
|
|
|
|
steps = []
|
|
for r in rows:
|
|
step_dict = dict(r)
|
|
|
|
# Strip raw binary bytes for API efficiency (clients fetch images by url/hash)
|
|
step_dict["pre_screenshot_bytes"] = None
|
|
step_dict["post_screenshot_bytes"] = None
|
|
|
|
# Parse JSON fields safely
|
|
for json_col in [
|
|
"action_taken",
|
|
"last_execution_result",
|
|
"extra_metadata",
|
|
]:
|
|
if step_dict.get(json_col):
|
|
try:
|
|
step_dict[json_col] = json.loads(step_dict[json_col])
|
|
except Exception:
|
|
pass
|
|
|
|
# If post_image_name is identical to pre_image_name, normalize to None
|
|
pre_img = step_dict.get("pre_image_name")
|
|
if pre_img and step_dict.get("post_image_name") == pre_img:
|
|
step_dict["post_image_name"] = None
|
|
|
|
# Fetch action and generic tool traces for this step
|
|
step_id = step_dict.get("step_id")
|
|
if step_id:
|
|
cursor.execute(
|
|
"SELECT t1.trace_id, t1.parent_trace_id, t1.step_id,"
|
|
" t1.type, t1.name, t1.timestamp, t1.duration, t1.status,"
|
|
" t1.payload, t2.name as agent_name FROM traces t1 LEFT"
|
|
" JOIN traces t2 ON t1.parent_trace_id = t2.trace_id WHERE"
|
|
" t1.step_id = ? AND (t1.type = 'tool' OR t1.type = 'action' OR (t1.type ="
|
|
" 'llm_call' AND t1.status IN ('failed', 'retrying')) OR (t1.type = 'log' AND"
|
|
" t1.payload LIKE '%LLM Error:%Pausing execution%' AND NOT EXISTS"
|
|
" (SELECT 1 FROM traces t3 WHERE t3.session_id = t1.session_id AND"
|
|
" t3.type = 'llm_call' AND t3.status = 'failed' AND"
|
|
" ABS(t3.timestamp - t1.timestamp) < 2))) ORDER BY"
|
|
" t1.timestamp ASC",
|
|
(step_id,),
|
|
)
|
|
trace_rows = cursor.fetchall()
|
|
generic_tools = []
|
|
for tr in trace_rows:
|
|
trace_dict = self._normalize_display_trace(dict(tr))
|
|
generic_tools.append(trace_dict)
|
|
step_dict["generic_tools"] = generic_tools
|
|
|
|
# Compute token usage for this step from all LLM call traces
|
|
cursor.execute(
|
|
"SELECT payload FROM traces WHERE step_id = ? AND type = 'llm_call'",
|
|
(step_id,),
|
|
)
|
|
llm_rows = cursor.fetchall()
|
|
step_p, step_c, step_t = 0, 0, 0
|
|
for lr in llm_rows:
|
|
p, c, t = self._extract_token_usage_from_trace(lr["payload"])
|
|
step_p += p
|
|
step_c += c
|
|
step_t += t
|
|
|
|
# If step_t is 0, also check extra_metadata for token_usage
|
|
if step_t == 0 and step_dict.get("extra_metadata"):
|
|
meta = (
|
|
step_dict["extra_metadata"]
|
|
if isinstance(step_dict["extra_metadata"], dict)
|
|
else {}
|
|
)
|
|
u = meta.get("token_usage")
|
|
if isinstance(u, dict):
|
|
step_p = int(u.get("prompt_tokens") or u.get("input_tokens") or 0)
|
|
step_c = int(u.get("completion_tokens") or u.get("output_tokens") or 0)
|
|
step_t = int(u.get("total_tokens") or (step_p + step_c))
|
|
|
|
step_dict["token_usage"] = {
|
|
"prompt_tokens": step_p,
|
|
"completion_tokens": step_c,
|
|
"total_tokens": step_t,
|
|
}
|
|
step_dict["total_tokens"] = step_t
|
|
else:
|
|
step_dict["generic_tools"] = []
|
|
step_dict["token_usage"] = {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0,
|
|
}
|
|
step_dict["total_tokens"] = 0
|
|
|
|
try:
|
|
from artemis.utils.coordinates import normalize_step_actions
|
|
|
|
step_dict = normalize_step_actions(step_dict)
|
|
except ImportError:
|
|
pass
|
|
|
|
steps.append(step_dict)
|
|
|
|
# Fetch step-less tool traces (e.g. planner tools)
|
|
cursor.execute(
|
|
"SELECT t1.trace_id, t1.parent_trace_id, t1.step_id, t1.type,"
|
|
" t1.name, t1.timestamp, t1.duration, t1.status, t1.payload,"
|
|
" t2.name as agent_name FROM traces t1 LEFT JOIN traces t2 ON"
|
|
" t1.parent_trace_id = t2.trace_id WHERE t1.session_id = ? AND"
|
|
" (t1.step_id IS NULL OR t1.step_id = '') AND (t1.type = 'tool' OR"
|
|
" (t1.type = 'llm_call' AND t1.status IN ('failed', 'retrying')) OR (t1.type = 'log'"
|
|
" AND t1.payload LIKE '%LLM Error:%Pausing execution%' AND NOT EXISTS"
|
|
" (SELECT 1 FROM traces t3 WHERE t3.session_id = t1.session_id AND"
|
|
" t3.type = 'llm_call' AND t3.status = 'failed' AND"
|
|
" ABS(t3.timestamp - t1.timestamp) < 2))) ORDER BY"
|
|
" t1.timestamp ASC",
|
|
(session_id,),
|
|
)
|
|
stepless_rows = cursor.fetchall()
|
|
if stepless_rows:
|
|
unassigned_tools = []
|
|
sorted_steps = sorted(steps, key=lambda s: s.get("timestamp") or 0)
|
|
|
|
for tr in stepless_rows:
|
|
trace_dict = self._normalize_display_trace(dict(tr))
|
|
|
|
tr_time = trace_dict.get("timestamp") or 0
|
|
target_step = None
|
|
for s in sorted_steps:
|
|
s_time = s.get("timestamp") or 0
|
|
if s_time <= tr_time:
|
|
target_step = s
|
|
else:
|
|
break
|
|
|
|
if target_step:
|
|
if "generic_tools" not in target_step or not target_step["generic_tools"]:
|
|
target_step["generic_tools"] = []
|
|
target_step["generic_tools"].append(trace_dict)
|
|
target_step["generic_tools"].sort(key=lambda t: t.get("timestamp") or 0)
|
|
else:
|
|
unassigned_tools.append(trace_dict)
|
|
|
|
if unassigned_tools:
|
|
unassigned_tools.sort(key=lambda t: t.get("timestamp") or 0)
|
|
first_timestamp = unassigned_tools[0].get("timestamp") or 0
|
|
|
|
# Calculate token usage ONLY for pre-planning LLM calls (before first step)
|
|
first_recorded_step_time = (
|
|
sorted_steps[0].get("timestamp") if sorted_steps else float("inf")
|
|
)
|
|
cursor.execute(
|
|
"SELECT payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type = 'llm_call'",
|
|
(session_id, first_recorded_step_time),
|
|
)
|
|
plan_llm_rows = cursor.fetchall()
|
|
plan_p, plan_c, plan_t = 0, 0, 0
|
|
for lr in plan_llm_rows:
|
|
p, c, t = self._extract_token_usage_from_trace(lr["payload"])
|
|
plan_p += p
|
|
plan_c += c
|
|
plan_t += t
|
|
|
|
# Retrieve planner thinking / thoughts
|
|
cursor.execute(
|
|
"SELECT type, payload FROM traces WHERE session_id = ? AND (step_id IS NULL OR step_id = '') AND timestamp < ? AND type IN ('raw_thinking', 'thinking') ORDER BY timestamp ASC",
|
|
(session_id, first_recorded_step_time),
|
|
)
|
|
think_rows = cursor.fetchall()
|
|
p_raw_th = []
|
|
p_nat_th = []
|
|
for thr in think_rows:
|
|
if thr["payload"]:
|
|
try:
|
|
th_obj = (
|
|
json.loads(thr["payload"])
|
|
if isinstance(thr["payload"], str)
|
|
else thr["payload"]
|
|
)
|
|
th_text = th_obj.get("thought") or th_obj.get("text") or ""
|
|
if thr["type"] == "raw_thinking" and th_text:
|
|
p_raw_th.append(str(th_text).strip())
|
|
elif thr["type"] == "thinking" and th_text:
|
|
p_nat_th.append(str(th_text).strip())
|
|
except Exception:
|
|
pass
|
|
|
|
virtual_step = {
|
|
"step_id": "pre-planning",
|
|
"step_type": "planning",
|
|
"session_id": session_id,
|
|
"step_number": 0,
|
|
"timestamp": first_timestamp,
|
|
"operator_native_thinking": "\n\n".join(p_nat_th) if p_nat_th else "",
|
|
"operator_raw_thinking": "\n\n".join(p_raw_th) if p_raw_th else "",
|
|
"action_taken": None,
|
|
"last_execution_result": None,
|
|
"extra_metadata": None,
|
|
"generic_tools": unassigned_tools,
|
|
"token_usage": {
|
|
"prompt_tokens": plan_p,
|
|
"completion_tokens": plan_c,
|
|
"total_tokens": plan_t,
|
|
},
|
|
"total_tokens": plan_t,
|
|
}
|
|
steps.append(virtual_step)
|
|
|
|
steps.sort(
|
|
key=lambda x: (
|
|
x.get("timestamp") or 0,
|
|
x.get("step_number") or 0,
|
|
)
|
|
)
|
|
|
|
return steps
|
|
|
|
|
|
step_repo = StepRepository()
|