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

Fix token estimation: include ToolUse arguments in text_length
This commit is contained in:
Jaber Jaber
2026-03-31 02:21:55 +03:00
committed by GitHub
9 changed files with 53 additions and 35 deletions
+1 -1
View File
@@ -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;
+2 -1
View File
@@ -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);
}
+1 -1
View File
@@ -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;
+3 -5
View File
@@ -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;
+6 -2
View File
@@ -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);
+10 -14
View File
@@ -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())
+8 -2
View File
@@ -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]
@@ -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]
+4 -3
View File
@@ -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(),
}