diff --git a/crates/goose-provider-types/src/thinking.rs b/crates/goose-provider-types/src/thinking.rs index e219549b3..20041b2ed 100644 --- a/crates/goose-provider-types/src/thinking.rs +++ b/crates/goose-provider-types/src/thinking.rs @@ -4,6 +4,7 @@ use regex::Regex; use serde::{Deserialize, Serialize}; pub const GEMINI_THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature"; +const MAX_BUFFERED_THINK_TAG_BYTES: usize = 8 * 1024; pub fn split_think_blocks(text: &str) -> (String, String) { let mut filter = ThinkFilter::new(); @@ -39,6 +40,9 @@ enum BufferEvent { end: usize, kind: ThinkTag, }, + OversizedTag { + end: usize, + }, Partial(usize), } @@ -98,6 +102,15 @@ impl ThinkFilter { ThinkTag::SelfClosing => {} } } + Some(BufferEvent::OversizedTag { end }) => { + let malformed = self.buffer.get(..end).unwrap_or_default().to_string(); + if self.inside_think { + out.thinking.push_str(&malformed); + } else { + out.content.push_str(&malformed); + } + self.buffer.drain(..end); + } Some(BufferEvent::Partial(pos)) => { if pos > 0 { let prefix = self.buffer.get(..pos).unwrap_or_default().to_string(); @@ -108,6 +121,14 @@ impl ThinkFilter { } self.buffer.drain(..pos); } + if self.buffer.len() > MAX_BUFFERED_THINK_TAG_BYTES { + let oversized = std::mem::take(&mut self.buffer); + if self.inside_think { + out.thinking.push_str(&oversized); + } else { + out.content.push_str(&oversized); + } + } break; } None => { @@ -124,6 +145,12 @@ impl ThinkFilter { } } + if self.buffer.capacity() > MAX_BUFFERED_THINK_TAG_BYTES { + let mut bounded = String::with_capacity(self.buffer.len()); + bounded.push_str(&self.buffer); + self.buffer = bounded; + } + out } } @@ -142,6 +169,9 @@ fn next_buffer_event(buffer: &str, inside_think: bool) -> Option { let suffix = buffer.get(pos..).unwrap_or_default(); if let Some((kind, end)) = parse_think_tag(buffer, pos) { + if end - pos > MAX_BUFFERED_THINK_TAG_BYTES { + return Some(BufferEvent::OversizedTag { end }); + } if inside_think || matches!(kind, ThinkTag::Open | ThinkTag::SelfClosing) { return Some(BufferEvent::Tag { pos, end, kind }); } @@ -583,4 +613,102 @@ mod tests { assert_eq!(out.content, "before visible"); assert_eq!(out.thinking, "hidden1 hidden2"); } + + #[test] + fn test_think_filter_bounds_unterminated_quoted_tag_candidates() { + for quote in ['"', '\''] { + let payload = format!( + "", + "a".repeat(MAX_BUFFERED_THINK_TAG_BYTES * 16) + ); + let mut filter = ThinkFilter::new(); + let out = filter.push(&payload); + + assert_eq!(out.content, payload); + assert!(out.thinking.is_empty()); + assert!(filter.buffer.is_empty()); + assert_eq!(filter.buffer.capacity(), 0); + } + + #[test] + fn test_think_filter_releases_oversized_prefix_capacity() { + let prefix = "a".repeat(MAX_BUFFERED_THINK_TAG_BYTES * 16); + let payload = format!("{prefix}hiddenvisible"); + let final_out = filter.finish(); + assert_eq!(second.content, "visible"); + assert_eq!(second.thinking, "hidden"); + assert_eq!(final_out, FilterOut::default()); + } }