@@ -5,17 +5,21 @@ use serde_json::Value;
55use std:: collections:: HashMap ;
66use std:: sync:: atomic:: { AtomicBool , Ordering } ;
77use std:: sync:: Arc ;
8+ use std:: path:: PathBuf ;
89use tokio:: process:: Command ;
9- use tokio:: sync:: { mpsc, broadcast, Mutex , oneshot:: Sender } ;
10+ use tokio:: sync:: { mpsc, broadcast, Mutex , oneshot:: Sender , RwLock } ;
1011use tokio:: io:: { self } ;
1112use tokio_util:: compat:: { TokioAsyncReadCompatExt , TokioAsyncWriteCompatExt } ;
1213use tracing:: { info, debug, error} ;
1314use anyhow:: { Result , anyhow} ;
1415use crate :: utils:: relative_to_current_dir;
16+ use crate :: acp_history:: AcpHistoryManager ;
1517
1618#[ derive( Debug , Clone , Serialize , Deserialize ) ]
1719pub struct AcpUserMessage {
1820 pub content : String ,
21+ #[ serde( skip_serializing_if = "Option::is_none" ) ]
22+ pub checkpoint_id : Option < String > ,
1923}
2024
2125#[ derive( Debug , Clone , Serialize , Deserialize ) ]
@@ -719,10 +723,19 @@ pub struct AcpAgent {
719723 io_handle : Option < tokio:: task:: JoinHandle < ( ) > > ,
720724 history : Arc < tokio:: sync:: Mutex < Vec < AcpMessage > > > ,
721725 pending_permissions : PendingPermissionsMap ,
726+ /// History manager for undo/redo support
727+ history_manager : Arc < RwLock < AcpHistoryManager > > ,
722728}
723729
724730impl AcpAgent {
725731 pub fn new ( agent_id : String , agent_name : String ) -> Self {
732+ // Initialize history manager with current working directory and agent ID
733+ let project_root = std:: env:: current_dir ( ) . unwrap_or_else ( |_| PathBuf :: from ( "." ) ) ;
734+ let mut history_manager = AcpHistoryManager :: new ( & project_root, & agent_id) ;
735+ if let Err ( e) = history_manager. init ( ) {
736+ error ! ( "Failed to initialize history manager: {}" , e) ;
737+ }
738+
726739 Self {
727740 agent_id,
728741 agent_name,
@@ -736,6 +749,7 @@ impl AcpAgent {
736749 io_handle : None ,
737750 history : Arc :: new ( tokio:: sync:: Mutex :: new ( Vec :: new ( ) ) ) ,
738751 pending_permissions : Arc :: new ( tokio:: sync:: Mutex :: new ( HashMap :: new ( ) ) ) ,
752+ history_manager : Arc :: new ( RwLock :: new ( history_manager) ) ,
739753 }
740754 }
741755
@@ -1104,13 +1118,27 @@ impl AcpAgent {
11041118 self . session_id = None ;
11051119 }
11061120
1107- pub async fn send_prompt ( & mut self , prompt : String ) -> Result < ( ) > {
1121+ pub async fn send_prompt ( & mut self , prompt : String ) -> Result < String > {
11081122 let prompt_tx = match & self . prompt_sender {
11091123 Some ( tx) => tx,
11101124 None => return Err ( anyhow ! ( "Prompt sender not initialized" ) )
11111125 } ;
11121126
1113- let user_message = AcpMessage :: User ( AcpUserMessage { content : prompt. clone ( ) } ) ;
1127+ // Create checkpoint before processing the message
1128+ let mut manager = self . history_manager . write ( ) . await ;
1129+ let checkpoint_id = match manager. create_checkpoint ( & prompt) {
1130+ Ok ( id) => Some ( id) ,
1131+ Err ( e) => {
1132+ error ! ( "Failed to create checkpoint: {}" , e) ;
1133+ None
1134+ }
1135+ } ;
1136+ drop ( manager) ;
1137+
1138+ let user_message = AcpMessage :: User ( AcpUserMessage {
1139+ content : prompt. clone ( ) ,
1140+ checkpoint_id,
1141+ } ) ;
11141142
11151143 // Save user message to history
11161144 let mut history = self . history . lock ( ) . await ;
@@ -1127,9 +1155,36 @@ impl AcpAgent {
11271155 // Send prompt to agent
11281156 prompt_tx. send ( prompt) . await ?;
11291157
1158+ Ok ( String :: new ( ) )
1159+ }
1160+
1161+ /// Restore project to state before a specific prompt was processed
1162+ pub async fn restore_to_prompt ( & self , prompt : & str ) -> Result < ( ) > {
1163+ let manager = self . history_manager . read ( ) . await ;
1164+ manager. restore_to_checkpoint ( prompt) ?;
1165+
1166+ info ! ( "Restored project to state before prompt" ) ;
1167+ Ok ( ( ) )
1168+ }
1169+
1170+ /// Restore project to state at a checkpoint id (commit hash)
1171+ pub async fn restore_to_checkpoint_id ( & self , checkpoint_id : & str ) -> Result < ( ) > {
1172+ let manager = self . history_manager . read ( ) . await ;
1173+ manager. restore_to_commit ( checkpoint_id) ?;
1174+
1175+ info ! ( "Restored project to checkpoint {}" , checkpoint_id) ;
11301176 Ok ( ( ) )
11311177 }
11321178
1179+ /// Get all available checkpoints
1180+ pub async fn get_checkpoints ( & self ) -> Vec < String > {
1181+ let manager = self . history_manager . read ( ) . await ;
1182+ manager. get_all_checkpoints ( )
1183+ . iter ( )
1184+ . map ( |cp| cp. prompt . clone ( ) )
1185+ . collect ( )
1186+ }
1187+
11331188 pub async fn cancel_prompt ( & self ) -> Result < ( ) > {
11341189 let cancel_sender_guard = self . cancel_sender . lock ( ) . await ;
11351190 let cancel_tx = match cancel_sender_guard. as_ref ( ) {
@@ -1189,7 +1244,6 @@ impl AcpManager {
11891244 pub async fn start_agent (
11901245 & mut self , agent_id : String , agent_name : String , cmd : & str , args : & [ String ] ,
11911246 ) -> Result < ( ) > {
1192-
11931247 if self . agents . contains_key ( & agent_id) {
11941248 return Err ( anyhow:: anyhow!( "Agent {} already running" , agent_id) ) ;
11951249 }
@@ -1250,4 +1304,28 @@ impl AcpManager {
12501304 . map ( |( id, agent) | ( id. clone ( ) , agent. agent_name ( ) . to_string ( ) ) )
12511305 . collect ( )
12521306 }
1307+
1308+ /// Restore agent's project to state before a specific prompt was processed
1309+ pub async fn restore_to_prompt ( & self , agent_id : & str , prompt : & str ) -> Result < ( ) > {
1310+ let agent = self . agents . get ( agent_id)
1311+ . ok_or_else ( || anyhow ! ( "Agent {} not found" , agent_id) ) ?;
1312+
1313+ agent. restore_to_prompt ( prompt) . await
1314+ }
1315+
1316+ /// Restore agent's project to state at a checkpoint id (commit hash)
1317+ pub async fn restore_to_checkpoint_id ( & self , agent_id : & str , checkpoint_id : & str ) -> Result < ( ) > {
1318+ let agent = self . agents . get ( agent_id)
1319+ . ok_or_else ( || anyhow ! ( "Agent {} not found" , agent_id) ) ?;
1320+
1321+ agent. restore_to_checkpoint_id ( checkpoint_id) . await
1322+ }
1323+
1324+ /// Get all checkpoints for an agent. Returns Vec of prompts
1325+ pub async fn get_checkpoints ( & self , agent_id : & str ) -> Result < Vec < String > > {
1326+ let agent = self . agents . get ( agent_id)
1327+ . ok_or_else ( || anyhow ! ( "Agent {} not found" , agent_id) ) ?;
1328+
1329+ Ok ( agent. get_checkpoints ( ) . await )
1330+ }
12531331}
0 commit comments