Update drp/diff.py

drp/diff.py

@@ -1,11 +1,10 @@
 import difflib
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
-from .bundle import Bundle, load_bundle
+from .bundle import load_bundle
 
 
 def _normalize_for_compare(x: Any) -> Any:
-    # Avoid false diffs from ordering
     if isinstance(x, dict):
         return {k: _normalize_for_compare(x[k]) for k in sorted(x.keys())}
     if isinstance(x, list):
@@ -13,11 +12,38 @@ def _normalize_for_compare(x: Any) -> Any:
     return x
 
 
+def _event_core(ev: Dict[str, Any]) -> Any:
+    return _normalize_for_compare({k: ev.get(k) for k in ("kind", "step", "payload")})
+
+
+def build_alignment(A_events: List[Dict[str, Any]], B_events: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    rows: List[Dict[str, Any]] = []
+    n = max(len(A_events), len(B_events))
+    for i in range(n):
+        a = A_events[i] if i < len(A_events) else None
+        b = B_events[i] if i < len(B_events) else None
+
+        if a is None:
+            status = "missing_in_A"
+        elif b is None:
+            status = "missing_in_B"
+        else:
+            status = "same" if _event_core(a) == _event_core(b) else "diff"
+
+        rows.append(
+            {
+                "i": i,
+                "status": status,
+                "kind_a": a.get("kind") if a else None,
+                "step_a": a.get("step") if a else None,
+                "kind_b": b.get("kind") if b else None,
+                "step_b": b.get("step") if b else None,
+            }
+        )
+    return rows
+
+
 def _json_diff(a: Any, b: Any, path: str = "") -> List[Dict[str, Any]]:
-    """
-    Small recursive diff (no heavy deps).
-    Emits list of {path, a, b, kind}.
-    """
     diffs: List[Dict[str, Any]] = []
 
     if type(a) != type(b):
@@ -36,7 +62,6 @@ def _json_diff(a: Any, b: Any, path: str = "") -> List[Dict[str, Any]]:
         return diffs
 
     if isinstance(a, list):
-        # list diff by index (simple)
         n = max(len(a), len(b))
         for i in range(n):
             pa = a[i] if i < len(a) else None
@@ -54,18 +79,16 @@ def _json_diff(a: Any, b: Any, path: str = "") -> List[Dict[str, Any]]:
     return diffs
 
 
-def _classify_divergence(ev_a: Dict[str, Any], ev_b: Dict[str, Any]) -> str:
-    ka = ev_a.get("kind")
-    kb = ev_b.get("kind")
-    if ka != kb:
+def _classify_divergence(kind_a: Optional[str], kind_b: Optional[str]) -> str:
+    if kind_a != kind_b:
         return "control-flow"
-    if ka in ("tool_call", "tool_result"):
+    if kind_a in ("tool_call", "tool_result"):
         return "tool"
-    if ka in ("memory_write", "memory_read"):
+    if kind_a in ("memory_write", "memory_read"):
         return "memory"
-    if ka in ("llm_sample", "llm_call"):
+    if kind_a in ("llm_sample", "llm_call"):
         return "sampling"
-    if ka in ("guardrail",):
+    if kind_a in ("guardrail",):
         return "governance"
     return "state"
 
@@ -77,48 +100,115 @@ def _text_delta(a: str, b: str) -> str:
     return "\n".join(diff)
 
 
+def _extract_final_reward(events: List[Dict[str, Any]]) -> Optional[float]:
+    """
+    Looks for last state_snapshot payload containing:
+      - payload.reward_total
+      - payload.metrics.reward_total
+    """
+    for ev in reversed(events):
+        if ev.get("kind") != "state_snapshot":
+            continue
+        p = ev.get("payload", {}) or {}
+        if isinstance(p, dict):
+            rt = p.get("reward_total")
+            if isinstance(rt, (int, float)):
+                return float(rt)
+            m = p.get("metrics")
+            if isinstance(m, dict):
+                rt2 = m.get("reward_total")
+                if isinstance(rt2, (int, float)):
+                    return float(rt2)
+    return None
+
+
+def _event_link(manifest: Dict[str, Any], i: int) -> Optional[str]:
+    """
+    Optional deep-link generation.
+    Supported:
+      - manifest.replay.base_url + manifest.replay.pattern with {run_id} and {i}
+      - manifest.run_url + ?i={i}
+    """
+    run_id = manifest.get("run_id")
+    replay = manifest.get("replay")
+
+    if isinstance(replay, dict):
+        base = replay.get("base_url")
+        pattern = replay.get("pattern", "")
+        if isinstance(base, str) and isinstance(pattern, str) and run_id:
+            try:
+                return base.rstrip("/") + pattern.format(run_id=run_id, i=i)
+            except Exception:
+                return None
+
+    run_url = manifest.get("run_url")
+    if isinstance(run_url, str) and run_url:
+        # append i in a minimal, non-destructive way
+        joiner = "&" if "?" in run_url else "?"
+        return f"{run_url}{joiner}i={i}"
+
+    return None
+
+
 def diff_bundles(zip_a: str, zip_b: str) -> Dict[str, Any]:
     A = load_bundle(zip_a)
     B = load_bundle(zip_b)
 
     ea = A.events
     eb = B.events
-    n = min(len(ea), len(eb))
 
+    alignment = build_alignment(ea, eb)
+
+    # first divergence index (including length mismatch)
     first_div: Optional[int] = None
-    per_event: List[Dict[str, Any]] = []
+    for row in alignment:
+        if row["status"] != "same":
+            first_div = row["i"]
+            break
 
+    # diff details (per index where both exist and differ)
+    per_event: List[Dict[str, Any]] = []
+    n = min(len(ea), len(eb))
     for i in range(n):
-        na = _normalize_for_compare({k: ea[i].get(k) for k in ("kind", "step", "payload")})
-        nb = _normalize_for_compare({k: eb[i].get(k) for k in ("kind", "step", "payload")})
-        if na != nb and first_div is None:
-            first_div = i
-
-        if na != nb:
-            diffs = _json_diff(na, nb)
-            item = {
-                "i": i,
-                "step_a": ea[i].get("step"),
-                "step_b": eb[i].get("step"),
-                "kind_a": ea[i].get("kind"),
-                "kind_b": eb[i].get("kind"),
-                "class": _classify_divergence(ea[i], eb[i]),
-                "diffs": diffs[:200],  # cap
-            }
-
-            # Optional friendly text diff if payload has 'text'
-            ta = ea[i].get("payload", {}).get("text")
-            tb = eb[i].get("payload", {}).get("text")
-            if isinstance(ta, str) and isinstance(tb, str) and ta != tb:
-                item["text_unified_diff"] = _text_delta(ta, tb)[:20000]
-
-            per_event.append(item)
-
-    # length mismatch
-    if len(ea) != len(eb):
-        first_div = first_div if first_div is not None else n
+        na = _event_core(ea[i])
+        nb = _event_core(eb[i])
+        if na == nb:
+            continue
+
+        diffs = _json_diff(na, nb)
+        item: Dict[str, Any] = {
+            "i": i,
+            "kind_a": ea[i].get("kind"),
+            "kind_b": eb[i].get("kind"),
+            "step_a": ea[i].get("step"),
+            "step_b": eb[i].get("step"),
+            "class": _classify_divergence(ea[i].get("kind"), eb[i].get("kind")),
+            "diffs": diffs[:200],
+            "link_a": _event_link(A.manifest, i),
+            "link_b": _event_link(B.manifest, i),
+        }
+
+        ta = (ea[i].get("payload", {}) or {}).get("text")
+        tb = (eb[i].get("payload", {}) or {}).get("text")
+        if isinstance(ta, str) and isinstance(tb, str) and ta != tb:
+            item["text_unified_diff"] = _text_delta(ta, tb)[:20000]
+
+        per_event.append(item)
+
+    diff_count = sum(1 for r in alignment if r["status"] == "diff")
+    missing_count = sum(1 for r in alignment if r["status"] in ("missing_in_A", "missing_in_B"))
+
+    ra = _extract_final_reward(ea)
+    rb = _extract_final_reward(eb)
+    reward_delta = (rb - ra) if (ra is not None and rb is not None) else None
+
+    # class counts
+    counts: Dict[str, int] = {}
+    for item in per_event:
+        c = item["class"]
+        counts[c] = counts.get(c, 0) + 1
 
-    summary = {
+    summary: Dict[str, Any] = {
         "run_a": A.manifest.get("run_id"),
         "run_b": B.manifest.get("run_id"),
         "framework_a": A.manifest.get("framework"),
@@ -128,16 +218,19 @@ def diff_bundles(zip_a: str, zip_b: str) -> Dict[str, Any]:
         "events_a": len(ea),
         "events_b": len(eb),
         "first_divergence_index": first_div,
+        "identical_until_index": first_div,  # same semantic, explicit name
+        "diff_event_count": diff_count,
+        "missing_event_count": missing_count,
+        "final_reward_a": ra,
+        "final_reward_b": rb,
+        "final_reward_delta": reward_delta,
+        "run_link_a": _event_link(A.manifest, 0),
+        "run_link_b": _event_link(B.manifest, 0),
     }
 
-    # simple counts by class
-    counts: Dict[str, int] = {}
-    for item in per_event:
-        counts[item["class"]] = counts.get(item["class"], 0) + 1
-
-    out = {
+    return {
         "summary": summary,
         "class_counts": counts,
-        "differences": per_event[:400],  # cap for UI
-    }
-    return out
+        "alignment": alignment,
+        "differences": per_event[:400],
+    }