| 2468 | } |
| 2469 | |
| 2470 | fn prune_after_loop(session: &mut Session) { |
| 2471 | let mut tool_name_by_call: HashMap<String, String> = HashMap::new(); |
| 2472 | for msg in &session.messages { |
| 2473 | for part in &msg.parts { |
| 2474 | if let PartType::ToolCall { id, name, .. } = &part.part_type { |
| 2475 | tool_name_by_call.insert(id.clone(), name.clone()); |
| 2476 | } |
| 2477 | } |
| 2478 | } |
| 2479 | |
| 2480 | let mut prune_messages: Vec<MessageForPrune> = session |
| 2481 | .messages |
| 2482 | .iter() |
| 2483 | .map(|m| { |
| 2484 | let parts: Vec<PruneToolPart> = m |
| 2485 | .parts |
| 2486 | .iter() |
| 2487 | .filter_map(|p| match &p.part_type { |
| 2488 | PartType::ToolResult { |
| 2489 | tool_call_id, |
| 2490 | content, |
| 2491 | is_error, |
| 2492 | } => Some(PruneToolPart { |
| 2493 | id: p.id.clone(), |
| 2494 | tool: tool_name_by_call |
| 2495 | .get(tool_call_id) |
| 2496 | .cloned() |
| 2497 | .unwrap_or_default(), |
| 2498 | output: content.clone(), |
| 2499 | status: if *is_error { |
| 2500 | ToolPartStatus::Error |
| 2501 | } else { |
| 2502 | ToolPartStatus::Completed |
| 2503 | }, |
| 2504 | compacted: None, |
| 2505 | }), |
| 2506 | _ => None, |
| 2507 | }) |
| 2508 | .collect(); |
| 2509 | MessageForPrune { |
| 2510 | role: match m.role { |
| 2511 | MessageRole::User => "user".to_string(), |
| 2512 | _ => "assistant".to_string(), |
| 2513 | }, |
| 2514 | parts, |
| 2515 | summary: false, |
| 2516 | } |
| 2517 | }) |
| 2518 | .collect(); |
| 2519 | |
| 2520 | let engine = CompactionEngine::new(CompactionConfig::default()); |
| 2521 | let pruned_ids = engine.prune(&mut prune_messages); |
| 2522 | if pruned_ids.is_empty() { |
| 2523 | return; |
| 2524 | } |
| 2525 | let pruned: HashSet<String> = pruned_ids.into_iter().collect(); |
| 2526 | for msg in &mut session.messages { |
| 2527 | for part in &mut msg.parts { |