use crate::history_cell::HistoryCell; use crate::history_cell::{self}; use codex_core::config::Config; use ratatui::text::Line; use super::StreamState; /// Controller that manages newline-gated streaming, header emission, and /// commit animation across streams. pub(crate) struct StreamController { config: Config, state: StreamState, finishing_after_drain: bool, header_emitted: bool, } impl StreamController { pub(crate) fn new(config: Config, width: Option) -> Self { Self { config, state: StreamState::new(width), finishing_after_drain: false, header_emitted: false, } } /// Push a delta; if it contains a newline, commit completed lines and start animation. pub(crate) fn push(&mut self, delta: &str) -> bool { let cfg = self.config.clone(); let state = &mut self.state; if !delta.is_empty() { state.has_seen_delta = true; } state.collector.push_delta(delta); if delta.contains('\n') { let newly_completed = state.collector.commit_complete_lines(&cfg); if !newly_completed.is_empty() { state.enqueue(newly_completed); return true; } } false } /// Finalize the active stream. Drain and emit now. pub(crate) fn finalize(&mut self) -> Option> { let cfg = self.config.clone(); // Finalize collector first. let remaining = { let state = &mut self.state; 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 state = &mut self.state; if !remaining.is_empty() { state.enqueue(remaining); } let step = state.drain_all(); out_lines.extend(step); } // 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>, bool) { let step = self.state.step(); (self.emit(step), self.state.is_idle()) } fn emit(&mut self, lines: Vec>) -> Option> { 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 }))) } } #[cfg(test)] mod tests { use super::*; use codex_core::config::Config; use codex_core::config::ConfigOverrides; fn test_config() -> Config { let overrides = ConfigOverrides { cwd: std::env::current_dir().ok(), ..Default::default() }; match Config::load_with_cli_overrides(vec![], overrides) { Ok(c) => c, Err(e) => panic!("load test config: {e}"), } } fn lines_to_plain_strings(lines: &[ratatui::text::Line<'_>]) -> Vec { lines .iter() .map(|l| { l.spans .iter() .map(|s| s.content.clone()) .collect::>() .join("") }) .collect() } #[test] fn controller_loose_vs_tight_with_commit_ticks_matches_full() { let cfg = test_config(); let mut ctrl = StreamController::new(cfg.clone(), None); let mut lines = Vec::new(); // Exact deltas from the session log (section: Loose vs. tight list items) let deltas = vec![ "\n\n", "Loose", " vs", ".", " tight", " list", " items", ":\n", "1", ".", " Tight", " item", "\n", "2", ".", " Another", " tight", " item", "\n\n", "1", ".", " Loose", " item", " with", " its", " own", " paragraph", ".\n\n", " ", " This", " paragraph", " belongs", " to", " the", " same", " list", " item", ".\n\n", "2", ".", " Second", " loose", " item", " with", " a", " nested", " list", " after", " a", " blank", " line", ".\n\n", " ", " -", " Nested", " bullet", " under", " a", " loose", " item", "\n", " ", " -", " Another", " nested", " bullet", "\n\n", ]; // 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; } } } // Finalize and flush remaining lines now. if let Some(cell) = ctrl.finalize() { lines.extend(cell.transcript_lines()); } let mut flat = lines; // 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); } if !flat.is_empty() { let s0 = lines_to_plain_strings(&[flat[0].clone()])[0].clone(); if s0 == "codex" { flat.remove(0); } } let streamed = lines_to_plain_strings(&flat); // Full render of the same source let source: String = deltas.iter().copied().collect(); let mut rendered: Vec> = Vec::new(); crate::markdown::append_markdown(&source, None, &mut rendered, &cfg); let rendered_strs = lines_to_plain_strings(&rendered); assert_eq!(streamed, rendered_strs); // Also assert exact expected plain strings for clarity. let expected = vec![ "Loose vs. tight list items:".to_string(), "".to_string(), "1. Tight item".to_string(), "2. Another tight item".to_string(), "3. Loose item with its own paragraph.".to_string(), "".to_string(), " This paragraph belongs to the same list item.".to_string(), "4. Second loose item with a nested list after a blank line.".to_string(), " - Nested bullet under a loose item".to_string(), " - Another nested bullet".to_string(), ]; assert_eq!( streamed, expected, "expected exact rendered lines for loose/tight section" ); } }