mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
Streaming markdown (#7233)
Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Douwe Osinga
Generated
+47
-12
@@ -1918,6 +1918,17 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "comfy-table"
|
||||
version = "7.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47"
|
||||
dependencies = [
|
||||
"crossterm",
|
||||
"unicode-segmentation",
|
||||
"unicode-width 0.2.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "compact_str"
|
||||
version = "0.7.1"
|
||||
@@ -2258,6 +2269,29 @@ version = "0.8.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
|
||||
|
||||
[[package]]
|
||||
name = "crossterm"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"crossterm_winapi",
|
||||
"document-features",
|
||||
"parking_lot",
|
||||
"rustix 1.1.3",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossterm_winapi"
|
||||
version = "0.9.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b"
|
||||
dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crunchy"
|
||||
version = "0.2.4"
|
||||
@@ -3031,7 +3065,7 @@ dependencies = [
|
||||
"libc",
|
||||
"option-ext",
|
||||
"redox_users 0.5.2",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3303,7 +3337,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4280,6 +4314,7 @@ dependencies = [
|
||||
"clap_complete",
|
||||
"clap_mangen",
|
||||
"cliclack",
|
||||
"comfy-table",
|
||||
"console 0.16.2",
|
||||
"dotenvy",
|
||||
"etcetera 0.11.0",
|
||||
@@ -4843,7 +4878,7 @@ dependencies = [
|
||||
"libc",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.2",
|
||||
"system-configuration 0.7.0",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
@@ -4863,7 +4898,7 @@ dependencies = [
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core 0.56.0",
|
||||
"windows-core 0.62.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5310,7 +5345,7 @@ dependencies = [
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"serde_core",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6039,7 +6074,7 @@ version = "0.50.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||
dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -7457,7 +7492,7 @@ dependencies = [
|
||||
"quinn-udp",
|
||||
"rustc-hash 2.1.1",
|
||||
"rustls 0.23.36",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.2",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
@@ -7494,9 +7529,9 @@ dependencies = [
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.2",
|
||||
"tracing",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.60.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8142,7 +8177,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys 0.11.0",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -9975,7 +10010,7 @@ dependencies = [
|
||||
"getrandom 0.4.1",
|
||||
"once_cell",
|
||||
"rustix 1.1.3",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -11493,7 +11528,7 @@ version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.48.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -61,6 +61,7 @@ open = "5.3.2"
|
||||
url = { workspace = true }
|
||||
urlencoding = { workspace = true }
|
||||
clap_complete = "4.5.62"
|
||||
comfy-table = "7.2.2"
|
||||
|
||||
[target.'cfg(target_os = "windows")'.dependencies]
|
||||
winapi = { version = "0.3", features = ["wincred"] }
|
||||
|
||||
@@ -5,6 +5,7 @@ mod elicitation;
|
||||
mod export;
|
||||
mod input;
|
||||
mod output;
|
||||
pub mod streaming_buffer;
|
||||
mod task_execution_display;
|
||||
mod thinking;
|
||||
|
||||
@@ -961,6 +962,7 @@ impl CliSession {
|
||||
|
||||
let mut progress_bars = output::McpSpinners::new();
|
||||
let cancel_token_clone = cancel_token.clone();
|
||||
let mut markdown_buffer = streaming_buffer::MarkdownBuffer::new();
|
||||
|
||||
use futures::StreamExt;
|
||||
loop {
|
||||
@@ -1033,7 +1035,7 @@ impl CliSession {
|
||||
if is_stream_json_mode {
|
||||
emit_stream_event(&StreamEvent::Message { message: message.clone() });
|
||||
} else if !is_json_mode {
|
||||
output::render_message(&message, self.debug);
|
||||
output::render_message_streaming(&message, &mut markdown_buffer, self.debug);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1087,6 +1089,10 @@ impl CliSession {
|
||||
}
|
||||
}
|
||||
|
||||
if !is_json_mode && !is_stream_json_mode {
|
||||
output::flush_markdown_buffer_current_theme(&mut markdown_buffer);
|
||||
}
|
||||
|
||||
if is_json_mode {
|
||||
let metadata = match self
|
||||
.agent
|
||||
|
||||
@@ -19,6 +19,8 @@ use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::streaming_buffer::MarkdownBuffer;
|
||||
|
||||
pub const DEFAULT_MIN_PRIORITY: f32 = 0.0;
|
||||
pub const DEFAULT_CLI_LIGHT_THEME: &str = "GitHub";
|
||||
pub const DEFAULT_CLI_DARK_THEME: &str = "zenburn";
|
||||
@@ -272,6 +274,94 @@ pub fn render_message(message: &Message, debug: bool) {
|
||||
let _ = std::io::stdout().flush();
|
||||
}
|
||||
|
||||
/// Render a streaming message, using a buffer to accumulate text content
|
||||
/// and only render when markdown constructs are complete.
|
||||
pub fn render_message_streaming(message: &Message, buffer: &mut MarkdownBuffer, debug: bool) {
|
||||
let theme = get_theme();
|
||||
|
||||
for content in &message.content {
|
||||
match content {
|
||||
MessageContent::Text(text) => {
|
||||
if let Some(safe_content) = buffer.push(&text.text) {
|
||||
print_markdown(&safe_content, theme);
|
||||
}
|
||||
}
|
||||
MessageContent::ToolRequest(req) => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
render_tool_request(req, theme, debug);
|
||||
}
|
||||
MessageContent::ToolResponse(resp) => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
render_tool_response(resp, theme, debug);
|
||||
}
|
||||
MessageContent::ActionRequired(action) => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
match &action.data {
|
||||
ActionRequiredData::ToolConfirmation { tool_name, .. } => {
|
||||
println!("action_required(tool_confirmation): {}", tool_name)
|
||||
}
|
||||
ActionRequiredData::Elicitation { message, .. } => {
|
||||
println!("action_required(elicitation): {}", message)
|
||||
}
|
||||
ActionRequiredData::ElicitationResponse { id, .. } => {
|
||||
println!("action_required(elicitation_response): {}", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
MessageContent::Image(image) => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
println!("Image: [data: {}, type: {}]", image.data, image.mime_type);
|
||||
}
|
||||
MessageContent::Thinking(thinking) => {
|
||||
if std::env::var("GOOSE_CLI_SHOW_THINKING").is_ok()
|
||||
&& std::io::stdout().is_terminal()
|
||||
{
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
println!("\n{}", style("Thinking:").dim().italic());
|
||||
print_markdown(&thinking.thinking, theme);
|
||||
}
|
||||
}
|
||||
MessageContent::RedactedThinking(_) => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
println!("\n{}", style("Thinking:").dim().italic());
|
||||
print_markdown("Thinking was redacted", theme);
|
||||
}
|
||||
MessageContent::SystemNotification(notification) => {
|
||||
use goose::conversation::message::SystemNotificationType;
|
||||
|
||||
match notification.notification_type {
|
||||
SystemNotificationType::ThinkingMessage => {
|
||||
show_thinking();
|
||||
set_thinking_message(¬ification.msg);
|
||||
}
|
||||
SystemNotificationType::InlineMessage => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
hide_thinking();
|
||||
println!("\n{}", style(¬ification.msg).yellow());
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
flush_markdown_buffer(buffer, theme);
|
||||
println!("WARNING: Message content type could not be rendered");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = std::io::stdout().flush();
|
||||
}
|
||||
|
||||
pub fn flush_markdown_buffer(buffer: &mut MarkdownBuffer, theme: Theme) {
|
||||
let remaining = buffer.flush();
|
||||
if !remaining.is_empty() {
|
||||
print_markdown(&remaining, theme);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn flush_markdown_buffer_current_theme(buffer: &mut MarkdownBuffer) {
|
||||
flush_markdown_buffer(buffer, get_theme());
|
||||
}
|
||||
|
||||
pub fn render_text(text: &str, color: Option<Color>, dim: bool) {
|
||||
render_text_no_newlines(format!("\n{}\n\n", text).as_str(), color, dim);
|
||||
}
|
||||
@@ -744,19 +834,199 @@ pub fn env_no_color() -> bool {
|
||||
|
||||
fn print_markdown(content: &str, theme: Theme) {
|
||||
if std::io::stdout().is_terminal() {
|
||||
bat::PrettyPrinter::new()
|
||||
.input(bat::Input::from_bytes(content.as_bytes()))
|
||||
.theme(theme.as_str())
|
||||
.colored_output(env_no_color())
|
||||
.language("Markdown")
|
||||
.wrapping_mode(WrappingMode::NoWrapping(true))
|
||||
.print()
|
||||
.unwrap();
|
||||
if let Some((before, table, after)) = extract_markdown_table(content) {
|
||||
if !before.is_empty() {
|
||||
print_markdown_raw(&before, theme);
|
||||
}
|
||||
print_table(&table, theme);
|
||||
if !after.is_empty() {
|
||||
print_markdown(after, theme);
|
||||
}
|
||||
} else {
|
||||
print_markdown_raw(content, theme);
|
||||
}
|
||||
} else {
|
||||
print!("{}", content);
|
||||
}
|
||||
}
|
||||
|
||||
/// Renders markdown content using bat (no table processing)
|
||||
fn print_markdown_raw(content: &str, theme: Theme) {
|
||||
bat::PrettyPrinter::new()
|
||||
.input(bat::Input::from_bytes(content.as_bytes()))
|
||||
.theme(theme.as_str())
|
||||
.colored_output(env_no_color())
|
||||
.language("Markdown")
|
||||
.wrapping_mode(WrappingMode::NoWrapping(true))
|
||||
.print()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn extract_markdown_table(content: &str) -> Option<(String, Vec<&str>, &str)> {
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
|
||||
// Track newline positions for safe slicing later
|
||||
let newline_indices: Vec<usize> = content
|
||||
.bytes()
|
||||
.enumerate()
|
||||
.filter_map(|(i, b)| if b == b'\n' { Some(i) } else { None })
|
||||
.collect();
|
||||
|
||||
// Skip tables inside code blocks
|
||||
let mut in_code_block = false;
|
||||
let mut table_start = None;
|
||||
let mut table_end = None;
|
||||
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
let trimmed = line.trim();
|
||||
|
||||
if trimmed.starts_with("```") || trimmed.starts_with("~~~") {
|
||||
in_code_block = !in_code_block;
|
||||
continue;
|
||||
}
|
||||
|
||||
if in_code_block {
|
||||
continue;
|
||||
}
|
||||
|
||||
if trimmed.starts_with('|') && trimmed.ends_with('|') {
|
||||
if table_start.is_none() {
|
||||
table_start = Some(i);
|
||||
}
|
||||
table_end = Some(i);
|
||||
} else if table_start.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let start = table_start?;
|
||||
let end = table_end?;
|
||||
|
||||
// Need at least header + separator (2 rows minimum)
|
||||
if end < start + 1 {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Require separator to be the second row with proper format
|
||||
let separator_line = lines.get(start + 1)?;
|
||||
let is_valid_separator = separator_line.trim().starts_with('|')
|
||||
&& separator_line.trim().ends_with('|')
|
||||
&& separator_line
|
||||
.trim()
|
||||
.trim_matches('|')
|
||||
.split('|')
|
||||
.all(|cell| {
|
||||
let t = cell.trim();
|
||||
!t.is_empty() && t.chars().all(|c| c == '-' || c == ':' || c == ' ')
|
||||
});
|
||||
|
||||
if !is_valid_separator {
|
||||
return None;
|
||||
}
|
||||
|
||||
let before = lines[..start].join("\n");
|
||||
let before = if before.is_empty() {
|
||||
before
|
||||
} else {
|
||||
before + "\n"
|
||||
};
|
||||
let table = lines[start..=end].to_vec();
|
||||
|
||||
let after = if end + 1 >= lines.len() {
|
||||
""
|
||||
} else if let Some(&newline_pos) = newline_indices.get(end) {
|
||||
content.get(newline_pos + 1..).unwrap_or("")
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
Some((before, table, after))
|
||||
}
|
||||
|
||||
fn print_table(table_lines: &[&str], theme: Theme) {
|
||||
use comfy_table::{presets, Cell, CellAlignment, ContentArrangement, Table};
|
||||
|
||||
let mut table = Table::new();
|
||||
table.set_content_arrangement(ContentArrangement::Dynamic);
|
||||
|
||||
table.load_preset(presets::ASCII_MARKDOWN);
|
||||
|
||||
let mut rows: Vec<Vec<String>> = Vec::new();
|
||||
let mut alignments: Vec<CellAlignment> = Vec::new();
|
||||
let mut separator_idx = None;
|
||||
|
||||
for (i, line) in table_lines.iter().enumerate() {
|
||||
let cells: Vec<String> = line
|
||||
.trim()
|
||||
.trim_matches('|')
|
||||
.split('|')
|
||||
.map(|s| s.trim().to_string())
|
||||
.collect();
|
||||
|
||||
let is_separator = cells.iter().all(|c| {
|
||||
let t = c.trim();
|
||||
t.chars().all(|ch| ch == '-' || ch == ':') && t.contains('-')
|
||||
});
|
||||
if is_separator {
|
||||
separator_idx = Some(i);
|
||||
alignments = cells
|
||||
.iter()
|
||||
.map(|c| {
|
||||
let t = c.trim();
|
||||
if t.starts_with(':') && t.ends_with(':') {
|
||||
CellAlignment::Center
|
||||
} else if t.ends_with(':') {
|
||||
CellAlignment::Right
|
||||
} else {
|
||||
CellAlignment::Left
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
} else {
|
||||
rows.push(cells);
|
||||
}
|
||||
}
|
||||
|
||||
if separator_idx.is_none() && !rows.is_empty() {
|
||||
alignments = vec![CellAlignment::Left; rows[0].len()];
|
||||
}
|
||||
|
||||
if let Some(header) = rows.first() {
|
||||
let header_cells: Vec<Cell> = header
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, text)| {
|
||||
let cell = Cell::new(text);
|
||||
if let Some(align) = alignments.get(i) {
|
||||
cell.set_alignment(*align)
|
||||
} else {
|
||||
cell
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
table.set_header(header_cells);
|
||||
}
|
||||
|
||||
for row in rows.iter().skip(1) {
|
||||
let cells: Vec<Cell> = row
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, text)| {
|
||||
let cell = Cell::new(text);
|
||||
if let Some(align) = alignments.get(i) {
|
||||
cell.set_alignment(*align)
|
||||
} else {
|
||||
cell
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
table.add_row(cells);
|
||||
}
|
||||
|
||||
let table_str = table.to_string();
|
||||
print_markdown_raw(&table_str, theme);
|
||||
}
|
||||
|
||||
const INDENT: &str = " ";
|
||||
|
||||
fn print_value_with_prefix(prefix: &String, value: &Value, debug: bool) {
|
||||
|
||||
@@ -0,0 +1,596 @@
|
||||
//! 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::sync::LazyLock;
|
||||
|
||||
/// 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.
|
||||
#[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.
|
||||
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(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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user