valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

simplify StreamController (#3928)

Browse Source 
no intended functional change, just simplifying the code.
This commit is contained in:
Jeremy Rose
2025-09-22 11:14:04 -07:00
committed by GitHub
parent 434eb4fd49
commit fa80bbb587
6 changed files with 86 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
codex-rs/tui/src/streaming/controller.rs
View File

@@ -3,7 +3,6 @@ use crate::history_cell::HistoryCell;
use codex_core::config::Config;
use ratatui::text::Line;
use super::HeaderEmitter;
use super::StreamState;
/// Sink for history insertions and animation control.
@@ -36,56 +35,25 @@ type Lines = Vec<Line<'static>>;
/// commit animation across streams.
pub(crate) struct StreamController {
config: Config,
header: HeaderEmitter,
state: StreamState,
active: bool,
finishing_after_drain: bool,
header_emitted: bool,
}
impl StreamController {
pub(crate) fn new(config: Config) -> Self {
Self {
config,
header: HeaderEmitter::new(),
state: StreamState::new(),
active: false,
finishing_after_drain: false,
header_emitted: false,
}
}
pub(crate) fn reset_headers_for_new_turn(&mut self) {
self.header.reset_for_new_turn();
}
pub(crate) fn is_write_cycle_active(&self) -> bool {
self.active
}
pub(crate) fn clear_all(&mut self) {
self.state.clear();
self.active = false;
self.finishing_after_drain = false;
// leave header state unchanged; caller decides when to reset
}
/// Begin an answer stream. Does not emit header yet; it is emitted on first commit.
pub(crate) fn begin(&mut self, _sink: &impl HistorySink) {
// Starting a new stream cancels any pending finish-from-previous-stream animation.
if !self.active {
self.header.reset_for_stream();
}
self.finishing_after_drain = false;
self.active = true;
}
/// Push a delta; if it contains a newline, commit completed lines and start animation.
pub(crate) fn push_and_maybe_commit(&mut self, delta: &str, sink: &impl HistorySink) {
if !self.active {
return;
}
let cfg = self.config.clone();
let state = &mut self.state;
// Record that at least one delta was received for this stream
if !delta.is_empty() {
state.has_seen_delta = true;
}
@@ -99,117 +67,54 @@ impl StreamController {
}
}
/// Finalize the active stream. If `flush_immediately` is true, drain and emit now.
pub(crate) fn finalize(&mut self, flush_immediately: bool, sink: &impl HistorySink) -> bool {
if !self.active {
return false;
}
/// Finalize the active stream. Drain and emit now.
pub(crate) fn finalize(&mut self, sink: &impl HistorySink) {
let cfg = self.config.clone();
// Finalize collector first.
let remaining = {
let state = &mut self.state;
state.collector.finalize_and_drain(&cfg)
};
if flush_immediately {
// Collect all output first to avoid emitting headers when there is no content.
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.history);
}
if !out_lines.is_empty() {
// Insert as a HistoryCell so display drops the header while transcript keeps it.
sink.insert_history_cell(Box::new(history_cell::AgentMessageCell::new(
out_lines,
self.header.maybe_emit_header(),
)));
}
// Cleanup
self.state.clear();
// Allow a subsequent block in this turn to emit its header.
self.header.allow_reemit_in_turn();
// Also clear the per-stream emitted flag so the header can render again.
self.header.reset_for_stream();
self.active = false;
self.finishing_after_drain = false;
true
} else {
// Collect all output first to avoid emitting headers when there is no content.
let mut out_lines: Lines = Vec::new();
{
let state = &mut self.state;
if !remaining.is_empty() {
let state = &mut self.state;
state.enqueue(remaining);
}
// Spacer animated out
self.state.enqueue(vec![Line::from("")]);
self.finishing_after_drain = true;
sink.start_commit_animation();
false
let step = state.drain_all();
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;
}
/// Step animation: commit at most one queued line and handle end-of-drain cleanup.
pub(crate) fn on_commit_tick(&mut self, sink: &impl HistorySink) -> bool {
if !self.active {
return false;
}
let step = { self.state.step() };
if !step.history.is_empty() {
sink.insert_history_cell(Box::new(history_cell::AgentMessageCell::new(
step.history,
self.header.maybe_emit_header(),
)));
self.emit(sink, step.history);
}
let is_idle = self.state.is_idle();
if is_idle {
sink.stop_commit_animation();
if self.finishing_after_drain {
// Reset and notify
self.state.clear();
// Allow a subsequent block in this turn to emit its header.
self.header.allow_reemit_in_turn();
// Also clear the per-stream emitted flag so the header can render again.
self.header.reset_for_stream();
self.active = false;
self.finishing_after_drain = false;
return true;
}
}
false
}
/// Apply a full final answer: replace queued content with only the remaining tail,
/// then finalize immediately and notify completion.
pub(crate) fn apply_final_answer(&mut self, message: &str, sink: &impl HistorySink) -> bool {
self.apply_full_final(message, sink)
}
fn apply_full_final(&mut self, message: &str, sink: &impl HistorySink) -> bool {
self.begin(sink);
{
let state = &mut self.state;
// Only inject the final full message if we have not seen any deltas for this stream.
// If deltas were received, rely on the collector's existing buffer to avoid duplication.
if !state.has_seen_delta && !message.is_empty() {
// normalize to end with newline
let mut msg = message.to_owned();
if !msg.ends_with('\n') {
msg.push('\n');
}
// replace while preserving already committed count
let committed = state.collector.committed_count();
state
.collector
.replace_with_and_mark_committed(&msg, committed);
}
}
self.finalize(true, sink)
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;
}
}
@@ -268,7 +173,6 @@ mod tests {
let cfg = test_config();
let mut ctrl = StreamController::new(cfg.clone());
let sink = TestSink::new();
ctrl.begin(&sink);
// Exact deltas from the session log (section: Loose vs. tight list items)
let deltas = vec![
@@ -347,7 +251,7 @@ mod tests {
let _ = ctrl.on_commit_tick(&sink);
}
// Finalize and flush remaining lines now.
let _ = ctrl.finalize(true, &sink);
ctrl.finalize(&sink);
// Flatten sink output and strip the header that the controller inserts (blank + "codex").
let mut flat: Vec<ratatui::text::Line<'static>> = Vec::new();