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