1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
// TODO(integrate): Wire streaming collector into TUI rendering pipeline
#![allow(dead_code)]
//! Markdown stream collector for newline-gated rendering.
//!
//! This module implements the pattern from codex-rs where:
//! - Streaming text is buffered until a newline is reached
//! - Only complete lines are committed to the UI
//! - This prevents visual flashing of partial words
//! - Final content is emitted when the stream ends
use ratatui::style::{Modifier, Style};
use ratatui::text::{Line, Span};
use unicode_width::UnicodeWidthStr;
use crate::palette;
/// Collects streaming text and commits complete lines.
#[derive(Debug, Clone)]
pub struct MarkdownStreamCollector {
/// Buffer for incoming text
buffer: String,
/// Number of lines already committed
committed_line_count: usize,
/// Terminal width for wrapping
width: Option<usize>,
/// Whether the stream is still active
is_streaming: bool,
/// Whether this is a thinking block
is_thinking: bool,
}
impl MarkdownStreamCollector {
/// Create a new collector
pub fn new(width: Option<usize>, is_thinking: bool) -> Self {
Self {
buffer: String::new(),
committed_line_count: 0,
width,
is_streaming: true,
is_thinking,
}
}
/// Push new content to the buffer
pub fn push(&mut self, content: &str) {
self.buffer.push_str(content);
}
/// Get the current buffer content (for display during streaming)
pub fn current_content(&self) -> &str {
&self.buffer
}
/// Check if there are complete lines to commit
pub fn has_complete_lines(&self) -> bool {
self.buffer.contains('\n')
}
/// Commit complete lines and return them.
/// Only lines ending with '\n' are committed.
/// Returns the newly committed lines since last call.
pub fn commit_complete_lines(&mut self) -> Vec<Line<'static>> {
let committed = self.commit_complete_text();
if committed.is_empty() {
return Vec::new();
}
self.render_lines(&committed)
}
/// Commit complete text chunks ending in a newline.
/// Returns the raw text that became visible since the last call.
pub fn commit_complete_text(&mut self) -> String {
if self.buffer.is_empty() {
return String::new();
}
// Find the last newline - only process up to there
let Some(last_newline_idx) = self.buffer.rfind('\n') else {
return String::new(); // No complete lines yet
};
// Extract the complete portion (up to and including last newline)
let complete_portion = self.buffer[..=last_newline_idx].to_string();
// Remove the committed portion from the buffer so finalize only emits the remainder
self.buffer = self.buffer[last_newline_idx + 1..].to_string();
self.committed_line_count = 0;
complete_portion
}
/// Finalize the stream and return any remaining content.
/// Call this when the stream ends to emit the final incomplete line.
pub fn finalize(&mut self) -> Vec<Line<'static>> {
let remaining = self.finalize_text();
if remaining.is_empty() {
return Vec::new();
}
self.render_lines(&remaining)
}
/// Finalize the stream and return any remaining raw text.
pub fn finalize_text(&mut self) -> String {
self.is_streaming = false;
if self.buffer.is_empty() {
return String::new();
}
let remaining = self.buffer.clone();
self.buffer.clear();
self.committed_line_count = 0;
remaining
}
/// Get all rendered lines (for final display after stream ends)
pub fn all_lines(&self) -> Vec<Line<'static>> {
self.render_lines(&self.buffer)
}
/// Render content into styled lines
fn render_lines(&self, content: &str) -> Vec<Line<'static>> {
let width = self.width.unwrap_or(80);
let style = if self.is_thinking {
Style::default()
.fg(palette::STATUS_WARNING)
.add_modifier(Modifier::DIM | Modifier::ITALIC)
} else {
Style::default()
};
let mut lines = Vec::new();
for line in content.lines() {
// Wrap long lines
let wrapped = wrap_line(line, width);
for wrapped_line in wrapped {
lines.push(Line::from(Span::styled(wrapped_line, style)));
}
}
// Handle trailing newline (add empty line)
if content.ends_with('\n') {
lines.push(Line::from(""));
}
lines
}
/// Check if the stream is still active
pub fn is_streaming(&self) -> bool {
self.is_streaming
}
/// Get the raw buffer length
pub fn buffer_len(&self) -> usize {
self.buffer.len()
}
/// Clear the buffer
pub fn clear(&mut self) {
self.buffer.clear();
self.committed_line_count = 0;
}
}
/// Wrap a single line to fit within the given width
fn wrap_line(line: &str, width: usize) -> Vec<String> {
if line.is_empty() {
return vec![String::new()];
}
let mut result = Vec::new();
let mut current_line = String::new();
let mut current_width = 0;
for word in line.split_whitespace() {
let word_width = word.width();
if current_width == 0 {
// First word on line
current_line = word.to_string();
current_width = word_width;
} else if current_width + 1 + word_width <= width {
// Word fits with space
current_line.push(' ');
current_line.push_str(word);
current_width += 1 + word_width;
} else {
// Word doesn't fit, start new line
result.push(current_line);
current_line = word.to_string();
current_width = word_width;
}
}
if !current_line.is_empty() {
result.push(current_line);
}
if result.is_empty() {
vec![String::new()]
} else {
result
}
}
/// State for managing multiple stream collectors (one per content block)
#[derive(Debug, Clone, Default)]
pub struct StreamingState {
/// Collectors for each content block by index
collectors: Vec<Option<MarkdownStreamCollector>>,
/// Whether any stream is currently active
pub is_active: bool,
/// Accumulated text for display
pub accumulated_text: String,
/// Accumulated thinking for display
pub accumulated_thinking: String,
}
impl StreamingState {
/// Create a new streaming state
pub fn new() -> Self {
Self::default()
}
/// Start a new text block
pub fn start_text(&mut self, index: usize, width: Option<usize>) {
self.ensure_capacity(index);
self.collectors[index] = Some(MarkdownStreamCollector::new(width, false));
self.is_active = true;
}
/// Start a new thinking block
pub fn start_thinking(&mut self, index: usize, width: Option<usize>) {
self.ensure_capacity(index);
self.collectors[index] = Some(MarkdownStreamCollector::new(width, true));
self.is_active = true;
}
/// Push content to a block
pub fn push_content(&mut self, index: usize, content: &str) {
if let Some(Some(collector)) = self.collectors.get_mut(index) {
collector.push(content);
// Update accumulated text
if collector.is_thinking {
self.accumulated_thinking.push_str(content);
} else {
self.accumulated_text.push_str(content);
}
}
}
/// Get newly committed lines from a block
pub fn commit_lines(&mut self, index: usize) -> Vec<Line<'static>> {
if let Some(Some(collector)) = self.collectors.get_mut(index) {
collector.commit_complete_lines()
} else {
Vec::new()
}
}
/// Get newly committed raw text from a block.
pub fn commit_text(&mut self, index: usize) -> String {
if let Some(Some(collector)) = self.collectors.get_mut(index) {
collector.commit_complete_text()
} else {
String::new()
}
}
/// Finalize a block and get remaining lines
pub fn finalize_block(&mut self, index: usize) -> Vec<Line<'static>> {
if let Some(Some(collector)) = self.collectors.get_mut(index) {
let lines = collector.finalize();
// Check if all blocks are done
self.check_active();
lines
} else {
Vec::new()
}
}
/// Finalize a block and get remaining raw text.
pub fn finalize_block_text(&mut self, index: usize) -> String {
if let Some(Some(collector)) = self.collectors.get_mut(index) {
let text = collector.finalize_text();
self.check_active();
text
} else {
String::new()
}
}
/// Finalize all blocks
pub fn finalize_all(&mut self) -> Vec<(usize, Vec<Line<'static>>)> {
let mut result = Vec::new();
for (i, collector) in self.collectors.iter_mut().enumerate() {
if let Some(c) = collector {
let lines = c.finalize();
if !lines.is_empty() {
result.push((i, lines));
}
}
}
self.is_active = false;
result
}
/// Check if any stream is still active
fn check_active(&mut self) {
self.is_active = self.collectors.iter().any(|c| {
c.as_ref()
.is_some_and(MarkdownStreamCollector::is_streaming)
});
}
/// Ensure capacity for the given index
fn ensure_capacity(&mut self, index: usize) {
while self.collectors.len() <= index {
self.collectors.push(None);
}
}
/// Reset the streaming state
pub fn reset(&mut self) {
self.collectors.clear();
self.is_active = false;
self.accumulated_text.clear();
self.accumulated_thinking.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_commit_complete_lines() {
let mut collector = MarkdownStreamCollector::new(Some(80), false);
// Push incomplete line
collector.push("Hello ");
let lines = collector.commit_complete_lines();
assert!(lines.is_empty()); // No complete lines yet
// Complete the line
collector.push("World\n");
let lines = collector.commit_complete_lines();
assert_eq!(lines.len(), 2); // "Hello World" + empty line from trailing \n
// Push more content
collector.push("Second line");
let lines = collector.commit_complete_lines();
assert!(lines.is_empty()); // No new complete lines
// Finalize
let lines = collector.finalize();
assert_eq!(lines.len(), 1); // "Second line"
}
#[test]
fn test_wrap_line() {
let result = wrap_line("This is a long line that should be wrapped", 20);
assert!(result.len() > 1);
}
}