diff --git a/src/agent/history.rs b/src/agent/history.rs index beabad29..74a0403c 100644 --- a/src/agent/history.rs +++ b/src/agent/history.rs @@ -21,6 +21,19 @@ pub(crate) fn floor_char_boundary(s: &str, i: usize) -> usize { pos } +/// Count the number of bytes a string would occupy when JSON-serialized +/// (excluding the surrounding quotes). This avoids allocating a throwaway +/// `serde_json::to_string` just to measure length. +fn json_escaped_len(s: &str) -> usize { + s.bytes() + .map(|b| match b { + b'"' | b'\\' | b'\n' | b'\r' | b'\t' => 2, + b if b < 0x20 => 6, // \u00XX + _ => 1, + }) + .sum() +} + /// Truncate a tool result to `max_chars`, keeping head (2/3) + tail (1/3) /// with a marker in the middle. Returns input unchanged if within limit or /// `max_chars == 0` (disabled). @@ -53,20 +66,51 @@ pub(crate) fn truncate_tool_message_content(content: &str, max_chars: usize) -> if value.get("tool_call_id").is_none() { return truncate_plain_text(content, max_chars); } - let Some(inner) = value.get("content").and_then(|v| v.as_str()) else { + // Clone inner content to release the immutable borrow on `value`, + // which we need to mutate below. + let Some(inner) = value + .get("content") + .and_then(|v| v.as_str()) + .map(String::from) + else { return truncate_plain_text(content, max_chars); }; - // inner.len() is raw bytes; content.len() is JSON-serialized (escaping may - // make the serialized form shorter than the raw value), so use saturating_sub. - let envelope_overhead = content.len().saturating_sub(inner.len()); - let inner_max = max_chars.saturating_sub(envelope_overhead); - if inner_max == 0 || inner.len() <= inner_max { - return truncate_plain_text(content, max_chars); + // Compute serialized length of the inner string without allocating: + // each byte that needs JSON escaping expands, plus 2 for the quotes. + let inner_serialized_len = json_escaped_len(&inner) + 2; + let envelope_overhead = content.len().saturating_sub(inner_serialized_len); + + let mut inner_max = max_chars.saturating_sub(envelope_overhead); + // When escaping makes the envelope estimate unreliable, use a + // conservative fraction of the budget as a starting point. + if inner_max == 0 { + inner_max = max_chars / 4; + } + if inner.len() <= inner_max { + return content.to_string(); + } + + // Shrink inner content until the re-serialized JSON fits the budget. + // Capped to prevent pathological inputs from hot-looping. + for _ in 0..8 { + let truncated = truncate_plain_text(&inner, inner_max); + value["content"] = serde_json::Value::String(truncated); + match serde_json::to_string(&value) { + Ok(result) if result.len() <= max_chars => return result, + Ok(result) => { + let overshoot = result.len() - max_chars; + inner_max = inner_max.saturating_sub(overshoot + 1); + if inner_max == 0 { + break; + } + } + Err(_) => return truncate_plain_text(content, max_chars), + } } - let truncated = truncate_plain_text(inner, inner_max); - value["content"] = serde_json::Value::String(truncated); + // Envelope alone exceeds budget or iterations exhausted; drop inner. + value["content"] = serde_json::Value::String("[content truncated]".into()); serde_json::to_string(&value).unwrap_or_else(|_| truncate_plain_text(content, max_chars)) } diff --git a/src/agent/loop_.rs b/src/agent/loop_.rs index 44e49610..942a73b4 100644 --- a/src/agent/loop_.rs +++ b/src/agent/loop_.rs @@ -4774,6 +4774,35 @@ mod tests { assert!(inner_content.contains("characters truncated")); } + #[test] + fn truncate_tool_message_content_preserves_json_with_heavy_escaping() { + // Simulate binary content (null bytes) that expands heavily under JSON escaping. + // Each \0 becomes \u0000 (6 chars) in JSON, so 5000 null bytes ≈ 30000 JSON chars. + let inner = "\0".repeat(5000); + let content = serde_json::json!({ + "tool_call_id": "call_binary_test", + "content": inner, + }) + .to_string(); + + // The JSON-serialized content is ~30k+ chars. Truncate to 2000. + let result = truncate_tool_message_content(&content, 2000); + + // Must remain valid JSON with tool_call_id preserved + let parsed: serde_json::Value = serde_json::from_str(&result) + .expect("truncated binary tool message must remain valid JSON"); + assert_eq!( + parsed.get("tool_call_id").and_then(|v| v.as_str()), + Some("call_binary_test"), + "tool_call_id must be preserved even with heavily-escaped content" + ); + assert!( + result.len() <= 2000, + "result len {} must be within budget 2000", + result.len() + ); + } + #[test] fn truncate_tool_message_content_plain_text_fallback() { let content = "x".repeat(5000); @@ -6652,7 +6681,11 @@ mod tests { .find(|msg| msg.role == "user" && msg.content.starts_with("[Tool results]")) .expect("tool results message should be present"); assert!(tool_results.content.contains("hello")); - assert!(!tool_results.content.contains("The operation was NOT performed.")); + assert!( + !tool_results + .content + .contains("The operation was NOT performed.") + ); } #[tokio::test] diff --git a/src/providers/openrouter.rs b/src/providers/openrouter.rs index 7cb655a9..938e85ab 100644 --- a/src/providers/openrouter.rs +++ b/src/providers/openrouter.rs @@ -316,6 +316,19 @@ impl OpenRouterProvider { reasoning_details: None, }; } + // JSON parse failed (e.g. truncation broke the wrapper). + // Try to salvage tool_call_id via substring search so + // providers that require it (Gemini) don't reject the + // entire request. + let tool_call_id = Self::extract_tool_call_id_from_broken_json(&m.content); + return NativeMessage { + role: "tool".to_string(), + content: Some(MessageContent::Text(m.content.clone())), + tool_call_id, + tool_calls: None, + reasoning_content: None, + reasoning_details: None, + }; } NativeMessage { @@ -387,6 +400,22 @@ impl OpenRouterProvider { } } + /// Best-effort extraction of `tool_call_id` from a tool message whose + /// JSON wrapper was broken by truncation. Falls back to `None` if the + /// field cannot be found. + fn extract_tool_call_id_from_broken_json(content: &str) -> Option { + // Look for "tool_call_id":"" pattern + let marker = "\"tool_call_id\":\""; + let start = content.find(marker)? + marker.len(); + let end = content[start..].find('"')? + start; + let id = &content[start..end]; + if id.is_empty() { + None + } else { + Some(id.to_string()) + } + } + fn compact_sanitized_body_snippet(body: &str) -> String { super::sanitize_api_error(body) .split_whitespace() @@ -1099,6 +1128,29 @@ mod tests { assert!(converted[0].tool_calls.is_none()); } + #[test] + fn convert_messages_salvages_tool_call_id_from_broken_json() { + // Simulate a tool message whose JSON wrapper was broken by truncation + let broken = r#"{"content":"SQLite\u0000\u0 + +[... 139762 characters truncated ...] + +\u0001","tool_call_id":"tool_file_read_abc123"}"#; + let messages = vec![ChatMessage { + role: "tool".into(), + content: broken.into(), + }]; + + let converted = OpenRouterProvider::convert_messages(&messages); + assert_eq!(converted.len(), 1); + assert_eq!(converted[0].role, "tool"); + assert_eq!( + converted[0].tool_call_id.as_deref(), + Some("tool_file_read_abc123"), + "must salvage tool_call_id from broken JSON" + ); + } + #[test] fn to_message_content_converts_image_markers_to_openai_parts() { let content = "Describe this\n\n[IMAGE:data:image/png;base64,abcd]"; diff --git a/src/providers/router.rs b/src/providers/router.rs index b655c63a..431cbeb5 100644 --- a/src/providers/router.rs +++ b/src/providers/router.rs @@ -288,9 +288,14 @@ impl Provider for RouterProvider { } fn supports_streaming_tool_events(&self) -> bool { + // Use `all()` not `any()`: the router doesn't know which provider + // will handle a given request, so it can only claim support when + // every routed provider supports it. With `any()`, the agent loop + // sees "supported", attempts streaming, but the resolved provider + // may not support it — causing a fallback on every tool-bearing call. self.providers .iter() - .any(|(_, provider)| provider.supports_streaming_tool_events()) + .all(|(_, provider)| provider.supports_streaming_tool_events()) } fn stream_chat_with_history( diff --git a/src/tools/file_read.rs b/src/tools/file_read.rs index 510d9515..ad145dc2 100644 --- a/src/tools/file_read.rs +++ b/src/tools/file_read.rs @@ -208,21 +208,36 @@ impl Tool for FileReadTool { }); } - // Lossy fallback — replaces invalid bytes with U+FFFD - let lossy = String::from_utf8_lossy(&bytes).into_owned(); - Ok(ToolResult { - success: true, - output: lossy, - error: None, - }) + let filename = resolved_path + .file_name() + .map(|n| n.to_string_lossy()) + .unwrap_or_default(); + + if is_pdf_magic(&bytes) { + return ToolResult::err(format!( + "'{filename}' is a PDF file but PDF text extraction \ + is not available. Use the pdf_read tool or web_fetch \ + if the content is available online." + )); + } + + ToolResult::err(format!( + "'{filename}' is a binary file and cannot be read as text. \ + If this is a database file (.db, .sqlite), use the \ + memory_recall tool or an appropriate query tool instead." + )) } } } } +fn is_pdf_magic(bytes: &[u8]) -> bool { + bytes.len() >= 5 && &bytes[..5] == b"%PDF-" +} + #[cfg(feature = "rag-pdf")] fn try_extract_pdf_text(bytes: &[u8]) -> Option { - if bytes.len() < 5 || &bytes[..5] != b"%PDF-" { + if !is_pdf_magic(bytes) { return None; } let text = pdf_extract::extract_text_from_mem(bytes).ok()?; @@ -638,6 +653,7 @@ mod tests { } /// PDF files should be readable via pdf-extract text extraction. + #[cfg(feature = "rag-pdf")] #[tokio::test] async fn file_read_extracts_pdf_text() { let dir = std::env::temp_dir().join("mentat_test_file_read_pdf"); @@ -667,14 +683,40 @@ mod tests { let _ = tokio::fs::remove_dir_all(&dir).await; } - /// Non-UTF-8 binary files should be read with lossy conversion. + /// Without rag-pdf feature, PDF files should be rejected with a helpful error. + #[cfg(not(feature = "rag-pdf"))] + #[tokio::test] + async fn file_read_rejects_pdf_without_feature() { + let dir = std::env::temp_dir().join("mentat_test_file_read_pdf_reject"); + let _ = tokio::fs::remove_dir_all(&dir).await; + tokio::fs::create_dir_all(&dir).await.unwrap(); + + let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/test_document.pdf"); + tokio::fs::copy(&fixture, dir.join("report.pdf")) + .await + .expect("copy PDF fixture"); + + let tool = FileReadTool::new(test_security(dir.clone())); + let result = tool.execute(json!({"path": "report.pdf"})).await.unwrap(); + + assert!(!result.success, "PDF read without feature must fail"); + let err = result.error.unwrap(); + assert!( + err.contains("PDF file"), + "error must mention PDF, got: {err}" + ); + + let _ = tokio::fs::remove_dir_all(&dir).await; + } + + /// Non-UTF-8 binary files must be rejected with a clear error. #[tokio::test] - async fn file_read_lossy_reads_binary_file() { - let dir = std::env::temp_dir().join("mentat_test_file_read_lossy"); + async fn file_read_rejects_binary_file() { + let dir = std::env::temp_dir().join("mentat_test_file_read_binary_reject"); let _ = tokio::fs::remove_dir_all(&dir).await; tokio::fs::create_dir_all(&dir).await.unwrap(); - // Write bytes that are not valid UTF-8 and not a PDF let binary_data: Vec = vec![0x00, 0x80, 0xFF, 0xFE, b'h', b'i', 0x80]; tokio::fs::write(dir.join("data.bin"), &binary_data) .await @@ -683,20 +725,15 @@ mod tests { let tool = FileReadTool::new(test_security(dir.clone())); let result = tool.execute(json!({"path": "data.bin"})).await.unwrap(); + assert!(!result.success, "binary read must fail"); + let err = result.error.unwrap(); assert!( - result.success, - "lossy read must succeed, error: {:?}", - result.error + err.contains("binary file"), + "error must mention binary file, got: {err}" ); assert!( - result.output.contains('\u{FFFD}'), - "lossy output must contain replacement character, got: {:?}", - result.output - ); - assert!( - result.output.contains("hi"), - "lossy output must preserve valid ASCII, got: {:?}", - result.output + err.contains("data.bin"), + "error must include filename, got: {err}" ); let _ = tokio::fs::remove_dir_all(&dir).await; @@ -782,6 +819,7 @@ mod tests { /// End-to-end test: scripted provider calls `file_read` on a real PDF /// fixture, the tool extracts text via pdf-extract, and the extracted /// content reaches the provider in the tool result message. + #[cfg(feature = "rag-pdf")] #[tokio::test] async fn e2e_agent_file_read_pdf_extraction() { use crate::agent::agent::Agent; @@ -880,16 +918,15 @@ mod tests { } /// End-to-end test: agent calls `file_read` on a binary file, gets - /// lossy UTF-8 output with replacement characters in the tool result. + /// a rejection error in the tool result. #[tokio::test] - async fn e2e_agent_file_read_lossy_binary() { + async fn e2e_agent_file_read_rejects_binary() { use crate::agent::agent::Agent; use crate::agent::dispatcher::NativeToolDispatcher; use crate::providers::{ChatResponse, Provider, ToolCall}; use e2e_helpers::*; - // ── Set up workspace with binary file ── - let workspace = std::env::temp_dir().join("mentat_test_e2e_file_read_lossy"); + let workspace = std::env::temp_dir().join("mentat_test_e2e_file_read_binary_reject"); let _ = tokio::fs::remove_dir_all(&workspace).await; tokio::fs::create_dir_all(&workspace).await.unwrap(); @@ -918,7 +955,7 @@ mod tests { provider_attrs: None, }, ChatResponse { - text: Some("The file appears to be binary data.".into()), + text: Some("The file is binary, I'll use memory_recall instead.".into()), tool_calls: vec![], usage: None, reasoning_content: None, @@ -939,11 +976,11 @@ mod tests { let response = agent.turn("Read data.bin").await.unwrap(); assert!( - response.contains("binary"), - "agent response must mention binary, got: {response}", + response.contains("binary") || response.contains("memory_recall"), + "agent response must acknowledge binary rejection, got: {response}", ); - // Verify tool result contains lossy output with replacement chars + // Verify tool result contains the rejection error { let all_requests = recorded.lock().unwrap(); assert!( @@ -958,13 +995,8 @@ mod tests { .expect("second request must contain a tool result message"); assert!( - tool_result_msg.content.contains("valid"), - "tool result must preserve valid ASCII from binary file, got: {}", - tool_result_msg.content, - ); - assert!( - tool_result_msg.content.contains('\u{FFFD}'), - "tool result must contain replacement character for invalid bytes, got: {}", + tool_result_msg.content.contains("binary file"), + "tool result must contain binary rejection, got: {}", tool_result_msg.content, ); }