diff --git a/src/harness/message/mod.rs b/src/harness/message/mod.rs index 679aabc..67ffe03 100644 --- a/src/harness/message/mod.rs +++ b/src/harness/message/mod.rs @@ -14,6 +14,16 @@ mod types; pub use types::*; +/// Approximate token-estimation weight of a single image content block, in +/// "characters" (the token estimator divides char weight by 4, so this is +/// ~1024 tokens per image). +/// +/// Vision models tokenize an image into a roughly fixed count that is +/// independent of the encoded byte length, so this is a flat conservative +/// estimate rather than the (potentially huge, e.g. a base64 `data:` URI) +/// [`ImageRef::url`] length, which would wildly over-count. +const IMAGE_CHAR_WEIGHT: usize = 4 * 1024; + impl ContentBlock { /// Returns the text of this block if it is a [`ContentBlock::Text`]. /// @@ -55,6 +65,29 @@ impl ContentBlock { ContentBlock::Thinking { .. } | ContentBlock::RedactedThinking { .. } ) } + + /// Approximate character weight of this block for token estimation. + /// + /// Unlike [`as_text`](Self::as_text) — which returns only visible + /// [`Text`](Self::Text) so reasoning never leaks into assistant output — + /// this accounts for *every* block that occupies model context: text, + /// structured JSON, reasoning ([`Thinking`](Self::Thinking) / + /// [`RedactedThinking`](Self::RedactedThinking)), provider extensions, and + /// a flat [`IMAGE_CHAR_WEIGHT`] per image. It is used only by token + /// budgeting / compaction gating, never by the visible-text accessors, so a + /// transcript dominated by images, large tool-result JSON, or model + /// reasoning no longer under-counts to near-zero and silently defeats + /// summarization. + pub fn estimated_char_weight(&self) -> usize { + match self { + ContentBlock::Text(text) => text.chars().count(), + ContentBlock::Json(value) => value.to_string().chars().count(), + ContentBlock::Image(_) => IMAGE_CHAR_WEIGHT, + ContentBlock::Thinking { text, .. } => text.chars().count(), + ContentBlock::RedactedThinking { data } => data.chars().count(), + ContentBlock::ProviderExtension(value) => value.to_string().chars().count(), + } + } } /// Concatenates the text of all [`ContentBlock::Text`] blocks in `content`. @@ -128,6 +161,28 @@ impl Message { .map(|t| t.chars().count()) .sum() } + + /// Approximate character weight of the message across *all* content blocks + /// (text, JSON, images, reasoning, provider extensions), for token + /// estimation and context-window gating. + /// + /// Distinct from [`char_len`](Self::char_len), which counts only visible + /// text: a transcript dominated by images, large tool-result JSON, or model + /// reasoning under-counts badly under `char_len`, so compaction/trim would + /// silently never trigger even as the real context window overflows. See + /// [`ContentBlock::estimated_char_weight`]. + pub fn estimated_char_weight(&self) -> usize { + let content = match self { + Message::System(m) => &m.content, + Message::User(m) => &m.content, + Message::Assistant(m) => &m.content, + Message::Tool(m) => &m.content, + }; + content + .iter() + .map(ContentBlock::estimated_char_weight) + .sum() + } } #[cfg(test)] diff --git a/src/harness/message/test.rs b/src/harness/message/test.rs index bf5d6dc..5db7b96 100644 --- a/src/harness/message/test.rs +++ b/src/harness/message/test.rs @@ -49,6 +49,46 @@ fn char_len_matches_text_char_count_including_multibyte() { ); } +#[test] +fn estimated_char_weight_counts_non_text_blocks() { + let text = "hello"; + let reasoning = "some private reasoning"; + let json_block = json!({"key": "value", "n": 42}); + let msg = Message::User(UserMessage { + content: vec![ + ContentBlock::Text(text.into()), + ContentBlock::Json(json_block.clone()), + ContentBlock::thinking(reasoning), + ], + }); + + // Visible-text accounting still ignores JSON and reasoning. + assert_eq!(msg.char_len(), text.chars().count()); + + // Token-estimation weight additionally counts JSON and reasoning, which + // occupy real context, so it is strictly larger than the visible-text len. + let expected = + text.chars().count() + json_block.to_string().chars().count() + reasoning.chars().count(); + assert_eq!(msg.estimated_char_weight(), expected); + assert!(msg.estimated_char_weight() > msg.char_len()); +} + +#[test] +fn estimated_char_weight_costs_images_that_carry_no_visible_text() { + let msg = Message::User(UserMessage { + content: vec![ContentBlock::Image(ImageRef { + url: "data:image/png;base64,AAAA".into(), + mime_type: Some("image/png".into()), + })], + }); + + // An image contributes no visible text but a flat, non-trivial weight, so a + // vision-heavy transcript can no longer under-count to zero and defeat + // compaction gating. + assert_eq!(msg.char_len(), 0); + assert_eq!(msg.estimated_char_weight(), super::IMAGE_CHAR_WEIGHT); +} + #[test] fn assistant_holds_tool_calls_and_usage() { let msg = Message::Assistant(AssistantMessage { diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index aa24cb1..ea3e6b8 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -753,10 +753,10 @@ impl StreamAccumulator { fn push_tool_chunk(&mut self, call_id: &str, content: &str, tool_name: Option<&str>) { if let Some(entry) = self.tool_chunks.iter_mut().find(|(id, ..)| id == call_id) { entry.1.push_str(content); - if entry.2.is_none() { - if let Some(name) = tool_name.filter(|n| !n.is_empty()) { - entry.2 = Some(name.to_string()); - } + if entry.2.is_none() + && let Some(name) = tool_name.filter(|n| !n.is_empty()) + { + entry.2 = Some(name.to_string()); } } else { self.tool_chunks.push(( diff --git a/src/harness/stream/mod.rs b/src/harness/stream/mod.rs index 14411f0..1712142 100644 --- a/src/harness/stream/mod.rs +++ b/src/harness/stream/mod.rs @@ -143,6 +143,5 @@ pub fn stream(chunks: &[StreamChunk], modes: &[StreamMode]) -> Vec .collect() } -#[cfg(test)] #[cfg(test)] mod test; diff --git a/src/harness/summarization/mod.rs b/src/harness/summarization/mod.rs index 3938d8e..74b0c6d 100644 --- a/src/harness/summarization/mod.rs +++ b/src/harness/summarization/mod.rs @@ -48,10 +48,16 @@ pub fn estimate_tokens(text: &str) -> u64 { } /// Estimate the total tokens for a [`Message`] using the same `chars / 4` -/// heuristic as [`estimate_tokens`], but counting characters directly over the -/// message's text blocks rather than allocating the concatenated string first. +/// heuristic as [`estimate_tokens`], counting weight directly over the message +/// content rather than allocating the concatenated string first. +/// +/// Uses [`Message::estimated_char_weight`] (all content blocks) rather than +/// `char_len` (visible text only): images, structured JSON, and model reasoning +/// occupy real context, so counting only text would under-estimate a +/// multimodal or tool-heavy transcript to near-zero and silently prevent +/// [`SummarizationPolicy::should_summarize`] from ever triggering. fn message_token_estimate(msg: &Message) -> u64 { - let chars = msg.char_len() as u64; + let chars = msg.estimated_char_weight() as u64; if chars == 0 { 0 } else { (chars / 4).max(1) } }