replay.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/replay.py
 
"""
Reconstruct an OC Chat session as a parent/child tree from its event NDJSON — M2.7.
 
Purpose
-------
M2.1–M2.6 produce and (in-session) read a per-session event log. This module is
the **offline** consumer: given an event directory, it reads every rotation slot,
reassembles the flat event list into a turn/tool-call/sub-agent tree, and renders
that tree as human-readable text or canonical ``replay/v0`` JSON. It never calls a
provider, never mutates the log, never touches the bus or writer — observability
only (the half-step between M2.6 ``/trace`` and the M7 packet).
 
Responsibilities
----------------
- :func:`read_events_from_dir` — directory-keyed twin of
  ``trace_view.read_session_events`` (ADR-4): walk ``events.NNNN.jsonl`` slots in
  sequence, parse to ``TurnEvent``, tolerate missing dirs and torn lines.
- :func:`build_replay_tree` — pure two-pass reduce: group by ``turn_id``, attach
  tool calls by ``tool_call_id`` and sub-agents by ``subagent_session_id``.
- :func:`render_tree_text` / :func:`render_tree_json` — the two output formats;
  JSON is canonical (byte-stable) for M2.9 golden fixtures.
 
Diagnostics
-----------
Domain: OC_CHAT.REPLAY
Levels:
    L2 — read/build lifecycle.
    L3 — slot/line counts and skipped malformed lines.
    L4 — module-load coverage marker.
 
Contracts
---------
- Reading never raises on missing dirs or torn lines (returns ``([], 0)`` / skips).
- ``build_replay_tree`` is pure; it does no I/O.
- ``render_tree_json`` is canonical and byte-stable; schema is ``"replay/v0"``.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- oc_chat.core.events (TurnEvent — parse persisted lines)
- option_c_agent_contracts (canonical_json — byte-stable JSON output)
- pydantic (BaseModel — the replay-tree node models)
"""
from __future__ import annotations
 
import re
from pathlib import Path
from typing import Any, Optional
 
from pydantic import BaseModel, Field
 
from oc_diagnostics import _dbg
 
from oc_chat.core.events import TurnEvent
from option_c_agent_contracts import canonical_json
 
_dbg("OC_CHAT.REPLAY", "replay", "module", 4)
 
REPLAY_SCHEMA = "replay/v0"
_SLOT_RE = re.compile(r"^events\.(\d+)\.jsonl$")
 
 
def _slot_seq(path: Path) -> int:
    """Return the rotation sequence encoded in an ``events.NNNN.jsonl`` filename."""
    match = _SLOT_RE.match(path.name)
    return int(match.group(1)) if match else 0
 
 
def read_events_from_dir(event_dir: Path) -> tuple[list[TurnEvent], int]:
    """Read every ``events.NNNN.jsonl`` slot under *event_dir* in sequence order.
 
    Returns ``(events, skipped)``. A missing dir yields ``([], 0)``; a torn or
    malformed line is skipped and counted; non-matching filenames are ignored.
    Deliberately duplicates the small slot-walk from ``trace_view`` (ADR-4) so
    that frozen module is left untouched; the on-disk layout is itself frozen
    (``event_writer.py``), so the two readers cannot drift.
    """
    func = "read_events_from_dir"
    base = Path(event_dir)
    if not base.exists():
        _dbg("OC_CHAT.REPLAY", func, f"no event dir at {base}", 3)
        return [], 0
    slots = sorted(
        (p for p in base.iterdir() if p.is_file() and _SLOT_RE.match(p.name)),
        key=_slot_seq,
    )
    events: list[TurnEvent] = []
    skipped = 0
    for slot in slots:
        for line in slot.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            try:
                events.append(TurnEvent.model_validate_json(line))
            except Exception:
                skipped += 1
    _dbg("OC_CHAT.REPLAY", func, f"read {len(events)} events from {len(slots)} slot(s); skipped={skipped}", 3)
    return events, skipped
 
 
class ToolCallNode(BaseModel):
    """One tool call within a turn, assembled from proposed/started/completed events."""
 
    tool_call_id: str
    tool_name: str = ""
    proposed_args: dict[str, Any] = Field(default_factory=dict)
    outcome: Optional[str] = None
    ts_proposed: Optional[str] = None
    ts_started: Optional[str] = None
    ts_completed: Optional[str] = None
 
 
class SubagentNode(BaseModel):
    """One sub-agent run nested under its parent turn (depth-1; ADR-5)."""
 
    subagent_session_id: str
    parent_turn_id: str
    outcome: Optional[str] = None
    ts_started: Optional[str] = None
    ts_completed: Optional[str] = None
 
 
class TurnNode(BaseModel):
    """A single turn: user input, tool calls, sub-agents, final answer, errors."""
 
    turn_id: str
    ts_first: Optional[str] = None
    ts_last: Optional[str] = None
    user_input: Optional[str] = None
    assistant_message: Optional[str] = None
    assistant_is_final: bool = False
    assistant_message_history: list[str] = Field(default_factory=list)
    tool_calls: list[ToolCallNode] = Field(default_factory=list)
    subagents: list[SubagentNode] = Field(default_factory=list)
    errors: list[str] = Field(default_factory=list)
    other_events: list[dict] = Field(default_factory=list)
 
 
class ReplayTree(BaseModel):
    """The reconstructed session: ordered turns + stream-level counts."""
 
    session_id: str = ""
    turn_count: int = 0
    event_count: int = 0
    skipped_count: int = 0
    source_dir: Optional[str] = None
    event_counts: dict[str, int] = Field(default_factory=dict)
    turns: list[TurnNode] = Field(default_factory=list)
 
 
_TOOL_KINDS = {"tool_call_proposed", "tool_call_started", "tool_call_completed", "tool_call_progress"}
_SUBAGENT_KINDS = {"subagent_started", "subagent_completed"}
_KNOWN_KINDS = (
    {"user_input_received", "assistant_message", "error"} | _TOOL_KINDS | _SUBAGENT_KINDS
)
 
 
def _build_turn_node(turn_id: str, events: list[TurnEvent]) -> TurnNode:
    """Assemble one :class:`TurnNode` from the events sharing *turn_id* (write order)."""
    node = TurnNode(turn_id=turn_id, ts_first=events[0].ts, ts_last=events[-1].ts)
    tool_nodes: dict[str, ToolCallNode] = {}
    proposed_q: list[Any] = []
    sub_nodes: dict[str, SubagentNode] = {}
 
    def _tool(tcid: str) -> ToolCallNode:
        if tcid not in tool_nodes:
            tn = ToolCallNode(tool_call_id=tcid)
            tool_nodes[tcid] = tn
            node.tool_calls.append(tn)
        return tool_nodes[tcid]
 
    for ev in events:
        p: Any = ev.payload
        kind = ev.kind
        if kind == "user_input_received":
            node.user_input = getattr(p, "text", None)
        elif kind == "assistant_message":
            text = getattr(p, "text", "")
            if getattr(p, "is_final", False):
                node.assistant_message = text
                node.assistant_is_final = True
            else:
                node.assistant_message_history.append(text)
        elif kind == "tool_call_proposed":
            proposed_q.append(p)
        elif kind == "tool_call_started":
            tn = _tool(ev.tool_call_id or getattr(p, "tool_call_id", "") or "?")
            tn.tool_name = getattr(p, "tool_name", "") or tn.tool_name
            tn.ts_started = ev.ts
            if proposed_q:
                prop = proposed_q.pop(0)
                tn.proposed_args = dict(getattr(prop, "requested_args", {}) or {})
                tn.ts_proposed = tn.ts_proposed or None
        elif kind == "tool_call_completed":
            tn = _tool(ev.tool_call_id or getattr(p, "tool_call_id", "") or "?")
            tn.outcome = getattr(p, "outcome", None)
            tn.ts_completed = ev.ts
        elif kind in _SUBAGENT_KINDS:
            sid = getattr(p, "subagent_session_id", "") or "?"
            sn = sub_nodes.get(sid)
            if sn is None:
                sn = SubagentNode(subagent_session_id=sid, parent_turn_id=turn_id)
                sub_nodes[sid] = sn
                node.subagents.append(sn)
            if kind == "subagent_started":
                sn.ts_started = ev.ts
            else:
                sn.outcome = getattr(p, "outcome", None)
                sn.ts_completed = ev.ts
        elif kind == "error":
            node.errors.append(getattr(p, "error_summary", "") or "")
        else:
            node.other_events.append(ev.model_dump(mode="json"))
    return node
 
 
def build_replay_tree(events: list[TurnEvent], *, skipped: int = 0) -> ReplayTree:
    """Reduce a flat event list into a :class:`ReplayTree` (pure; two-pass, O(N))."""
    session_id = ""
    order: list[str] = []
    buckets: dict[str, list[TurnEvent]] = {}
    counts: dict[str, int] = {}
    for ev in events:
        counts[ev.kind] = counts.get(ev.kind, 0) + 1
        if ev.session_id and not session_id:
            session_id = ev.session_id
        tid = ev.turn_id
        if not tid:
            continue  # session_started / session_ended — no turn bucket
        if tid not in buckets:
            buckets[tid] = []
            order.append(tid)
        buckets[tid].append(ev)
    nodes = [(_build_turn_node(tid, buckets[tid]), idx) for idx, tid in enumerate(order)]
    nodes.sort(key=lambda pair: (pair[0].ts_first or "", pair[1]))
    return ReplayTree(
        session_id=session_id,
        turn_count=len(nodes),
        event_count=len(events),
        skipped_count=skipped,
        event_counts=counts,
        turns=[node for node, _ in nodes],
    )
 
 
def _filter_turns(tree: ReplayTree, single_turn_id: Optional[str]) -> list[TurnNode]:
    """Return all turns, or only the one whose id matches *single_turn_id* exactly."""
    if single_turn_id is None:
        return list(tree.turns)
    return [t for t in tree.turns if t.turn_id == single_turn_id]
 
 
def render_tree_text(tree: ReplayTree, *, single_turn_id: Optional[str] = None) -> str:
    """Render *tree* as an indented ASCII tree + a per-kind count summary."""
    if tree.event_count == 0:
        return "── Replay: empty session (0 events) ──"
    turns = _filter_turns(tree, single_turn_id)
    head = (
        f"── Replay: session {(tree.session_id or '-')[:8]} ──\n"
        f"  events: {tree.event_count}  turns: {tree.turn_count}  "
        f"skipped: {tree.skipped_count}  source: {tree.source_dir or '-'}"
    )
    lines = [head]
    for ti, turn in enumerate(turns):
        last_turn = ti == len(turns) - 1
        tbranch = "└──" if last_turn else "├──"
        gutter = "    " if last_turn else "│   "
        lines.append(f"{tbranch} turn {turn.turn_id[:4]} [{turn.ts_first} → {turn.ts_last}]")
        if turn.user_input is not None:
            lines.append(f"{gutter}  user: {turn.user_input!r}")
        for tc in turn.tool_calls:
            lines.append(f"{gutter}  ├── tool {tc.tool_name} id={tc.tool_call_id[:8]} outcome={tc.outcome or '-'}")
        for sa in turn.subagents:
            lines.append(f"{gutter}  ├── subagent {sa.subagent_session_id[:4]} outcome={sa.outcome or '-'}")
        for err in turn.errors:
            lines.append(f"{gutter}  ├── error: {err!r}")
        if turn.assistant_message is not None:
            final = "final" if turn.assistant_is_final else "intermediate"
            lines.append(f"{gutter}  └── assistant ({final}): {turn.assistant_message!r}")
        if turn.other_events:
            lines.append(f"{gutter}  ({len(turn.other_events)} other event(s))")
    lines.append("── Summary ──")
    for kind in sorted(tree.event_counts):
        lines.append(f"  {kind:<24} {tree.event_counts[kind]}")
    return "\n".join(lines)
 
 
def render_tree_json(tree: ReplayTree, *, single_turn_id: Optional[str] = None) -> str:
    """Render *tree* (optionally one turn) as canonical, byte-stable ``replay/v0`` JSON."""
    turns = _filter_turns(tree, single_turn_id)
    payload = {
        "schema": REPLAY_SCHEMA,
        "session_id": tree.session_id,
        "event_count": tree.event_count,
        "skipped_count": tree.skipped_count,
        "turn_count": tree.turn_count,
        "source_dir": tree.source_dir,
        "event_counts": tree.event_counts,
        "turns": [t.model_dump(mode="json") for t in turns],
    }
    return canonical_json(payload)