diff --git a/crates/tui/src/tui/ui.rs b/crates/tui/src/tui/ui.rs index 55f93aa3..95f7219f 100644 --- a/crates/tui/src/tui/ui.rs +++ b/crates/tui/src/tui/ui.rs @@ -1335,7 +1335,7 @@ async fn run_event_loop( } let tool_content = match &result { Ok(output) => sanitize_stream_chunk( - &tool_result_content_for_api_message(app, &id, &name, output), + &tool_result_content_for_api_message(app, &id, &name, output).await, ), Err(err) => sanitize_stream_chunk(&format!("Error: {err}")), }; @@ -4088,7 +4088,7 @@ fn push_assistant_message( } } -fn tool_result_content_for_api_message( +async fn tool_result_content_for_api_message( app: &App, id: &str, name: &str, @@ -4100,31 +4100,71 @@ fn tool_result_content_for_api_message( } if raw.chars().count() > crate::tool_output_receipts::RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS { - let mut messages = app.api_messages.clone(); - messages.push(Message { - role: "user".to_string(), - content: vec![ContentBlock::ToolResult { - tool_use_id: id.to_string(), - content: raw.to_string(), - is_error: Some(!output.success), - content_blocks: None, - }], - }); - let (compacted, stats) = crate::tool_output_receipts::compact_messages_for_persistence( - &messages, - &app.session_artifacts, - ); - if stats.compacted_count > 0 - && let Some(Message { content, .. }) = compacted.last() - && let Some(ContentBlock::ToolResult { content, .. }) = content.first() + let messages = live_tool_receipt_messages(app, id, raw, output.success); + let artifacts = app.session_artifacts.clone(); + let raw = raw.to_string(); + match tokio::task::spawn_blocking(move || { + compact_live_tool_receipt(messages, artifacts, raw) + }) + .await { - return content.clone(); + Ok(Some(receipt)) => return receipt, + Ok(None) => {} + Err(err) => { + crate::logging::warn(format!("live tool-output receipt compaction failed: {err}")); + } } } crate::core::engine::compact_tool_result_for_context(&app.model, name, output) } +fn live_tool_receipt_messages(app: &App, id: &str, raw: &str, success: bool) -> Vec { + let mut messages = Vec::with_capacity(2); + if let Some(tool_use_msg) = app.api_messages.iter().rev().find(|message| { + message.content.iter().any(|block| { + matches!(block, ContentBlock::ToolUse { id: tool_use_id, .. } if tool_use_id == id) + }) + }) { + messages.push(tool_use_msg.clone()); + } + messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: id.to_string(), + content: raw.to_string(), + is_error: Some(!success), + content_blocks: None, + }], + }); + messages +} + +fn compact_live_tool_receipt( + messages: Vec, + artifacts: Vec, + raw: String, +) -> Option { + let (compacted, _) = + crate::tool_output_receipts::compact_messages_for_persistence(&messages, &artifacts); + let content = compacted + .last() + .and_then(|message| message.content.first()) + .and_then(|block| match block { + ContentBlock::ToolResult { content, .. } => Some(content), + _ => None, + })?; + if content != &raw && live_tool_content_is_receipt(content) { + Some(content.clone()) + } else { + None + } +} + +fn live_tool_content_is_receipt(content: &str) -> bool { + content.trim_start().starts_with("[TOOL_OUTPUT_RECEIPT]") +} + fn replace_matching_assistant_text( app: &mut App, original_text: &str, diff --git a/crates/tui/src/tui/ui/tests.rs b/crates/tui/src/tui/ui/tests.rs index a4ed5ae7..b8dfa26e 100644 --- a/crates/tui/src/tui/ui/tests.rs +++ b/crates/tui/src/tui/ui/tests.rs @@ -1359,8 +1359,8 @@ fn create_test_options() -> TuiOptions { } } -#[test] -fn tool_result_api_content_receipts_large_live_output() { +#[tokio::test] +async fn tool_result_api_content_receipts_large_live_output() { let _guard = crate::tools::truncate::TEST_SPILLOVER_GUARD .lock() .unwrap_or_else(|err| err.into_inner()); @@ -1389,7 +1389,8 @@ fn tool_result_api_content_receipts_large_live_output() { let raw = "LIVE_RAW_SENTINEL\n".repeat(900); let output = crate::tools::spec::ToolResult::success(raw.clone()); - let content = tool_result_content_for_api_message(&app, "call-live-big", "exec_shell", &output); + let content = + tool_result_content_for_api_message(&app, "call-live-big", "exec_shell", &output).await; assert!(content.contains("[TOOL_OUTPUT_RECEIPT]")); assert!(content.contains("tool: exec_shell")); @@ -1403,6 +1404,51 @@ fn tool_result_api_content_receipts_large_live_output() { ); } +#[test] +fn live_tool_receipt_messages_clones_only_matching_tool_use() { + let mut app = App::new(create_test_options(), &Config::default()); + app.api_messages.push(Message { + role: "assistant".to_string(), + content: vec![ContentBlock::ToolUse { + id: "call-old".to_string(), + name: "exec_shell".to_string(), + input: serde_json::json!({"command": "old"}), + caller: None, + }], + }); + app.api_messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: "call-old".to_string(), + content: "OLD_RAW\n".repeat(2_000), + is_error: None, + content_blocks: None, + }], + }); + app.api_messages.push(Message { + role: "assistant".to_string(), + content: vec![ContentBlock::ToolUse { + id: "call-new".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({"path": "src/main.rs"}), + caller: None, + }], + }); + + let messages = live_tool_receipt_messages(&app, "call-new", "NEW_RAW", true); + + assert_eq!(messages.len(), 2); + assert!(matches!( + &messages[0].content[0], + ContentBlock::ToolUse { id, name, .. } if id == "call-new" && name == "read_file" + )); + assert!(matches!( + &messages[1].content[0], + ContentBlock::ToolResult { tool_use_id, content, .. } + if tool_use_id == "call-new" && content == "NEW_RAW" + )); +} + fn text_message(role: &str, text: &str) -> Message { Message { role: role.to_string(),