Douwe Osinga
672 lines
21 KiB
Rust
672 lines
21 KiB
Rust
//! Streaming markdown buffer for safe incremental rendering.
|
|
//!
|
|
//! This module provides a buffer that accumulates streaming markdown chunks
|
|
//! and determines safe points to flush content for rendering. It tracks
|
|
//! open markdown constructs (code blocks, bold, links, etc.) to ensure
|
|
//! we only output complete, well-formed markdown.
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```
|
|
//! use goose_cli::session::streaming_buffer::MarkdownBuffer;
|
|
//!
|
|
//! let mut buf = MarkdownBuffer::new();
|
|
//!
|
|
//! // Partial bold - buffers until closed
|
|
//! assert_eq!(buf.push("Hello **wor"), Some("Hello ".to_string()));
|
|
//! assert_eq!(buf.push("ld**!"), Some("**world**!".to_string()));
|
|
//!
|
|
//! // At end of stream, flush remaining content
|
|
//! let remaining = buf.flush();
|
|
//! ```
|
|
|
|
use regex::Regex;
|
|
use std::io::Write;
|
|
use std::sync::LazyLock;
|
|
|
|
const MAX_CODE_BLOCK_LINES: usize = 50;
|
|
const TRUNCATED_SHOW_LINES: usize = 20;
|
|
|
|
fn truncate_code_blocks(content: &str) -> String {
|
|
let (open_pos, fence) = match (content.find("```"), content.find("~~~")) {
|
|
(Some(a), Some(b)) if a <= b => (a, "```"),
|
|
(Some(a), None) => (a, "```"),
|
|
(None, Some(b)) => (b, "~~~"),
|
|
(Some(_), Some(b)) => (b, "~~~"),
|
|
(None, None) => return content.to_string(),
|
|
};
|
|
|
|
let Some(after_open) = content.get(open_pos + 3..) else {
|
|
return content.to_string();
|
|
};
|
|
let Some(newline_pos) = after_open.find('\n') else {
|
|
return content.to_string();
|
|
};
|
|
let code_start = open_pos + 3 + newline_pos + 1;
|
|
|
|
let Some(code_region) = content.get(code_start..) else {
|
|
return content.to_string();
|
|
};
|
|
let close_pattern = format!("\n{}", fence);
|
|
let Some(close_offset) = code_region.find(&close_pattern) else {
|
|
return content.to_string();
|
|
};
|
|
|
|
let Some(code_content) = code_region.get(..close_offset) else {
|
|
return content.to_string();
|
|
};
|
|
let lines: Vec<&str> = code_content.lines().collect();
|
|
|
|
if lines.len() <= MAX_CODE_BLOCK_LINES {
|
|
return content.to_string();
|
|
}
|
|
|
|
let truncated: String = lines
|
|
.iter()
|
|
.take(TRUNCATED_SHOW_LINES)
|
|
.copied()
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
let remaining = lines.len() - TRUNCATED_SHOW_LINES;
|
|
|
|
let file_msg = save_to_temp_file(code_content)
|
|
.map(|p| format!(" → {}", p))
|
|
.unwrap_or_default();
|
|
|
|
let close_pos = code_start + close_offset + 1; // +1 to include the \n
|
|
let prefix = content.get(..code_start).unwrap_or("");
|
|
let suffix = content.get(close_pos..).unwrap_or("");
|
|
format!(
|
|
"{}{}\n... ({} more lines{})\n{}",
|
|
prefix, truncated, remaining, file_msg, suffix
|
|
)
|
|
}
|
|
|
|
fn save_to_temp_file(content: &str) -> Option<String> {
|
|
let mut file = tempfile::Builder::new()
|
|
.prefix("goose-")
|
|
.suffix(".txt")
|
|
.tempfile()
|
|
.ok()?;
|
|
|
|
file.write_all(content.as_bytes()).ok()?;
|
|
|
|
// Keep the file (don't delete on drop) and get the path
|
|
let (_, path) = file.keep().ok()?;
|
|
Some(path.display().to_string())
|
|
}
|
|
|
|
/// Regex that tokenizes markdown inline elements.
|
|
/// Order matters: longer/more-specific patterns first.
|
|
static INLINE_TOKEN_RE: LazyLock<Regex> = LazyLock::new(|| {
|
|
Regex::new(concat!(
|
|
r"(",
|
|
r"\\.", // Escaped char (highest priority)
|
|
r"|`+", // Inline code (variable length backticks)
|
|
r"|\*\*\*", // Bold+italic
|
|
r"|\*\*", // Bold
|
|
r"|\*", // Italic
|
|
r"|___", // Bold+italic (underscore)
|
|
r"|__", // Bold (underscore)
|
|
r"|_", // Italic (underscore)
|
|
r"|~~", // Strikethrough
|
|
r"|\!\[", // Image start
|
|
r"|\]\(", // Link URL start
|
|
r"|\[", // Link text start
|
|
r"|\]", // Bracket close (without following paren)
|
|
r"|\)", // Link URL end
|
|
r"|[^\\\*_`~\[\]!()]+", // Plain text (no special chars)
|
|
r"|.", // Any other single char
|
|
r")"
|
|
))
|
|
.unwrap()
|
|
});
|
|
|
|
/// A streaming markdown buffer that tracks open constructs.
|
|
///
|
|
/// Accumulates chunks and returns content that is safe to render,
|
|
/// holding back any incomplete markdown constructs. Large code blocks
|
|
/// are automatically truncated with full content saved to a temp file.
|
|
#[derive(Default)]
|
|
pub struct MarkdownBuffer {
|
|
buffer: String,
|
|
}
|
|
|
|
/// Tracks the current parsing state for markdown constructs.
|
|
#[derive(Default, Debug, Clone, PartialEq)]
|
|
struct ParseState {
|
|
in_code_block: bool,
|
|
code_fence_char: char,
|
|
code_fence_len: usize,
|
|
in_table: bool,
|
|
pending_heading: bool,
|
|
in_inline_code: bool,
|
|
inline_code_len: usize,
|
|
in_bold: bool,
|
|
in_italic: bool,
|
|
in_strikethrough: bool,
|
|
in_link_text: bool,
|
|
in_link_url: bool,
|
|
in_image_alt: bool,
|
|
}
|
|
|
|
impl ParseState {
|
|
/// Returns true if no markdown constructs are currently open.
|
|
fn is_clean(&self) -> bool {
|
|
!self.in_code_block
|
|
&& !self.in_table
|
|
&& !self.pending_heading
|
|
&& !self.in_inline_code
|
|
&& !self.in_bold
|
|
&& !self.in_italic
|
|
&& !self.in_strikethrough
|
|
&& !self.in_link_text
|
|
&& !self.in_link_url
|
|
&& !self.in_image_alt
|
|
}
|
|
}
|
|
|
|
// SAFETY: All string slicing in this impl is safe because:
|
|
// - We only slice at positions derived from ASCII characters (newlines, #, |, etc.)
|
|
// - The regex tokenizer operates on valid UTF-8 and returns byte positions at char boundaries
|
|
// - Code fence detection uses chars().take_while() which respects UTF-8
|
|
#[allow(clippy::string_slice)]
|
|
impl MarkdownBuffer {
|
|
/// Create a new empty buffer.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Add a chunk of markdown text to the buffer.
|
|
///
|
|
/// Returns any content that is safe to render, or None if the buffer
|
|
/// contains only incomplete constructs. Large code blocks are automatically
|
|
/// truncated with full content saved to a temp file.
|
|
pub fn push(&mut self, chunk: &str) -> Option<String> {
|
|
self.buffer.push_str(chunk);
|
|
let safe_end = self.find_safe_end();
|
|
|
|
if safe_end > 0 {
|
|
// SAFETY: safe_end is always at a valid UTF-8 char boundary because:
|
|
// - We only set it after processing complete regex tokens (which match
|
|
// valid UTF-8 sequences) or at newline positions (ASCII, single byte)
|
|
// - The regex tokenizer operates on &str which guarantees UTF-8
|
|
let to_render = self.buffer[..safe_end].to_string();
|
|
self.buffer = self.buffer[safe_end..].to_string();
|
|
Some(truncate_code_blocks(&to_render))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Flush any remaining content from the buffer.
|
|
///
|
|
/// Call this at the end of a stream to get any buffered content,
|
|
/// even if markdown constructs are unclosed.
|
|
pub fn flush(&mut self) -> String {
|
|
std::mem::take(&mut self.buffer)
|
|
}
|
|
|
|
/// Find the last byte position where the parse state is "clean".
|
|
fn find_safe_end(&self) -> usize {
|
|
let mut state = ParseState::default();
|
|
let mut last_safe: usize = 0;
|
|
let bytes = self.buffer.as_bytes();
|
|
let len = bytes.len();
|
|
let mut pos: usize = 0;
|
|
|
|
while pos < len {
|
|
let at_line_start = pos == 0 || bytes[pos - 1] == b'\n';
|
|
|
|
if at_line_start {
|
|
if let Some(new_pos) = self.process_line_start(&mut state, pos) {
|
|
pos = new_pos;
|
|
if state.is_clean() {
|
|
last_safe = pos;
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
|
|
if state.in_code_block {
|
|
while pos < len && bytes[pos] != b'\n' {
|
|
pos += 1;
|
|
}
|
|
if pos < len {
|
|
pos += 1;
|
|
}
|
|
continue;
|
|
}
|
|
|
|
let remaining = &self.buffer[pos..];
|
|
let line_end = remaining.find('\n').map(|i| pos + i + 1).unwrap_or(len);
|
|
let line_content = &self.buffer[pos..line_end];
|
|
|
|
for cap in INLINE_TOKEN_RE.find_iter(line_content) {
|
|
let token = cap.as_str();
|
|
let token_end = pos + cap.end();
|
|
|
|
self.process_inline_token(&mut state, token);
|
|
|
|
if state.is_clean() {
|
|
last_safe = token_end;
|
|
}
|
|
}
|
|
|
|
if line_end <= len && line_end > pos && bytes[line_end - 1] == b'\n' {
|
|
state.pending_heading = false;
|
|
if state.is_clean() {
|
|
last_safe = line_end;
|
|
}
|
|
}
|
|
|
|
pos = line_end;
|
|
}
|
|
|
|
last_safe
|
|
}
|
|
|
|
/// Process block-level constructs at the start of a line.
|
|
///
|
|
/// Returns the new position after processing, or None if no block construct found.
|
|
fn process_line_start(&self, state: &mut ParseState, pos: usize) -> Option<usize> {
|
|
let remaining = &self.buffer[pos..];
|
|
|
|
if state.pending_heading {
|
|
state.pending_heading = false;
|
|
}
|
|
|
|
if let Some(fence_result) = self.check_code_fence(remaining, state) {
|
|
return Some(pos + fence_result);
|
|
}
|
|
|
|
if state.in_code_block {
|
|
return None;
|
|
}
|
|
|
|
if remaining.starts_with('#') {
|
|
let hashes = remaining.chars().take_while(|&c| c == '#').count();
|
|
if hashes <= 6 {
|
|
let after_hashes = &remaining[hashes..];
|
|
if after_hashes.is_empty()
|
|
|| after_hashes.starts_with(' ')
|
|
|| after_hashes.starts_with('\n')
|
|
{
|
|
state.pending_heading = true;
|
|
return None;
|
|
}
|
|
}
|
|
}
|
|
|
|
if remaining.starts_with('|') {
|
|
state.in_table = true;
|
|
return None;
|
|
}
|
|
|
|
if (remaining.starts_with('\n') || remaining.is_empty()) && state.in_table {
|
|
state.in_table = false;
|
|
return Some(pos + 1);
|
|
}
|
|
|
|
if state.in_table && !remaining.starts_with('|') {
|
|
state.in_table = false;
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
/// Check for a code fence and update state accordingly.
|
|
///
|
|
/// Returns the position after the fence line if found, None otherwise.
|
|
fn check_code_fence(&self, line: &str, state: &mut ParseState) -> Option<usize> {
|
|
let trimmed = line.trim_start();
|
|
|
|
let fence_char = trimmed.chars().next()?;
|
|
if fence_char != '`' && fence_char != '~' {
|
|
return None;
|
|
}
|
|
|
|
let fence_len = trimmed.chars().take_while(|&c| c == fence_char).count();
|
|
if fence_len < 3 {
|
|
return None;
|
|
}
|
|
|
|
let after_fence = &trimmed[fence_len..];
|
|
|
|
if state.in_code_block {
|
|
if fence_char == state.code_fence_char
|
|
&& fence_len >= state.code_fence_len
|
|
&& (after_fence.is_empty()
|
|
|| after_fence.starts_with('\n')
|
|
|| after_fence.trim().is_empty())
|
|
{
|
|
state.in_code_block = false;
|
|
state.code_fence_char = '\0';
|
|
state.code_fence_len = 0;
|
|
|
|
if let Some(newline_pos) = line.find('\n') {
|
|
return Some(newline_pos + 1);
|
|
} else {
|
|
return Some(line.len());
|
|
}
|
|
}
|
|
} else {
|
|
state.in_code_block = true;
|
|
state.code_fence_char = fence_char;
|
|
state.code_fence_len = fence_len;
|
|
|
|
if let Some(newline_pos) = line.find('\n') {
|
|
return Some(newline_pos + 1);
|
|
} else {
|
|
return Some(line.len());
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
/// Process an inline token and update state.
|
|
fn process_inline_token(&self, state: &mut ParseState, token: &str) {
|
|
if token.starts_with('\\') && token.len() == 2 {
|
|
return;
|
|
}
|
|
|
|
if token.starts_with('`') {
|
|
let tick_count = token.len();
|
|
if state.in_inline_code {
|
|
if tick_count == state.inline_code_len {
|
|
state.in_inline_code = false;
|
|
state.inline_code_len = 0;
|
|
}
|
|
} else {
|
|
state.in_inline_code = true;
|
|
state.inline_code_len = tick_count;
|
|
}
|
|
return;
|
|
}
|
|
|
|
if state.in_inline_code {
|
|
return;
|
|
}
|
|
|
|
match token {
|
|
"***" | "___" => {
|
|
if state.in_bold && state.in_italic {
|
|
state.in_bold = false;
|
|
state.in_italic = false;
|
|
} else if state.in_bold {
|
|
state.in_italic = !state.in_italic;
|
|
} else if state.in_italic {
|
|
state.in_bold = !state.in_bold;
|
|
} else {
|
|
state.in_bold = true;
|
|
state.in_italic = true;
|
|
}
|
|
}
|
|
"**" | "__" => {
|
|
state.in_bold = !state.in_bold;
|
|
}
|
|
"*" | "_" => {
|
|
state.in_italic = !state.in_italic;
|
|
}
|
|
"~~" => {
|
|
state.in_strikethrough = !state.in_strikethrough;
|
|
}
|
|
"![" => {
|
|
state.in_image_alt = true;
|
|
}
|
|
"[" => {
|
|
if !state.in_link_text && !state.in_image_alt {
|
|
state.in_link_text = true;
|
|
}
|
|
}
|
|
"](" => {
|
|
if state.in_link_text {
|
|
state.in_link_text = false;
|
|
state.in_link_url = true;
|
|
} else if state.in_image_alt {
|
|
state.in_image_alt = false;
|
|
state.in_link_url = true;
|
|
}
|
|
}
|
|
"]" => {}
|
|
")" => {
|
|
if state.in_link_url {
|
|
state.in_link_url = false;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use test_case::test_case;
|
|
|
|
/// Process chunks through the buffer and return all outputs (skipping None, including flush)
|
|
fn stream(chunks: &[&str]) -> Vec<String> {
|
|
let mut buf = MarkdownBuffer::new();
|
|
let mut results: Vec<String> = chunks.iter().filter_map(|chunk| buf.push(chunk)).collect();
|
|
let remaining = buf.flush();
|
|
if !remaining.is_empty() {
|
|
results.push(remaining);
|
|
}
|
|
results
|
|
}
|
|
|
|
// ===========================================
|
|
// Realistic LLM streaming scenarios
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["I'll", " help", " you", " with", " that", "!"],
|
|
&["I'll", " help", " you", " with", " that", "!"]
|
|
; "simple sentence streams through immediately without markdown"
|
|
)]
|
|
#[test_case(
|
|
&["Here's the **important", "** part."],
|
|
&["Here's the ", "**important** part."]
|
|
; "bold split mid-word"
|
|
)]
|
|
#[test_case(
|
|
&["Use the `println!", "` macro."],
|
|
&["Use the ", "`println!` macro."]
|
|
; "inline code split"
|
|
)]
|
|
#[test_case(
|
|
&["Check [the docs](https://doc", "s.rs) for more."],
|
|
&["Check ", "[the docs](https://docs.rs) for more."]
|
|
; "link url split"
|
|
)]
|
|
fn test_inline_streaming(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Code blocks (most important for bat rendering)
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["```rust\n", "fn main() {\n", " println!(\"hello\");\n", "}\n", "```\n"],
|
|
&["```rust\nfn main() {\n println!(\"hello\");\n}\n```\n"]
|
|
; "rust code block streamed line by line"
|
|
)]
|
|
#[test_case(
|
|
&["Here's an exa", "mple:\n\n```python\nprint(\"``", "`nested```\")\n```\n\nNice!"],
|
|
&["Here's an exa", "mple:\n", "\n```python\nprint(\"```nested```\")\n```\n\nNice!"]
|
|
; "code block with backticks in string literal"
|
|
)]
|
|
#[test_case(
|
|
&["````md\n", "```\ninner\n```\n", "````\n"],
|
|
&["````md\n```\ninner\n```\n````\n"]
|
|
; "nested code fence with longer outer fence"
|
|
)]
|
|
#[test_case(
|
|
&["~~~bash\n", "echo 'hello'\n", "~", "~~\n"],
|
|
&["~~~bash\necho 'hello'\n~~~\n"]
|
|
; "tilde code fence"
|
|
)]
|
|
#[test_case(
|
|
&["```\ncode"],
|
|
&["```\ncode"]
|
|
; "unclosed code block flushes at end"
|
|
)]
|
|
fn test_code_blocks(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Headings
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["# Getting St", "arted\n\nFirst, install..."],
|
|
&["# Getting Started\n\nFirst, install..."]
|
|
; "heading split mid-word"
|
|
)]
|
|
#[test_case(
|
|
&["## API Reference\n\n###", " Methods\n\n"],
|
|
&["## API Reference\n\n", "### Methods\n\n"]
|
|
; "multiple headings in one chunk"
|
|
)]
|
|
fn test_headings(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Tables
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["| Name | Value |\n", "|------|-------|\n", "| foo | 42 |\n", "\nMore text"],
|
|
&["| Name | Value |\n|------|-------|\n| foo | 42 |\n\nMore text"]
|
|
; "table streamed row by row"
|
|
)]
|
|
#[test_case(
|
|
&["| A | B |\n|---|---|\n| 1 | 2 |\n\n"],
|
|
&["| A | B |\n|---|---|\n| 1 | 2 |\n\n"]
|
|
; "table followed by blank line"
|
|
)]
|
|
fn test_tables(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Mixed formatting (realistic assistant responses)
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&[
|
|
"Here's how to do it:\n\n",
|
|
"1. First, run `cargo", " build`\n",
|
|
"2. Then check the **out", "put**\n\n",
|
|
"```rust\n",
|
|
"fn main() {}\n",
|
|
"```\n"
|
|
],
|
|
&[
|
|
"Here's how to do it:\n\n",
|
|
"1. First, run ",
|
|
"`cargo build`\n",
|
|
"2. Then check the ",
|
|
"**output**\n\n",
|
|
"```rust\nfn main() {}\n```\n"
|
|
]
|
|
; "typical assistant response with list code and formatting"
|
|
)]
|
|
#[test_case(
|
|
&[
|
|
"See the [**Rust Book**](https://doc.rust-l",
|
|
"ang.org/book/) for more info.\n\n",
|
|
"Key points:\n- Use `Result` for errors\n- Prefer `Option` over null"
|
|
],
|
|
&[
|
|
"See the ",
|
|
"[**Rust Book**](https://doc.rust-lang.org/book/) for more info.\n\n",
|
|
"Key points:\n- Use `Result` for errors\n- Prefer `Option` over null"
|
|
]
|
|
; "link with nested bold and list"
|
|
)]
|
|
#[test_case(
|
|
&[
|
|
"![screenshot](./img/sc",
|
|
"reen.png)\n\nAs shown above..."
|
|
],
|
|
&[
|
|
"\n\nAs shown above..."
|
|
]
|
|
; "image with split url"
|
|
)]
|
|
fn test_mixed_content(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Edge cases and escapes
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["Use \\* for bullet points, not \\`code\\`"],
|
|
&["Use \\* for bullet points, not \\`code\\`"]
|
|
; "escaped markdown characters"
|
|
)]
|
|
#[test_case(
|
|
&["Price: $100 * 2 = $200"],
|
|
&["Price: $100 ", "* 2 = $200"]
|
|
; "asterisk in math context treated as italic marker"
|
|
)]
|
|
#[test_case(
|
|
&[""],
|
|
&[] as &[&str]
|
|
; "empty input produces no output"
|
|
)]
|
|
#[test_case(
|
|
&["Hello 世界! Here's some **太字** text."],
|
|
&["Hello 世界! Here's some **太字** text."]
|
|
; "unicode content"
|
|
)]
|
|
#[test_case(
|
|
&["**bold *and italic* together**"],
|
|
&["**bold *and italic* together**"]
|
|
; "nested bold and italic"
|
|
)]
|
|
#[test_case(
|
|
&["***bold italic***"],
|
|
&["***bold italic***"]
|
|
; "combined bold italic marker"
|
|
)]
|
|
#[test_case(
|
|
&["~~stri", "ke~~ and **bo", "ld**"],
|
|
&["~~strike~~ and ", "**bold**"]
|
|
; "strikethrough and bold split"
|
|
)]
|
|
fn test_edge_cases(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
|
|
// ===========================================
|
|
// Incomplete constructs at stream end
|
|
// ===========================================
|
|
|
|
#[test_case(
|
|
&["This is **incomplete bold"],
|
|
&["This is ", "**incomplete bold"]
|
|
; "unclosed bold flushes"
|
|
)]
|
|
#[test_case(
|
|
&["Check [broken link](http://"],
|
|
&["Check ", "[broken link](http://"]
|
|
; "unclosed link flushes"
|
|
)]
|
|
#[test_case(
|
|
&["Start of `code"],
|
|
&["Start of ", "`code"]
|
|
; "unclosed inline code flushes"
|
|
)]
|
|
fn test_incomplete_constructs(chunks: &[&str], expected: &[&str]) {
|
|
assert_eq!(stream(chunks), expected);
|
|
}
|
|
}
|