Merge pull request #881 from pbranchu/fix/token-estimation-tooluse

Fix token estimation: include ToolUse arguments in text_length

Author: jaberjaber23

crates/openfang-api/src/channel_bridge.rs

@@ -49,8 +49,8 @@ use openfang_channels::discourse::DiscourseAdapter;
 use openfang_channels::gitter::GitterAdapter;
 use openfang_channels::gotify::GotifyAdapter;
 use openfang_channels::linkedin::LinkedInAdapter;
-use openfang_channels::mumble::MumbleAdapter;
 use openfang_channels::mqtt::MqttAdapter;
+use openfang_channels::mumble::MumbleAdapter;
 use openfang_channels::ntfy::NtfyAdapter;
 use openfang_channels::webhook::WebhookAdapter;
 use openfang_channels::wecom::WeComAdapter;

crates/openfang-api/src/routes.rs

@@ -579,7 +579,8 @@ pub async fn get_agent_session(
                                         msg.get_mut("tools").and_then(|v| v.as_array_mut())
                                     {
                                         if let Some(tool_obj) = tools_arr.get_mut(tool_idx) {
-                                            tool_obj["result"] = serde_json::Value::String(result.clone());
+                                            tool_obj["result"] =
+                                                serde_json::Value::String(result.clone());
                                             tool_obj["is_error"] =
                                                 serde_json::Value::Bool(*is_error);
                                         }

crates/openfang-channels/src/lib.rs

@@ -48,8 +48,8 @@ pub mod discourse;
 pub mod gitter;
 pub mod gotify;
 pub mod linkedin;
-pub mod mumble;
 pub mod mqtt;
+pub mod mumble;
 pub mod ntfy;
 pub mod webhook;
 pub mod wecom;

crates/openfang-channels/src/line.rs

@@ -108,7 +108,7 @@ impl LineAdapter {
             diff |= a ^ b;
         }
         if diff != 0 {
-            let computed = base64::engine::general_purpose::STANDARD.encode(&result);
+            let computed = base64::engine::general_purpose::STANDARD.encode(result);
             // Log first/last 4 chars of each signature for debugging without leaking full HMAC
             let comp_redacted = format!(
                 "{}...{}",
@@ -381,8 +381,7 @@ impl ChannelAdapter for LineAdapter {
                 axum::routing::post({
                     let secret = Arc::clone(&channel_secret);
                     let tx = Arc::clone(&tx);
-                    move |headers: axum::http::HeaderMap,
-                          body: axum::body::Bytes| {
+                    move |headers: axum::http::HeaderMap, body: axum::body::Bytes| {
                         let secret = Arc::clone(&secret);
                         let tx = Arc::clone(&tx);
                         async move {
@@ -404,8 +403,7 @@ impl ChannelAdapter for LineAdapter {
                                 shutdown_rx: watch::channel(false).1,
                             };
 
-                            if !signature.is_empty()
-                                && !adapter.verify_signature(&body, signature)
+                            if !signature.is_empty() && !adapter.verify_signature(&body, signature)
                             {
                                 warn!("LINE: invalid webhook signature");
                                 return axum::http::StatusCode::UNAUTHORIZED;

crates/openfang-channels/src/mqtt.rs

@@ -152,7 +152,10 @@ impl MqttAdapter {
     }
 
     /// Parse host:port string.
-    fn parse_host_port(s: &str, default_port: u16) -> Result<(String, u16), Box<dyn std::error::Error>> {
+    fn parse_host_port(
+        s: &str,
+        default_port: u16,
+    ) -> Result<(String, u16), Box<dyn std::error::Error>> {
         let s = s.trim();
         if let Some(colon_pos) = s.rfind(':') {
             let host = s[..colon_pos].to_string();
@@ -239,7 +242,8 @@ impl ChannelAdapter for MqttAdapter {
 
     async fn start(
         &self,
-    ) -> Result<Pin<Box<dyn Stream<Item = ChannelMessage> + Send>>, Box<dyn std::error::Error>> {
+    ) -> Result<Pin<Box<dyn Stream<Item = ChannelMessage> + Send>>, Box<dyn std::error::Error>>
+    {
         let options = self.build_mqtt_options()?;
         let (client, mut eventloop) = AsyncClient::new(options, 10);
 

crates/openfang-kernel/src/kernel.rs

@@ -2930,20 +2930,16 @@ impl OpenFangKernel {
         model: &str,
         explicit_provider: Option<&str>,
     ) -> KernelResult<()> {
-        let catalog_entry = self
-            .model_catalog
-            .read()
-            .ok()
-            .and_then(|catalog| {
-                // When the caller specifies a provider, use provider-aware lookup
-                // so we resolve the model on the correct provider — not a builtin
-                // from a different provider that happens to share the same name (#833).
-                if let Some(ep) = explicit_provider {
-                    catalog.find_model_for_provider(model, ep).cloned()
-                } else {
-                    catalog.find_model(model).cloned()
-                }
-            });
+        let catalog_entry = self.model_catalog.read().ok().and_then(|catalog| {
+            // When the caller specifies a provider, use provider-aware lookup
+            // so we resolve the model on the correct provider — not a builtin
+            // from a different provider that happens to share the same name (#833).
+            if let Some(ep) = explicit_provider {
+                catalog.find_model_for_provider(model, ep).cloned()
+            } else {
+                catalog.find_model(model).cloned()
+            }
+        });
         let provider = if let Some(ep) = explicit_provider {
             // User explicitly set the provider — use it as-is
             Some(ep.to_string())

crates/openfang-runtime/src/compactor.rs

@@ -1478,7 +1478,10 @@ mod tests {
             Message::assistant("Done reading."),
         ];
         let adjusted = adjust_split_for_tool_pairs(&messages, 2);
-        assert_eq!(adjusted, 1, "Should pull back split to keep ToolUse + ToolResult together");
+        assert_eq!(
+            adjusted, 1,
+            "Should pull back split to keep ToolUse + ToolResult together"
+        );
     }
 
     #[test]
@@ -1489,7 +1492,10 @@ mod tests {
             Message::user("c"),
         ];
         let adjusted = adjust_split_for_tool_pairs(&messages, 1);
-        assert_eq!(adjusted, 1, "Should not change split for plain text messages");
+        assert_eq!(
+            adjusted, 1,
+            "Should not change split for plain text messages"
+        );
     }
 
     #[test]

crates/openfang-runtime/src/context_overflow.rs

@@ -32,10 +32,14 @@ fn safe_drain_boundary(messages: &[Message], mut boundary: usize) -> usize {
     // is in the last drained message (boundary - 1).  Pull boundary back by 1.
     if messages[boundary].role == Role::User {
         if let MessageContent::Blocks(blocks) = &messages[boundary].content {
-            let has_tool_result = blocks.iter().any(|b| matches!(b, ContentBlock::ToolResult { .. }));
+            let has_tool_result = blocks
+                .iter()
+                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
             if has_tool_result && boundary > 0 && messages[boundary - 1].role == Role::Assistant {
                 if let MessageContent::Blocks(asst_blocks) = &messages[boundary - 1].content {
-                    let has_tool_use = asst_blocks.iter().any(|b| matches!(b, ContentBlock::ToolUse { .. }));
+                    let has_tool_use = asst_blocks
+                        .iter()
+                        .any(|b| matches!(b, ContentBlock::ToolUse { .. }));
                     if has_tool_use {
                         boundary -= 1;
                         debug!(
@@ -135,7 +139,8 @@ pub fn recover_from_overflow(
             debug!(
                 estimated_tokens = estimated,
                 removing = remove,
-                "Stage 1: moderate trim to last {} messages", messages.len() - remove
+                "Stage 1: moderate trim to last {} messages",
+                messages.len() - remove
             );
             messages.drain(..remove);
             // Re-check after trim
@@ -156,7 +161,8 @@ pub fn recover_from_overflow(
             warn!(
                 estimated_tokens = estimate_tokens(messages, system_prompt, tools),
                 removing = remove,
-                "Stage 2: aggressive overflow compaction to last {} messages", messages.len() - remove
+                "Stage 2: aggressive overflow compaction to last {} messages",
+                messages.len() - remove
             );
             let summary = Message::user(format!(
                 "[System: {} earlier messages were removed due to context overflow. \
@@ -373,7 +379,10 @@ mod tests {
         ];
         // Boundary 2 would cut between the assistant(ToolUse) at [1] and user(ToolResult) at [2].
         let adjusted = safe_drain_boundary(&msgs, 2);
-        assert_eq!(adjusted, 1, "Should pull boundary back to keep the ToolUse/ToolResult pair together");
+        assert_eq!(
+            adjusted, 1,
+            "Should pull boundary back to keep the ToolUse/ToolResult pair together"
+        );
     }
 
     #[test]
@@ -385,7 +394,10 @@ mod tests {
             Message::assistant("d"),
         ];
         let adjusted = safe_drain_boundary(&msgs, 2);
-        assert_eq!(adjusted, 2, "Should not change boundary for plain text messages");
+        assert_eq!(
+            adjusted, 2,
+            "Should not change boundary for plain text messages"
+        );
     }
 
     #[test]

crates/openfang-types/src/message.rs

@@ -142,9 +142,10 @@ impl MessageContent {
                     ContentBlock::Text { text, .. } => text.len(),
                     ContentBlock::ToolResult { content, .. } => content.len(),
                     ContentBlock::Thinking { thinking } => thinking.len(),
-                    ContentBlock::ToolUse { .. }
-                    | ContentBlock::Image { .. }
-                    | ContentBlock::Unknown => 0,
+                    ContentBlock::ToolUse { name, input, .. } => {
+                        name.len() + input.to_string().len()
+                    }
+                    ContentBlock::Image { .. } | ContentBlock::Unknown => 0,
                 })
                 .sum(),
         }