use crate::ui::tui::{InlineMessageKind, InlineSegment};
#[derive(Clone, Debug)]
pub struct StreamConfig {
pub batch_size: usize,
pub max_buffer_bytes: usize,
}
impl Default for StreamConfig {
fn default() -> Self {
Self {
batch_size: 20, max_buffer_bytes: 65536, }
}
}
#[derive(Debug)]
pub struct StreamBuffer {
lines: Vec<Vec<InlineSegment>>,
config: StreamConfig,
approximate_size: usize,
}
impl StreamBuffer {
pub fn new() -> Self {
Self::with_config(StreamConfig::default())
}
pub fn with_config(config: StreamConfig) -> Self {
Self {
lines: Vec::with_capacity(config.batch_size),
config,
approximate_size: 0,
}
}
pub fn append_line(&mut self, segments: Vec<InlineSegment>) -> bool {
let line_size: usize = segments.iter().map(|s| s.text.len()).sum();
self.approximate_size += line_size;
self.lines.push(segments);
self.should_flush()
}
fn should_flush(&self) -> bool {
self.lines.len() >= self.config.batch_size
|| self.approximate_size >= self.config.max_buffer_bytes
}
pub fn flush(&mut self) -> Vec<Vec<InlineSegment>> {
self.approximate_size = 0;
std::mem::take(&mut self.lines)
}
pub fn len(&self) -> usize {
self.lines.len()
}
pub fn is_empty(&self) -> bool {
self.lines.is_empty()
}
pub fn approximate_bytes(&self) -> usize {
self.approximate_size
}
pub fn force_flush(&mut self) -> Vec<Vec<InlineSegment>> {
self.flush()
}
pub fn clear(&mut self) {
self.lines.clear();
self.approximate_size = 0;
}
}
impl Default for StreamBuffer {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct StreamingContext {
pub kind: InlineMessageKind,
pub buffer: StreamBuffer,
pub total_lines: usize,
}
impl StreamingContext {
pub fn new(kind: InlineMessageKind) -> Self {
Self {
kind,
buffer: StreamBuffer::new(),
total_lines: 0,
}
}
pub fn with_config(kind: InlineMessageKind, config: StreamConfig) -> Self {
Self {
kind,
buffer: StreamBuffer::with_config(config),
total_lines: 0,
}
}
pub fn append(&mut self, segments: Vec<InlineSegment>) -> bool {
let should_flush = self.buffer.append_line(segments);
self.total_lines += 1;
should_flush
}
pub fn flush(&mut self) -> Vec<Vec<InlineSegment>> {
self.buffer.flush()
}
}
pub struct AllocationPredictor {
bytes_per_line: usize,
}
impl AllocationPredictor {
pub fn new() -> Self {
Self {
bytes_per_line: 120, }
}
pub fn estimate_total_bytes(&self, line_count: usize) -> usize {
line_count * self.bytes_per_line
}
pub fn optimal_batch_size(&self, _total_bytes: usize) -> usize {
let target_batch_bytes = 8192;
let batch_lines = (target_batch_bytes / self.bytes_per_line).max(5);
batch_lines.min(50) }
pub fn pre_allocation_capacity(&self, estimated_lines: usize) -> usize {
(estimated_lines as f64 * 1.2) as usize }
}
impl Default for AllocationPredictor {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_stream_buffer_creation() {
let buffer = StreamBuffer::new();
assert!(buffer.is_empty());
assert_eq!(buffer.len(), 0);
}
#[test]
fn test_stream_buffer_append() {
let mut buffer = StreamBuffer::new();
let segment = InlineSegment {
text: "test".to_string(),
style: std::sync::Arc::new(Default::default()),
};
let should_flush = buffer.append_line(vec![segment]);
assert!(!should_flush); assert_eq!(buffer.len(), 1);
}
#[test]
fn test_stream_buffer_batch_flush() {
let mut buffer = StreamBuffer::with_config(StreamConfig {
batch_size: 5,
max_buffer_bytes: usize::MAX,
});
for i in 0..5 {
let segment = InlineSegment {
text: format!("line {}", i),
style: std::sync::Arc::new(Default::default()),
};
let should_flush = buffer.append_line(vec![segment]);
if i < 4 {
assert!(!should_flush);
} else {
assert!(should_flush);
}
}
assert_eq!(buffer.len(), 5);
}
#[test]
fn test_stream_buffer_byte_limit_flush() {
let mut buffer = StreamBuffer::with_config(StreamConfig {
batch_size: 100,
max_buffer_bytes: 50,
});
let segment = InlineSegment {
text: "x".repeat(60),
style: std::sync::Arc::new(Default::default()),
};
let should_flush = buffer.append_line(vec![segment]);
assert!(should_flush);
}
#[test]
fn test_stream_buffer_flush_returns_lines() {
let mut buffer = StreamBuffer::new();
let segment = InlineSegment {
text: "test".to_string(),
style: std::sync::Arc::new(Default::default()),
};
buffer.append_line(vec![segment]);
let flushed = buffer.flush();
assert_eq!(flushed.len(), 1);
assert!(buffer.is_empty());
}
#[test]
fn test_streaming_context() {
let mut ctx = StreamingContext::new(InlineMessageKind::Agent);
assert_eq!(ctx.total_lines, 0);
let segment = InlineSegment {
text: "test".to_string(),
style: std::sync::Arc::new(Default::default()),
};
ctx.append(vec![segment]);
assert_eq!(ctx.total_lines, 1);
}
#[test]
fn test_allocation_predictor() {
let predictor = AllocationPredictor::new();
let estimate = predictor.estimate_total_bytes(100);
assert!(estimate > 0);
let batch = predictor.optimal_batch_size(10000);
assert!(batch > 0 && batch <= 50);
}
#[test]
fn test_stream_config_defaults() {
let config = StreamConfig::default();
assert_eq!(config.batch_size, 20);
assert_eq!(config.max_buffer_bytes, 65536);
}
#[test]
fn test_pre_allocation_capacity() {
let predictor = AllocationPredictor::new();
let capacity = predictor.pre_allocation_capacity(100);
assert!(capacity >= 100); assert!(capacity <= 120); }
}