| 1 | //! Compact receipts for oversized tool outputs in saved session history. |
| 2 | |
| 3 | use crate::artifacts::{ArtifactKind, ArtifactRecord}; |
| 4 | |
| 5 | use serde_json::Value; |
| 6 | |
| 7 | use crate::fast_hash::FastHashMap; |
| 8 | use codewhale_localization::{Locale, MessageId, tr}; |
| 9 | use codewhale_models::{ContentBlock, Message}; |
| 10 | |
| 11 | /// Match the provider-wire budget so persisted/resumed history does not keep a |
| 12 | /// larger raw body than the model would receive on a fresh request. |
| 13 | pub const RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS: usize = 12_000; |
| 14 | |
| 15 | #[derive(Debug, Clone, Default, PartialEq, Eq)] |
| 16 | pub struct ToolOutputReceiptStats { |
| 17 | pub compacted_count: usize, |
| 18 | pub artifact_receipts: usize, |
| 19 | pub sha_receipts: usize, |
| 20 | pub unavailable_receipts: usize, |
| 21 | pub original_chars: usize, |
| 22 | } |
| 23 | |
| 24 | #[derive(Debug, Clone, Default, PartialEq, Eq)] |
| 25 | pub struct ToolOutputStatus { |
| 26 | pub raw_large_count: usize, |
| 27 | pub raw_large_chars: usize, |
| 28 | pub receipt_count: usize, |
| 29 | pub artifact_count: usize, |
| 30 | pub artifact_bytes: u64, |
| 31 | } |
| 32 | |
| 33 | #[derive(Debug, Clone)] |
| 34 | struct ToolUseInfo { |
| 35 | name: String, |
| 36 | input: Value, |
| 37 | } |
| 38 | |
| 39 | #[derive(Debug, Clone)] |
| 40 | enum DetailHandle { |
| 41 | Artifact(ArtifactRecord), |
| 42 | /// Raw legacy result with no session-owned artifact. Compacted |
| 43 | /// truthfully, but never given a retrieval handle: a process-wide |
| 44 | /// SHA store cannot prove which session owns the bytes. |
| 45 | Unavailable, |
| 46 | } |
| 47 | |
| 48 | /// Return a copy of `messages` with oversized raw tool-result bodies replaced |
| 49 | /// by compact receipts. Full output is kept behind existing session artifacts |
| 50 | /// when available. Raw legacy results without a session-owned artifact are |
| 51 | /// compacted truthfully, but never receive a retrieval handle: a process-wide |
| 52 | /// SHA store cannot prove which session owns the bytes. |
| 53 | pub fn compact_messages_for_persistence( |
| 54 | messages: &[Message], |
| 55 | artifacts: &[ArtifactRecord], |
| 56 | ) -> (Vec<Message>, ToolOutputReceiptStats) { |
| 57 | // Tool-call IDs here come from engine transcript blocks and artifact records, |
| 58 | // making this save/resume bookkeeping a safe FastHashMap target. |
| 59 | let artifacts_by_call = artifacts_by_tool_call(artifacts); |
| 60 | let mut tool_uses: FastHashMap<String, ToolUseInfo> = FastHashMap::default(); |
| 61 | let mut stats = ToolOutputReceiptStats::default(); |
| 62 | let mut compacted = Vec::with_capacity(messages.len()); |
| 63 | |
| 64 | for message in messages { |
| 65 | let mut next = message.clone(); |
| 66 | for block in &mut next.content { |
| 67 | match block { |
| 68 | ContentBlock::ToolUse { |
| 69 | id, name, input, .. |
| 70 | } => { |
| 71 | tool_uses.insert( |
| 72 | id.clone(), |
| 73 | ToolUseInfo { |
| 74 | name: name.clone(), |
| 75 | input: input.clone(), |
| 76 | }, |
| 77 | ); |
| 78 | } |
| 79 | ContentBlock::ToolResult { |
| 80 | tool_use_id, |
| 81 | content, |
| 82 | is_error, |
| 83 | .. |
| 84 | } => { |
| 85 | let char_count = content.chars().count(); |
| 86 | if char_count <= RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS |
| 87 | || looks_like_receipt(content) |
| 88 | { |
| 89 | continue; |
| 90 | } |
| 91 | |
| 92 | let tool_info = tool_uses.get(tool_use_id); |
| 93 | let handle = artifacts_by_call |
| 94 | .get(tool_use_id.as_str()) |
| 95 | .cloned() |
| 96 | .map(|artifact| DetailHandle::Artifact((*artifact).clone())) |
| 97 | .unwrap_or(DetailHandle::Unavailable); |
| 98 | let source = match &handle { |
| 99 | DetailHandle::Artifact(_) => ReceiptSource::Artifact, |
| 100 | DetailHandle::Unavailable => ReceiptSource::Unavailable, |
| 101 | }; |
| 102 | |
| 103 | *content = render_tool_output_receipt( |
| 104 | tool_use_id, |
| 105 | tool_info, |
| 106 | content, |
| 107 | *is_error, |
| 108 | &handle, |
| 109 | ); |
| 110 | stats.compacted_count += 1; |
| 111 | stats.original_chars = stats.original_chars.saturating_add(char_count); |
| 112 | match source { |
| 113 | ReceiptSource::Artifact => stats.artifact_receipts += 1, |
| 114 | ReceiptSource::Unavailable => stats.unavailable_receipts += 1, |
| 115 | } |
| 116 | } |
| 117 | _ => {} |
| 118 | } |
| 119 | } |
| 120 | compacted.push(next); |
| 121 | } |
| 122 | |
| 123 | (compacted, stats) |
| 124 | } |
| 125 | |
| 126 | pub fn tool_output_status(messages: &[Message], artifacts: &[ArtifactRecord]) -> ToolOutputStatus { |
| 127 | let mut status = ToolOutputStatus { |
| 128 | artifact_count: artifacts.len(), |
| 129 | artifact_bytes: artifacts |
| 130 | .iter() |
| 131 | .map(|artifact| artifact.byte_size) |
| 132 | .sum::<u64>(), |
| 133 | ..ToolOutputStatus::default() |
| 134 | }; |
| 135 | |
| 136 | for message in messages { |
| 137 | for block in &message.content { |
| 138 | if let ContentBlock::ToolResult { content, .. } = block { |
| 139 | if looks_like_receipt(content) { |
| 140 | status.receipt_count += 1; |
| 141 | } else { |
| 142 | let chars = content.chars().count(); |
| 143 | if chars > RAW_TOOL_OUTPUT_RECEIPT_THRESHOLD_CHARS { |
| 144 | status.raw_large_count += 1; |
| 145 | status.raw_large_chars = status.raw_large_chars.saturating_add(chars); |
| 146 | } |
| 147 | } |
| 148 | } |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | status |
| 153 | } |
| 154 | |
| 155 | pub fn format_tool_output_status(status: &ToolOutputStatus, locale: Locale) -> String { |
| 156 | let mut parts = Vec::new(); |
| 157 | if status.raw_large_count > 0 { |
| 158 | parts.push( |
| 159 | tr(locale, MessageId::StatusToolRawPressure) |
| 160 | .replace("{count}", &status.raw_large_count.to_string()) |
| 161 | .replace("{chars}", &format_count(status.raw_large_chars)), |
| 162 | ); |
| 163 | } |
| 164 | if status.receipt_count > 0 { |
| 165 | parts.push( |
| 166 | tr(locale, MessageId::StatusToolCompactReceipts) |
| 167 | .replace("{count}", &status.receipt_count.to_string()), |
| 168 | ); |
| 169 | } |
| 170 | if status.artifact_count > 0 { |
| 171 | parts.push( |
| 172 | tr(locale, MessageId::StatusToolArtifacts) |
| 173 | .replace("{count}", &status.artifact_count.to_string()) |
| 174 | .replace( |
| 175 | "{bytes}", |
| 176 | &crate::artifacts::format_byte_size(status.artifact_bytes), |
| 177 | ), |
| 178 | ); |
| 179 | } |
| 180 | if parts.is_empty() { |
| 181 | tr(locale, MessageId::StatusToolNone).into_owned() |
| 182 | } else { |
| 183 | parts.join("; ") |
| 184 | } |
| 185 | } |
| 186 | |
| 187 | fn artifacts_by_tool_call(artifacts: &[ArtifactRecord]) -> FastHashMap<&str, &ArtifactRecord> { |
| 188 | artifacts |
| 189 | .iter() |
| 190 | .filter(|artifact| artifact.kind == ArtifactKind::ToolOutput) |
| 191 | .map(|artifact| (artifact.tool_call_id.as_str(), artifact)) |
| 192 | .collect() |
| 193 | } |
| 194 | |
| 195 | #[derive(Debug, Clone, Copy)] |
| 196 | enum ReceiptSource { |
| 197 | Artifact, |
| 198 | Unavailable, |
| 199 | } |
| 200 | |
| 201 | fn render_tool_output_receipt( |
| 202 | tool_call_id: &str, |
| 203 | tool_info: Option<&ToolUseInfo>, |
| 204 | original_content: &str, |
| 205 | is_error: Option<bool>, |
| 206 | handle: &DetailHandle, |
| 207 | ) -> String { |
| 208 | let original_chars = original_content.chars().count(); |
| 209 | let original_bytes = original_content.len() as u64; |
| 210 | let tool_name = match handle { |
| 211 | DetailHandle::Artifact(record) if !record.tool_name.trim().is_empty() => { |
| 212 | record.tool_name.as_str() |
| 213 | } |
| 214 | _ => tool_info |
| 215 | .map(|info| info.name.as_str()) |
| 216 | .filter(|name| !name.trim().is_empty()) |
| 217 | .unwrap_or("unknown"), |
| 218 | }; |
| 219 | let command_or_query = tool_info |
| 220 | .map(|info| summarize_input(&info.input, 300)) |
| 221 | .unwrap_or_else(|| "unknown".to_string()); |
| 222 | let status = if is_error.unwrap_or(false) { |
| 223 | "error" |
| 224 | } else { |
| 225 | "success" |
| 226 | }; |
| 227 | let exit_status = infer_exit_status(original_content).unwrap_or_else(|| "unknown".to_string()); |
| 228 | let preview = preview_for_receipt(handle, original_content); |
| 229 | |
| 230 | format!( |
| 231 | "[TOOL_OUTPUT_RECEIPT]\n\ |
| 232 | tool: {tool_name}\n\ |
| 233 | tool_call_id: {tool_call_id}\n\ |
| 234 | status: {status}\n\ |
| 235 | exit_status: {exit_status}\n\ |
| 236 | elapsed: unknown\n\ |
| 237 | output: {bytes} ({chars} chars, ~{tokens} tokens)\n\ |
| 238 | truncation: raw output omitted — full output in the tool details view\n\ |
| 239 | command_or_query: {command_or_query}\n\ |
| 240 | preview: {preview}\n\ |
| 241 | [/TOOL_OUTPUT_RECEIPT]", |
| 242 | bytes = crate::artifacts::format_byte_size(original_bytes), |
| 243 | chars = format_count(original_chars), |
| 244 | tokens = format_count(approx_tokens(original_chars)), |
| 245 | ) |
| 246 | } |
| 247 | |
| 248 | fn preview_for_receipt(handle: &DetailHandle, original_content: &str) -> String { |
| 249 | let preview = match handle { |
| 250 | DetailHandle::Artifact(record) if !record.preview.trim().is_empty() => { |
| 251 | record.preview.as_str() |
| 252 | } |
| 253 | _ => original_content, |
| 254 | }; |
| 255 | summarize_text(preview, 240) |
| 256 | } |
| 257 | |
| 258 | fn looks_like_receipt(content: &str) -> bool { |
| 259 | let trimmed = content.trim_start(); |
| 260 | trimmed.starts_with("[TOOL_OUTPUT_RECEIPT]") |
| 261 | || trimmed.starts_with("[artifact:") |
| 262 | || trimmed.starts_with("[TOOL_RESULT_TRUNCATED]") |
| 263 | || trimmed.starts_with("<TOOL_RESULT_REF") |
| 264 | } |
| 265 | |
| 266 | fn infer_exit_status(content: &str) -> Option<String> { |
| 267 | if let Ok(value) = serde_json::from_str::<Value>(content) { |
| 268 | for key in ["exit_code", "exit_status", "status", "code"] { |
| 269 | if let Some(value) = value.get(key) { |
| 270 | return Some(summarize_input(value, 120)); |
| 271 | } |
| 272 | } |
| 273 | } |
| 274 | |
| 275 | for line in content.lines().take(40) { |
| 276 | let trimmed = line.trim(); |
| 277 | for prefix in ["Exit code:", "exit code:", "Exit status:", "exit status:"] { |
| 278 | if let Some(value) = trimmed.strip_prefix(prefix) { |
| 279 | return Some(summarize_text(value.trim(), 120)); |
| 280 | } |
| 281 | } |
| 282 | } |
| 283 | None |
| 284 | } |
| 285 | |
| 286 | fn summarize_input(value: &Value, max_chars: usize) -> String { |
| 287 | let raw = value |
| 288 | .as_str() |
| 289 | .map(str::to_string) |
| 290 | .unwrap_or_else(|| value.to_string()); |
| 291 | summarize_text(&raw, max_chars) |
| 292 | } |
| 293 | |
| 294 | fn summarize_text(text: &str, max_chars: usize) -> String { |
| 295 | let escaped = text.replace('\n', "\\n"); |
| 296 | let mut summary: String = escaped.chars().take(max_chars).collect(); |
| 297 | if escaped.chars().count() > max_chars { |
| 298 | summary.push_str("..."); |
| 299 | } |
| 300 | summary |
| 301 | } |
| 302 | |
| 303 | fn approx_tokens(chars: usize) -> usize { |
| 304 | chars.div_ceil(4) |
| 305 | } |
| 306 | |
| 307 | fn format_count(value: usize) -> String { |
| 308 | value.to_string() |
| 309 | } |
| 310 | |
| 311 | #[cfg(test)] |
| 312 | mod tests { |
| 313 | use codewhale_models::Role; |
| 314 | use std::path::{Path, PathBuf}; |
| 315 | |
| 316 | use super::*; |
| 317 | use chrono::Utc; |
| 318 | use serde_json::json; |
| 319 | |
| 320 | fn tool_use_message(id: &str, name: &str, input: Value) -> Message { |
| 321 | Message { |
| 322 | role: Role::Assistant, |
| 323 | content: vec![ContentBlock::ToolUse { |
| 324 | id: id.to_string(), |
| 325 | name: name.to_string(), |
| 326 | input, |
| 327 | caller: None, |
| 328 | thought_signature: None, |
| 329 | }], |
| 330 | } |
| 331 | } |
| 332 | |
| 333 | fn tool_result_message(id: &str, content: &str) -> Message { |
| 334 | Message { |
| 335 | role: Role::User, |
| 336 | content: vec![ContentBlock::ToolResult { |
| 337 | tool_use_id: id.to_string(), |
| 338 | content: content.to_string(), |
| 339 | is_error: None, |
| 340 | content_blocks: None, |
| 341 | }], |
| 342 | } |
| 343 | } |
| 344 | |
| 345 | fn artifact_record(tool_call_id: &str, raw: &str) -> ArtifactRecord { |
| 346 | ArtifactRecord { |
| 347 | id: crate::artifacts::artifact_id_for_tool_call(tool_call_id), |
| 348 | kind: ArtifactKind::ToolOutput, |
| 349 | session_id: "session-123".to_string(), |
| 350 | tool_call_id: tool_call_id.to_string(), |
| 351 | tool_name: "exec_shell".to_string(), |
| 352 | created_at: Utc::now(), |
| 353 | byte_size: raw.len() as u64, |
| 354 | preview: "checking crate ... error[E0425]".to_string(), |
| 355 | storage_path: PathBuf::from("artifacts").join("art_call-big.txt"), |
| 356 | } |
| 357 | } |
| 358 | |
| 359 | #[test] |
| 360 | fn compacts_large_tool_result_to_artifact_receipt() { |
| 361 | let raw = "RAW_SENTINEL\n".repeat(2_000); |
| 362 | let messages = vec![ |
| 363 | tool_use_message( |
| 364 | "call-big", |
| 365 | "exec_shell", |
| 366 | json!({"command": "cargo test -p codewhale-tui"}), |
| 367 | ), |
| 368 | tool_result_message("call-big", &raw), |
| 369 | ]; |
| 370 | let artifacts = vec![artifact_record("call-big", &raw)]; |
| 371 | |
| 372 | let (compacted, stats) = compact_messages_for_persistence(&messages, &artifacts); |
| 373 | let ContentBlock::ToolResult { content, .. } = &compacted[1].content[0] else { |
| 374 | panic!("expected tool result"); |
| 375 | }; |
| 376 | |
| 377 | assert_eq!(stats.compacted_count, 1); |
| 378 | assert_eq!(stats.artifact_receipts, 1); |
| 379 | assert!(!content.contains("RAW_SENTINEL")); |
| 380 | assert!(content.contains("[TOOL_OUTPUT_RECEIPT]")); |
| 381 | assert!(content.contains("tool: exec_shell")); |
| 382 | assert!(!content.contains("detail_handle")); |
| 383 | assert!(!content.contains("retrieve_tool_result")); |
| 384 | assert!(content.contains("full output in the tool details view")); |
| 385 | assert!( |
| 386 | content.contains("command_or_query: {\"command\":\"cargo test -p codewhale-tui\"}") |
| 387 | ); |
| 388 | } |
| 389 | |
| 390 | #[test] |
| 391 | fn compacts_unowned_large_tool_result_without_false_storage_claims() { |
| 392 | let raw = format!("{}\n{}", "H".repeat(320), "NO_ARTIFACT_RAW\n".repeat(2_000)); |
| 393 | let messages = vec![ |
| 394 | tool_use_message("call-big", "grep_files", json!({"pattern": "TODO"})), |
| 395 | tool_result_message("call-big", &raw), |
| 396 | ]; |
| 397 | |
| 398 | let (compacted, stats) = compact_messages_for_persistence(&messages, &[]); |
| 399 | let ContentBlock::ToolResult { content, .. } = &compacted[1].content[0] else { |
| 400 | panic!("expected tool result"); |
| 401 | }; |
| 402 | |
| 403 | assert_eq!(stats.compacted_count, 1); |
| 404 | assert_eq!(stats.sha_receipts, 0); |
| 405 | assert_eq!(stats.unavailable_receipts, 1); |
| 406 | assert!(!content.contains("NO_ARTIFACT_RAW")); |
| 407 | assert!(!content.contains("detail_handle")); |
| 408 | assert!(!content.contains("storage:")); |
| 409 | assert!(!content.contains("retrieve_tool_result")); |
| 410 | assert!(content.contains("full output in the tool details view")); |
| 411 | } |
| 412 | |
| 413 | #[test] |
| 414 | fn small_tool_results_remain_inline() { |
| 415 | let messages = vec![ |
| 416 | tool_use_message("call-small", "exec_shell", json!({"command": "pwd"})), |
| 417 | tool_result_message("call-small", "ok"), |
| 418 | ]; |
| 419 | |
| 420 | let (compacted, stats) = compact_messages_for_persistence(&messages, &[]); |
| 421 | let ContentBlock::ToolResult { content, .. } = &compacted[1].content[0] else { |
| 422 | panic!("expected tool result"); |
| 423 | }; |
| 424 | |
| 425 | assert_eq!(content, "ok"); |
| 426 | assert_eq!(stats.compacted_count, 0); |
| 427 | } |
| 428 | |
| 429 | #[test] |
| 430 | fn status_reports_raw_large_receipts_and_artifacts() { |
| 431 | let raw = "RAW_STATUS\n".repeat(2_000); |
| 432 | let receipt = "[TOOL_OUTPUT_RECEIPT]\ntruncation: raw output omitted — full output in the tool details view"; |
| 433 | let messages = vec![ |
| 434 | tool_result_message("call-raw", &raw), |
| 435 | tool_result_message("call-receipt", receipt), |
| 436 | ]; |
| 437 | let artifacts = vec![ArtifactRecord { |
| 438 | storage_path: Path::new("artifacts/art_call-big.txt").to_path_buf(), |
| 439 | ..artifact_record("call-big", &raw) |
| 440 | }]; |
| 441 | |
| 442 | let status = tool_output_status(&messages, &artifacts); |
| 443 | assert_eq!(status.raw_large_count, 1); |
| 444 | assert_eq!(status.receipt_count, 1); |
| 445 | assert_eq!(status.artifact_count, 1); |
| 446 | |
| 447 | let rendered = format_tool_output_status(&status, Locale::En); |
| 448 | assert!(rendered.contains("raw over cap")); |
| 449 | assert!(rendered.contains("compact receipt")); |
| 450 | assert!(rendered.contains("artifact")); |
| 451 | } |
| 452 | } |
| 453 |