use std::time::Duration;
use std::time::Instant;
pub const DEFAULT_STREAM_COMMIT_INTERVAL: Duration = Duration::from_millis(16);
pub const CATCH_UP_QUEUE_DEPTH: usize = 160;
pub const CATCH_UP_OLDEST_AGE: Duration = Duration::from_millis(1_200);
#[derive(Debug, Clone)]
pub struct StreamDisplayClock {
interval: Duration,
pending: bool,
next_due_at: Option<Instant>,
last_commit_at: Option<Instant>,
commit_count: u64,
catch_up_count: u64,
allow_catch_up: bool,
}
impl Default for StreamDisplayClock {
fn default() -> Self {
Self::new(DEFAULT_STREAM_COMMIT_INTERVAL)
}
}
impl StreamDisplayClock {
pub fn new(interval: Duration) -> Self {
Self {
interval,
pending: false,
next_due_at: None,
last_commit_at: None,
commit_count: 0,
catch_up_count: 0,
allow_catch_up: true,
}
}
pub fn set_allow_catch_up(&mut self, allow: bool) {
self.allow_catch_up = allow;
}
pub fn note_delta(&mut self, now: Instant) {
self.note_delta_with_backlog(now, 1, None);
}
pub fn note_delta_with_backlog(
&mut self,
now: Instant,
queued: usize,
oldest_age: Option<Duration>,
) {
self.pending = true;
let catch_up = self.allow_catch_up
&& (queued >= CATCH_UP_QUEUE_DEPTH
|| oldest_age.is_some_and(|age| age >= CATCH_UP_OLDEST_AGE));
if catch_up {
self.next_due_at = Some(now);
self.catch_up_count = self.catch_up_count.saturating_add(1);
return;
}
if self.next_due_at.is_some() {
return;
}
self.next_due_at = Some(match self.last_commit_at {
Some(last) => last.checked_add(self.interval).unwrap_or(now).max(now),
None => now,
});
}
pub fn due_in(&self, now: Instant) -> Option<Duration> {
let due = self.next_due_at?;
Some(due.saturating_duration_since(now))
}
pub fn take_due(&mut self, now: Instant) -> bool {
if !self.pending {
self.next_due_at = None;
return false;
}
let Some(due) = self.next_due_at else {
return false;
};
if now < due {
return false;
}
self.pending = false;
self.next_due_at = None;
self.last_commit_at = Some(now);
self.commit_count = self.commit_count.saturating_add(1);
true
}
pub fn flush_now(&mut self, now: Instant) -> bool {
let had_pending = self.pending;
self.pending = false;
self.next_due_at = None;
if had_pending {
self.last_commit_at = Some(now);
self.commit_count = self.commit_count.saturating_add(1);
}
had_pending
}
pub fn reset(&mut self) {
self.pending = false;
self.next_due_at = None;
self.last_commit_at = None;
self.commit_count = 0;
self.catch_up_count = 0;
}
#[cfg(test)]
pub fn commit_count(&self) -> u64 {
self.commit_count
}
#[cfg(test)]
pub fn catch_up_count(&self) -> u64 {
self.catch_up_count
}
}
#[derive(Debug, Default, Clone)]
pub struct StreamBuffer {
pending: String,
}
impl StreamBuffer {
pub fn new() -> Self {
Self::default()
}
pub fn push_delta(&mut self, delta: &str) {
self.pending.push_str(delta);
}
pub fn has_pending(&self) -> bool {
!self.pending.is_empty()
}
pub fn take(&mut self) -> String {
std::mem::take(&mut self.pending)
}
}
#[derive(Debug, Default)]
struct BlockState {
is_thinking: bool,
is_streaming: bool,
buffer: StreamBuffer,
}
#[derive(Debug, Default)]
pub struct StreamingState {
blocks: Vec<Option<BlockState>>,
pub is_active: bool,
pub accumulated_text: String,
pub accumulated_thinking: String,
}
impl StreamingState {
pub fn new() -> Self {
Self::default()
}
pub fn start_text(&mut self, index: usize) {
self.start_block(index, false);
}
pub fn start_thinking(&mut self, index: usize) {
self.start_block(index, true);
}
fn start_block(&mut self, index: usize, is_thinking: bool) {
self.ensure_capacity(index);
self.blocks[index] = Some(BlockState {
is_thinking,
is_streaming: true,
buffer: StreamBuffer::new(),
});
self.is_active = true;
}
pub fn push_content(&mut self, index: usize, content: &str) {
if let Some(Some(block)) = self.blocks.get_mut(index) {
if block.is_thinking {
self.accumulated_thinking.push_str(content);
} else {
self.accumulated_text.push_str(content);
}
block.buffer.push_delta(content);
}
}
pub fn commit_text(&mut self, index: usize) -> String {
match self.blocks.get_mut(index) {
Some(Some(block)) => block.buffer.take(),
_ => String::new(),
}
}
pub fn has_pending_stream_text(&self, index: usize) -> bool {
self.blocks
.get(index)
.and_then(|b| b.as_ref())
.is_some_and(|b| b.buffer.has_pending())
}
pub fn finalize_block_text(&mut self, index: usize) -> String {
let out = match self.blocks.get_mut(index) {
Some(Some(block)) => {
block.is_streaming = false;
block.buffer.take()
}
_ => return String::new(),
};
self.check_active();
out
}
fn check_active(&mut self) {
self.is_active = self.blocks.iter().flatten().any(|b| b.is_streaming);
}
fn ensure_capacity(&mut self, index: usize) {
while self.blocks.len() <= index {
self.blocks.push(None);
}
}
pub fn reset(&mut self) {
self.blocks.clear();
self.is_active = false;
self.accumulated_text.clear();
self.accumulated_thinking.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn assistant_text_streams_before_newline() {
let mut state = StreamingState::new();
state.start_text(0);
state.push_content(0, "hello world");
assert_eq!(state.commit_text(0), "hello world");
assert!(!state.has_pending_stream_text(0));
}
#[test]
fn thinking_text_streams_before_newline() {
let mut state = StreamingState::new();
state.start_thinking(0);
state.push_content(0, "thinking deeply");
assert_eq!(state.commit_text(0), "thinking deeply");
assert!(!state.has_pending_stream_text(0));
}
#[test]
fn commit_beat_drains_everything_received_since_the_previous_beat() {
let mut state = StreamingState::new();
state.start_text(0);
let burst = "abcdefghijklmnopqrstuvwxyz".repeat(8);
state.push_content(0, &burst);
assert_eq!(state.commit_text(0), burst);
assert_eq!(state.commit_text(0), "");
}
#[test]
fn combining_marks_stay_with_their_base_letter() {
let mut state = StreamingState::new();
state.start_text(0);
state.push_content(0, "e\u{301}x");
assert_eq!(state.commit_text(0), "e\u{301}x");
}
#[test]
fn finalize_drains_partial_tail() {
let mut state = StreamingState::new();
state.start_text(0);
state.push_content(0, "done\nno-newline-here");
assert_eq!(state.finalize_block_text(0), "done\nno-newline-here");
assert!(!state.is_active);
}
#[test]
fn finalize_after_commit_has_nothing_left() {
let mut state = StreamingState::new();
state.start_text(0);
state.push_content(0, "abc");
assert_eq!(state.commit_text(0), "abc");
assert_eq!(state.finalize_block_text(0), "");
}
#[test]
fn accumulators_track_raw_stream_by_block_kind() {
let mut state = StreamingState::new();
state.start_thinking(0);
state.push_content(0, "reasoning");
state.start_text(1);
state.push_content(1, "answer");
assert_eq!(state.accumulated_thinking, "reasoning");
assert_eq!(state.accumulated_text, "answer");
}
#[test]
fn bursty_stream_state_has_no_text_loss_after_coalesced_flushes() {
let mut state = StreamingState::new();
state.start_text(0);
let mut expected = String::new();
for idx in 0..250 {
let chunk = format!("{idx}.");
expected.push_str(&chunk);
state.push_content(0, &chunk);
}
let first_flush = state.commit_text(0);
assert_eq!(first_flush, expected);
assert_eq!(state.finalize_block_text(0), "");
}
#[test]
fn stream_display_clock_coalesces_bursty_tiny_deltas() {
let interval = Duration::from_millis(33);
let mut clock = StreamDisplayClock::new(interval);
let t0 = Instant::now();
for _ in 0..100 {
clock.note_delta(t0);
}
assert_eq!(clock.due_in(t0), Some(Duration::ZERO));
assert!(clock.take_due(t0));
assert_eq!(clock.commit_count(), 1);
for _ in 0..25 {
clock.note_delta(t0 + Duration::from_millis(5));
}
assert!(!clock.take_due(t0 + Duration::from_millis(5)));
assert_eq!(
clock.due_in(t0 + Duration::from_millis(5)),
Some(Duration::from_millis(28))
);
assert!(clock.take_due(t0 + interval));
assert_eq!(clock.commit_count(), 2);
}
#[test]
fn stream_display_clock_bounds_long_reasoning_commit_count() {
let interval = Duration::from_millis(33);
let mut clock = StreamDisplayClock::new(interval);
let t0 = Instant::now();
let mut commits = 0u64;
for millis in 0..300 {
let now = t0 + Duration::from_millis(millis);
clock.note_delta(now);
if clock.take_due(now) {
commits += 1;
}
}
assert!(commits > 1, "long streams should keep advancing visibly");
assert!(
commits <= 11,
"300 one-ms deltas should not commit on provider cadence: {commits}"
);
assert_eq!(commits, clock.commit_count());
}
#[test]
fn stream_display_clock_final_flush_consumes_pending_delta() {
let mut clock = StreamDisplayClock::new(Duration::from_millis(33));
let t0 = Instant::now();
clock.note_delta(t0);
assert!(clock.take_due(t0));
clock.note_delta(t0 + Duration::from_millis(4));
assert!(!clock.take_due(t0 + Duration::from_millis(4)));
assert!(clock.flush_now(t0 + Duration::from_millis(5)));
assert_eq!(clock.due_in(t0 + Duration::from_millis(5)), None);
assert!(!clock.take_due(t0 + Duration::from_millis(33)));
assert_eq!(clock.commit_count(), 2);
}
#[test]
fn normal_clock_catch_up_only_when_backlog_crosses_threshold() {
let interval = Duration::from_millis(33);
let mut clock = StreamDisplayClock::new(interval);
let t0 = Instant::now();
clock.note_delta(t0);
assert!(clock.take_due(t0));
clock.note_delta_with_backlog(
t0 + Duration::from_millis(1),
3,
Some(Duration::from_millis(5)),
);
assert!(!clock.take_due(t0 + Duration::from_millis(1)));
assert_eq!(clock.catch_up_count(), 0);
clock.note_delta_with_backlog(
t0 + Duration::from_millis(2),
CATCH_UP_QUEUE_DEPTH,
Some(Duration::from_millis(10)),
);
assert!(clock.take_due(t0 + Duration::from_millis(2)));
assert!(clock.catch_up_count() >= 1);
}
#[test]
fn reduced_motion_keeps_steady_clock_without_catch_up_or_typewriter() {
let interval = Duration::from_millis(33);
let mut clock = StreamDisplayClock::new(interval);
clock.set_allow_catch_up(false);
let t0 = Instant::now();
clock.note_delta(t0);
assert!(clock.take_due(t0));
clock.note_delta_with_backlog(
t0 + Duration::from_millis(1),
CATCH_UP_QUEUE_DEPTH * 2,
Some(CATCH_UP_OLDEST_AGE),
);
assert!(!clock.take_due(t0 + Duration::from_millis(1)));
assert_eq!(clock.catch_up_count(), 0);
assert_eq!(
clock.due_in(t0 + Duration::from_millis(1)),
Some(Duration::from_millis(32))
);
assert!(clock.take_due(t0 + interval));
assert_eq!(clock.commit_count(), 2);
}
}