//! ACD (Aggressive Context Dehydration) tool: rehydrate. //! //! This tool allows the LLM to restore dehydrated conversation history //! from a previous context segment. use anyhow::Result; use tracing::{debug, warn}; use crate::acd::Fragment; use crate::ui_writer::UiWriter; use crate::ToolCall; use super::executor::ToolContext; /// Execute the rehydrate tool. /// Loads a fragment from disk and returns its contents for the LLM to review. pub async fn execute_rehydrate( tool_call: &ToolCall, ctx: &mut ToolContext<'_, W>, ) -> Result { let fragment_id = tool_call .args .get("fragment_id") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required 'fragment_id' parameter"))?; // Get session ID from context let session_id = ctx .session_id .as_ref() .ok_or_else(|| anyhow::anyhow!("No session ID available - cannot rehydrate fragment"))?; debug!("Rehydrating fragment {} for session {}", fragment_id, session_id); // Load the fragment let fragment = match Fragment::load(session_id, fragment_id) { Ok(f) => f, Err(e) => { warn!("Failed to load fragment {}: {}", fragment_id, e); return Ok(format!( "❌ Failed to rehydrate fragment '{}': {}\n\nThe fragment may have been deleted or the ID may be incorrect.", fragment_id, e )); } }; // Check if rehydration would be useful (warn if context is nearly full) let context_percentage = (ctx.context_used_tokens as f64 / ctx.context_total_tokens as f64) * 100.0; let fragment_tokens = fragment.estimated_tokens; let available_tokens = ctx.context_total_tokens.saturating_sub(ctx.context_used_tokens); if fragment_tokens > available_tokens { return Ok(format!( "⚠️ Cannot rehydrate fragment '{}': it contains ~{} tokens but only {} tokens are available in context.\n\n\ Consider compacting the context first with /compact, or continue without the full history.", fragment_id, fragment_tokens, available_tokens )); } if context_percentage > 70.0 && ctx.context_total_tokens > 0 { ctx.ui_writer.println(&format!( "⚠️ Warning: Context is at {:.0}% capacity. Rehydrating {} tokens may trigger compaction soon.", context_percentage, fragment_tokens )); } // Format the rehydrated content let mut output = String::new(); output.push_str(&format!( "✅ Rehydrated fragment '{}' ({} messages, ~{} tokens)\n\n", fragment_id, fragment.message_count, fragment.estimated_tokens )); // Add fragment metadata output.push_str("## Fragment Metadata\n"); output.push_str(&format!("- Created: {}\n", fragment.created_at)); if let Some(ref preceding) = fragment.preceding_fragment_id { output.push_str(&format!("- Preceding fragment: {}\n", preceding)); } if !fragment.topics.is_empty() { output.push_str(&format!("- Topics: {}\n", fragment.topics.join(", "))); } output.push_str("\n"); // Add the conversation history output.push_str("## Restored Conversation\n\n"); for (i, msg) in fragment.messages.iter().enumerate() { let role_str = match msg.role { g3_providers::MessageRole::User => "**User**", g3_providers::MessageRole::Assistant => "**Assistant**", g3_providers::MessageRole::System => "**System**", }; // Truncate very long messages for readability let content = if msg.content.chars().count() > 2000 { let chars: String = msg.content.chars().take(2000).collect(); format!("{}... [truncated, {} chars total]", chars, msg.content.chars().count()) } else { msg.content.clone() }; output.push_str(&format!("### Message {} - {}\n{}\n\n", i + 1, role_str, content)); } // Add note about preceding fragments if fragment.preceding_fragment_id.is_some() { output.push_str(&format!( "---\n💡 This fragment has a preceding fragment. To see earlier history, call: rehydrate(fragment_id: \"{}\")\n", fragment.preceding_fragment_id.as_ref().unwrap() )); } Ok(output) } #[cfg(test)] mod tests { use super::*; use crate::acd::Fragment; use crate::ui_writer::NullUiWriter; use crate::background_process::BackgroundProcessManager; use crate::webdriver_session::WebDriverSession; use g3_providers::{Message, MessageRole}; use std::sync::Arc; use tokio::sync::RwLock; use serde_json::json; struct TestContext { ui_writer: NullUiWriter, webdriver_session: Arc>>>>, webdriver_process: Arc>>, background_process_manager: Arc, todo_content: Arc>, pending_images: Vec, config: g3_config::Config, } impl TestContext { fn new() -> Self { Self { ui_writer: NullUiWriter, webdriver_session: Arc::new(RwLock::new(None)), webdriver_process: Arc::new(RwLock::new(None)), background_process_manager: Arc::new(BackgroundProcessManager::new(std::path::PathBuf::from("/tmp"))), todo_content: Arc::new(RwLock::new(String::new())), pending_images: Vec::new(), config: g3_config::Config::default(), } } } #[tokio::test] async fn test_rehydrate_missing_fragment_id() { let mut test_ctx = TestContext::new(); let mut ctx = ToolContext { working_dir: None, session_id: Some("test-session"), ui_writer: &test_ctx.ui_writer, config: &test_ctx.config, computer_controller: None, webdriver_session: &test_ctx.webdriver_session, webdriver_process: &test_ctx.webdriver_process, background_process_manager: &test_ctx.background_process_manager, todo_content: &test_ctx.todo_content, pending_images: &mut test_ctx.pending_images, is_autonomous: false, requirements_sha: None, context_total_tokens: 100000, context_used_tokens: 10000, }; let tool_call = ToolCall { tool: "rehydrate".to_string(), args: json!({}), }; let result = execute_rehydrate(&tool_call, &mut ctx).await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Missing required")); } #[tokio::test] async fn test_rehydrate_no_session_id() { let mut test_ctx = TestContext::new(); let mut ctx = ToolContext { working_dir: None, session_id: None, ui_writer: &test_ctx.ui_writer, config: &test_ctx.config, computer_controller: None, webdriver_session: &test_ctx.webdriver_session, webdriver_process: &test_ctx.webdriver_process, background_process_manager: &test_ctx.background_process_manager, todo_content: &test_ctx.todo_content, pending_images: &mut test_ctx.pending_images, is_autonomous: false, requirements_sha: None, context_total_tokens: 100000, context_used_tokens: 10000, }; let tool_call = ToolCall { tool: "rehydrate".to_string(), args: json!({"fragment_id": "test-fragment"}), }; let result = execute_rehydrate(&tool_call, &mut ctx).await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("No session ID")); } #[tokio::test] async fn test_rehydrate_nonexistent_fragment() { let mut test_ctx = TestContext::new(); let mut ctx = ToolContext { working_dir: None, session_id: Some("nonexistent-session"), ui_writer: &test_ctx.ui_writer, config: &test_ctx.config, computer_controller: None, webdriver_session: &test_ctx.webdriver_session, webdriver_process: &test_ctx.webdriver_process, background_process_manager: &test_ctx.background_process_manager, todo_content: &test_ctx.todo_content, pending_images: &mut test_ctx.pending_images, is_autonomous: false, requirements_sha: None, context_total_tokens: 100000, context_used_tokens: 10000, }; let tool_call = ToolCall { tool: "rehydrate".to_string(), args: json!({"fragment_id": "nonexistent-fragment"}), }; let result = execute_rehydrate(&tool_call, &mut ctx).await; assert!(result.is_ok()); let output = result.unwrap(); assert!(output.contains("Failed to rehydrate")); assert!(output.contains("nonexistent-fragment")); } #[tokio::test] async fn test_rehydrate_success() { // Create a temporary fragment let test_session_id = format!("test_rehydrate_{}", std::process::id()); let messages = vec![ Message::new(MessageRole::User, "Test user message".to_string()), Message::new(MessageRole::Assistant, "Test assistant response".to_string()), ]; let fragment = Fragment::new(messages, None); let fragment_id = fragment.fragment_id.clone(); // Save fragment using the Fragment::save method let save_result = fragment.save(&test_session_id); assert!(save_result.is_ok()); let file_path = save_result.unwrap(); assert!(file_path.exists(), "Fragment file should exist after save"); // Verify we can load it back let loaded = Fragment::load(&test_session_id, &fragment_id); assert!(loaded.is_ok()); let loaded_fragment = loaded.unwrap(); assert_eq!(loaded_fragment.fragment_id, fragment_id); assert_eq!(loaded_fragment.message_count, 2); // Cleanup let _ = std::fs::remove_file(&file_path); let _ = std::fs::remove_dir(file_path.parent().unwrap()); } }