| 1 | """Event Stacker — record agent events as a structured event stream.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import json |
| 6 | import threading |
| 7 | from contextlib import contextmanager |
| 8 | from datetime import datetime, timezone |
| 9 | from pathlib import Path |
| 10 | from typing import Any, Generator |
| 11 | |
| 12 | |
| 13 | _STACKER_DIR = ".event_stacker" |
| 14 | |
| 15 | |
| 16 | class EventStacker: |
| 17 | _instance: EventStacker | None = None |
| 18 | _lock = threading.Lock() |
| 19 | |
| 20 | def __init__(self, workspace: Path, *, max_traces: int = 100) -> None: |
| 21 | self._workspace = workspace |
| 22 | self._max_traces = max_traces |
| 23 | self._dir = workspace / _STACKER_DIR |
| 24 | self._dir.mkdir(parents=True, exist_ok=True) |
| 25 | self._tls = threading.local() |
| 26 | self._turn_counter = self._recover_max_turn_id() |
| 27 | self._event_counter = 0 |
| 28 | |
| 29 | def _recover_max_turn_id(self) -> int: |
| 30 | max_id = 0 |
| 31 | try: |
| 32 | for f in self._dir.glob("*.jsonl"): |
| 33 | with open(f, "rb") as fh: |
| 34 | fh.seek(max(0, f.stat().st_size - 8192)) |
| 35 | tail = fh.read().decode("utf-8", errors="replace") |
| 36 | for line in reversed(tail.strip().split("\n")): |
| 37 | line = line.strip() |
| 38 | if not line: |
| 39 | continue |
| 40 | try: |
| 41 | rec = json.loads(line) |
| 42 | tid = rec.get("turn_id", "") |
| 43 | if tid.startswith("turn_"): |
| 44 | num = int(tid.split("_", 1)[1]) |
| 45 | max_id = max(max_id, num) |
| 46 | break |
| 47 | except (json.JSONDecodeError, ValueError): |
| 48 | continue |
| 49 | except OSError: |
| 50 | pass |
| 51 | return max_id |
| 52 | |
| 53 | @classmethod |
| 54 | def init(cls, workspace: Path, *, enabled: bool = False, max_traces: int = 100) -> None: |
| 55 | with cls._lock: |
| 56 | cls._instance = cls(workspace, max_traces=max_traces) if enabled else None |
| 57 | |
| 58 | @classmethod |
| 59 | def begin_turn(cls, session_key: str | None, model: str | None) -> None: |
| 60 | inst = cls._instance |
| 61 | if inst is None: |
| 62 | return |
| 63 | key = session_key or "unknown" |
| 64 | inst._tls.session_key = key |
| 65 | with inst._lock: |
| 66 | inst._turn_counter += 1 |
| 67 | inst._tls.turn_id = f"turn_{inst._turn_counter:04d}" |
| 68 | inst._tls.phase = "turn_context" |
| 69 | cls.emit("turn_start", {"model": model or ""}) |
| 70 | |
| 71 | @classmethod |
| 72 | def set_phase(cls, phase: str | None) -> None: |
| 73 | inst = cls._instance |
| 74 | if inst is None: |
| 75 | return |
| 76 | inst._tls.phase = phase |
| 77 | |
| 78 | @classmethod |
| 79 | @contextmanager |
| 80 | def phase(cls, name: str) -> Generator[None, None, None]: |
| 81 | """Context manager: set phase for the duration, restore previous on exit.""" |
| 82 | inst = cls._instance |
| 83 | if inst is None: |
| 84 | yield |
| 85 | return |
| 86 | prev = getattr(inst._tls, "phase", None) |
| 87 | inst._tls.phase = name |
| 88 | try: |
| 89 | yield |
| 90 | finally: |
| 91 | inst._tls.phase = prev |
| 92 | |
| 93 | @classmethod |
| 94 | def log(cls, label: str, content: Any) -> None: |
| 95 | """Shorthand for emit('context_part', ...) — mirrors PromptStacker.log API.""" |
| 96 | inst = cls._instance |
| 97 | if inst is None: |
| 98 | return |
| 99 | current_phase = getattr(inst._tls, "phase", None) |
| 100 | if current_phase is None: |
| 101 | return |
| 102 | text = content if isinstance(content, str) else json.dumps( |
| 103 | content, ensure_ascii=False, default=str, |
| 104 | ) |
| 105 | cls.emit("context_part", { |
| 106 | "label": label, |
| 107 | "phase": current_phase, |
| 108 | "content": text, |
| 109 | "char_count": len(text), |
| 110 | }) |
| 111 | |
| 112 | @classmethod |
| 113 | def emit(cls, event_type: str, data: dict[str, Any] | None = None) -> None: |
| 114 | inst = cls._instance |
| 115 | if inst is None: |
| 116 | return |
| 117 | key = getattr(inst._tls, "session_key", None) |
| 118 | if key is None: |
| 119 | return |
| 120 | turn_id = getattr(inst._tls, "turn_id", None) |
| 121 | |
| 122 | with inst._lock: |
| 123 | inst._event_counter += 1 |
| 124 | seq = inst._event_counter |
| 125 | |
| 126 | record = { |
| 127 | "seq": seq, |
| 128 | "type": event_type, |
| 129 | "turn_id": turn_id or "", |
| 130 | "timestamp": datetime.now(timezone.utc).isoformat(), |
| 131 | "session_key": key, |
| 132 | "data": data or {}, |
| 133 | } |
| 134 | |
| 135 | safe_key = key.replace(":", "_").replace("/", "_") |
| 136 | trace_file = inst._dir / f"{safe_key}.jsonl" |
| 137 | with open(trace_file, "a", encoding="utf-8") as f: |
| 138 | f.write(json.dumps(record, ensure_ascii=False, default=str) + "\n") |
| 139 | |
| 140 | @classmethod |
| 141 | def end_turn( |
| 142 | cls, |
| 143 | stop_reason: str | None = None, |
| 144 | usage: dict[str, Any] | None = None, |
| 145 | ) -> None: |
| 146 | cls.emit("turn_end", { |
| 147 | "stop_reason": stop_reason or "", |
| 148 | "usage": usage or {}, |
| 149 | }) |
| 150 | inst = cls._instance |
| 151 | if inst is not None: |
| 152 | inst._tls.phase = None |
| 153 | inst._gc_traces() |
| 154 | |
| 155 | def _gc_traces(self) -> None: |
| 156 | try: |
| 157 | traces = sorted(self._dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime) |
| 158 | while len(traces) > self._max_traces: |
| 159 | traces.pop(0).unlink(missing_ok=True) |
| 160 | except OSError: |
| 161 | pass |
| 162 | |
| 163 | @classmethod |
| 164 | def get_sessions(cls) -> list[dict[str, Any]]: |
| 165 | inst = cls._instance |
| 166 | if inst is None: |
| 167 | return [] |
| 168 | results = [] |
| 169 | for f in sorted(inst._dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True): |
| 170 | stat = f.stat() |
| 171 | event_count = 0 |
| 172 | last_record: dict[str, Any] = {} |
| 173 | try: |
| 174 | with open(f, "r", encoding="utf-8") as fh: |
| 175 | for line in fh: |
| 176 | line = line.strip() |
| 177 | if line: |
| 178 | event_count += 1 |
| 179 | try: |
| 180 | last_record = json.loads(line) |
| 181 | except json.JSONDecodeError: |
| 182 | pass |
| 183 | except OSError: |
| 184 | pass |
| 185 | results.append({ |
| 186 | "id": f.stem, |
| 187 | "filename": f.name, |
| 188 | "event_count": event_count, |
| 189 | "size_bytes": stat.st_size, |
| 190 | "modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat(), |
| 191 | "last_turn_id": last_record.get("turn_id", ""), |
| 192 | "last_session_key": last_record.get("session_key", ""), |
| 193 | }) |
| 194 | return results |
| 195 | |
| 196 | @classmethod |
| 197 | def get_trace(cls, session_id: str) -> list[dict[str, Any]]: |
| 198 | inst = cls._instance |
| 199 | if inst is None: |
| 200 | return [] |
| 201 | trace_file = inst._dir / f"{session_id}.jsonl" |
| 202 | if not trace_file.is_file(): |
| 203 | return [] |
| 204 | try: |
| 205 | trace_file.relative_to(inst._dir) |
| 206 | except ValueError: |
| 207 | return [] |
| 208 | records = [] |
| 209 | with open(trace_file, "r", encoding="utf-8") as f: |
| 210 | for line in f: |
| 211 | line = line.strip() |
| 212 | if line: |
| 213 | try: |
| 214 | records.append(json.loads(line)) |
| 215 | except json.JSONDecodeError: |
| 216 | continue |
| 217 | return records |
| 218 |