From c28f347979549ee49bd3c27ad68a263b3695a36a Mon Sep 17 00:00:00 2001 From: M3gA-Mind Date: Tue, 7 Jul 2026 19:13:06 +0530 Subject: [PATCH 1/2] harness: size token estimation over all content blocks, not just text MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Images, structured JSON (large tool results), and model reasoning occupy real context, but char_len counts only Text blocks — so a multimodal or tool-heavy transcript under-counts and should_summarize / MaxTokens trim silently never fire. Add Message::estimated_char_weight over all content blocks and route the token estimator through it; the visible-text accessors (as_text/text/char_len) are left unchanged so reasoning still never leaks into visible output. --- src/harness/message/mod.rs | 55 ++++++++++++++++++++++++++++++++ src/harness/message/test.rs | 40 +++++++++++++++++++++++ src/harness/summarization/mod.rs | 12 +++++-- 3 files changed, 104 insertions(+), 3 deletions(-) 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/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) } } From 94b0d996734db216461bebb2a7f4df3c309daf0d Mon Sep 17 00:00:00 2001 From: M3gA-Mind Date: Tue, 7 Jul 2026 19:38:09 +0530 Subject: [PATCH 2/2] harness: fix clippy lints surfaced by stable 1.96 collapsible_if in StreamAccumulator::push_tool_chunk and a duplicated #[cfg(test)] on the stream test module both fail CI's cargo clippy --all-targets -- -D warnings under rust 1.96. Behavior-neutral. --- src/harness/model/mod.rs | 8 ++++---- src/harness/stream/mod.rs | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) 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;