返回 DeepSeek-TUI-2026
mod.rs
根目录 / crates / tui / src / tui / streaming / mod.rs
1 #![allow(dead_code)]
2
3 //! Markdown stream collector for live micro-chunk rendering.
4 //!
5 //! This module implements the pattern from codex-rs where:
6 //! - Streaming text is split into small grapheme-aligned chunks
7 //! - Commit ticks drip chunks into the transcript between provider deltas
8 //! - Final content is emitted when the stream ends
9
10 use ratatui::style::{Modifier, Style};
11 use ratatui::text::{Line, Span};
12 use std::time::Instant;
13 use unicode_width::UnicodeWidthStr;
14
15 use crate::palette;
16
17 pub mod chunking;
18 pub mod commit_tick;
19 pub mod line_buffer;
20
21 pub use chunking::{AdaptiveChunkingPolicy, ChunkingMode};
22 pub use commit_tick::{StreamChunker, run_commit_tick};
23 pub use line_buffer::LineBuffer;
24 /// Collects streaming text and commits complete lines.
25 #[derive(Debug, Clone)]
26 pub struct MarkdownStreamCollector {
27 /// Buffer for incoming text
28 buffer: String,
29 /// Number of lines already committed
30 committed_line_count: usize,
31 /// Terminal width for wrapping
32 width: Option<usize>,
33 /// Whether the stream is still active
34 is_streaming: bool,
35 /// Whether this is a thinking block
36 is_thinking: bool,
37 }
38
39 impl Default for MarkdownStreamCollector {
40 fn default() -> Self {
41 // `is_streaming: true` matches `MarkdownStreamCollector::new` so a
42 // freshly-default block behaves like a freshly-started stream.
43 Self::new(None, false)
44 }
45 }
46
47 impl MarkdownStreamCollector {
48 /// Create a new collector
49 pub fn new(width: Option<usize>, is_thinking: bool) -> Self {
50 Self {
51 buffer: String::new(),
52 committed_line_count: 0,
53 width,
54 is_streaming: true,
55 is_thinking,
56 }
57 }
58
59 /// Push new content to the buffer
60 pub fn push(&mut self, content: &str) {
61 self.buffer.push_str(content);
62 }
63
64 /// Get the current buffer content (for display during streaming)
65 pub fn current_content(&self) -> &str {
66 &self.buffer
67 }
68
69 /// Check if there are complete lines to commit
70 pub fn has_complete_lines(&self) -> bool {
71 self.buffer.contains('\n')
72 }
73
74 /// Commit complete lines and return them.
75 /// Only lines ending with '\n' are committed.
76 /// Returns the newly committed lines since last call.
77 pub fn commit_complete_lines(&mut self) -> Vec<Line<'static>> {
78 let committed = self.commit_complete_text();
79 if committed.is_empty() {
80 return Vec::new();
81 }
82 self.render_lines(&committed)
83 }
84
85 /// Commit complete text chunks ending in a newline.
86 /// Returns the raw text that became visible since the last call.
87 pub fn commit_complete_text(&mut self) -> String {
88 if self.buffer.is_empty() {
89 return String::new();
90 }
91
92 // Find the last newline - only process up to there
93 let Some(last_newline_idx) = self.buffer.rfind('\n') else {
94 return String::new(); // No complete lines yet
95 };
96
97 // Extract the complete portion (up to and including last newline)
98 let complete_portion = self.buffer[..=last_newline_idx].to_string();
99
100 // Remove the committed portion from the buffer so finalize only emits the remainder
101 self.buffer = self.buffer[last_newline_idx + 1..].to_string();
102 self.committed_line_count = 0;
103
104 complete_portion
105 }
106
107 /// Finalize the stream and return any remaining content.
108 /// Call this when the stream ends to emit the final incomplete line.
109 pub fn finalize(&mut self) -> Vec<Line<'static>> {
110 let remaining = self.finalize_text();
111 if remaining.is_empty() {
112 return Vec::new();
113 }
114 self.render_lines(&remaining)
115 }
116
117 /// Finalize the stream and return any remaining raw text.
118 pub fn finalize_text(&mut self) -> String {
119 self.is_streaming = false;
120
121 if self.buffer.is_empty() {
122 return String::new();
123 }
124
125 let remaining = self.buffer.clone();
126 self.buffer.clear();
127 self.committed_line_count = 0;
128 remaining
129 }
130
131 /// Get all rendered lines (for final display after stream ends)
132 pub fn all_lines(&self) -> Vec<Line<'static>> {
133 self.render_lines(&self.buffer)
134 }
135
136 /// Render content into styled lines
137 fn render_lines(&self, content: &str) -> Vec<Line<'static>> {
138 let width = self.width.unwrap_or(80);
139 let style = if self.is_thinking {
140 Style::default()
141 .fg(palette::STATUS_WARNING)
142 .add_modifier(Modifier::DIM | Modifier::ITALIC)
143 } else {
144 Style::default()
145 };
146
147 let mut lines = Vec::new();
148
149 for line in content.lines() {
150 // Wrap long lines
151 let wrapped = wrap_line(line, width);
152 for wrapped_line in wrapped {
153 lines.push(Line::from(Span::styled(wrapped_line, style)));
154 }
155 }
156
157 // Handle trailing newline (add empty line)
158 if content.ends_with('\n') {
159 lines.push(Line::from(""));
160 }
161
162 lines
163 }
164
165 /// Check if the stream is still active
166 pub fn is_streaming(&self) -> bool {
167 self.is_streaming
168 }
169
170 /// Get the raw buffer length
171 pub fn buffer_len(&self) -> usize {
172 self.buffer.len()
173 }
174
175 /// Clear the buffer
176 pub fn clear(&mut self) {
177 self.buffer.clear();
178 self.committed_line_count = 0;
179 }
180 }
181
182 /// Wrap a single line to fit within the given width
183 fn wrap_line(line: &str, width: usize) -> Vec<String> {
184 if line.is_empty() {
185 return vec![String::new()];
186 }
187
188 let mut result = Vec::new();
189 let mut current_line = String::new();
190 let mut current_width = 0;
191
192 for word in line.split_whitespace() {
193 let word_width = word.width();
194
195 if current_width == 0 {
196 // First word on line
197 current_line = word.to_string();
198 current_width = word_width;
199 } else if current_width + 1 + word_width <= width {
200 // Word fits with space
201 current_line.push(' ');
202 current_line.push_str(word);
203 current_width += 1 + word_width;
204 } else {
205 // Word doesn't fit, start new line
206 result.push(current_line);
207 current_line = word.to_string();
208 current_width = word_width;
209 }
210 }
211
212 if !current_line.is_empty() {
213 result.push(current_line);
214 }
215
216 if result.is_empty() {
217 vec![String::new()]
218 } else {
219 result
220 }
221 }
222
223 /// Per-block streaming substate: optional line-buffer feeding a collector +
224 /// chunker/policy for two-gear pacing.
225 ///
226 /// Pipeline:
227 /// ```text
228 /// raw delta -> LineBuffer.push -> take_committable -> collector + chunker -> commit tick
229 /// ```
230 ///
231 /// The [`LineBuffer`] remains available for line-sensitive modes. Normal
232 /// assistant prose and thinking blocks bypass it so text can stream in live
233 /// micro-chunks instead of waiting for newline boundaries.
234 #[derive(Debug, Default)]
235 struct BlockState {
236 /// Newline gate: holds back trailing partial-line text between deltas.
237 /// Bypassed when `bypass_gate` is true (thinking blocks).
238 line_buffer: LineBuffer,
239 /// Whether to bypass the [`LineBuffer`] (thinking blocks stream live).
240 bypass_gate: bool,
241 collector: MarkdownStreamCollector,
242 chunker: StreamChunker,
243 policy: AdaptiveChunkingPolicy,
244 }
245
246 /// State for managing multiple stream collectors (one per content block)
247 #[derive(Debug, Default)]
248 pub struct StreamingState {
249 /// Per-block state by index (collector + chunker + policy).
250 blocks: Vec<Option<BlockState>>,
251 /// Whether any stream is currently active
252 pub is_active: bool,
253 /// Accumulated text for display
254 pub accumulated_text: String,
255 /// Accumulated thinking for display
256 pub accumulated_thinking: String,
257 }
258
259 impl StreamingState {
260 /// Create a new streaming state
261 pub fn new() -> Self {
262 Self::default()
263 }
264
265 /// Start a new text block. Assistant prose streams live in micro-chunks so
266 /// users can visually track the answer as it forms instead of waiting for
267 /// a newline-terminated line.
268 pub fn start_text(&mut self, index: usize, width: Option<usize>) {
269 self.ensure_capacity(index);
270 self.blocks[index] = Some(BlockState {
271 line_buffer: LineBuffer::new(),
272 bypass_gate: true,
273 collector: MarkdownStreamCollector::new(width, false),
274 chunker: StreamChunker::new(),
275 policy: AdaptiveChunkingPolicy::new(),
276 });
277 self.is_active = true;
278 }
279
280 /// Start a new thinking block. Thinking deltas bypass the newline gate so
281 /// they remain visually live — long reasoning often arrives as a single
282 /// paragraph without intermediate newlines, and gating it would create
283 /// long pauses where the user sees nothing.
284 pub fn start_thinking(&mut self, index: usize, width: Option<usize>) {
285 self.ensure_capacity(index);
286 self.blocks[index] = Some(BlockState {
287 line_buffer: LineBuffer::new(),
288 bypass_gate: true,
289 collector: MarkdownStreamCollector::new(width, true),
290 chunker: StreamChunker::new(),
291 policy: AdaptiveChunkingPolicy::new(),
292 });
293 self.is_active = true;
294 }
295
296 /// Push content to a block. Routing depends on the block kind:
297 ///
298 /// - Assistant text blocks: incoming bytes normally bypass [`LineBuffer`]
299 /// and are split into small display chunks downstream.
300 /// - Thinking blocks: bytes bypass the gate and go straight to the
301 /// collector/chunker so reasoning stays visually live (long thoughts
302 /// often have no intermediate newlines).
303 ///
304 /// `accumulated_text` / `accumulated_thinking` always track the full raw
305 /// stream so callers building API messages or doing retries see exactly
306 /// what the model emitted, regardless of UI gating.
307 pub fn push_content(&mut self, index: usize, content: &str) {
308 if let Some(Some(block)) = self.blocks.get_mut(index) {
309 // Always update the raw accumulator first — UI gating must not
310 // affect what we send back to the model on retry/continuation.
311 if block.collector.is_thinking {
312 self.accumulated_thinking.push_str(content);
313 } else {
314 self.accumulated_text.push_str(content);
315 }
316
317 // Determine what bytes are safe to expose downstream on this push.
318 let downstream: String = if block.bypass_gate {
319 // Thinking: forward verbatim to collector + chunker.
320 content.to_string()
321 } else {
322 // Assistant text: gate at the last-newline boundary.
323 block.line_buffer.push(content);
324 block.line_buffer.take_committable()
325 };
326
327 if downstream.is_empty() {
328 return;
329 }
330
331 if block.bypass_gate {
332 block.chunker.push_delta(&downstream);
333 } else {
334 block.collector.push(&downstream);
335 let committed = block.collector.commit_complete_text();
336 if !committed.is_empty() {
337 block.chunker.push_delta(&committed);
338 }
339 }
340 }
341 }
342
343 /// Get newly committed lines from a block. (Legacy entry point that maps
344 /// onto the chunker.)
345 pub fn commit_lines(&mut self, index: usize) -> Vec<Line<'static>> {
346 let text = self.commit_text(index);
347 if text.is_empty() {
348 return Vec::new();
349 }
350 // Re-render the text through the same path the collector used.
351 let style = if self
352 .blocks
353 .get(index)
354 .and_then(|b| b.as_ref())
355 .is_some_and(|b| b.collector.is_thinking)
356 {
357 Style::default()
358 .fg(palette::STATUS_WARNING)
359 .add_modifier(Modifier::DIM | Modifier::ITALIC)
360 } else {
361 Style::default()
362 };
363 let mut lines = Vec::new();
364 for line in text.lines() {
365 lines.push(Line::from(Span::styled(line.to_string(), style)));
366 }
367 if text.ends_with('\n') {
368 lines.push(Line::from(""));
369 }
370 lines
371 }
372
373 /// Run one commit-tick of the chunker policy and return any text safe to
374 /// flush to the transcript on this tick. May be empty (Smooth-mode tick
375 /// against an empty queue) or contain anywhere from one line up to the
376 /// full backlog (CatchUp-mode burst drain).
377 pub fn commit_text(&mut self, index: usize) -> String {
378 if let Some(Some(block)) = self.blocks.get_mut(index) {
379 let now = Instant::now();
380 let out = run_commit_tick(&mut block.policy, &mut block.chunker, now);
381 out.committed_text
382 } else {
383 String::new()
384 }
385 }
386
387 /// Inspect the current chunking mode for a block (testing/observability).
388 pub fn chunking_mode(&self, index: usize) -> Option<ChunkingMode> {
389 self.blocks
390 .get(index)
391 .and_then(|b| b.as_ref())
392 .map(|b| b.policy.mode())
393 }
394
395 /// Whether the chunker has queued content waiting to be flushed by the
396 /// next commit tick. Useful for callers that want to drive an extra tick
397 /// while the queue drains under Smooth-mode pacing.
398 pub fn has_pending_chunker_lines(&self, index: usize) -> bool {
399 self.blocks
400 .get(index)
401 .and_then(|b| b.as_ref())
402 .is_some_and(|b| b.chunker.queued_lines() > 0)
403 }
404
405 /// Finalize a block and get remaining lines
406 pub fn finalize_block(&mut self, index: usize) -> Vec<Line<'static>> {
407 let text = self.finalize_block_text(index);
408 if text.is_empty() {
409 return Vec::new();
410 }
411 let style = if self
412 .blocks
413 .get(index)
414 .and_then(|b| b.as_ref())
415 .is_some_and(|b| b.collector.is_thinking)
416 {
417 Style::default()
418 .fg(palette::STATUS_WARNING)
419 .add_modifier(Modifier::DIM | Modifier::ITALIC)
420 } else {
421 Style::default()
422 };
423 let mut lines = Vec::new();
424 for line in text.lines() {
425 lines.push(Line::from(Span::styled(line.to_string(), style)));
426 }
427 if text.ends_with('\n') {
428 lines.push(Line::from(""));
429 }
430 lines
431 }
432
433 /// Finalize a block and get remaining raw text. Drains the full pipeline
434 /// in upstream-to-downstream order:
435 ///
436 /// 1. [`LineBuffer::flush`] returns any post-newline tail held by the gate.
437 /// For gated blocks this is critical — without it, a final partial
438 /// line (e.g. text the model emitted without a trailing newline before
439 /// the turn ended) would otherwise be stranded in the gate.
440 /// 2. The collector's `finalize_text` releases any partial line it still
441 /// holds (relevant for the bypass path where the collector receives
442 /// raw deltas directly).
443 /// 3. The chunker's `drain_remaining` releases queued whole-line text
444 /// that the policy hadn't yet committed.
445 pub fn finalize_block_text(&mut self, index: usize) -> String {
446 if let Some(Some(block)) = self.blocks.get_mut(index) {
447 // Flush the gate first so any held tail rejoins the stream
448 // before the collector/chunker drain. For thinking blocks the
449 // gate is unused, so this is a no-op.
450 let gate_tail = block.line_buffer.flush();
451 if !gate_tail.is_empty() {
452 block.collector.push(&gate_tail);
453 }
454 // Any newly committable text after the gate flush feeds the
455 // chunker so drain order remains "queued-lines, then partial-tail".
456 let post_flush = block.collector.commit_complete_text();
457 if !post_flush.is_empty() {
458 block.chunker.push_delta(&post_flush);
459 }
460 // Any unterminated tail still in the collector is returned raw.
461 let tail = block.collector.finalize_text();
462 // Any whole-line text held by the chunker is safe to emit now.
463 let mut out = block.chunker.drain_remaining();
464 if !tail.is_empty() {
465 out.push_str(&tail);
466 }
467 self.check_active();
468 out
469 } else {
470 String::new()
471 }
472 }
473
474 /// Finalize all blocks
475 pub fn finalize_all(&mut self) -> Vec<(usize, Vec<Line<'static>>)> {
476 let mut result = Vec::new();
477 let len = self.blocks.len();
478 for i in 0..len {
479 let lines = self.finalize_block(i);
480 if !lines.is_empty() {
481 result.push((i, lines));
482 }
483 }
484 self.is_active = false;
485 result
486 }
487
488 /// Propagate the low-motion flag to every block's chunking policy.
489 /// When true, all policies stay in `Smooth` regardless of queue pressure,
490 /// preventing CatchUp burst drains that would create sudden visual jumps.
491 pub fn set_low_motion(&mut self, low_motion: bool) {
492 for block in self.blocks.iter_mut().flatten() {
493 block.policy.set_low_motion(low_motion);
494 }
495 }
496
497 /// Check if any stream is still active
498 fn check_active(&mut self) {
499 self.is_active = self.blocks.iter().any(|b| {
500 b.as_ref()
501 .is_some_and(|state| state.collector.is_streaming())
502 });
503 }
504
505 /// Ensure capacity for the given index
506 fn ensure_capacity(&mut self, index: usize) {
507 while self.blocks.len() <= index {
508 self.blocks.push(None);
509 }
510 }
511
512 /// Reset the streaming state
513 pub fn reset(&mut self) {
514 self.blocks.clear();
515 self.is_active = false;
516 self.accumulated_text.clear();
517 self.accumulated_thinking.clear();
518 }
519 }
520
521 #[cfg(test)]
522 mod tests {
523 use super::*;
524
525 #[test]
526 fn test_commit_complete_lines() {
527 let mut collector = MarkdownStreamCollector::new(Some(80), false);
528
529 // Push incomplete line
530 collector.push("Hello ");
531 let lines = collector.commit_complete_lines();
532 assert!(lines.is_empty()); // No complete lines yet
533
534 // Complete the line
535 collector.push("World\n");
536 let lines = collector.commit_complete_lines();
537 assert_eq!(lines.len(), 2); // "Hello World" + empty line from trailing \n
538
539 // Push more content
540 collector.push("Second line");
541 let lines = collector.commit_complete_lines();
542 assert!(lines.is_empty()); // No new complete lines
543
544 // Finalize
545 let lines = collector.finalize();
546 assert_eq!(lines.len(), 1); // "Second line"
547 }
548
549 #[test]
550 fn test_wrap_line() {
551 let result = wrap_line("This is a long line that should be wrapped", 20);
552 assert!(result.len() > 1);
553 }
554
555 #[test]
556 fn assistant_text_streams_before_newline() {
557 let mut state = StreamingState::new();
558 state.start_text(0, None);
559 state.push_content(0, "hello world");
560
561 assert_eq!(state.commit_text(0), "h");
562 assert_eq!(state.commit_text(0), "e");
563 assert!(state.has_pending_chunker_lines(0));
564 }
565
566 #[test]
567 fn thinking_text_streams_before_newline() {
568 let mut state = StreamingState::new();
569 state.start_thinking(0, None);
570 state.push_content(0, "thinking deeply");
571
572 assert_eq!(state.commit_text(0), "t");
573 assert_eq!(state.commit_text(0), "h");
574 assert!(state.has_pending_chunker_lines(0));
575 }
576
577 #[test]
578 fn finalize_preserves_uncommitted_micro_chunks() {
579 let mut state = StreamingState::new();
580 state.start_text(0, None);
581 state.push_content(0, "abc");
582 assert_eq!(state.commit_text(0), "a");
583
584 assert_eq!(state.finalize_block_text(0), "bc");
585 }
586 }
587
587 lines RUST