#!/usr/bin/env python3 """Build a session-closed replay window from an already materialized trace. This is deliberately a *window-scoped fallback* for when the raw trace span needed by ``prepare_trace_windows.py`` is unavailable. It closes every known parent/child edge within the supplied window and groups siblings whose common parent lies outside the window. It cannot recover turns outside the window; the manifest records that residual rather than calling the result fully session-coherent. """ from __future__ import annotations import argparse import hashlib import json import math from collections import Counter from pathlib import Path from typing import Any SCORE_ALGORITHM = "collectivespec-window-cc-v1" def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def typed_id(value: Any) -> str: """Encode IDs without conflating values such as ``1`` and ``"1"``.""" return json.dumps( {"type": type(value).__name__, "value": value}, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ) def session_uniform(*, seed: int, window_id: str, component_key: str) -> float: payload = json.dumps( { "algorithm": SCORE_ALGORITHM, "seed": seed, "window_id": window_id, "component_key": component_key, }, sort_keys=True, separators=(",", ":"), ).encode("utf-8") digest = hashlib.blake2b(payload, digest_size=8).digest() return int.from_bytes(digest, byteorder="big", signed=False) / float(1 << 64) def resolve_trace_path(windows_path: Path, trace_file: str) -> Path: candidate = Path(trace_file) if candidate.is_absolute(): return candidate return (windows_path.parent / candidate).resolve() def is_root_parent(value: Any) -> bool: return value is None or ( isinstance(value, (int, float)) and not isinstance(value, bool) and int(value) < 0 ) class UnionFind: def __init__(self) -> None: self.parent: dict[str, str] = {} def add(self, item: str) -> None: self.parent.setdefault(item, item) def find(self, item: str) -> str: parent = self.parent[item] if parent != item: self.parent[item] = self.find(parent) return self.parent[item] def union(self, left: str, right: str) -> None: left_root, right_root = self.find(left), self.find(right) if left_root != right_root: # A canonical representative makes the component key independent # of trace order and union order. self.parent[max(left_root, right_root)] = min(left_root, right_root) def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--windows-path", type=Path, required=True) parser.add_argument("--window-id", required=True) parser.add_argument("--sample-seed", type=int, default=20260325) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--output-window-id", required=True) args = parser.parse_args() source_payload = json.loads(args.windows_path.read_text(encoding="utf-8")) source_windows = source_payload.get("windows") if isinstance(source_payload, dict) else source_payload if not isinstance(source_windows, list): raise SystemExit("windows payload must contain a list") source_window = next( (item for item in source_windows if item.get("window_id") == args.window_id), None, ) if not isinstance(source_window, dict): raise SystemExit(f"window id not found: {args.window_id}") trace_file = source_window.get("trace_file") if not isinstance(trace_file, str) or not trace_file: raise SystemExit(f"window {args.window_id} has no trace_file") source_trace = resolve_trace_path(args.windows_path, trace_file) if not source_trace.is_file(): raise SystemExit(f"source trace does not exist: {source_trace}") rows: list[dict[str, Any]] = [] id_to_index: dict[str, int] = {} with source_trace.open("r", encoding="utf-8") as source_handle: for row_index, raw in enumerate(source_handle): if not raw.strip(): continue row = json.loads(raw) if not isinstance(row, dict): raise SystemExit(f"row {row_index}: expected an object") if "chat_id" not in row or "parent_chat_id" not in row: raise SystemExit(f"row {row_index}: chat_id and parent_chat_id are required") if "timestamp" not in row: raise SystemExit(f"row {row_index}: timestamp is required") try: timestamp = float(row["timestamp"]) except (TypeError, ValueError) as error: raise SystemExit(f"row {row_index}: timestamp must be numeric") from error if not math.isfinite(timestamp): raise SystemExit(f"row {row_index}: timestamp must be finite") key = typed_id(row["chat_id"]) if key in id_to_index: raise SystemExit(f"row {row_index}: duplicate chat_id {row['chat_id']!r}") id_to_index[key] = len(rows) rows.append(row) if not rows: raise SystemExit("source trace has no rows") union_find = UnionFind() for chat_key in id_to_index: union_find.add(chat_key) known_parent_edges = 0 absent_parent_edges = 0 root_request_count = 0 for row in rows: chat_key = typed_id(row["chat_id"]) parent = row["parent_chat_id"] if is_root_parent(parent): root_request_count += 1 continue parent_key = typed_id(parent) union_find.add(parent_key) if parent_key in id_to_index: union_find.union(chat_key, parent_key) known_parent_edges += 1 else: # Keep the absent parent as a virtual node. This groups siblings # with a common boundary parent without pretending that the parent # turn itself is in this window. union_find.union(chat_key, parent_key) absent_parent_edges += 1 component_members: dict[str, list[str]] = {} for node in union_find.parent: component_members.setdefault(union_find.find(node), []).append(node) component_key_by_chat: dict[str, str] = {} for members in component_members.values(): # Virtual boundary parents intentionally participate: they ensure all # children of an unavailable parent share a deterministic component. component_key = min(members) for chat_key in (member for member in members if member in id_to_index): component_key_by_chat[chat_key] = component_key for row in rows: parent = row["parent_chat_id"] if is_root_parent(parent): continue parent_key = typed_id(parent) if parent_key in id_to_index and union_find.find(typed_id(row["chat_id"])) != union_find.find(parent_key): raise SystemExit("same-window parent edge did not close into one component") args.output_dir.mkdir(parents=True, exist_ok=True) trace_dir = args.output_dir / "traces" trace_dir.mkdir(exist_ok=True) output_trace = trace_dir / f"{args.output_window_id}.jsonl" score_values: list[float] = [] timestamps: list[float] = [] component_counts: Counter[str] = Counter() with output_trace.open("w", encoding="utf-8") as handle: for row in rows: chat_key = typed_id(row["chat_id"]) component_key = component_key_by_chat[chat_key] component_id = f"window-component:{component_key}" out = dict(row) out["collectivespec_original_sampling_u"] = out.get("sampling_u") out["session_root"] = component_id out["sampling_u"] = session_uniform( seed=args.sample_seed, window_id=args.window_id, component_key=component_key, ) handle.write(json.dumps(out, ensure_ascii=False, separators=(",", ":")) + "\n") score_values.append(float(out["sampling_u"])) timestamps.append(float(out["timestamp"])) component_counts[component_id] += 1 component_sizes = list(component_counts.values()) output_window = dict(source_window) output_window.update( { "window_id": args.output_window_id, "trace_file": str(output_trace.resolve()), "num_requests": len(rows), "sampling_strategy": "window_session_closed_uniform_score", "sampling_seed": args.sample_seed, "collectivespec_sampling_algorithm": SCORE_ALGORITHM, "collectivespec_source_window_id": args.window_id, "collectivespec_source_trace_path": str(source_trace), "collectivespec_source_trace_sha256": sha256(source_trace), "collectivespec_session_closure_scope": "window_only", "collectivespec_known_parent_edges": known_parent_edges, "collectivespec_absent_parent_edges": absent_parent_edges, } ) output_windows = { "kind": "collectivespec_window_session_closed_fallback", "source_windows_path": str(args.windows_path.resolve()), "source_window_id": args.window_id, "sample_seed": args.sample_seed, "windows": [output_window], } output_manifest = args.output_dir / "windows.json" output_manifest.write_text( json.dumps(output_windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) sanity = { "n": len(rows), "timestamp": { "min": min(timestamps), "max": max(timestamps), "non_negative": all(value >= 0 for value in timestamps), "nondecreasing": all( timestamps[index] >= timestamps[index - 1] for index in range(1, len(timestamps)) ), }, "sampling_u": { "min": min(score_values), "max": max(score_values), "distinct_value_count": len(set(score_values)), "within_0_1": all(0 <= value < 1 for value in score_values), }, "window_session_closure": { "component_count": len(component_sizes), "component_size_min": min(component_sizes), "component_size_max": max(component_sizes), "component_size_distinct_value_count": len(set(component_sizes)), "known_parent_edges": known_parent_edges, "absent_parent_edges": absent_parent_edges, "root_request_count": root_request_count, "all_known_parent_edges_closed": True, "scope": "window_only", }, "source_trace_sha256": sha256(source_trace), "output_trace_sha256": sha256(output_trace), } print( json.dumps( { "windows_path": str(output_manifest), "trace_path": str(output_trace), "data_sanity": sanity, }, ensure_ascii=False, indent=2, sort_keys=True, ) ) return 0 if __name__ == "__main__": raise SystemExit(main())