Revamp status UI with rate limit reset times

codex-rs/tui/src/streaming/controller.rs

@@ -1,10 +1,36 @@
+use crate::history_cell;
 use crate::history_cell::HistoryCell;
-use crate::history_cell::{self};
 use codex_core::config::Config;
 use ratatui::text::Line;
 
 use super::StreamState;
 
+/// Sink for history insertions and animation control.
+pub(crate) trait HistorySink {
+    fn insert_history_cell(&self, cell: Box<dyn HistoryCell>);
+    fn start_commit_animation(&self);
+    fn stop_commit_animation(&self);
+}
+
+/// Concrete sink backed by `AppEventSender`.
+pub(crate) struct AppEventHistorySink(pub(crate) crate::app_event_sender::AppEventSender);
+
+impl HistorySink for AppEventHistorySink {
+    fn insert_history_cell(&self, cell: Box<dyn crate::history_cell::HistoryCell>) {
+        self.0
+            .send(crate::app_event::AppEvent::InsertHistoryCell(cell))
+    }
+    fn start_commit_animation(&self) {
+        self.0
+            .send(crate::app_event::AppEvent::StartCommitAnimation)
+    }
+    fn stop_commit_animation(&self) {
+        self.0.send(crate::app_event::AppEvent::StopCommitAnimation)
+    }
+}
+
+type Lines = Vec<Line<'static>>;
+
 /// Controller that manages newline-gated streaming, header emission, and
 /// commit animation across streams.
 pub(crate) struct StreamController {
@@ -25,7 +51,7 @@ impl StreamController {
     }
 
     /// Push a delta; if it contains a newline, commit completed lines and start animation.
-    pub(crate) fn push(&mut self, delta: &str) -> bool {
+    pub(crate) fn push_and_maybe_commit(&mut self, delta: &str, sink: &impl HistorySink) {
         let cfg = self.config.clone();
         let state = &mut self.state;
         if !delta.is_empty() {
@@ -36,14 +62,13 @@ impl StreamController {
             let newly_completed = state.collector.commit_complete_lines(&cfg);
             if !newly_completed.is_empty() {
                 state.enqueue(newly_completed);
-                return true;
+                sink.start_commit_animation();
             }
         }
-        false
     }
 
     /// Finalize the active stream. Drain and emit now.
-    pub(crate) fn finalize(&mut self) -> Option<Box<dyn HistoryCell>> {
+    pub(crate) fn finalize(&mut self, sink: &impl HistorySink) {
         let cfg = self.config.clone();
         // Finalize collector first.
         let remaining = {
@@ -51,37 +76,45 @@ impl StreamController {
             state.collector.finalize_and_drain(&cfg)
         };
         // Collect all output first to avoid emitting headers when there is no content.
-        let mut out_lines = Vec::new();
+        let mut out_lines: Lines = Vec::new();
         {
             let state = &mut self.state;
             if !remaining.is_empty() {
                 state.enqueue(remaining);
             }
             let step = state.drain_all();
-            out_lines.extend(step);
+            out_lines.extend(step.history);
+        }
+        if !out_lines.is_empty() {
+            // Insert as a HistoryCell so display drops the header while transcript keeps it.
+            self.emit(sink, out_lines);
         }
 
         // Cleanup
         self.state.clear();
         self.finishing_after_drain = false;
-        self.emit(out_lines)
     }
 
     /// Step animation: commit at most one queued line and handle end-of-drain cleanup.
-    pub(crate) fn on_commit_tick(&mut self) -> (Option<Box<dyn HistoryCell>>, bool) {
-        let step = self.state.step();
-        (self.emit(step), self.state.is_idle())
+    pub(crate) fn on_commit_tick(&mut self, sink: &impl HistorySink) -> bool {
+        let step = { self.state.step() };
+        if !step.history.is_empty() {
+            self.emit(sink, step.history);
+        }
+
+        let is_idle = self.state.is_idle();
+        if is_idle {
+            sink.stop_commit_animation();
+        }
+        false
     }
 
-    fn emit(&mut self, lines: Vec<Line<'static>>) -> Option<Box<dyn HistoryCell>> {
-        if lines.is_empty() {
-            return None;
-        }
-        Some(Box::new(history_cell::AgentMessageCell::new(lines, {
-            let header_emitted = self.header_emitted;
-            self.header_emitted = true;
-            !header_emitted
-        })))
+    fn emit(&mut self, sink: &impl HistorySink, lines: Vec<Line<'static>>) {
+        sink.insert_history_cell(Box::new(history_cell::AgentMessageCell::new(
+            lines,
+            !self.header_emitted,
+        )));
+        self.header_emitted = true;
     }
 }
 
@@ -90,6 +123,7 @@ mod tests {
     use super::*;
     use codex_core::config::Config;
     use codex_core::config::ConfigOverrides;
+    use std::cell::RefCell;
 
     fn test_config() -> Config {
         let overrides = ConfigOverrides {
@@ -102,6 +136,25 @@ mod tests {
         }
     }
 
+    struct TestSink {
+        pub lines: RefCell<Vec<Vec<Line<'static>>>>,
+    }
+    impl TestSink {
+        fn new() -> Self {
+            Self {
+                lines: RefCell::new(Vec::new()),
+            }
+        }
+    }
+    impl HistorySink for TestSink {
+        fn insert_history_cell(&self, cell: Box<dyn crate::history_cell::HistoryCell>) {
+            // For tests, store the transcript representation of the cell.
+            self.lines.borrow_mut().push(cell.transcript_lines());
+        }
+        fn start_commit_animation(&self) {}
+        fn stop_commit_animation(&self) {}
+    }
+
     fn lines_to_plain_strings(lines: &[ratatui::text::Line<'_>]) -> Vec<String> {
         lines
             .iter()
@@ -119,7 +172,7 @@ mod tests {
     fn controller_loose_vs_tight_with_commit_ticks_matches_full() {
         let cfg = test_config();
         let mut ctrl = StreamController::new(cfg.clone());
-        let mut lines = Vec::new();
+        let sink = TestSink::new();
 
         // Exact deltas from the session log (section: Loose vs. tight list items)
         let deltas = vec![
@@ -193,21 +246,20 @@ mod tests {
         ];
 
         // Simulate streaming with a commit tick attempt after each delta.
-        for d in deltas.iter() {
-            ctrl.push(d);
-            while let (Some(cell), idle) = ctrl.on_commit_tick() {
-                lines.extend(cell.transcript_lines());
-                if idle {
-                    break;
-                }
-            }
+        for d in &deltas {
+            ctrl.push_and_maybe_commit(d, &sink);
+            let _ = ctrl.on_commit_tick(&sink);
         }
         // Finalize and flush remaining lines now.
-        if let Some(cell) = ctrl.finalize() {
-            lines.extend(cell.transcript_lines());
-        }
+        ctrl.finalize(&sink);
 
-        let mut flat = lines;
+        // Flatten sink output and strip the header that the controller inserts (blank + "codex").
+        let mut flat: Vec<ratatui::text::Line<'static>> = Vec::new();
+        for batch in sink.lines.borrow().iter() {
+            for l in batch {
+                flat.push(l.clone());
+            }
+        }
         // Drop leading blank and header line if present.
         if !flat.is_empty() && lines_to_plain_strings(&[flat[0].clone()])[0].is_empty() {
             flat.remove(0);