rustant_core/
replay.rs

1//! Session replay engine — step-by-step playback of execution traces.
2//!
3//! Enables reviewing past agent executions by stepping through recorded
4//! trace events, providing context at each step about what happened
5//! and why.
6
7use crate::audit::{AuditStore, ExecutionTrace, TraceEvent, TraceEventKind};
8use crate::types::{CostEstimate, TokenUsage};
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11use std::collections::HashSet;
12use uuid::Uuid;
13
14// ---------------------------------------------------------------------------
15// ReplayError
16// ---------------------------------------------------------------------------
17
18/// Errors that can occur during replay operations.
19#[derive(Debug, Clone, thiserror::Error)]
20pub enum ReplayError {
21    #[error("trace not found: {0}")]
22    TraceNotFound(Uuid),
23    #[error("position {position} out of bounds (total: {total})")]
24    OutOfBounds { position: usize, total: usize },
25    #[error("bookmark index {0} out of bounds")]
26    BookmarkNotFound(usize),
27    #[error("empty trace: no events to replay")]
28    EmptyTrace,
29}
30
31// ---------------------------------------------------------------------------
32// Bookmark
33// ---------------------------------------------------------------------------
34
35/// A saved position in the replay.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct Bookmark {
38    pub position: usize,
39    pub label: String,
40    pub created_at: DateTime<Utc>,
41}
42
43// ---------------------------------------------------------------------------
44// ReplaySnapshot
45// ---------------------------------------------------------------------------
46
47/// A point-in-time snapshot of the replay state.
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ReplaySnapshot {
50    pub trace_id: Uuid,
51    pub position: usize,
52    pub total_events: usize,
53    /// Progress percentage from 0.0 to 100.0.
54    pub progress_pct: f64,
55    pub current_event: Option<TraceEvent>,
56    /// Milliseconds elapsed from the trace start to the current event.
57    pub elapsed_from_start: Option<u64>,
58    pub cumulative_usage: TokenUsage,
59    pub cumulative_cost: CostEstimate,
60    pub tools_executed_so_far: Vec<String>,
61    pub errors_so_far: usize,
62}
63
64// ---------------------------------------------------------------------------
65// TimelineEntry
66// ---------------------------------------------------------------------------
67
68/// One entry in the timeline view.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct TimelineEntry {
71    pub sequence: usize,
72    pub timestamp: DateTime<Utc>,
73    /// Milliseconds elapsed from the trace start.
74    pub elapsed_ms: u64,
75    pub description: String,
76    /// Whether this entry is the current replay position.
77    pub is_current: bool,
78    /// Whether this entry has a bookmark.
79    pub is_bookmarked: bool,
80}
81
82// ---------------------------------------------------------------------------
83// ReplayEngine
84// ---------------------------------------------------------------------------
85
86/// The main replay controller — provides step-by-step playback of an
87/// execution trace.
88pub struct ReplayEngine {
89    trace: ExecutionTrace,
90    /// Current event index (0-based).
91    position: usize,
92    bookmarks: Vec<Bookmark>,
93}
94
95impl ReplayEngine {
96    /// Create a new replay from an execution trace.
97    pub fn new(trace: ExecutionTrace) -> Self {
98        Self {
99            trace,
100            position: 0,
101            bookmarks: Vec::new(),
102        }
103    }
104
105    /// Create a replay engine from a trace in the audit store.
106    pub fn from_store(store: &AuditStore, trace_id: Uuid) -> Result<Self, ReplayError> {
107        let trace = store
108            .get_trace(trace_id)
109            .ok_or(ReplayError::TraceNotFound(trace_id))?;
110        Ok(Self::new(trace.clone()))
111    }
112
113    /// Get the current position (0-based event index).
114    pub fn position(&self) -> usize {
115        self.position
116    }
117
118    /// Get total number of events.
119    pub fn total_events(&self) -> usize {
120        self.trace.events.len()
121    }
122
123    /// Is the replay at the beginning?
124    pub fn is_at_start(&self) -> bool {
125        self.position == 0
126    }
127
128    /// Is the replay at the end?
129    pub fn is_at_end(&self) -> bool {
130        self.trace.events.is_empty() || self.position >= self.trace.events.len() - 1
131    }
132
133    /// Step forward one event. Returns the event at the new position, or
134    /// `None` if already at the end.
135    pub fn step_forward(&mut self) -> Option<&TraceEvent> {
136        if self.position + 1 < self.trace.events.len() {
137            self.position += 1;
138            self.trace.events.get(self.position)
139        } else {
140            None
141        }
142    }
143
144    /// Step backward one event. Returns the event at the new position, or
145    /// `None` if already at the start.
146    pub fn step_backward(&mut self) -> Option<&TraceEvent> {
147        if self.position > 0 {
148            self.position -= 1;
149            self.trace.events.get(self.position)
150        } else {
151            None
152        }
153    }
154
155    /// Jump to a specific position.
156    pub fn seek(&mut self, position: usize) -> Result<&TraceEvent, ReplayError> {
157        if position >= self.trace.events.len() {
158            return Err(ReplayError::OutOfBounds {
159                position,
160                total: self.trace.events.len(),
161            });
162        }
163        self.position = position;
164        Ok(&self.trace.events[self.position])
165    }
166
167    /// Go to the start.
168    pub fn rewind(&mut self) {
169        self.position = 0;
170    }
171
172    /// Go to the end.
173    pub fn fast_forward(&mut self) {
174        if !self.trace.events.is_empty() {
175            self.position = self.trace.events.len() - 1;
176        }
177    }
178
179    /// Get the current event.
180    pub fn current_event(&self) -> Option<&TraceEvent> {
181        self.trace.events.get(self.position)
182    }
183
184    /// Get a snapshot of the current replay state.
185    pub fn snapshot(&self) -> ReplaySnapshot {
186        let total_events = self.trace.events.len();
187        let current_event = self.trace.events.get(self.position).cloned();
188
189        let elapsed_from_start = current_event.as_ref().map(|e| {
190            (e.timestamp - self.trace.started_at)
191                .num_milliseconds()
192                .max(0) as u64
193        });
194
195        let progress_pct = if total_events == 0 {
196            0.0
197        } else if total_events == 1 {
198            100.0
199        } else {
200            (self.position as f64 / (total_events - 1) as f64) * 100.0
201        };
202
203        let end = if total_events == 0 {
204            0
205        } else {
206            self.position + 1
207        };
208
209        let tools_executed_so_far: Vec<String> = self
210            .trace
211            .events
212            .iter()
213            .take(end)
214            .filter_map(|e| {
215                if let TraceEventKind::ToolExecuted { ref tool, .. } = e.kind {
216                    Some(tool.clone())
217                } else {
218                    None
219                }
220            })
221            .collect();
222
223        let errors_so_far = self
224            .trace
225            .events
226            .iter()
227            .take(end)
228            .filter(|e| matches!(&e.kind, TraceEventKind::Error { .. }))
229            .count();
230
231        ReplaySnapshot {
232            trace_id: self.trace.trace_id,
233            position: self.position,
234            total_events,
235            progress_pct,
236            current_event,
237            elapsed_from_start,
238            cumulative_usage: self.cumulative_usage(),
239            cumulative_cost: self.cumulative_cost(),
240            tools_executed_so_far,
241            errors_so_far,
242        }
243    }
244
245    /// Get the current step as a formatted description.
246    pub fn describe_current(&self) -> String {
247        match self.current_event() {
248            Some(event) => format!(
249                "[{}/{}] {}",
250                self.position + 1,
251                self.total_events(),
252                describe_event(&event.kind)
253            ),
254            None => "No events".to_string(),
255        }
256    }
257
258    /// Get a full timeline description of all events.
259    pub fn timeline(&self) -> Vec<TimelineEntry> {
260        let bookmark_positions: HashSet<usize> =
261            self.bookmarks.iter().map(|b| b.position).collect();
262
263        self.trace
264            .events
265            .iter()
266            .map(|event| {
267                let elapsed_ms = (event.timestamp - self.trace.started_at)
268                    .num_milliseconds()
269                    .max(0) as u64;
270
271                TimelineEntry {
272                    sequence: event.sequence,
273                    timestamp: event.timestamp,
274                    elapsed_ms,
275                    description: describe_event(&event.kind),
276                    is_current: event.sequence == self.position,
277                    is_bookmarked: bookmark_positions.contains(&event.sequence),
278                }
279            })
280            .collect()
281    }
282
283    /// Add a bookmark at the current position.
284    pub fn add_bookmark(&mut self, label: impl Into<String>) {
285        self.bookmarks.push(Bookmark {
286            position: self.position,
287            label: label.into(),
288            created_at: Utc::now(),
289        });
290    }
291
292    /// Get all bookmarks.
293    pub fn bookmarks(&self) -> &[Bookmark] {
294        &self.bookmarks
295    }
296
297    /// Jump to a bookmark by index.
298    pub fn goto_bookmark(&mut self, index: usize) -> Result<&TraceEvent, ReplayError> {
299        let position = self
300            .bookmarks
301            .get(index)
302            .ok_or(ReplayError::BookmarkNotFound(index))?
303            .position;
304        self.seek(position)
305    }
306
307    /// Get the trace being replayed.
308    pub fn trace(&self) -> &ExecutionTrace {
309        &self.trace
310    }
311
312    /// Skip forward to the next event of a tool-related kind
313    /// (`ToolRequested`, `ToolApproved`, `ToolDenied`, or `ToolExecuted`).
314    pub fn skip_to_next_tool_event(&mut self) -> Option<&TraceEvent> {
315        let start = self.position + 1;
316        for i in start..self.trace.events.len() {
317            match &self.trace.events[i].kind {
318                TraceEventKind::ToolRequested { .. }
319                | TraceEventKind::ToolApproved { .. }
320                | TraceEventKind::ToolDenied { .. }
321                | TraceEventKind::ToolExecuted { .. } => {
322                    self.position = i;
323                    return self.trace.events.get(i);
324                }
325                _ => continue,
326            }
327        }
328        None
329    }
330
331    /// Get cumulative token usage up to and including the current position.
332    pub fn cumulative_usage(&self) -> TokenUsage {
333        let mut usage = TokenUsage::default();
334        let end = self.position + 1;
335        for event in self.trace.events.iter().take(end) {
336            if let TraceEventKind::LlmCall {
337                input_tokens,
338                output_tokens,
339                ..
340            } = &event.kind
341            {
342                usage.input_tokens += input_tokens;
343                usage.output_tokens += output_tokens;
344            }
345        }
346        usage
347    }
348
349    /// Get cumulative cost up to and including the current position.
350    ///
351    /// The per-call cost is split proportionally between input and output
352    /// based on token counts.
353    pub fn cumulative_cost(&self) -> CostEstimate {
354        let mut estimate = CostEstimate::default();
355        let end = self.position + 1;
356        for event in self.trace.events.iter().take(end) {
357            if let TraceEventKind::LlmCall {
358                cost,
359                input_tokens,
360                output_tokens,
361                ..
362            } = &event.kind
363            {
364                let total_tokens = input_tokens + output_tokens;
365                if total_tokens > 0 {
366                    estimate.input_cost += cost * (*input_tokens as f64 / total_tokens as f64);
367                    estimate.output_cost += cost * (*output_tokens as f64 / total_tokens as f64);
368                }
369            }
370        }
371        estimate
372    }
373}
374
375// ---------------------------------------------------------------------------
376// describe_event helper
377// ---------------------------------------------------------------------------
378
379/// Format a [`TraceEventKind`] as a human-readable description.
380pub fn describe_event(kind: &TraceEventKind) -> String {
381    match kind {
382        TraceEventKind::TaskStarted { goal, .. } => {
383            format!("Task started: {}", goal)
384        }
385        TraceEventKind::TaskCompleted {
386            success,
387            iterations,
388            ..
389        } => {
390            format!(
391                "Task {} after {} iterations",
392                if *success {
393                    "completed successfully"
394                } else {
395                    "failed"
396                },
397                iterations
398            )
399        }
400        TraceEventKind::ToolRequested {
401            tool, risk_level, ..
402        } => {
403            format!("Tool requested: {} (risk: {:?})", tool, risk_level)
404        }
405        TraceEventKind::ToolApproved { tool } => {
406            format!("Tool approved: {}", tool)
407        }
408        TraceEventKind::ToolDenied { tool, reason } => {
409            format!("Tool denied: {} - {}", tool, reason)
410        }
411        TraceEventKind::ApprovalRequested { tool, .. } => {
412            format!("Approval requested for: {}", tool)
413        }
414        TraceEventKind::ApprovalDecision { tool, approved } => {
415            format!(
416                "Approval {}: {}",
417                if *approved { "granted" } else { "rejected" },
418                tool
419            )
420        }
421        TraceEventKind::ToolExecuted {
422            tool,
423            success,
424            duration_ms,
425            ..
426        } => {
427            format!(
428                "Tool executed: {} ({}, {}ms)",
429                tool,
430                if *success { "ok" } else { "failed" },
431                duration_ms
432            )
433        }
434        TraceEventKind::LlmCall {
435            model,
436            input_tokens,
437            output_tokens,
438            ..
439        } => {
440            format!(
441                "LLM call: {} ({}/{} tokens)",
442                model, input_tokens, output_tokens
443            )
444        }
445        TraceEventKind::StatusChange { from, to } => {
446            format!("Status: {} -> {}", from, to)
447        }
448        TraceEventKind::Error { message } => {
449            format!("Error: {}", message)
450        }
451    }
452}
453
454// ---------------------------------------------------------------------------
455// ReplaySession
456// ---------------------------------------------------------------------------
457
458/// Manages multiple replay engines, with one active at a time.
459pub struct ReplaySession {
460    engines: Vec<ReplayEngine>,
461    active_index: Option<usize>,
462}
463
464impl ReplaySession {
465    /// Create a new empty replay session.
466    pub fn new() -> Self {
467        Self {
468            engines: Vec::new(),
469            active_index: None,
470        }
471    }
472
473    /// Add a replay for the given execution trace. Returns its index.
474    /// The first added replay is automatically set as active.
475    pub fn add_replay(&mut self, trace: ExecutionTrace) -> usize {
476        let index = self.engines.len();
477        self.engines.push(ReplayEngine::new(trace));
478        if self.active_index.is_none() {
479            self.active_index = Some(index);
480        }
481        index
482    }
483
484    /// Set the active replay by index.
485    pub fn set_active(&mut self, index: usize) -> Result<(), ReplayError> {
486        if index >= self.engines.len() {
487            return Err(ReplayError::OutOfBounds {
488                position: index,
489                total: self.engines.len(),
490            });
491        }
492        self.active_index = Some(index);
493        Ok(())
494    }
495
496    /// Get a reference to the active replay engine.
497    pub fn active(&self) -> Option<&ReplayEngine> {
498        self.active_index.and_then(|i| self.engines.get(i))
499    }
500
501    /// Get a mutable reference to the active replay engine.
502    pub fn active_mut(&mut self) -> Option<&mut ReplayEngine> {
503        self.active_index.and_then(|i| self.engines.get_mut(i))
504    }
505
506    /// List all replays with summary information.
507    pub fn list_replays(&self) -> Vec<ReplaySummary> {
508        self.engines
509            .iter()
510            .enumerate()
511            .map(|(i, engine)| ReplaySummary {
512                index: i,
513                trace_id: engine.trace().trace_id,
514                goal: engine.trace().goal.clone(),
515                event_count: engine.total_events(),
516                is_active: self.active_index == Some(i),
517            })
518            .collect()
519    }
520
521    /// Get the number of replays in the session.
522    pub fn len(&self) -> usize {
523        self.engines.len()
524    }
525
526    /// Check whether the session has no replays.
527    pub fn is_empty(&self) -> bool {
528        self.engines.is_empty()
529    }
530}
531
532impl Default for ReplaySession {
533    fn default() -> Self {
534        Self::new()
535    }
536}
537
538/// Summary information about a replay in a session.
539#[derive(Debug, Clone, Serialize, Deserialize)]
540pub struct ReplaySummary {
541    pub index: usize,
542    pub trace_id: Uuid,
543    pub goal: String,
544    pub event_count: usize,
545    pub is_active: bool,
546}
547
548// ---------------------------------------------------------------------------
549// Tests
550// ---------------------------------------------------------------------------
551
552#[cfg(test)]
553mod tests {
554    use super::*;
555    use crate::audit::{AuditStore, ExecutionTrace, TraceEventKind};
556    use crate::types::RiskLevel;
557    use uuid::Uuid;
558
559    /// Build a sample execution trace with multiple event types.
560    ///
561    /// Events (0-based index):
562    ///   0: TaskStarted          (auto-pushed by ExecutionTrace::new)
563    ///   1: ToolRequested         (file_read, ReadOnly)
564    ///   2: ToolApproved          (file_read)
565    ///   3: ToolExecuted          (file_read, success, 42ms)
566    ///   4: LlmCall               (gpt-4, 500/200 tokens, $0.021)
567    ///   5: ToolRequested         (file_write, Write)
568    ///   6: ToolDenied            (file_write, "path denied")
569    ///   7: Error                 ("write denied")
570    ///   8: TaskCompleted         (auto-pushed by complete())
571    fn sample_trace() -> ExecutionTrace {
572        let session_id = Uuid::new_v4();
573        let task_id = Uuid::new_v4();
574        let mut trace = ExecutionTrace::new(session_id, task_id, "test task");
575
576        // Events 1..7
577        trace.push_event(TraceEventKind::ToolRequested {
578            tool: "file_read".into(),
579            risk_level: RiskLevel::ReadOnly,
580            args_summary: "path=/src/main.rs".into(),
581        });
582        trace.push_event(TraceEventKind::ToolApproved {
583            tool: "file_read".into(),
584        });
585        trace.push_event(TraceEventKind::ToolExecuted {
586            tool: "file_read".into(),
587            success: true,
588            duration_ms: 42,
589            output_preview: "fn main() {...}".into(),
590        });
591        trace.push_event(TraceEventKind::LlmCall {
592            model: "gpt-4".into(),
593            input_tokens: 500,
594            output_tokens: 200,
595            cost: 0.021,
596        });
597        trace.push_event(TraceEventKind::ToolRequested {
598            tool: "file_write".into(),
599            risk_level: RiskLevel::Write,
600            args_summary: "path=/src/lib.rs".into(),
601        });
602        trace.push_event(TraceEventKind::ToolDenied {
603            tool: "file_write".into(),
604            reason: "path denied".into(),
605        });
606        trace.push_event(TraceEventKind::Error {
607            message: "write denied".into(),
608        });
609
610        // Event 8 (auto-pushed by complete)
611        trace.iterations = 3;
612        trace.complete(true);
613        trace
614    }
615
616    // -----------------------------------------------------------------------
617    // 1. test_replay_engine_new
618    // -----------------------------------------------------------------------
619    #[test]
620    fn test_replay_engine_new() {
621        let trace = sample_trace();
622        let engine = ReplayEngine::new(trace);
623        assert_eq!(engine.position(), 0);
624        assert_eq!(engine.total_events(), 9);
625        assert!(engine.bookmarks().is_empty());
626    }
627
628    // -----------------------------------------------------------------------
629    // 2. test_replay_engine_step_forward
630    // -----------------------------------------------------------------------
631    #[test]
632    fn test_replay_engine_step_forward() {
633        let mut engine = ReplayEngine::new(sample_trace());
634        assert_eq!(engine.position(), 0);
635
636        let event = engine.step_forward().unwrap();
637        assert_eq!(event.sequence, 1);
638        assert_eq!(engine.position(), 1);
639
640        let event = engine.step_forward().unwrap();
641        assert_eq!(event.sequence, 2);
642        assert_eq!(engine.position(), 2);
643    }
644
645    // -----------------------------------------------------------------------
646    // 3. test_replay_engine_step_backward
647    // -----------------------------------------------------------------------
648    #[test]
649    fn test_replay_engine_step_backward() {
650        let mut engine = ReplayEngine::new(sample_trace());
651        engine.seek(3).unwrap();
652        assert_eq!(engine.position(), 3);
653
654        let event = engine.step_backward().unwrap();
655        assert_eq!(event.sequence, 2);
656        assert_eq!(engine.position(), 2);
657
658        let event = engine.step_backward().unwrap();
659        assert_eq!(event.sequence, 1);
660        assert_eq!(engine.position(), 1);
661    }
662
663    // -----------------------------------------------------------------------
664    // 4. test_replay_engine_at_boundaries
665    // -----------------------------------------------------------------------
666    #[test]
667    fn test_replay_engine_at_boundaries() {
668        let trace = sample_trace();
669        let total = trace.events.len();
670        let mut engine = ReplayEngine::new(trace);
671
672        // At start
673        assert!(engine.is_at_start());
674        assert!(!engine.is_at_end());
675
676        // step_backward at start returns None and position stays 0
677        assert!(engine.step_backward().is_none());
678        assert_eq!(engine.position(), 0);
679
680        // Fast forward to end
681        engine.fast_forward();
682        assert_eq!(engine.position(), total - 1);
683        assert!(!engine.is_at_start());
684        assert!(engine.is_at_end());
685
686        // step_forward at end returns None and position stays at end
687        assert!(engine.step_forward().is_none());
688        assert_eq!(engine.position(), total - 1);
689    }
690
691    // -----------------------------------------------------------------------
692    // 5. test_replay_engine_seek
693    // -----------------------------------------------------------------------
694    #[test]
695    fn test_replay_engine_seek() {
696        let mut engine = ReplayEngine::new(sample_trace());
697        let event = engine.seek(4).unwrap();
698        assert_eq!(event.sequence, 4);
699        assert_eq!(engine.position(), 4);
700
701        let event = engine.seek(0).unwrap();
702        assert_eq!(event.sequence, 0);
703        assert_eq!(engine.position(), 0);
704
705        let event = engine.seek(8).unwrap();
706        assert_eq!(event.sequence, 8);
707        assert_eq!(engine.position(), 8);
708    }
709
710    // -----------------------------------------------------------------------
711    // 6. test_replay_engine_seek_out_of_bounds
712    // -----------------------------------------------------------------------
713    #[test]
714    fn test_replay_engine_seek_out_of_bounds() {
715        let mut engine = ReplayEngine::new(sample_trace());
716        let result = engine.seek(100);
717        assert!(matches!(
718            result,
719            Err(ReplayError::OutOfBounds {
720                position: 100,
721                total: 9
722            })
723        ));
724
725        // Seeking to exactly total_events is also out of bounds
726        let result = engine.seek(9);
727        assert!(matches!(result, Err(ReplayError::OutOfBounds { .. })));
728    }
729
730    // -----------------------------------------------------------------------
731    // 7. test_replay_engine_rewind
732    // -----------------------------------------------------------------------
733    #[test]
734    fn test_replay_engine_rewind() {
735        let mut engine = ReplayEngine::new(sample_trace());
736        engine.seek(5).unwrap();
737        assert_eq!(engine.position(), 5);
738
739        engine.rewind();
740        assert_eq!(engine.position(), 0);
741        assert!(engine.is_at_start());
742    }
743
744    // -----------------------------------------------------------------------
745    // 8. test_replay_engine_fast_forward
746    // -----------------------------------------------------------------------
747    #[test]
748    fn test_replay_engine_fast_forward() {
749        let trace = sample_trace();
750        let total = trace.events.len();
751        let mut engine = ReplayEngine::new(trace);
752        assert_eq!(engine.position(), 0);
753
754        engine.fast_forward();
755        assert_eq!(engine.position(), total - 1);
756        assert!(engine.is_at_end());
757    }
758
759    // -----------------------------------------------------------------------
760    // 9. test_replay_engine_current_event
761    // -----------------------------------------------------------------------
762    #[test]
763    fn test_replay_engine_current_event() {
764        let engine = ReplayEngine::new(sample_trace());
765        let event = engine.current_event().unwrap();
766        assert_eq!(event.sequence, 0);
767        assert!(matches!(
768            &event.kind,
769            TraceEventKind::TaskStarted { goal, .. } if goal == "test task"
770        ));
771    }
772
773    // -----------------------------------------------------------------------
774    // 10. test_replay_engine_snapshot
775    // -----------------------------------------------------------------------
776    #[test]
777    fn test_replay_engine_snapshot() {
778        let mut engine = ReplayEngine::new(sample_trace());
779
780        // At position 0 — no tools executed, no errors
781        let snap = engine.snapshot();
782        assert_eq!(snap.trace_id, engine.trace().trace_id);
783        assert_eq!(snap.position, 0);
784        assert_eq!(snap.total_events, 9);
785        assert!(snap.current_event.is_some());
786        assert!(snap.elapsed_from_start.is_some());
787        assert!(snap.tools_executed_so_far.is_empty());
788        assert_eq!(snap.errors_so_far, 0);
789
790        // Advance past the ToolExecuted event (index 3)
791        engine.seek(3).unwrap();
792        let snap = engine.snapshot();
793        assert_eq!(snap.position, 3);
794        assert_eq!(snap.tools_executed_so_far.len(), 1);
795        assert_eq!(snap.tools_executed_so_far[0], "file_read");
796
797        // Advance past the Error event (index 7)
798        engine.seek(7).unwrap();
799        let snap = engine.snapshot();
800        assert_eq!(snap.errors_so_far, 1);
801    }
802
803    // -----------------------------------------------------------------------
804    // 11. test_replay_engine_describe_current
805    // -----------------------------------------------------------------------
806    #[test]
807    fn test_replay_engine_describe_current() {
808        let mut engine = ReplayEngine::new(sample_trace());
809        let desc = engine.describe_current();
810        assert!(desc.contains("[1/9]"));
811        assert!(desc.contains("Task started"));
812        assert!(desc.contains("test task"));
813
814        engine.seek(4).unwrap();
815        let desc = engine.describe_current();
816        assert!(desc.contains("[5/9]"));
817        assert!(desc.contains("LLM call"));
818        assert!(desc.contains("gpt-4"));
819    }
820
821    // -----------------------------------------------------------------------
822    // 12. test_replay_engine_timeline
823    // -----------------------------------------------------------------------
824    #[test]
825    fn test_replay_engine_timeline() {
826        let trace = sample_trace();
827        let total = trace.events.len();
828        let mut engine = ReplayEngine::new(trace);
829
830        let timeline = engine.timeline();
831        assert_eq!(timeline.len(), total);
832
833        // First entry should be current
834        assert!(timeline[0].is_current);
835        assert!(!timeline[1].is_current);
836
837        // Step forward and verify current marker moves
838        engine.step_forward();
839        let timeline = engine.timeline();
840        assert!(!timeline[0].is_current);
841        assert!(timeline[1].is_current);
842
843        // All entries should have descriptions
844        for entry in &timeline {
845            assert!(!entry.description.is_empty());
846        }
847    }
848
849    // -----------------------------------------------------------------------
850    // 13. test_replay_engine_bookmarks
851    // -----------------------------------------------------------------------
852    #[test]
853    fn test_replay_engine_bookmarks() {
854        let mut engine = ReplayEngine::new(sample_trace());
855
856        engine.add_bookmark("start");
857        engine.step_forward();
858        engine.step_forward();
859        engine.add_bookmark("after two steps");
860
861        assert_eq!(engine.bookmarks().len(), 2);
862        assert_eq!(engine.bookmarks()[0].position, 0);
863        assert_eq!(engine.bookmarks()[0].label, "start");
864        assert_eq!(engine.bookmarks()[1].position, 2);
865        assert_eq!(engine.bookmarks()[1].label, "after two steps");
866
867        // Jump back to the first bookmark
868        let event = engine.goto_bookmark(0).unwrap();
869        assert_eq!(event.sequence, 0);
870        assert_eq!(engine.position(), 0);
871
872        // Timeline should reflect bookmarks
873        let timeline = engine.timeline();
874        assert!(timeline[0].is_bookmarked);
875        assert!(!timeline[1].is_bookmarked);
876        assert!(timeline[2].is_bookmarked);
877    }
878
879    // -----------------------------------------------------------------------
880    // 14. test_replay_engine_bookmark_out_of_bounds
881    // -----------------------------------------------------------------------
882    #[test]
883    fn test_replay_engine_bookmark_out_of_bounds() {
884        let mut engine = ReplayEngine::new(sample_trace());
885        let result = engine.goto_bookmark(0);
886        assert!(matches!(result, Err(ReplayError::BookmarkNotFound(0))));
887
888        engine.add_bookmark("only one");
889        let result = engine.goto_bookmark(5);
890        assert!(matches!(result, Err(ReplayError::BookmarkNotFound(5))));
891    }
892
893    // -----------------------------------------------------------------------
894    // 15. test_replay_engine_skip_to_tool
895    // -----------------------------------------------------------------------
896    #[test]
897    fn test_replay_engine_skip_to_tool() {
898        let mut engine = ReplayEngine::new(sample_trace());
899
900        // From position 0 (TaskStarted), skip to first tool event (index 1)
901        let event = engine.skip_to_next_tool_event().unwrap();
902        assert_eq!(event.sequence, 1);
903        assert!(matches!(
904            &event.kind,
905            TraceEventKind::ToolRequested { tool, .. } if tool == "file_read"
906        ));
907        assert_eq!(engine.position(), 1);
908
909        // Skip again — next tool event is ToolApproved at index 2
910        let event = engine.skip_to_next_tool_event().unwrap();
911        assert_eq!(event.sequence, 2);
912        assert!(matches!(
913            &event.kind,
914            TraceEventKind::ToolApproved { tool } if tool == "file_read"
915        ));
916
917        // Skip to ToolExecuted at index 3
918        let event = engine.skip_to_next_tool_event().unwrap();
919        assert_eq!(event.sequence, 3);
920
921        // Skip to second ToolRequested at index 5
922        let event = engine.skip_to_next_tool_event().unwrap();
923        assert_eq!(event.sequence, 5);
924
925        // Skip to ToolDenied at index 6
926        let event = engine.skip_to_next_tool_event().unwrap();
927        assert_eq!(event.sequence, 6);
928
929        // No more tool events after this
930        assert!(engine.skip_to_next_tool_event().is_none());
931    }
932
933    // -----------------------------------------------------------------------
934    // 16. test_replay_engine_cumulative_usage
935    // -----------------------------------------------------------------------
936    #[test]
937    fn test_replay_engine_cumulative_usage() {
938        let mut engine = ReplayEngine::new(sample_trace());
939
940        // At position 0 (TaskStarted), no LLM calls yet
941        let usage = engine.cumulative_usage();
942        assert_eq!(usage.input_tokens, 0);
943        assert_eq!(usage.output_tokens, 0);
944
945        // Seek to position 4 (LlmCall event)
946        engine.seek(4).unwrap();
947        let usage = engine.cumulative_usage();
948        assert_eq!(usage.input_tokens, 500);
949        assert_eq!(usage.output_tokens, 200);
950        assert_eq!(usage.total(), 700);
951
952        // At end, should still be 500/200 (only one LLM call in the trace)
953        engine.fast_forward();
954        let usage = engine.cumulative_usage();
955        assert_eq!(usage.input_tokens, 500);
956        assert_eq!(usage.output_tokens, 200);
957    }
958
959    // -----------------------------------------------------------------------
960    // 17. test_replay_engine_cumulative_cost
961    // -----------------------------------------------------------------------
962    #[test]
963    fn test_replay_engine_cumulative_cost() {
964        let mut engine = ReplayEngine::new(sample_trace());
965
966        // At start, no cost
967        let cost = engine.cumulative_cost();
968        assert!((cost.total() - 0.0).abs() < f64::EPSILON);
969
970        // After LlmCall event at position 4
971        engine.seek(4).unwrap();
972        let cost = engine.cumulative_cost();
973        assert!((cost.total() - 0.021).abs() < 0.001);
974        // input_cost and output_cost should be proportional to tokens
975        assert!(cost.input_cost > 0.0);
976        assert!(cost.output_cost > 0.0);
977
978        // At end, total cost remains the same
979        engine.fast_forward();
980        let cost = engine.cumulative_cost();
981        assert!((cost.total() - 0.021).abs() < 0.001);
982    }
983
984    // -----------------------------------------------------------------------
985    // 18. test_replay_engine_from_store
986    // -----------------------------------------------------------------------
987    #[test]
988    fn test_replay_engine_from_store() {
989        let trace = sample_trace();
990        let trace_id = trace.trace_id;
991        let mut store = AuditStore::new();
992        store.add_trace(trace);
993
994        let engine = ReplayEngine::from_store(&store, trace_id).unwrap();
995        assert_eq!(engine.trace().trace_id, trace_id);
996        assert_eq!(engine.position(), 0);
997        assert_eq!(engine.total_events(), 9);
998    }
999
1000    // -----------------------------------------------------------------------
1001    // 19. test_replay_engine_from_store_not_found
1002    // -----------------------------------------------------------------------
1003    #[test]
1004    fn test_replay_engine_from_store_not_found() {
1005        let store = AuditStore::new();
1006        let missing_id = Uuid::new_v4();
1007        let result = ReplayEngine::from_store(&store, missing_id);
1008        assert!(matches!(result, Err(ReplayError::TraceNotFound(id)) if id == missing_id));
1009    }
1010
1011    // -----------------------------------------------------------------------
1012    // 20. test_replay_session_new
1013    // -----------------------------------------------------------------------
1014    #[test]
1015    fn test_replay_session_new() {
1016        let session = ReplaySession::new();
1017        assert!(session.is_empty());
1018        assert_eq!(session.len(), 0);
1019        assert!(session.active().is_none());
1020        assert!(session.list_replays().is_empty());
1021    }
1022
1023    // -----------------------------------------------------------------------
1024    // 21. test_replay_session_add_and_activate
1025    // -----------------------------------------------------------------------
1026    #[test]
1027    fn test_replay_session_add_and_activate() {
1028        let mut session = ReplaySession::new();
1029
1030        let idx0 = session.add_replay(sample_trace());
1031        assert_eq!(idx0, 0);
1032        assert_eq!(session.len(), 1);
1033        assert!(!session.is_empty());
1034
1035        // First replay is auto-activated
1036        assert!(session.active().is_some());
1037        assert_eq!(session.active().unwrap().position(), 0);
1038
1039        let idx1 = session.add_replay(sample_trace());
1040        assert_eq!(idx1, 1);
1041        assert_eq!(session.len(), 2);
1042
1043        // Active is still the first replay
1044        let summaries = session.list_replays();
1045        assert!(summaries[0].is_active);
1046        assert!(!summaries[1].is_active);
1047
1048        // Switch to second replay
1049        session.set_active(1).unwrap();
1050        let summaries = session.list_replays();
1051        assert!(!summaries[0].is_active);
1052        assert!(summaries[1].is_active);
1053
1054        // Invalid index
1055        assert!(session.set_active(10).is_err());
1056    }
1057
1058    // -----------------------------------------------------------------------
1059    // 22. test_replay_session_list
1060    // -----------------------------------------------------------------------
1061    #[test]
1062    fn test_replay_session_list() {
1063        let mut session = ReplaySession::new();
1064
1065        let t1 = sample_trace();
1066        let id1 = t1.trace_id;
1067        session.add_replay(t1);
1068
1069        let t2 = sample_trace();
1070        let id2 = t2.trace_id;
1071        session.add_replay(t2);
1072
1073        let list = session.list_replays();
1074        assert_eq!(list.len(), 2);
1075        assert_eq!(list[0].index, 0);
1076        assert_eq!(list[0].trace_id, id1);
1077        assert_eq!(list[0].goal, "test task");
1078        assert_eq!(list[0].event_count, 9);
1079        assert!(list[0].is_active);
1080
1081        assert_eq!(list[1].index, 1);
1082        assert_eq!(list[1].trace_id, id2);
1083        assert!(!list[1].is_active);
1084    }
1085
1086    // -----------------------------------------------------------------------
1087    // 23. test_describe_event_variants
1088    // -----------------------------------------------------------------------
1089    #[test]
1090    fn test_describe_event_variants() {
1091        // TaskStarted
1092        let desc = describe_event(&TraceEventKind::TaskStarted {
1093            task_id: Uuid::new_v4(),
1094            goal: "refactor auth".into(),
1095        });
1096        assert!(desc.contains("Task started"));
1097        assert!(desc.contains("refactor auth"));
1098
1099        // TaskCompleted — success
1100        let desc = describe_event(&TraceEventKind::TaskCompleted {
1101            task_id: Uuid::new_v4(),
1102            success: true,
1103            iterations: 5,
1104        });
1105        assert!(desc.contains("completed successfully"));
1106        assert!(desc.contains("5 iterations"));
1107
1108        // TaskCompleted — failure
1109        let desc = describe_event(&TraceEventKind::TaskCompleted {
1110            task_id: Uuid::new_v4(),
1111            success: false,
1112            iterations: 3,
1113        });
1114        assert!(desc.contains("failed"));
1115        assert!(desc.contains("3 iterations"));
1116
1117        // ToolRequested
1118        let desc = describe_event(&TraceEventKind::ToolRequested {
1119            tool: "file_read".into(),
1120            risk_level: RiskLevel::ReadOnly,
1121            args_summary: "".into(),
1122        });
1123        assert!(desc.contains("Tool requested"));
1124        assert!(desc.contains("file_read"));
1125        assert!(desc.contains("ReadOnly"));
1126
1127        // ToolApproved
1128        let desc = describe_event(&TraceEventKind::ToolApproved {
1129            tool: "shell_exec".into(),
1130        });
1131        assert!(desc.contains("Tool approved"));
1132        assert!(desc.contains("shell_exec"));
1133
1134        // ToolDenied
1135        let desc = describe_event(&TraceEventKind::ToolDenied {
1136            tool: "file_write".into(),
1137            reason: "not allowed".into(),
1138        });
1139        assert!(desc.contains("Tool denied"));
1140        assert!(desc.contains("file_write"));
1141        assert!(desc.contains("not allowed"));
1142
1143        // ApprovalRequested
1144        let desc = describe_event(&TraceEventKind::ApprovalRequested {
1145            tool: "deploy".into(),
1146            context: "production".into(),
1147        });
1148        assert!(desc.contains("Approval requested for"));
1149        assert!(desc.contains("deploy"));
1150
1151        // ApprovalDecision — granted
1152        let desc = describe_event(&TraceEventKind::ApprovalDecision {
1153            tool: "deploy".into(),
1154            approved: true,
1155        });
1156        assert!(desc.contains("granted"));
1157        assert!(desc.contains("deploy"));
1158
1159        // ApprovalDecision — rejected
1160        let desc = describe_event(&TraceEventKind::ApprovalDecision {
1161            tool: "deploy".into(),
1162            approved: false,
1163        });
1164        assert!(desc.contains("rejected"));
1165        assert!(desc.contains("deploy"));
1166
1167        // ToolExecuted
1168        let desc = describe_event(&TraceEventKind::ToolExecuted {
1169            tool: "grep".into(),
1170            success: true,
1171            duration_ms: 42,
1172            output_preview: "match found".into(),
1173        });
1174        assert!(desc.contains("Tool executed"));
1175        assert!(desc.contains("grep"));
1176        assert!(desc.contains("ok"));
1177        assert!(desc.contains("42ms"));
1178
1179        // ToolExecuted — failure
1180        let desc = describe_event(&TraceEventKind::ToolExecuted {
1181            tool: "grep".into(),
1182            success: false,
1183            duration_ms: 100,
1184            output_preview: "".into(),
1185        });
1186        assert!(desc.contains("failed"));
1187        assert!(desc.contains("100ms"));
1188
1189        // LlmCall
1190        let desc = describe_event(&TraceEventKind::LlmCall {
1191            model: "gpt-4".into(),
1192            input_tokens: 1000,
1193            output_tokens: 500,
1194            cost: 0.05,
1195        });
1196        assert!(desc.contains("LLM call"));
1197        assert!(desc.contains("gpt-4"));
1198        assert!(desc.contains("1000/500 tokens"));
1199
1200        // StatusChange
1201        let desc = describe_event(&TraceEventKind::StatusChange {
1202            from: "idle".into(),
1203            to: "thinking".into(),
1204        });
1205        assert!(desc.contains("Status"));
1206        assert!(desc.contains("idle"));
1207        assert!(desc.contains("thinking"));
1208
1209        // Error
1210        let desc = describe_event(&TraceEventKind::Error {
1211            message: "something failed".into(),
1212        });
1213        assert!(desc.contains("Error"));
1214        assert!(desc.contains("something failed"));
1215    }
1216
1217    // -----------------------------------------------------------------------
1218    // 24. test_replay_error_display
1219    // -----------------------------------------------------------------------
1220    #[test]
1221    fn test_replay_error_display() {
1222        let err = ReplayError::TraceNotFound(Uuid::nil());
1223        assert!(err.to_string().contains("trace not found"));
1224
1225        let err = ReplayError::OutOfBounds {
1226            position: 10,
1227            total: 5,
1228        };
1229        let msg = err.to_string();
1230        assert!(msg.contains("10"));
1231        assert!(msg.contains("5"));
1232
1233        let err = ReplayError::BookmarkNotFound(3);
1234        assert!(err.to_string().contains("3"));
1235
1236        let err = ReplayError::EmptyTrace;
1237        assert!(err.to_string().contains("empty trace"));
1238    }
1239
1240    // -----------------------------------------------------------------------
1241    // 25. test_replay_snapshot_progress
1242    // -----------------------------------------------------------------------
1243    #[test]
1244    fn test_replay_snapshot_progress() {
1245        let mut engine = ReplayEngine::new(sample_trace());
1246
1247        // At start: 0%
1248        let snap = engine.snapshot();
1249        assert!((snap.progress_pct - 0.0).abs() < f64::EPSILON);
1250
1251        // At midpoint (position 4 of 9 events = 4/8 = 50%)
1252        engine.seek(4).unwrap();
1253        let snap = engine.snapshot();
1254        assert!((snap.progress_pct - 50.0).abs() < f64::EPSILON);
1255
1256        // At end: 100%
1257        engine.fast_forward();
1258        let snap = engine.snapshot();
1259        assert!((snap.progress_pct - 100.0).abs() < f64::EPSILON);
1260    }
1261
1262    // -----------------------------------------------------------------------
1263    // 26. test_replay_session_default
1264    // -----------------------------------------------------------------------
1265    #[test]
1266    fn test_replay_session_default() {
1267        let session = ReplaySession::default();
1268        assert!(session.is_empty());
1269        assert_eq!(session.len(), 0);
1270        assert!(session.active().is_none());
1271    }
1272
1273    // -----------------------------------------------------------------------
1274    // Bonus tests
1275    // -----------------------------------------------------------------------
1276
1277    #[test]
1278    fn test_replay_session_active_mut() {
1279        let mut session = ReplaySession::new();
1280        session.add_replay(sample_trace());
1281
1282        // Use active_mut to advance the engine
1283        let engine = session.active_mut().unwrap();
1284        engine.step_forward();
1285        assert_eq!(engine.position(), 1);
1286
1287        // Verify position is persisted
1288        assert_eq!(session.active().unwrap().position(), 1);
1289    }
1290
1291    #[test]
1292    fn test_replay_engine_empty_trace() {
1293        // Create a trace and manually clear its events to test empty behavior
1294        let mut trace = ExecutionTrace::new(Uuid::new_v4(), Uuid::new_v4(), "empty");
1295        trace.events.clear();
1296
1297        let mut engine = ReplayEngine::new(trace);
1298        assert_eq!(engine.position(), 0);
1299        assert_eq!(engine.total_events(), 0);
1300        assert!(engine.is_at_start());
1301        assert!(engine.is_at_end());
1302        assert!(engine.current_event().is_none());
1303        assert!(engine.step_forward().is_none());
1304        assert!(engine.step_backward().is_none());
1305        assert!(engine.skip_to_next_tool_event().is_none());
1306        assert_eq!(engine.describe_current(), "No events");
1307
1308        let snap = engine.snapshot();
1309        assert_eq!(snap.total_events, 0);
1310        assert!((snap.progress_pct - 0.0).abs() < f64::EPSILON);
1311        assert!(snap.current_event.is_none());
1312        assert!(snap.elapsed_from_start.is_none());
1313
1314        let timeline = engine.timeline();
1315        assert!(timeline.is_empty());
1316    }
1317
1318    #[test]
1319    fn test_replay_engine_walk_all_events() {
1320        let mut engine = ReplayEngine::new(sample_trace());
1321        let total = engine.total_events();
1322        let mut count = 1; // already at position 0
1323
1324        while engine.step_forward().is_some() {
1325            count += 1;
1326        }
1327
1328        assert_eq!(count, total);
1329        assert!(engine.is_at_end());
1330
1331        // Walk backwards
1332        count = 1;
1333        while engine.step_backward().is_some() {
1334            count += 1;
1335        }
1336
1337        assert_eq!(count, total);
1338        assert!(engine.is_at_start());
1339    }
1340
1341    #[test]
1342    fn test_replay_snapshot_serialization() {
1343        let engine = ReplayEngine::new(sample_trace());
1344        let snap = engine.snapshot();
1345
1346        let json = serde_json::to_string(&snap).unwrap();
1347        let restored: ReplaySnapshot = serde_json::from_str(&json).unwrap();
1348
1349        assert_eq!(restored.trace_id, snap.trace_id);
1350        assert_eq!(restored.position, snap.position);
1351        assert_eq!(restored.total_events, snap.total_events);
1352        assert_eq!(restored.errors_so_far, snap.errors_so_far);
1353    }
1354
1355    #[test]
1356    fn test_timeline_entry_serialization() {
1357        let engine = ReplayEngine::new(sample_trace());
1358        let timeline = engine.timeline();
1359
1360        let json = serde_json::to_string(&timeline).unwrap();
1361        let restored: Vec<TimelineEntry> = serde_json::from_str(&json).unwrap();
1362        assert_eq!(restored.len(), timeline.len());
1363        assert_eq!(restored[0].sequence, 0);
1364    }
1365
1366    #[test]
1367    fn test_bookmark_serialization() {
1368        let bookmark = Bookmark {
1369            position: 5,
1370            label: "important point".into(),
1371            created_at: Utc::now(),
1372        };
1373
1374        let json = serde_json::to_string(&bookmark).unwrap();
1375        let restored: Bookmark = serde_json::from_str(&json).unwrap();
1376        assert_eq!(restored.position, 5);
1377        assert_eq!(restored.label, "important point");
1378    }
1379
1380    #[test]
1381    fn test_replay_summary_serialization() {
1382        let summary = ReplaySummary {
1383            index: 0,
1384            trace_id: Uuid::new_v4(),
1385            goal: "test goal".into(),
1386            event_count: 42,
1387            is_active: true,
1388        };
1389
1390        let json = serde_json::to_string(&summary).unwrap();
1391        let restored: ReplaySummary = serde_json::from_str(&json).unwrap();
1392        assert_eq!(restored.index, 0);
1393        assert_eq!(restored.goal, "test goal");
1394        assert_eq!(restored.event_count, 42);
1395        assert!(restored.is_active);
1396    }
1397}
rustant_core/replay.rs

rustant_core/
replay.rs
