tokio_prompt_orchestrator/
trace_ui.rs

1#![allow(dead_code)]
2//! # Module: Request Lifecycle Tracer TUI Panel
3//!
4//! ## Responsibility
5//! Provides a Ratatui panel that renders the full lifecycle of a single request
6//! as a Gantt-chart timeline. Activated with the `t` key in the TUI; the user
7//! types a `request_id` to trace, and each pipeline stage is rendered as a
8//! colour-coded bar proportional to time spent.
9//!
10//! ## Guarantees
11//! - No panics in any rendering or state-update path
12//! - All arithmetic is saturating / checked; no integer overflow
13//! - Zero allocations in the hot render path (only on stage record insertion)
14//!
15//! ## Feature Gate
16//! This module is compiled unconditionally but [`TracePanel`] only renders
17//! content when the `tui` feature is active (the binary wires it in).
18
19use std::collections::HashMap;
20use std::time::{Duration, Instant};
21
22use serde::{Deserialize, Serialize};
23
24// ─── Stage names ─────────────────────────────────────────────────────────────
25
26/// Ordered pipeline stage labels.
27pub const LIFECYCLE_STAGES: [&str; 5] = ["RAG", "Assemble", "Inference", "Post", "Stream"];
28
29// ─── Stage record ────────────────────────────────────────────────────────────
30
31/// Timing record for one pipeline stage within a single request.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct StageRecord {
34    /// Stage index (0 = RAG … 4 = Stream).
35    pub stage_index: usize,
36    /// When the stage began processing this request (milliseconds since request entry).
37    pub start_ms: u64,
38    /// When the stage finished (milliseconds since request entry). `None` if still running.
39    pub end_ms: Option<u64>,
40    /// Whether the stage completed without error.
41    pub success: bool,
42}
43
44impl StageRecord {
45    /// Duration in milliseconds, if the stage has completed.
46    #[must_use]
47    pub fn duration_ms(&self) -> Option<u64> {
48        self.end_ms.map(|e| e.saturating_sub(self.start_ms))
49    }
50}
51
52// ─── Request trace ───────────────────────────────────────────────────────────
53
54/// Complete lifecycle record for one request.
55#[derive(Debug, Clone)]
56pub struct RequestTrace {
57    /// Unique identifier for the request.
58    pub request_id: String,
59    /// Wall-clock time when the request entered the pipeline.
60    pub entry_time: Instant,
61    /// Per-stage records, indexed by stage index (0-4).
62    pub stages: [Option<StageRecord>; 5],
63    /// Whether the entire request has completed.
64    pub complete: bool,
65}
66
67impl RequestTrace {
68    /// Create a new, empty trace for the given request id.
69    #[must_use]
70    pub fn new(request_id: impl Into<String>) -> Self {
71        Self {
72            request_id: request_id.into(),
73            entry_time: Instant::now(),
74            stages: [None, None, None, None, None],
75            complete: false,
76        }
77    }
78
79    /// Record the start of a stage.
80    ///
81    /// Idempotent: if the stage was already started, this is a no-op.
82    pub fn stage_started(&mut self, stage_index: usize) {
83        if stage_index >= 5 {
84            return;
85        }
86        if self.stages[stage_index].is_none() {
87            let start_ms = self
88                .entry_time
89                .elapsed()
90                .as_millis()
91                .try_into()
92                .unwrap_or(u64::MAX);
93            self.stages[stage_index] = Some(StageRecord {
94                stage_index,
95                start_ms,
96                end_ms: None,
97                success: false,
98            });
99        }
100    }
101
102    /// Record the completion of a stage.
103    pub fn stage_finished(&mut self, stage_index: usize, success: bool) {
104        if stage_index >= 5 {
105            return;
106        }
107        // Ensure a start record exists (handle missing start gracefully)
108        if self.stages[stage_index].is_none() {
109            self.stage_started(stage_index);
110        }
111        if let Some(ref mut rec) = self.stages[stage_index] {
112            let now_ms = self
113                .entry_time
114                .elapsed()
115                .as_millis()
116                .try_into()
117                .unwrap_or(u64::MAX);
118            rec.end_ms = Some(now_ms);
119            rec.success = success;
120        }
121        // Mark overall completion when the final stage (Stream) finishes
122        if stage_index == 4 {
123            self.complete = true;
124        }
125    }
126
127    /// Total elapsed time from entry to now (or until Stream completed).
128    #[must_use]
129    pub fn total_elapsed(&self) -> Duration {
130        if let Some(Some(stream)) = self.stages.get(4) {
131            if let Some(end_ms) = stream.end_ms {
132                return Duration::from_millis(end_ms);
133            }
134        }
135        self.entry_time.elapsed()
136    }
137}
138
139// ─── Trace store ─────────────────────────────────────────────────────────────
140
141/// Thread-safe store of recent request traces.
142///
143/// Bounded to `capacity` entries; oldest entries are evicted when the cap is
144/// reached. All public methods take `&self` to allow sharing across tasks via
145/// `Arc<TraceStore>`.
146#[derive(Debug)]
147pub struct TraceStore {
148    inner: std::sync::Mutex<TraceStoreInner>,
149}
150
151#[derive(Debug)]
152struct TraceStoreInner {
153    traces: HashMap<String, RequestTrace>,
154    insertion_order: std::collections::VecDeque<String>,
155    capacity: usize,
156}
157
158impl TraceStore {
159    /// Create a store with the given capacity.
160    #[must_use]
161    pub fn new(capacity: usize) -> Self {
162        Self {
163            inner: std::sync::Mutex::new(TraceStoreInner {
164                traces: HashMap::new(),
165                insertion_order: std::collections::VecDeque::new(),
166                capacity: capacity.max(1),
167            }),
168        }
169    }
170
171    /// Insert or update a trace. Evicts the oldest entry when at capacity.
172    pub fn upsert(&self, trace: RequestTrace) {
173        let Ok(mut inner) = self.inner.lock() else {
174            return;
175        };
176        let id = trace.request_id.clone();
177        if !inner.traces.contains_key(&id) {
178            // Evict oldest if at cap
179            if inner.insertion_order.len() >= inner.capacity {
180                if let Some(oldest) = inner.insertion_order.pop_front() {
181                    inner.traces.remove(&oldest);
182                }
183            }
184            inner.insertion_order.push_back(id.clone());
185        }
186        inner.traces.insert(id, trace);
187    }
188
189    /// Record a stage start event.
190    pub fn stage_started(&self, request_id: &str, stage_index: usize) {
191        let Ok(mut inner) = self.inner.lock() else {
192            return;
193        };
194        if let Some(trace) = inner.traces.get_mut(request_id) {
195            trace.stage_started(stage_index);
196        } else {
197            // Create a new trace if we haven't seen this request yet
198            drop(inner);
199            let mut trace = RequestTrace::new(request_id);
200            trace.stage_started(stage_index);
201            self.upsert(trace);
202        }
203    }
204
205    /// Record a stage finish event.
206    pub fn stage_finished(&self, request_id: &str, stage_index: usize, success: bool) {
207        let Ok(mut inner) = self.inner.lock() else {
208            return;
209        };
210        if let Some(trace) = inner.traces.get_mut(request_id) {
211            trace.stage_finished(stage_index, success);
212        }
213    }
214
215    /// Look up a trace by request id.
216    #[must_use]
217    pub fn get(&self, request_id: &str) -> Option<RequestTrace> {
218        let Ok(inner) = self.inner.lock() else {
219            return None;
220        };
221        inner.traces.get(request_id).cloned()
222    }
223
224    /// Return a list of all known request ids (newest last).
225    #[must_use]
226    pub fn known_ids(&self) -> Vec<String> {
227        let Ok(inner) = self.inner.lock() else {
228            return vec![];
229        };
230        inner.insertion_order.iter().cloned().collect()
231    }
232}
233
234impl Default for TraceStore {
235    fn default() -> Self {
236        Self::new(1000)
237    }
238}
239
240// ─── TUI panel state ─────────────────────────────────────────────────────────
241
242/// Interaction mode for the trace panel.
243#[derive(Debug, Clone, Copy, PartialEq, Eq)]
244pub enum TracePanelMode {
245    /// Panel is hidden; user presses `t` to activate.
246    Hidden,
247    /// User is typing a request_id to look up.
248    Searching,
249    /// Displaying the timeline for a resolved request.
250    Viewing,
251}
252
253/// State for the request-lifecycle tracer TUI panel.
254#[derive(Debug)]
255pub struct TracePanel {
256    /// Current interaction mode.
257    pub mode: TracePanelMode,
258    /// Text the user has typed so far.
259    pub search_input: String,
260    /// The currently displayed trace (if any).
261    pub current_trace: Option<RequestTrace>,
262    /// Error message to display (e.g. "not found").
263    pub error_msg: Option<String>,
264    /// Shared trace store.
265    pub store: std::sync::Arc<TraceStore>,
266}
267
268impl TracePanel {
269    /// Create a new panel backed by the given store.
270    #[must_use]
271    pub fn new(store: std::sync::Arc<TraceStore>) -> Self {
272        Self {
273            mode: TracePanelMode::Hidden,
274            search_input: String::new(),
275            current_trace: None,
276            error_msg: None,
277            store,
278        }
279    }
280
281    /// Toggle panel visibility (bound to `t` key).
282    pub fn toggle(&mut self) {
283        match self.mode {
284            TracePanelMode::Hidden => {
285                self.mode = TracePanelMode::Searching;
286                self.search_input.clear();
287                self.error_msg = None;
288            }
289            _ => {
290                self.mode = TracePanelMode::Hidden;
291            }
292        }
293    }
294
295    /// Feed a character from keyboard input (while in Searching mode).
296    pub fn push_char(&mut self, c: char) {
297        if self.mode == TracePanelMode::Searching {
298            self.search_input.push(c);
299        }
300    }
301
302    /// Delete last character (Backspace) while searching.
303    pub fn pop_char(&mut self) {
304        if self.mode == TracePanelMode::Searching {
305            self.search_input.pop();
306        }
307    }
308
309    /// Submit the current search input (Enter key).
310    pub fn submit(&mut self) {
311        if self.mode != TracePanelMode::Searching {
312            return;
313        }
314        let id = self.search_input.trim().to_string();
315        if id.is_empty() {
316            self.error_msg = Some("Enter a request_id".to_string());
317            return;
318        }
319        match self.store.get(&id) {
320            Some(trace) => {
321                self.current_trace = Some(trace);
322                self.error_msg = None;
323                self.mode = TracePanelMode::Viewing;
324            }
325            None => {
326                self.error_msg = Some(format!("request_id '{id}' not found"));
327            }
328        }
329    }
330
331    /// Refresh the displayed trace from the store (call on each tick while Viewing).
332    pub fn refresh(&mut self) {
333        if self.mode == TracePanelMode::Viewing {
334            if let Some(ref trace) = self.current_trace.clone() {
335                self.current_trace = self.store.get(&trace.request_id);
336            }
337        }
338    }
339}
340
341// ─── Ratatui rendering ───────────────────────────────────────────────────────
342
343#[cfg(feature = "tui")]
344pub mod render {
345    use super::{TracePanel, TracePanelMode, LIFECYCLE_STAGES};
346    use ratatui::layout::{Alignment, Constraint, Direction, Layout, Rect};
347    use ratatui::style::{Color, Modifier, Style};
348    use ratatui::text::{Line, Span};
349    use ratatui::widgets::{Block, Borders, Clear, Paragraph};
350    use ratatui::Frame;
351
352    /// Draw the trace panel as a centred overlay.
353    ///
354    /// Does nothing if `panel.mode == Hidden`.
355    pub fn draw_trace_panel(f: &mut Frame, panel: &TracePanel) {
356        if panel.mode == TracePanelMode::Hidden {
357            return;
358        }
359
360        let area = f.area();
361        // Centre a 70%-wide, 60%-tall popup
362        let popup = centred_rect(70, 60, area);
363        f.render_widget(Clear, popup);
364
365        let block = Block::default()
366            .title(" REQUEST LIFECYCLE TRACER  [t] to close ")
367            .borders(Borders::ALL)
368            .border_style(Style::default().fg(Color::Cyan));
369        let inner = block.inner(popup);
370        f.render_widget(block, popup);
371
372        let chunks = Layout::default()
373            .direction(Direction::Vertical)
374            .constraints([
375                Constraint::Length(3), // search bar
376                Constraint::Min(4),    // timeline / message
377            ])
378            .split(inner);
379
380        // Search bar
381        draw_search_bar(f, chunks[0], panel);
382
383        // Timeline or status
384        match panel.mode {
385            TracePanelMode::Searching => {
386                if let Some(ref msg) = panel.error_msg {
387                    let p = Paragraph::new(msg.as_str())
388                        .style(Style::default().fg(Color::Red))
389                        .alignment(Alignment::Center);
390                    f.render_widget(p, chunks[1]);
391                } else {
392                    let ids = panel.store.known_ids();
393                    let hint = if ids.is_empty() {
394                        "No traces recorded yet.".to_string()
395                    } else {
396                        format!(
397                            "Known IDs (newest last): {}",
398                            ids.iter()
399                                .rev()
400                                .take(5)
401                                .cloned()
402                                .collect::<Vec<_>>()
403                                .join(", ")
404                        )
405                    };
406                    let p = Paragraph::new(hint)
407                        .style(Style::default().fg(Color::DarkGray))
408                        .alignment(Alignment::Center);
409                    f.render_widget(p, chunks[1]);
410                }
411            }
412            TracePanelMode::Viewing => {
413                if let Some(ref trace) = panel.current_trace {
414                    draw_timeline(f, chunks[1], trace);
415                }
416            }
417            TracePanelMode::Hidden => {}
418        }
419    }
420
421    fn draw_search_bar(f: &mut Frame, area: Rect, panel: &TracePanel) {
422        let label = if panel.mode == TracePanelMode::Searching {
423            format!(" Search request_id: {}|", panel.search_input)
424        } else {
425            format!(" Tracing: {}", panel.search_input)
426        };
427        let p = Paragraph::new(label)
428            .block(
429                Block::default()
430                    .borders(Borders::ALL)
431                    .border_style(Style::default().fg(Color::Yellow)),
432            )
433            .style(Style::default().fg(Color::White));
434        f.render_widget(p, area);
435    }
436
437    fn draw_timeline(
438        f: &mut Frame,
439        area: Rect,
440        trace: &super::RequestTrace,
441    ) {
442        if area.width < 20 || area.height < 3 {
443            return;
444        }
445
446        let total_ms = trace.total_elapsed().as_millis() as f64;
447        let label_w: u16 = 10;
448        let bar_w = area.width.saturating_sub(label_w + 2);
449
450        // Header
451        let header = Line::from(vec![
452            Span::raw(format!("{:<10}", "Stage")),
453            Span::styled(
454                format!(" {:>w$}", format!("{total_ms:.0} ms total"), w = bar_w as usize),
455                Style::default()
456                    .fg(Color::DarkGray)
457                    .add_modifier(Modifier::DIM),
458            ),
459        ]);
460
461        let mut lines: Vec<Line> = vec![header, Line::raw("")];
462
463        for (idx, name) in LIFECYCLE_STAGES.iter().enumerate() {
464            let line = if let Some(Some(rec)) = trace.stages.get(idx) {
465                let start_frac = rec.start_ms as f64 / total_ms.max(1.0);
466                let dur_ms = rec
467                    .end_ms
468                    .map(|e| e.saturating_sub(rec.start_ms))
469                    .unwrap_or(0);
470                let dur_frac = dur_ms as f64 / total_ms.max(1.0);
471
472                let offset = ((start_frac * bar_w as f64) as u16).min(bar_w);
473                let width = ((dur_frac * bar_w as f64) as u16)
474                    .max(1)
475                    .min(bar_w.saturating_sub(offset));
476
477                let color = if rec.end_ms.is_none() {
478                    Color::Yellow // still running
479                } else if rec.success {
480                    // Color by proportion: green < 30%, yellow < 60%, red otherwise
481                    if dur_frac < 0.30 {
482                        Color::Green
483                    } else if dur_frac < 0.60 {
484                        Color::Yellow
485                    } else {
486                        Color::Red
487                    }
488                } else {
489                    Color::Red
490                };
491
492                let prefix = " ".repeat(offset as usize);
493                let bar_str = "\u{2588}".repeat(width as usize); // ██
494                let suffix_w = bar_w.saturating_sub(offset + width) as usize;
495                let suffix = " ".repeat(suffix_w);
496                let timing = format!(
497                    " {dur_ms}ms",
498                );
499
500                Line::from(vec![
501                    Span::styled(
502                        format!("{name:<10}"),
503                        Style::default()
504                            .fg(Color::White)
505                            .add_modifier(Modifier::BOLD),
506                    ),
507                    Span::raw(prefix),
508                    Span::styled(bar_str, Style::default().fg(color)),
509                    Span::raw(suffix),
510                    Span::styled(timing, Style::default().fg(Color::DarkGray)),
511                ])
512            } else {
513                // Stage not started
514                Line::from(vec![
515                    Span::styled(
516                        format!("{name:<10}"),
517                        Style::default().fg(Color::DarkGray),
518                    ),
519                    Span::styled(
520                        format!("{:-<w$}", "", w = bar_w as usize),
521                        Style::default().fg(Color::DarkGray),
522                    ),
523                ])
524            };
525            lines.push(line);
526        }
527
528        // Summary line
529        lines.push(Line::raw(""));
530        let status = if trace.complete {
531            Span::styled(
532                format!(" COMPLETE  ({:.0} ms)", total_ms),
533                Style::default()
534                    .fg(Color::Green)
535                    .add_modifier(Modifier::BOLD),
536            )
537        } else {
538            Span::styled(
539                format!(" IN FLIGHT ({:.0} ms elapsed)", total_ms),
540                Style::default()
541                    .fg(Color::Yellow)
542                    .add_modifier(Modifier::BOLD),
543            )
544        };
545        lines.push(Line::from(vec![status]));
546
547        let p = Paragraph::new(lines);
548        f.render_widget(p, area);
549    }
550
551    /// Returns a centred [`Rect`] occupying `percent_x` / `percent_y` of `r`.
552    fn centred_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect {
553        let layout = Layout::default()
554            .direction(Direction::Vertical)
555            .constraints([
556                Constraint::Percentage((100 - percent_y) / 2),
557                Constraint::Percentage(percent_y),
558                Constraint::Percentage((100 - percent_y) / 2),
559            ])
560            .split(r);
561
562        Layout::default()
563            .direction(Direction::Horizontal)
564            .constraints([
565                Constraint::Percentage((100 - percent_x) / 2),
566                Constraint::Percentage(percent_x),
567                Constraint::Percentage((100 - percent_x) / 2),
568            ])
569            .split(layout[1])[1]
570    }
571}
572
573// ─── Tests ───────────────────────────────────────────────────────────────────
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578
579    #[test]
580    fn trace_records_stages() {
581        let mut t = RequestTrace::new("req-1");
582        t.stage_started(0);
583        t.stage_finished(0, true);
584        t.stage_started(4);
585        t.stage_finished(4, true);
586        assert!(t.complete);
587        let rec = t.stages[0].as_ref().expect("stage 0 should be recorded");
588        assert!(rec.duration_ms().is_some());
589    }
590
591    #[test]
592    fn store_evicts_at_capacity() {
593        let store = TraceStore::new(2);
594        for i in 0..4u32 {
595            store.upsert(RequestTrace::new(format!("req-{i}")));
596        }
597        let ids = store.known_ids();
598        assert_eq!(ids.len(), 2);
599        assert_eq!(ids[0], "req-2");
600        assert_eq!(ids[1], "req-3");
601    }
602
603    #[test]
604    fn panel_toggle_hides() {
605        let store = std::sync::Arc::new(TraceStore::default());
606        let mut panel = TracePanel::new(store);
607        assert_eq!(panel.mode, TracePanelMode::Hidden);
608        panel.toggle();
609        assert_eq!(panel.mode, TracePanelMode::Searching);
610        panel.toggle();
611        assert_eq!(panel.mode, TracePanelMode::Hidden);
612    }
613
614    #[test]
615    fn panel_submit_not_found() {
616        let store = std::sync::Arc::new(TraceStore::default());
617        let mut panel = TracePanel::new(store);
618        panel.toggle();
619        panel.push_char('x');
620        panel.submit();
621        assert!(panel.error_msg.is_some());
622        assert_eq!(panel.mode, TracePanelMode::Searching);
623    }
624}
tokio_prompt_orchestrator/trace_ui.rs

tokio_prompt_orchestrator/
trace_ui.rs
