chendelian
633ecff4e2
Keep the last in-template copy on main so history can restore it; product continues in jiangchang-ui-recorder. Co-authored-by: Cursor <cursoragent@cursor.com>
401 lines
15 KiB
Python
401 lines
15 KiB
Python
"""Snapshot service: wrap SurfaceDriver.dump into deduped UIState assets."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import asdict, is_dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from models.observation import UIObservation
|
|
from models.serialize import to_dict, write_model_json
|
|
from models.state import UIState, compute_state_fingerprint
|
|
from redaction import (
|
|
redact_locatable_rows,
|
|
sanitize_artifact,
|
|
screenshot_sensitivity_meta,
|
|
)
|
|
|
|
|
|
def _dump_to_parts(dump: Any) -> dict[str, Any]:
|
|
if dump is None:
|
|
return {}
|
|
if isinstance(dump, dict):
|
|
return dict(dump)
|
|
# DumpResult dataclass or duck-typed dump objects (tests / drivers).
|
|
if is_dataclass(dump) and not isinstance(dump, type) or hasattr(dump, "surface"):
|
|
return {
|
|
"surface": getattr(dump, "surface", "") or "",
|
|
"driver": getattr(dump, "driver", "") or "",
|
|
"window": dict(getattr(dump, "window", None) or {}),
|
|
"tree": getattr(dump, "tree", None) or {},
|
|
"index": list(getattr(dump, "index", None) or []),
|
|
"locatable": list(getattr(dump, "locatable", None) or []),
|
|
"stats": dict(getattr(dump, "stats", None) or {}),
|
|
"selection": dict(getattr(dump, "selection", None) or {}),
|
|
"url": getattr(dump, "url", "") or "",
|
|
"page_title": getattr(dump, "page_title", "") or "",
|
|
"screenshot_bytes": getattr(dump, "screenshot_bytes", None),
|
|
"use_hwnd_screenshot": bool(getattr(dump, "use_hwnd_screenshot", False)),
|
|
"locatable_text": getattr(dump, "locatable_text", "") or "",
|
|
"tree_text": getattr(dump, "tree_text", "") or "",
|
|
}
|
|
return {}
|
|
|
|
|
|
class SnapshotService:
|
|
"""
|
|
Creates UIState records with fingerprint dedupe.
|
|
|
|
Reuses SurfaceDriver.dump(); does not reimplement browser/desktop dump.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
flow_dir: Path,
|
|
*,
|
|
save_sensitive_screenshots: bool = False,
|
|
persist_evidence: bool = True,
|
|
) -> None:
|
|
self.flow_dir = Path(flow_dir)
|
|
self.states_dir = self.flow_dir / "states"
|
|
self.states_dir.mkdir(parents=True, exist_ok=True)
|
|
self.observations_dir = self.flow_dir / "observations"
|
|
self.observations_dir.mkdir(parents=True, exist_ok=True)
|
|
self.save_sensitive_screenshots = save_sensitive_screenshots
|
|
self.persist_evidence = persist_evidence
|
|
self._by_fp: dict[str, UIState] = {}
|
|
self._order: list[str] = []
|
|
self.latest: Optional[UIState] = None
|
|
self.latest_observation: Optional[UIObservation] = None
|
|
self._observation_order: list[str] = []
|
|
|
|
def load_existing(self) -> None:
|
|
"""Index existing state JSON files (best-effort)."""
|
|
if not self.states_dir.exists():
|
|
return
|
|
for path in sorted(self.states_dir.glob("*.json")):
|
|
try:
|
|
import json
|
|
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
st = UIState(
|
|
id=data.get("id") or path.stem,
|
|
surface=data.get("surface") or "windows",
|
|
driver=data.get("driver") or "",
|
|
captured_at=data.get("captured_at") or "",
|
|
fingerprint=data.get("fingerprint") or "",
|
|
previous_state_id=data.get("previous_state_id"),
|
|
change_reason=data.get("change_reason") or "",
|
|
app=data.get("app") or {},
|
|
window=data.get("window") or {},
|
|
url=data.get("url") or "",
|
|
page_title=data.get("page_title") or "",
|
|
tree_path=data.get("tree_path") or "",
|
|
index_path=data.get("index_path") or "",
|
|
locatable_path=data.get("locatable_path") or "",
|
|
screenshot_path=data.get("screenshot_path") or "",
|
|
stats=data.get("stats") or {},
|
|
selection=data.get("selection") or {},
|
|
duplicate_of=data.get("duplicate_of"),
|
|
observation_count=int(data.get("observation_count") or 1),
|
|
may_contain_sensitive_data=bool(
|
|
data.get("may_contain_sensitive_data")
|
|
),
|
|
)
|
|
if st.fingerprint:
|
|
self._by_fp[st.fingerprint] = st
|
|
if st.id not in self._order:
|
|
self._order.append(st.id)
|
|
self.latest = st
|
|
except Exception:
|
|
continue
|
|
for path in sorted(self.observations_dir.glob("*/meta.json")):
|
|
try:
|
|
import json
|
|
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
obs = UIObservation(
|
|
**{k: v for k, v in data.items() if k in UIObservation.__dataclass_fields__}
|
|
)
|
|
self._observation_order.append(obs.id)
|
|
self.latest_observation = obs
|
|
except Exception:
|
|
continue
|
|
|
|
def capture_from_dump(
|
|
self,
|
|
dump: Any,
|
|
*,
|
|
change_reason: str = "",
|
|
force_persist: bool = False,
|
|
previous_state_id: Optional[str] = None,
|
|
redact: bool = True,
|
|
) -> UIState:
|
|
parts = _dump_to_parts(dump)
|
|
locatable = list(parts.get("locatable") or [])
|
|
sensitive = False
|
|
if redact:
|
|
locatable, sensitive = redact_locatable_rows(locatable)
|
|
|
|
window = dict(parts.get("window") or {})
|
|
url = parts.get("url") or ""
|
|
page_title = parts.get("page_title") or ""
|
|
surface = parts.get("surface") or "windows"
|
|
driver = parts.get("driver") or ""
|
|
fp = compute_state_fingerprint(
|
|
surface=surface,
|
|
driver=driver,
|
|
url=url,
|
|
page_title=page_title,
|
|
window=window,
|
|
locatable=locatable,
|
|
app={
|
|
"title": window.get("title"),
|
|
"class_name": window.get("class_name"),
|
|
"process_name": window.get("process_name"),
|
|
},
|
|
)
|
|
|
|
existing = self._by_fp.get(fp)
|
|
if existing is not None:
|
|
existing.bump_observation()
|
|
existing.duplicate_of = None
|
|
# Update meta path only — do not rewrite full tree evidence.
|
|
write_model_json(self.states_dir / f"{existing.id}.json", existing)
|
|
self.latest = existing
|
|
self._persist_observation(
|
|
existing,
|
|
parts,
|
|
locatable=locatable,
|
|
sensitive=sensitive,
|
|
change_reason=change_reason or "capture",
|
|
)
|
|
return existing
|
|
|
|
prev = previous_state_id
|
|
if prev is None and self.latest is not None:
|
|
prev = self.latest.id
|
|
|
|
state = UIState(
|
|
surface=surface,
|
|
driver=driver,
|
|
fingerprint=fp,
|
|
previous_state_id=prev,
|
|
change_reason=change_reason or "capture",
|
|
app={
|
|
"title": window.get("title"),
|
|
"class_name": window.get("class_name"),
|
|
"process_name": window.get("process_name"),
|
|
"pid": window.get("pid"),
|
|
"hwnd": window.get("hwnd"),
|
|
},
|
|
window=window,
|
|
url=url,
|
|
page_title=page_title,
|
|
stats=dict(parts.get("stats") or {}),
|
|
selection=dict(parts.get("selection") or {}),
|
|
locatable=locatable,
|
|
may_contain_sensitive_data=sensitive,
|
|
)
|
|
state.ensure_fingerprint()
|
|
|
|
should_write_evidence = force_persist or bool(change_reason) or self.persist_evidence
|
|
if should_write_evidence:
|
|
self._persist_evidence(state, parts, sensitive=sensitive)
|
|
|
|
write_model_json(self.states_dir / f"{state.id}.json", state)
|
|
self._by_fp[fp] = state
|
|
self._order.append(state.id)
|
|
self.latest = state
|
|
self._persist_observation(
|
|
state,
|
|
parts,
|
|
locatable=locatable,
|
|
sensitive=sensitive,
|
|
change_reason=change_reason or "capture",
|
|
)
|
|
return state
|
|
|
|
def _persist_observation(
|
|
self,
|
|
state: UIState,
|
|
parts: dict[str, Any],
|
|
*,
|
|
locatable: list[dict[str, Any]],
|
|
sensitive: bool,
|
|
change_reason: str,
|
|
) -> UIObservation:
|
|
stats = dict(parts.get("stats") or {})
|
|
selection = dict(parts.get("selection") or {})
|
|
active_page = dict(stats.get("active_page") or {})
|
|
browser_session = dict(stats.get("browser_session") or {})
|
|
obs = UIObservation(
|
|
state_id=state.id,
|
|
previous_observation_id=(
|
|
self.latest_observation.id if self.latest_observation else None
|
|
),
|
|
change_reason=change_reason,
|
|
surface=state.surface,
|
|
driver=state.driver,
|
|
page_id=str(
|
|
active_page.get("page_id")
|
|
or browser_session.get("page_id")
|
|
or selection.get("page_id")
|
|
or ""
|
|
),
|
|
document_id=str(selection.get("document_id") or ""),
|
|
frame_id=str(selection.get("frame_id") or ""),
|
|
url=state.url,
|
|
page_title=state.page_title,
|
|
stats=stats,
|
|
selection=selection,
|
|
may_contain_sensitive_data=sensitive,
|
|
)
|
|
evidence_dir = self.observations_dir / obs.id
|
|
evidence_dir.mkdir(parents=True, exist_ok=True)
|
|
import artifacts as art
|
|
|
|
art.write_json(
|
|
evidence_dir / "tree.json",
|
|
sanitize_artifact(parts.get("tree") or {}),
|
|
)
|
|
clean_index = sanitize_artifact(list(parts.get("index") or []))
|
|
art.write_jsonl(evidence_dir / "index.jsonl", list(clean_index or []))
|
|
art.write_json(evidence_dir / "locatable.json", list(locatable or []))
|
|
obs.tree_path = f"observations/{obs.id}/tree.json"
|
|
obs.index_path = f"observations/{obs.id}/index.jsonl"
|
|
obs.locatable_path = f"observations/{obs.id}/locatable.json"
|
|
|
|
sensitivity_meta = screenshot_sensitivity_meta(
|
|
may_contain_sensitive=sensitive,
|
|
save_sensitive_screenshots=self.save_sensitive_screenshots,
|
|
)
|
|
art.write_json(evidence_dir / "sensitivity.json", sensitivity_meta)
|
|
if (not sensitive) or self.save_sensitive_screenshots:
|
|
png = parts.get("screenshot_bytes")
|
|
if png:
|
|
try:
|
|
(evidence_dir / "screen.png").write_bytes(png)
|
|
obs.screenshot_path = f"observations/{obs.id}/screen.png"
|
|
except Exception:
|
|
pass
|
|
elif parts.get("use_hwnd_screenshot"):
|
|
hwnd = int((parts.get("window") or {}).get("hwnd") or 0)
|
|
if hwnd and art.capture_window_png(hwnd, evidence_dir / "screen.png"):
|
|
obs.screenshot_path = f"observations/{obs.id}/screen.png"
|
|
|
|
write_model_json(evidence_dir / "meta.json", obs)
|
|
self.latest_observation = obs
|
|
self._observation_order.append(obs.id)
|
|
return obs
|
|
|
|
def _persist_evidence(
|
|
self,
|
|
state: UIState,
|
|
parts: dict[str, Any],
|
|
*,
|
|
sensitive: bool,
|
|
) -> None:
|
|
evidence_dir = self.states_dir / state.id
|
|
evidence_dir.mkdir(parents=True, exist_ok=True)
|
|
import artifacts as art
|
|
|
|
art.write_json(
|
|
evidence_dir / "tree.json",
|
|
sanitize_artifact(parts.get("tree") or {}),
|
|
)
|
|
clean_index = sanitize_artifact(list(parts.get("index") or []))
|
|
art.write_jsonl(evidence_dir / "index.jsonl", list(clean_index or []))
|
|
art.write_json(evidence_dir / "locatable.json", list(state.locatable or []))
|
|
state.tree_path = f"states/{state.id}/tree.json"
|
|
state.index_path = f"states/{state.id}/index.jsonl"
|
|
state.locatable_path = f"states/{state.id}/locatable.json"
|
|
|
|
sens_meta = screenshot_sensitivity_meta(
|
|
may_contain_sensitive=sensitive,
|
|
save_sensitive_screenshots=self.save_sensitive_screenshots,
|
|
)
|
|
art.write_json(evidence_dir / "sensitivity.json", sens_meta)
|
|
|
|
save_png = (not sensitive) or self.save_sensitive_screenshots
|
|
if save_png:
|
|
png = parts.get("screenshot_bytes")
|
|
if png:
|
|
try:
|
|
(evidence_dir / "screen.png").write_bytes(png)
|
|
state.screenshot_path = f"states/{state.id}/screen.png"
|
|
except Exception:
|
|
pass
|
|
elif parts.get("use_hwnd_screenshot"):
|
|
hwnd = int((parts.get("window") or {}).get("hwnd") or 0)
|
|
if hwnd:
|
|
ok = art.capture_window_png(hwnd, evidence_dir / "screen.png")
|
|
if ok:
|
|
state.screenshot_path = f"states/{state.id}/screen.png"
|
|
elif sensitive:
|
|
state.may_contain_sensitive_data = True
|
|
|
|
def wait_until_stable(
|
|
self,
|
|
dump_fn,
|
|
*,
|
|
timeout_sec: float = 3.0,
|
|
poll_sec: float = 0.25,
|
|
change_reason: str = "after_action",
|
|
) -> UIState:
|
|
"""
|
|
Poll dump_fn until two consecutive fingerprints match or timeout.
|
|
dump_fn() -> DumpResult / dump-like object
|
|
"""
|
|
deadline = time.monotonic() + max(0.2, float(timeout_sec))
|
|
last_fp: Optional[str] = None
|
|
last_dump = None
|
|
while True:
|
|
last_dump = dump_fn()
|
|
parts = _dump_to_parts(last_dump)
|
|
locatable, _ = redact_locatable_rows(list(parts.get("locatable") or []))
|
|
fp = compute_state_fingerprint(
|
|
surface=parts.get("surface") or "windows",
|
|
driver=parts.get("driver") or "",
|
|
url=parts.get("url") or "",
|
|
page_title=parts.get("page_title") or "",
|
|
window=parts.get("window") or {},
|
|
locatable=locatable,
|
|
)
|
|
if last_fp is not None and fp == last_fp:
|
|
break
|
|
last_fp = fp
|
|
if time.monotonic() >= deadline:
|
|
break
|
|
time.sleep(max(0.05, float(poll_sec)))
|
|
return self.capture_from_dump(
|
|
last_dump,
|
|
change_reason=change_reason,
|
|
force_persist=True,
|
|
)
|
|
|
|
def state_ids(self) -> list[str]:
|
|
return list(self._order)
|
|
|
|
def observation_ids(self) -> list[str]:
|
|
return list(self._observation_order)
|
|
|
|
def get(self, state_id: str) -> Optional[UIState]:
|
|
for st in self._by_fp.values():
|
|
if st.id == state_id:
|
|
return st
|
|
path = self.states_dir / f"{state_id}.json"
|
|
if path.exists():
|
|
import json
|
|
from models.state import UIState as S
|
|
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
return S(**{k: v for k, v in data.items() if k in S.__dataclass_fields__})
|
|
return None
|
|
|
|
def to_summary_dict(self, state: UIState) -> dict[str, Any]:
|
|
d = to_dict(state)
|
|
d["locatable"] = list(state.locatable or [])
|
|
return d
|