From 867a83cfc761f152ba14b900bfe9017688abddd8 Mon Sep 17 00:00:00 2001 From: Jasper Date: Wed, 26 Aug 2026 21:02:35 +0000 Subject: [PATCH] fix(local-inference): keep nested execute fences inert (#11117) --- Cargo.lock | 12 + crates/goose-local-inference/Cargo.toml | 1 + crates/goose-local-inference/src/lib.rs | 1 - .../src/llamacpp/inference_emulated_tools.rs | 414 +--------- .../src/tool_emulation.rs | 770 +++++++++++------- 5 files changed, 492 insertions(+), 706 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9bed6c90f..82af1ed7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5458,6 +5458,7 @@ dependencies = [ "llama-cpp-2", "llama-cpp-sys-2", "minijinja", + "pulldown-cmark", "regex", "reqwest 0.13.4", "rmcp 3.1.4", @@ -9601,6 +9602,17 @@ dependencies = [ "psl-types", ] +[[package]] +name = "pulldown-cmark" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f068eba8e7071c5f9511831b44f32c740d5adf574e990f946ddb53db2f314e" +dependencies = [ + "bitflags 2.13.0", + "memchr", + "unicase", +] + [[package]] name = "pulp" version = "0.21.5" diff --git a/crates/goose-local-inference/Cargo.toml b/crates/goose-local-inference/Cargo.toml index 5d197cda2..3da359e6e 100644 --- a/crates/goose-local-inference/Cargo.toml +++ b/crates/goose-local-inference/Cargo.toml @@ -35,6 +35,7 @@ include_dir = { workspace = true } llama-cpp-2 = { workspace = true } llama-cpp-sys-2 = { workspace = true } minijinja = { version = "2.18", default-features = false, features = ["loader", "multi_template", "serde"] } +pulldown-cmark = { version = "0.13.4", default-features = false } reqwest = { workspace = true, features = ["json", "stream"] } regex = { workspace = true } rmcp = { workspace = true } diff --git a/crates/goose-local-inference/src/lib.rs b/crates/goose-local-inference/src/lib.rs index bb231575c..46e436323 100644 --- a/crates/goose-local-inference/src/lib.rs +++ b/crates/goose-local-inference/src/lib.rs @@ -15,7 +15,6 @@ pub(crate) mod multimodal; #[cfg(feature = "mlx")] mod native_tool_parsing; pub(crate) mod thinking_output; -#[cfg(feature = "mlx")] mod tool_emulation; mod tool_parsing; diff --git a/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs index 27f7d4d0c..0eaf4a66f 100644 --- a/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs @@ -1,21 +1,21 @@ //! Tool call emulation for models without native tool-calling support. //! //! The model is prompted to emit shell commands as `$ command` on a new line and -//! code blocks as `` ```execute `` fenced blocks. A streaming parser detects these +//! code blocks as `` ```execute_typescript `` fenced blocks. A streaming parser detects these //! patterns and converts them into tool-call messages. //! //! # Known false-positive scenarios //! -//! Because detection is purely text-based, the parser can misinterpret model output: +//! Top-level protocol markers remain ambiguous when a model emits them as explanation: //! //! - **`$` at line start in explanatory text.** If the model writes a line starting //! with `$` as an example (e.g. "$ is the jQuery selector"), it will be treated as //! a shell command. Mid-sentence `$` (e.g. "costs $50") is safe — only `\n$` or //! `$` at the very start of output triggers command detection. //! -//! - **`` ```execute `` in explanatory code fences.** If the model uses this exact -//! fence tag in prose, the content will be executed. Standard `` ```js `` or -//! `` ```python `` fences are not affected. +//! - **`` ```execute_typescript `` as a top-level explanatory fence.** If the model +//! uses this exact fence tag in prose, the content will be executed. Nested Markdown +//! examples and fences with other info strings remain text. //! //! These are inherent to text-based tool emulation. Models with native tool-calling //! support should use the `inference_native_tools` path instead. @@ -31,13 +31,11 @@ use super::super::{finalize_usage, thinking_output::ThinkingOutputFilter, Stream use super::inference_engine::{ generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, TokenAction, }; +use crate::tool_emulation::{EmulatorAction, StreamingEmulatorParser}; const SHELL_TOOL: &str = "developer__shell"; const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript"; -const HOLD_BACK_CODE_MODE: usize = " ```execute_typescript\n".len(); -const HOLD_BACK_SHELL_ONLY: usize = "\n$".len(); - pub(super) fn load_tiny_model_prompt() -> String { use std::env; @@ -138,167 +136,6 @@ pub(super) fn build_emulator_tool_description(tools: &[Tool], code_mode_enabled: tool_desc } -enum EmulatorAction { - Text(String), - ShellCommand(String), - ExecuteCode(String), -} - -enum ParserState { - Normal, - InCommand, - InExecuteBlock, -} - -struct StreamingEmulatorParser { - buffer: String, - state: ParserState, - code_mode_enabled: bool, -} - -impl StreamingEmulatorParser { - fn new(code_mode_enabled: bool) -> Self { - Self { - buffer: String::new(), - state: ParserState::Normal, - code_mode_enabled, - } - } - - fn process_chunk(&mut self, chunk: &str) -> Vec { - self.buffer.push_str(chunk); - let mut results = Vec::new(); - - loop { - match self.state { - ParserState::InCommand => { - if let Some((command_line, rest)) = self.buffer.split_once('\n') { - if let Some(command) = command_line.strip_prefix('$') { - let command = command.trim(); - if !command.is_empty() { - results.push(EmulatorAction::ShellCommand(command.to_string())); - } - } - self.buffer = rest.to_string(); - self.state = ParserState::Normal; - } else { - break; - } - } - ParserState::InExecuteBlock => { - // Look for closing ``` to end the execute block - if let Some(end_idx) = self.buffer.find("\n```") { - #[allow(clippy::string_slice)] - let code = self.buffer[..end_idx].to_string(); - // Skip past the closing ``` and any trailing newline - #[allow(clippy::string_slice)] - let rest = &self.buffer[end_idx + 4..]; - let rest = rest.strip_prefix('\n').unwrap_or(rest); - self.buffer = rest.to_string(); - self.state = ParserState::Normal; - if !code.trim().is_empty() { - results.push(EmulatorAction::ExecuteCode(code)); - } - } else { - // Still accumulating code — wait for closing fence - break; - } - } - ParserState::Normal => { - // Check for ```execute block (code mode) - if self.code_mode_enabled { - if let Some((before, after)) = - self.buffer.split_once("```execute_typescript\n") - { - if !before.trim().is_empty() { - results.push(EmulatorAction::Text(before.to_string())); - } - self.buffer = after.to_string(); - self.state = ParserState::InExecuteBlock; - continue; - } - // Also handle without newline after tag (accumulating) - if self.buffer.ends_with("```execute_typescript") { - let before = self.buffer.trim_end_matches("```execute_typescript"); - if !before.trim().is_empty() { - results.push(EmulatorAction::Text(before.to_string())); - } - self.buffer.clear(); - self.state = ParserState::InExecuteBlock; - continue; - } - } - - // Check for $ command - if let Some((before_dollar, from_dollar)) = self.buffer.split_once("\n$") { - let text = format!("{}\n", before_dollar); - if !text.trim().is_empty() { - results.push(EmulatorAction::Text(text)); - } - self.buffer = format!("${}", from_dollar); - self.state = ParserState::InCommand; - } else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() { - self.state = ParserState::InCommand; - } else { - let hold_back = if self.code_mode_enabled { - HOLD_BACK_CODE_MODE - } else { - HOLD_BACK_SHELL_ONLY - }; - let char_count = self.buffer.chars().count(); - if char_count > hold_back && !self.buffer.ends_with('\n') { - let mut chars = self.buffer.chars(); - let emit_count = char_count - hold_back; - let emit_text: String = chars.by_ref().take(emit_count).collect(); - let keep_text: String = chars.collect(); - if !emit_text.is_empty() { - results.push(EmulatorAction::Text(emit_text)); - } - self.buffer = keep_text; - } - break; - } - } - } - } - - results - } - - fn flush(&mut self) -> Vec { - let mut results = Vec::new(); - - if !self.buffer.is_empty() { - match self.state { - ParserState::InCommand => { - let command_line = self.buffer.trim(); - if let Some(command) = command_line.strip_prefix('$') { - let command = command.trim(); - if !command.is_empty() { - results.push(EmulatorAction::ShellCommand(command.to_string())); - } - } else if !command_line.is_empty() { - results.push(EmulatorAction::Text(self.buffer.clone())); - } - } - ParserState::InExecuteBlock => { - let code = self.buffer.trim(); - if !code.is_empty() { - results.push(EmulatorAction::ExecuteCode(code.to_string())); - } - } - ParserState::Normal => { - results.push(EmulatorAction::Text(self.buffer.clone())); - } - } - self.buffer.clear(); - self.state = ParserState::Normal; - } - - results - } -} - fn send_emulator_action( action: &EmulatorAction, message_id: &str, @@ -465,21 +302,6 @@ pub(super) fn generate_with_emulated_tools( mod tests { use super::*; - /// Collect all actions from feeding chunks through the parser, then flushing. - fn parse_chunks(chunks: &[&str], code_mode: bool) -> Vec { - let mut parser = StreamingEmulatorParser::new(code_mode); - let mut actions = Vec::new(); - for chunk in chunks { - actions.extend(parser.process_chunk(chunk)); - } - actions.extend(parser.flush()); - actions - } - - fn parse_all(input: &str, code_mode: bool) -> Vec { - parse_chunks(&[input], code_mode) - } - fn trim_chunks(chunks: &[&str], stops: &[String]) -> (String, bool) { let mut trimmer = StopSuffixTrimmer::new(stops); let mut output = String::new(); @@ -524,39 +346,6 @@ mod tests { (thinking, actions) } - fn assert_text(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::Text(t) => assert_eq!(t.trim(), expected.trim(), "text mismatch"), - other => panic!("expected Text, got {:?}", action_label(other)), - } - } - - fn assert_shell(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::ShellCommand(cmd) => { - assert_eq!(cmd, expected, "shell command mismatch") - } - other => panic!("expected ShellCommand, got {:?}", action_label(other)), - } - } - - fn assert_execute(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::ExecuteCode(code) => { - assert_eq!(code.trim(), expected.trim(), "execute code mismatch") - } - other => panic!("expected ExecuteCode, got {:?}", action_label(other)), - } - } - - fn action_label(a: &EmulatorAction) -> &'static str { - match a { - EmulatorAction::Text(_) => "Text", - EmulatorAction::ShellCommand(_) => "ShellCommand", - EmulatorAction::ExecuteCode(_) => "ExecuteCode", - } - } - #[test] fn stop_suffix_trimmer_strips_split_stop() { let stops = vec!["<|eom_id|>".to_string()]; @@ -575,166 +364,6 @@ mod tests { assert_eq!(content, "Use the < symbol"); } - #[test] - fn plain_text_no_tools() { - let actions = parse_all("Hello, world!", false); - // Hold-back may split text across actions; concatenate all text - let all_text: String = actions - .iter() - .map(|a| match a { - EmulatorAction::Text(t) => t.as_str(), - _ => panic!("expected only Text actions"), - }) - .collect(); - assert_eq!(all_text.trim(), "Hello, world!"); - } - - #[test] - fn single_shell_command() { - let actions = parse_all("$ ls -la\n", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "ls -la"); - } - - #[test] - fn text_then_shell_command() { - let actions = parse_all("Let me check:\n$ ls -la\n", false); - assert!(actions.len() >= 2); - assert_text(&actions[0], "Let me check:"); - assert_shell(&actions[actions.len() - 1], "ls -la"); - } - - #[test] - fn shell_command_at_start_of_output() { - let actions = parse_all("$ whoami\n", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "whoami"); - } - - #[test] - fn shell_command_without_trailing_newline() { - // Flush should handle unterminated command - let actions = parse_all("$ whoami", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "whoami"); - } - - #[test] - fn dollar_sign_mid_sentence_is_not_command() { - let actions = parse_all("It costs $50 per month", false); - for action in &actions { - assert!( - matches!(action, EmulatorAction::Text(_)), - "mid-sentence $ should not trigger a shell command" - ); - } - let all_text: String = actions - .iter() - .filter_map(|a| match a { - EmulatorAction::Text(t) => Some(t.as_str()), - _ => None, - }) - .collect(); - assert_eq!(all_text.trim(), "It costs $50 per month"); - } - - #[test] - fn execute_block() { - let input = "Here's the code:\n```execute_typescript\nconsole.log('hi');\n```\n"; - let actions = parse_all(input, true); - assert!(actions.len() >= 2); - assert_text(&actions[0], "Here's the code:"); - assert_execute(&actions[actions.len() - 1], "console.log('hi');"); - } - - #[test] - fn execute_block_not_detected_without_code_mode() { - let input = "```execute_typescript\nconsole.log('hi');\n```\n"; - let actions = parse_all(input, false); - // Should be treated as plain text - for action in &actions { - assert!(matches!(action, EmulatorAction::Text(_))); - } - } - - #[test] - fn dollar_split_across_chunks() { - // The \n and $ arrive in separate chunks - let actions = parse_chunks(&["Let me check\n", "$ ls -la\n"], false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 1); - assert_shell(shells[0], "ls -la"); - } - - #[test] - fn execute_fence_split_across_chunks() { - let actions = parse_chunks( - &["Here:\n```ex", "ecute_typescript\nlet x = 1;\n", "```\n"], - true, - ); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - assert_execute(executes[0], "let x = 1;"); - } - - #[test] - fn multiple_commands_on_separate_lines() { - // In practice, generation stops after the first tool call. But the - // parser should detect commands separated by \n$ when fed as chunks. - let actions = parse_chunks(&["Here:\n$ cd /tmp\n", "Done.\n$ ls\n"], false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 2); - assert_shell(shells[0], "cd /tmp"); - assert_shell(shells[1], "ls"); - } - - #[test] - fn regular_code_fence_not_treated_as_execute() { - let input = "```python\nprint('hi')\n```\n"; - let actions = parse_all(input, true); - for action in &actions { - assert!( - matches!(action, EmulatorAction::Text(_)), - "regular code fence should be text" - ); - } - } - - #[test] - fn empty_command_ignored() { - let actions = parse_all("$\n", false); - // Empty command after $ should not produce a ShellCommand - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 0); - } - - #[test] - fn token_by_token_streaming() { - // Simulate LLM generating one token at a time - let input = "$ echo hello\n"; - let chars: Vec = input.chars().map(|c| c.to_string()).collect(); - let chunks: Vec<&str> = chars.iter().map(|s| s.as_str()).collect(); - let actions = parse_chunks(&chunks, false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 1); - assert_shell(shells[0], "echo hello"); - } - #[test] fn thinking_seeded_from_generation_prompt_is_not_emulated_text() { let (thinking, actions) = @@ -753,35 +382,4 @@ mod tests { .collect(); assert_eq!(text.trim(), "The answer."); } - - #[test] - fn execute_block_with_multiline_code() { - let input = "```execute_typescript\nasync function run() {\n const r = await Developer.shell({ command: \"ls\" });\n return r;\n}\n```\n"; - let actions = parse_all(input, true); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - match executes[0] { - EmulatorAction::ExecuteCode(code) => { - assert!(code.contains("async function run()")); - assert!(code.contains("Developer.shell")); - } - _ => unreachable!(), - } - } - - #[test] - fn unclosed_execute_block_flushed() { - // Model stops generating mid-block - let input = "```execute_typescript\nlet x = 1;"; - let actions = parse_all(input, true); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - assert_execute(executes[0], "let x = 1;"); - } } diff --git a/crates/goose-local-inference/src/tool_emulation.rs b/crates/goose-local-inference/src/tool_emulation.rs index ac528be23..4e9046f7f 100644 --- a/crates/goose-local-inference/src/tool_emulation.rs +++ b/crates/goose-local-inference/src/tool_emulation.rs @@ -4,18 +4,25 @@ //! as `$ command` on a new line and code blocks as ```execute_typescript fenced blocks. //! The parser converts those patterns into Goose tool-call messages. +use pulldown_cmark::{CodeBlockKind, Event, Parser, Tag}; + +#[cfg(feature = "mlx")] use goose_provider_types::conversation::message::{Message, MessageContent}; +#[cfg(feature = "mlx")] use rmcp::model::{CallToolRequestParams, Tool}; +#[cfg(feature = "mlx")] use serde_json::json; +#[cfg(feature = "mlx")] use std::borrow::Cow; +#[cfg(feature = "mlx")] use uuid::Uuid; +#[cfg(feature = "mlx")] pub(crate) const SHELL_TOOL: &str = "developer__shell"; +#[cfg(feature = "mlx")] pub(crate) const CODE_EXECUTION_TOOL: &str = "code_execution__execute_typescript"; -const HOLD_BACK_CODE_MODE: usize = " ```execute_typescript\n".len(); -const HOLD_BACK_SHELL_ONLY: usize = "\n$".len(); - +#[cfg(feature = "mlx")] pub(crate) fn load_tiny_model_prompt() -> String { use std::env; @@ -48,6 +55,7 @@ pub(crate) fn load_tiny_model_prompt() -> String { }) } +#[cfg(feature = "mlx")] pub(crate) fn build_emulator_tool_description(tools: &[Tool], code_mode_enabled: bool) -> String { let mut tool_desc = String::new(); @@ -124,114 +132,271 @@ pub(crate) enum EmulatorAction { ExecuteCode(String), } +// pulldown-cmark parses complete documents, so fail closed once reparsing would become expensive. +const MAX_MARKDOWN_CONTEXT_BYTES: usize = 256 * 1024; +const MAX_MARKDOWN_PARSE_BYTES: usize = 1024 * 1024; + +#[derive(Clone, Copy)] enum ParserState { Normal, - InCommand, - InExecuteBlock, + InExecuteBlock { fence_len: usize }, } pub(crate) struct StreamingEmulatorParser { buffer: String, + document: String, + markdown_context_available: bool, + markdown_parse_bytes_remaining: usize, state: ParserState, code_mode_enabled: bool, } +fn execute_fence_len(line: &str) -> Option { + let line = line.strip_suffix('\r').unwrap_or(line); + let indent = line.bytes().take_while(|byte| *byte == b' ').count(); + if indent > 3 { + return None; + } + + let rest = line.get(indent..)?; + let fence_len = rest.bytes().take_while(|byte| *byte == b'`').count(); + if fence_len < 3 { + return None; + } + + rest.get(fence_len..)? + .trim_matches([' ', '\t']) + .eq("execute_typescript") + .then_some(fence_len) +} + +fn is_closing_fence(line: &str, minimum_len: usize) -> bool { + let line = line.strip_suffix('\r').unwrap_or(line); + let indent = line.bytes().take_while(|byte| *byte == b' ').count(); + if indent > 3 { + return false; + } + + let Some(rest) = line.get(indent..) else { + return false; + }; + let fence_len = rest.bytes().take_while(|byte| *byte == b'`').count(); + fence_len >= minimum_len + && rest + .get(fence_len..) + .is_some_and(|suffix| suffix.bytes().all(|byte| matches!(byte, b' ' | b'\t'))) +} + +fn is_top_level_execute_fence(markdown: &str, line_start: usize) -> bool { + let mut depth = 0; + for (event, range) in Parser::new(markdown).into_offset_iter() { + match event { + Event::Start(tag) => { + if depth == 0 + && range.start == line_start + && matches!( + tag, + Tag::CodeBlock(CodeBlockKind::Fenced(ref info)) + if info.as_ref() == "execute_typescript" + ) + { + return true; + } + depth += 1; + } + Event::End(_) => depth -= 1, + _ => {} + } + } + false +} + +fn is_top_level_paragraph_line(markdown: &str, line_start: usize) -> bool { + let mut depth = 0; + for (event, range) in Parser::new(markdown).into_offset_iter() { + match event { + Event::Start(tag) => { + if depth == 0 + && matches!(tag, Tag::Paragraph) + && range.start <= line_start + && line_start < range.end + { + return true; + } + depth += 1; + } + Event::End(_) => depth -= 1, + _ => {} + } + } + false +} + +fn closing_fence_range( + input: &str, + minimum_len: usize, + allow_end_of_stream: bool, +) -> Option<(usize, usize)> { + let mut line_start = 0; + loop { + let remaining = input + .get(line_start..) + .expect("line start must be a character boundary"); + let newline_offset = remaining.find('\n'); + if newline_offset.is_none() && !allow_end_of_stream { + return None; + } + let line_end = newline_offset + .map(|offset| line_start + offset) + .unwrap_or(input.len()); + if is_closing_fence( + input + .get(line_start..line_end) + .expect("line range must be on character boundaries"), + minimum_len, + ) { + let consumed = if line_end < input.len() { + line_end + 1 + } else { + line_end + }; + return Some((line_start, consumed)); + } + if line_end == input.len() { + return None; + } + line_start = line_end + 1; + } +} + impl StreamingEmulatorParser { pub(crate) fn new(code_mode_enabled: bool) -> Self { Self { buffer: String::new(), + document: String::new(), + markdown_context_available: true, + markdown_parse_bytes_remaining: MAX_MARKDOWN_PARSE_BYTES, state: ParserState::Normal, code_mode_enabled, } } + fn append_markdown_context(&mut self, chunk: &str) { + if !self.markdown_context_available { + return; + } + + if self.document.len().saturating_add(chunk.len()) > MAX_MARKDOWN_CONTEXT_BYTES { + self.disable_markdown_emulation(); + } else { + self.document.push_str(chunk); + } + } + + fn consume_markdown_parse_budget(&mut self, bytes: usize) -> bool { + if bytes > self.markdown_parse_bytes_remaining { + self.disable_markdown_emulation(); + false + } else { + self.markdown_parse_bytes_remaining -= bytes; + true + } + } + + fn markdown_matches( + &mut self, + range: Option<(usize, usize)>, + predicate: fn(&str, usize) -> bool, + ) -> bool { + let Some((line_start, markdown_end)) = range else { + return false; + }; + if !self.consume_markdown_parse_budget(markdown_end) { + return false; + } + + let markdown = self + .document + .get(..markdown_end) + .expect("markdown end must be a character boundary"); + predicate(markdown, line_start) + } + + fn disable_markdown_emulation(&mut self) { + self.document.clear(); + self.markdown_context_available = false; + self.markdown_parse_bytes_remaining = 0; + } + pub(crate) fn process_chunk(&mut self, chunk: &str) -> Vec { self.buffer.push_str(chunk); + self.append_markdown_context(chunk); let mut results = Vec::new(); loop { match self.state { - ParserState::InCommand => { - if let Some((command_line, rest)) = self.buffer.split_once('\n') { - if let Some(command) = command_line.strip_prefix('$') { + ParserState::InExecuteBlock { fence_len } => { + let Some((closing_start, consumed)) = + closing_fence_range(&self.buffer, fence_len, false) + else { + break; + }; + let code_end = closing_start + .checked_sub(1) + .filter(|index| self.buffer.as_bytes()[*index] == b'\n') + .unwrap_or(closing_start); + let code = self + .buffer + .get(..code_end) + .expect("code boundary must be a character boundary"); + let code = code.strip_suffix('\r').unwrap_or(code).to_string(); + self.buffer.replace_range(..consumed, ""); + self.state = ParserState::Normal; + if !code.trim().is_empty() { + results.push(EmulatorAction::ExecuteCode(code)); + } + } + ParserState::Normal => { + let Some(line_end) = self.buffer.find('\n') else { + break; + }; + let markdown_range = self.markdown_context_available.then(|| { + let line_start = self.document.len() - self.buffer.len(); + (line_start, line_start + line_end + 1) + }); + let line = self + .buffer + .get(..line_end) + .expect("line end must be a character boundary") + .to_string(); + let line_with_newline = self + .buffer + .get(..=line_end) + .expect("newline must be a character boundary") + .to_string(); + self.buffer.replace_range(..=line_end, ""); + + if self.code_mode_enabled { + if let Some(fence_len) = execute_fence_len(&line) { + if self.markdown_matches(markdown_range, is_top_level_execute_fence) { + self.state = ParserState::InExecuteBlock { fence_len }; + continue; + } + } + } + + let line_without_cr = line.strip_suffix('\r').unwrap_or(&line); + if let Some(command) = line_without_cr.strip_prefix('$') { + if self.markdown_matches(markdown_range, is_top_level_paragraph_line) { let command = command.trim(); if !command.is_empty() { results.push(EmulatorAction::ShellCommand(command.to_string())); } - } - self.buffer = rest.to_string(); - self.state = ParserState::Normal; - } else { - break; - } - } - ParserState::InExecuteBlock => { - if let Some(end_idx) = self.buffer.find("\n```") { - #[allow(clippy::string_slice)] - let code = self.buffer[..end_idx].to_string(); - #[allow(clippy::string_slice)] - let rest = &self.buffer[end_idx + 4..]; - let rest = rest.strip_prefix('\n').unwrap_or(rest); - self.buffer = rest.to_string(); - self.state = ParserState::Normal; - if !code.trim().is_empty() { - results.push(EmulatorAction::ExecuteCode(code)); - } - } else { - break; - } - } - ParserState::Normal => { - if self.code_mode_enabled { - if let Some((before, after)) = - self.buffer.split_once("```execute_typescript\n") - { - if !before.trim().is_empty() { - results.push(EmulatorAction::Text(before.to_string())); - } - self.buffer = after.to_string(); - self.state = ParserState::InExecuteBlock; - continue; - } - if self.buffer.ends_with("```execute_typescript") { - let before = self.buffer.trim_end_matches("```execute_typescript"); - if !before.trim().is_empty() { - results.push(EmulatorAction::Text(before.to_string())); - } - self.buffer.clear(); - self.state = ParserState::InExecuteBlock; continue; } } - if let Some((before_dollar, from_dollar)) = self.buffer.split_once("\n$") { - let text = format!("{}\n", before_dollar); - if !text.trim().is_empty() { - results.push(EmulatorAction::Text(text)); - } - self.buffer = format!("${}", from_dollar); - self.state = ParserState::InCommand; - } else if self.buffer.starts_with('$') && self.buffer.len() == chunk.len() { - self.state = ParserState::InCommand; - } else { - let hold_back = if self.code_mode_enabled { - HOLD_BACK_CODE_MODE - } else { - HOLD_BACK_SHELL_ONLY - }; - let char_count = self.buffer.chars().count(); - if char_count > hold_back && !self.buffer.ends_with('\n') { - let mut chars = self.buffer.chars(); - let emit_count = char_count - hold_back; - let emit_text: String = chars.by_ref().take(emit_count).collect(); - let keep_text: String = chars.collect(); - if !emit_text.is_empty() { - results.push(EmulatorAction::Text(emit_text)); - } - self.buffer = keep_text; - } - break; - } + results.push(EmulatorAction::Text(line_with_newline)); } } } @@ -240,39 +405,68 @@ impl StreamingEmulatorParser { } pub(crate) fn flush(&mut self) -> Vec { - let mut results = Vec::new(); + let mut results = self.process_chunk(""); - if !self.buffer.is_empty() { - match self.state { - ParserState::InCommand => { - let command_line = self.buffer.trim(); - if let Some(command) = command_line.strip_prefix('$') { + match self.state { + ParserState::InExecuteBlock { fence_len } => { + let code_end = closing_fence_range(&self.buffer, fence_len, true) + .map(|(closing_start, _)| { + closing_start + .checked_sub(1) + .filter(|index| self.buffer.as_bytes()[*index] == b'\n') + .unwrap_or(closing_start) + }) + .unwrap_or(self.buffer.len()); + let code = self + .buffer + .get(..code_end) + .expect("code boundary must be a character boundary"); + let code = code.strip_suffix('\r').unwrap_or(code).trim(); + if !code.is_empty() { + results.push(EmulatorAction::ExecuteCode(code.to_string())); + } + } + ParserState::Normal if !self.buffer.is_empty() => { + let line = self + .buffer + .strip_suffix('\r') + .unwrap_or(&self.buffer) + .to_string(); + let markdown_range = self.markdown_context_available.then(|| { + let line_start = self.document.len() - self.buffer.len(); + (line_start, self.document.len()) + }); + + if self.code_mode_enabled + && execute_fence_len(&line).is_some() + && self.markdown_matches(markdown_range, is_top_level_execute_fence) + { + } else if let Some(command) = line.strip_prefix('$') { + if self.markdown_matches(markdown_range, is_top_level_paragraph_line) { let command = command.trim(); if !command.is_empty() { results.push(EmulatorAction::ShellCommand(command.to_string())); } - } else if !command_line.is_empty() { + } else { results.push(EmulatorAction::Text(self.buffer.clone())); } - } - ParserState::InExecuteBlock => { - let code = self.buffer.trim(); - if !code.is_empty() { - results.push(EmulatorAction::ExecuteCode(code.to_string())); - } - } - ParserState::Normal => { + } else { results.push(EmulatorAction::Text(self.buffer.clone())); } } - self.buffer.clear(); - self.state = ParserState::Normal; + ParserState::Normal => {} } + self.buffer.clear(); + self.document.clear(); + self.markdown_context_available = true; + self.markdown_parse_bytes_remaining = MAX_MARKDOWN_PARSE_BYTES; + self.state = ParserState::Normal; results } } +#[cfg(feature = "mlx")] pub(crate) fn message_for_emulator_action( action: &EmulatorAction, message_id: &str, @@ -335,221 +529,203 @@ mod tests { parse_chunks(&[input], code_mode) } - fn assert_text(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::Text(t) => assert_eq!(t.trim(), expected.trim(), "text mismatch"), - other => panic!("expected Text, got {:?}", action_label(other)), - } - } - - fn assert_shell(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::ShellCommand(cmd) => { - assert_eq!(cmd, expected, "shell command mismatch") - } - other => panic!("expected ShellCommand, got {:?}", action_label(other)), - } - } - - fn assert_execute(action: &EmulatorAction, expected: &str) { - match action { - EmulatorAction::ExecuteCode(code) => { - assert_eq!(code.trim(), expected.trim(), "execute code mismatch") - } - other => panic!("expected ExecuteCode, got {:?}", action_label(other)), - } - } - - fn action_label(a: &EmulatorAction) -> &'static str { - match a { - EmulatorAction::Text(_) => "Text", - EmulatorAction::ShellCommand(_) => "ShellCommand", - EmulatorAction::ExecuteCode(_) => "ExecuteCode", - } - } - - #[test] - fn plain_text_no_tools() { - let actions = parse_all("Hello, world!", false); - let all_text: String = actions + fn text(actions: &[EmulatorAction]) -> String { + actions .iter() - .map(|a| match a { - EmulatorAction::Text(t) => t.as_str(), - _ => panic!("expected only Text actions"), - }) - .collect(); - assert_eq!(all_text.trim(), "Hello, world!"); - } - - #[test] - fn single_shell_command() { - let actions = parse_all("$ ls -la\n", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "ls -la"); - } - - #[test] - fn text_then_shell_command() { - let actions = parse_all("Let me check:\n$ ls -la\n", false); - assert!(actions.len() >= 2); - assert_text(&actions[0], "Let me check:"); - assert_shell(&actions[actions.len() - 1], "ls -la"); - } - - #[test] - fn shell_command_at_start_of_output() { - let actions = parse_all("$ whoami\n", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "whoami"); - } - - #[test] - fn shell_command_without_trailing_newline() { - let actions = parse_all("$ whoami", false); - assert_eq!(actions.len(), 1); - assert_shell(&actions[0], "whoami"); - } - - #[test] - fn dollar_sign_mid_sentence_is_not_command() { - let actions = parse_all("It costs $50 per month", false); - for action in &actions { - assert!(matches!(action, EmulatorAction::Text(_))); - } - let all_text: String = actions - .iter() - .filter_map(|a| match a { - EmulatorAction::Text(t) => Some(t.as_str()), + .filter_map(|action| match action { + EmulatorAction::Text(text) => Some(text.as_str()), _ => None, }) - .collect(); - assert_eq!(all_text.trim(), "It costs $50 per month"); + .collect() + } + + fn shell_commands(actions: &[EmulatorAction]) -> Vec<&str> { + actions + .iter() + .filter_map(|action| match action { + EmulatorAction::ShellCommand(command) => Some(command.as_str()), + _ => None, + }) + .collect() + } + + fn execute_blocks(actions: &[EmulatorAction]) -> Vec<&str> { + actions + .iter() + .filter_map(|action| match action { + EmulatorAction::ExecuteCode(code) => Some(code.as_str()), + _ => None, + }) + .collect() } #[test] - fn execute_block() { + fn plain_text_is_preserved() { + let input = "Hello, world!\n"; + let actions = parse_all(input, false); + + assert_eq!(text(&actions), input); + assert!(shell_commands(&actions).is_empty()); + assert!(execute_blocks(&actions).is_empty()); + } + + #[test] + fn top_level_shell_commands_are_detected() { + let actions = parse_chunks(&["Let me check:\n", "$ who", "ami\n$ pwd"], false); + + assert_eq!(shell_commands(&actions), ["whoami", "pwd"]); + assert_eq!(text(&actions), "Let me check:\n"); + } + + #[test] + fn dollar_sign_mid_sentence_is_text() { + let input = "It costs $50 per month"; + let actions = parse_all(input, false); + + assert_eq!(text(&actions), input); + assert!(shell_commands(&actions).is_empty()); + } + + #[test] + fn markdown_contained_shell_commands_remain_text() { + let cases = [ + "````markdown\n$ inert\n````\n", + "> $ inert\n", + "- $ inert\n", + "
\n$ inert\n\n", + ]; + + for input in cases { + let actions = parse_all(input, true); + assert_eq!(text(&actions), input, "input: {input:?}"); + assert!(shell_commands(&actions).is_empty(), "input: {input:?}"); + } + } + + #[test] + fn top_level_execute_block_after_text_is_detected() { let input = "Here's the code:\n```execute_typescript\nconsole.log('hi');\n```\n"; let actions = parse_all(input, true); - assert!(actions.len() >= 2); - assert_text(&actions[0], "Here's the code:"); - assert_execute(&actions[actions.len() - 1], "console.log('hi');"); + + assert_eq!(text(&actions), "Here's the code:\n"); + assert_eq!(execute_blocks(&actions), ["console.log('hi');"]); } #[test] - fn tool_description_uses_parser_execute_fence() { + fn execute_fence_can_be_split_across_chunks() { + let input = "```execute_typescript\nlet x = 1;\n```\n"; + let chunks: Vec = input.chars().map(|ch| ch.to_string()).collect(); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let actions = parse_chunks(&chunk_refs, true); + + assert_eq!(execute_blocks(&actions), ["let x = 1;"]); + } + + #[test] + fn execute_blocks_are_disabled_without_code_mode() { + let input = "```execute_typescript\nlet x = 1;\n```\n"; + let actions = parse_all(input, false); + + assert_eq!(text(&actions), input); + assert!(execute_blocks(&actions).is_empty()); + } + + #[test] + fn nested_execute_fences_remain_text() { + let cases = [ + "````markdown\n```execute_typescript\ninert();\n```\n````\n", + "~~~markdown\n```execute_typescript\ninert();\n```\n~~~\n", + "> ```execute_typescript\n> inert();\n> ```\n", + "- ```execute_typescript\n inert();\n ```\n", + "
\n```execute_typescript\ninert();\n```\n\n", + ]; + + for input in cases { + let actions = parse_all(input, true); + assert_eq!(text(&actions), input, "input: {input:?}"); + assert!(execute_blocks(&actions).is_empty(), "input: {input:?}"); + } + } + + #[test] + fn invalid_fence_info_does_not_hide_following_execute() { + let input = "```execute_typescript```\n```execute_typescript\nsafe();\n```\n"; + let actions = parse_all(input, true); + + assert_eq!(execute_blocks(&actions), ["safe();"]); + assert!(text(&actions).contains("```execute_typescript```")); + } + + #[test] + fn longer_execute_fence_requires_matching_close() { + let input = "````execute_typescript\nlet before = 1;\n```\nlet after = 2;\n````\n"; + let actions = parse_all(input, true); + + assert_eq!( + execute_blocks(&actions), + ["let before = 1;\n```\nlet after = 2;"] + ); + } + + #[test] + fn closing_fence_waits_for_complete_line() { + let mut parser = StreamingEmulatorParser::new(true); + + assert!(parser + .process_chunk("```execute_typescript\nlet x = 1;\n```") + .iter() + .all(|action| !matches!(action, EmulatorAction::ExecuteCode(_)))); + assert!(parser + .process_chunk("not-a-close") + .iter() + .all(|action| !matches!(action, EmulatorAction::ExecuteCode(_)))); + let actions = parser.process_chunk("\n```\n"); + + assert_eq!(execute_blocks(&actions), ["let x = 1;\n```not-a-close"]); + } + + #[test] + fn unicode_whitespace_does_not_close_execute_fence() { + let input = "```execute_typescript\nlet before = 1;\n```\u{a0}\nlet after = 2;\n```\n"; + let actions = parse_all(input, true); + + assert_eq!( + execute_blocks(&actions), + ["let before = 1;\n```\u{a0}\nlet after = 2;"] + ); + } + + #[test] + fn markdown_parse_work_is_bounded() { + let mut input = String::from("````markdown\n"); + for _ in 0..1000 { + input.push_str("$ inert\n"); + } + input.push_str("````\n$ still inert\n"); + + let actions = parse_all(&input, true); + + assert_eq!(text(&actions), input); + assert!(shell_commands(&actions).is_empty()); + assert!(execute_blocks(&actions).is_empty()); + } + + #[test] + fn eof_flushes_closed_and_unclosed_execute_blocks() { + for (input, expected) in [ + ("```execute_typescript\nsafe();\n```", "safe();"), + ("```execute_typescript\nsafe();", "safe();"), + ] { + let actions = parse_all(input, true); + assert_eq!(execute_blocks(&actions), [expected]); + } + } + + #[cfg(feature = "mlx")] + #[test] + fn tool_description_uses_execute_typescript_fence() { let description = build_emulator_tool_description(&[], true); assert!(description.contains("```execute_typescript")); - assert!(!description.contains("```execute block")); - assert!(!description.contains("Use ```execute for tool calls")); - } - - #[test] - fn execute_block_not_detected_without_code_mode() { - let input = "```execute_typescript\nconsole.log('hi');\n```\n"; - let actions = parse_all(input, false); - for action in &actions { - assert!(matches!(action, EmulatorAction::Text(_))); - } - } - - #[test] - fn dollar_split_across_chunks() { - let actions = parse_chunks(&["Let me check\n", "$ ls -la\n"], false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 1); - assert_shell(shells[0], "ls -la"); - } - - #[test] - fn execute_fence_split_across_chunks() { - let actions = parse_chunks( - &["Here:\n```ex", "ecute_typescript\nlet x = 1;\n", "```\n"], - true, - ); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - assert_execute(executes[0], "let x = 1;"); - } - - #[test] - fn multiple_commands_on_separate_lines() { - let actions = parse_chunks(&["Here:\n$ cd /tmp\n", "Done.\n$ ls\n"], false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 2); - assert_shell(shells[0], "cd /tmp"); - assert_shell(shells[1], "ls"); - } - - #[test] - fn regular_code_fence_not_treated_as_execute() { - let input = "```python\nprint('hi')\n```\n"; - let actions = parse_all(input, true); - for action in &actions { - assert!(matches!(action, EmulatorAction::Text(_))); - } - } - - #[test] - fn empty_command_ignored() { - let actions = parse_all("$\n", false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 0); - } - - #[test] - fn token_by_token_streaming() { - let input = "$ echo hello\n"; - let chars: Vec = input.chars().map(|c| c.to_string()).collect(); - let chunks: Vec<&str> = chars.iter().map(|s| s.as_str()).collect(); - let actions = parse_chunks(&chunks, false); - let shells: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ShellCommand(_))) - .collect(); - assert_eq!(shells.len(), 1); - assert_shell(shells[0], "echo hello"); - } - - #[test] - fn execute_block_with_multiline_code() { - let input = "```execute_typescript\nasync function run() {\n const r = await Developer.shell({ command: \"ls\" });\n return r;\n}\n```\n"; - let actions = parse_all(input, true); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - match executes[0] { - EmulatorAction::ExecuteCode(code) => { - assert!(code.contains("async function run()")); - assert!(code.contains("Developer.shell")); - } - _ => unreachable!(), - } - } - - #[test] - fn unclosed_execute_block_flushed() { - let input = "```execute_typescript\nlet x = 1;"; - let actions = parse_all(input, true); - let executes: Vec<_> = actions - .iter() - .filter(|a| matches!(a, EmulatorAction::ExecuteCode(_))) - .collect(); - assert_eq!(executes.len(), 1); - assert_execute(executes[0], "let x = 1;"); } }