Refine replay UX and similar failure search

Author: acailic

api/entity_routes.py

@@ -40,6 +40,7 @@ class EntityListResponse(BaseModel):
 class EntitySummaryResponse(BaseModel):
     """Response schema for entity summary statistics."""
 
+    agent_name_count: int
     tool_name_count: int
     error_type_count: int
     model_count: int
@@ -156,6 +157,7 @@ async def get_entity_summary(
     summary = await repo.get_entity_summary()
 
     return EntitySummaryResponse(
+        agent_name_count=summary.get(EntityType.AGENT_NAME, 0),
         tool_name_count=summary.get(EntityType.TOOL_NAME, 0),
         error_type_count=summary.get(EntityType.ERROR_TYPE, 0),
         model_count=summary.get(EntityType.MODEL, 0),

api/replay_routes.py

@@ -26,6 +26,10 @@
 router = APIRouter(tags=["replay"])
 
 
+def _split_csv_param(value: str | None) -> set[str]:
+    return {item.strip() for item in (value or "").split(",") if item.strip()}
+
+
 @router.get("/api/sessions/{session_id}/replay", response_model=ReplayResponse)
 async def replay_session(
     session_id: str,
@@ -46,6 +50,8 @@ async def replay_session(
     # Extract the default value when the raw Query object is passed through.
     if hasattr(collapse_threshold, "default"):
         collapse_threshold = float(collapse_threshold.default)
+    if hasattr(stop_at_breakpoint, "default"):
+        stop_at_breakpoint = bool(stop_at_breakpoint.default)
 
     # Record analytics event (fire-and-forget)
     record_event("replay_started", session_id=session_id, properties={"mode": mode})
@@ -72,10 +78,10 @@ async def replay_session(
         checkpoints,
         mode=mode,
         focus_event_id=focus_event_id,
-        breakpoint_event_types={item for item in (breakpoint_event_types or "").split(",") if item},
-        breakpoint_tool_names={item for item in (breakpoint_tool_names or "").split(",") if item},
+        breakpoint_event_types=_split_csv_param(breakpoint_event_types),
+        breakpoint_tool_names=_split_csv_param(breakpoint_tool_names),
         breakpoint_confidence_below=breakpoint_confidence_below,
-        breakpoint_safety_outcomes={item for item in (breakpoint_safety_outcomes or "").split(",") if item},
+        breakpoint_safety_outcomes=_split_csv_param(breakpoint_safety_outcomes),
     )
 
     # Handle segment collapsing for highlights mode
@@ -98,8 +104,10 @@ async def replay_session(
         stopped_at_breakpoint = True
         # Build O(1) event_id -> index map for efficient breakpoint lookup
         event_id_to_index = {event.get("id"): i for i, event in enumerate(replay_data["events"])}
-        first_breakpoint_id = replay_data["breakpoints"][0].get("id")
-        stopped_at_index = event_id_to_index.get(first_breakpoint_id)
+        for breakpoint_event in replay_data["breakpoints"]:
+            stopped_at_index = event_id_to_index.get(breakpoint_event.get("id"))
+            if stopped_at_index is not None:
+                break
 
     return ReplayResponse(
         session_id=session_id,

api/schemas.py

@@ -70,6 +70,11 @@ class TraceEventSchema(BaseModel):
     outcome: SafetyOutcome | None = None
     risk_level: RiskLevel | None = None
     rationale: str | None = None
+    attempted_fix: str | None = None
+    validation_result: str | None = None
+    repair_outcome: str | None = None
+    repair_sequence_id: str | None = None
+    repair_diff: str | None = None
     blocked_action: str | None = None
     reason: str | None = None
     safe_alternative: str | None = None

api/services.py

@@ -7,6 +7,7 @@
 import logging
 from typing import Any
 
+from sqlalchemy import String, cast, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
 
 from agent_debugger_sdk.core.events import Checkpoint, EventType, Session, SessionStatus, TraceEvent
@@ -17,9 +18,12 @@
 from collector.intelligence.facade import TraceIntelligence
 from redaction.pipeline import RedactionPipeline
 from storage import TraceRepository
+from storage.converters import orm_to_event, orm_to_session
+from storage.models import EventModel, SessionModel
 
 logger = logging.getLogger(__name__)
 SESSION_ANALYSIS_CAP = 100
+FAILURE_SIMILARITY_THRESHOLD = 0.5
 
 
 def normalize_session(
@@ -400,69 +404,97 @@ async def find_similar_failures(
     # Get the failure event
     failure_event = await repo.get_event(failure_event_id)
     if not failure_event:
-        return []
+        raise NotFoundError(f"Failure event {failure_event_id} not found")
+    if failure_event.session_id != session_id:
+        raise NotFoundError(
+            f"Failure event {failure_event_id} was not found in session {session_id}"
+        )
 
     # Determine failure characteristics
-    error_text = failure_event.error or failure_event.error_message or failure_event.name or ""
-    error_type = failure_event.error_type or ""
-
-    # Get all sessions with failures
-    all_sessions = await repo.list_sessions(limit=500, offset=0, sort_by="started_at")
-
-    similar_failures: list[dict[str, Any]] = []
-
-    for session in all_sessions:
-        # Skip the current session
-        if session.id == session_id:
-            continue
-
-        # Skip sessions without errors
-        if session.errors == 0:
-            continue
-
-        # Get events from this session to find matching failures
-        try:
-            session_events = await repo.list_events(session.id, limit=1000)
-        except Exception:
+    error_text = _event_error_text(failure_event)
+    error_type = _event_error_type(failure_event)
+    candidate_failures = await _load_candidate_failure_events(repo, failure_event, session_id)
+
+    best_match_by_session: dict[str, dict[str, Any]] = {}
+
+    for event, session in candidate_failures:
+        similarity = _calculate_failure_similarity(
+            failure_event,
+            event,
+            error_text,
+            error_type,
+        )
+        if similarity < FAILURE_SIMILARITY_THRESHOLD:
             continue
 
-        # Find failure events in this session
-        for event in session_events:
-            if not _is_failure_event(event):
-                continue
-
-            # Calculate similarity score
-            similarity = _calculate_failure_similarity(
-                failure_event,
-                event,
-                error_text,
-                error_type,
-            )
-
-            # Only include reasonably similar failures
-            if similarity >= 0.3:
-                # Derive failure mode and root cause
-                failure_mode = _derive_failure_mode(event)
-                root_cause = _derive_root_cause(event)
-
-                similar_failures.append({
-                    "session_id": session.id,
-                    "agent_name": session.agent_name,
-                    "framework": session.framework,
-                    "started_at": session.started_at,
-                    "failure_type": str(event.event_type),
-                    "failure_mode": failure_mode,
-                    "root_cause": root_cause,
-                    "similarity": similarity,
-                    "fix_note": session.fix_note,
-                })
-                break  # Only add one failure per session
+        failure_summary = {
+            "session_id": session.id,
+            "agent_name": session.agent_name,
+            "framework": session.framework,
+            "started_at": session.started_at,
+            "failure_type": str(event.event_type),
+            "failure_mode": _derive_failure_mode(event),
+            "root_cause": _derive_root_cause(event),
+            "similarity": similarity,
+            "fix_note": session.fix_note,
+        }
+        existing = best_match_by_session.get(session.id)
+        if existing is None or failure_summary["similarity"] > existing["similarity"]:
+            best_match_by_session[session.id] = failure_summary
 
     # Sort by similarity and limit
+    similar_failures = list(best_match_by_session.values())
     similar_failures.sort(key=lambda x: x["similarity"], reverse=True)
     return similar_failures[:limit]
 
 
+async def _load_candidate_failure_events(
+    repo: TraceRepository,
+    failure_event: TraceEvent,
+    session_id: str,
+) -> list[tuple[TraceEvent, Session]]:
+    """Load tenant-scoped failure candidates without per-session N+1 queries."""
+    failure_event_types = [
+        str(EventType.ERROR),
+        str(EventType.REFUSAL),
+        str(EventType.POLICY_VIOLATION),
+        str(EventType.BEHAVIOR_ALERT),
+        str(EventType.TOOL_RESULT),
+        str(EventType.SAFETY_CHECK),
+    ]
+
+    source_clues = [EventModel.event_type == str(failure_event.event_type)]
+    source_error_type = _event_error_type(failure_event)
+    if source_error_type:
+        source_clues.append(cast(EventModel.data, String).ilike(f"%{source_error_type}%"))
+    source_tool_name = getattr(failure_event, "tool_name", None)
+    if source_tool_name:
+        source_clues.append(cast(EventModel.data, String).ilike(f"%{source_tool_name}%"))
+
+    stmt = (
+        select(EventModel, SessionModel)
+        .join(SessionModel, EventModel.session_id == SessionModel.id)
+        .where(
+            SessionModel.tenant_id == repo.tenant_id,
+            EventModel.tenant_id == repo.tenant_id,
+            SessionModel.id != session_id,
+            SessionModel.errors > 0,
+            EventModel.event_type.in_(failure_event_types),
+            or_(*source_clues),
+        )
+        .order_by(SessionModel.started_at.desc(), EventModel.timestamp.desc())
+    )
+    result = await repo.session.execute(stmt)
+
+    candidates: list[tuple[TraceEvent, Session]] = []
+    for db_event, db_session in result.all():
+        event = orm_to_event(db_event)
+        if not _is_failure_event(event):
+            continue
+        candidates.append((event, orm_to_session(db_session)))
+    return candidates
+
+
 def _is_failure_event(event: TraceEvent) -> bool:
     """Check if an event represents a failure."""
     return (
@@ -475,6 +507,28 @@ def _is_failure_event(event: TraceEvent) -> bool:
     )
 
 
+def _event_error_text(event: TraceEvent) -> str:
+    """Return the most useful error-like text available on an event."""
+    return (
+        getattr(event, "error", None)
+        or getattr(event, "error_message", None)
+        or getattr(event, "reason", None)
+        or event.name
+        or ""
+    )
+
+
+def _event_error_type(event: TraceEvent) -> str:
+    """Return the most useful error-like type available on an event."""
+    return (
+        getattr(event, "error_type", None)
+        or getattr(event, "violation_type", None)
+        or getattr(event, "alert_type", None)
+        or ""
+    )
+
+
+
 def _calculate_failure_similarity(
     source_event: TraceEvent,
     candidate_event: TraceEvent,
@@ -492,13 +546,13 @@ def _calculate_failure_similarity(
         score += 0.4
 
     # Error type match
-    candidate_error_type = candidate_event.error_type or ""
+    candidate_error_type = _event_error_type(candidate_event)
     if source_error_type and candidate_error_type:
         if source_error_type.lower() == candidate_error_type.lower():
             score += 0.3
 
     # Error text similarity (simple keyword overlap)
-    candidate_error_text = candidate_event.error or candidate_event.error_message or candidate_event.name or ""
+    candidate_error_text = _event_error_text(candidate_event)
     if source_error_text and candidate_error_text:
         source_words = set(source_error_text.lower().split())
         candidate_words = set(candidate_error_text.lower().split())

collector/replay.py

@@ -243,9 +243,10 @@ def build_replay(
     breakpoint_tool_names = breakpoint_tool_names or set()
     breakpoint_safety_outcomes = breakpoint_safety_outcomes or set()
 
+    breakpoint_source_events = replay_events if mode == "focus" else replay_window_events
     breakpoints = [
         event.to_dict()
-        for event in replay_window_events
+        for event in breakpoint_source_events
         if matches_breakpoint(
             event,
             event_types=breakpoint_event_types,