diff --git a/src/apps/desktop/src/api/snapshot_service.rs b/src/apps/desktop/src/api/snapshot_service.rs index eb1700f2..0a88e360 100644 --- a/src/apps/desktop/src/api/snapshot_service.rs +++ b/src/apps/desktop/src/api/snapshot_service.rs @@ -7,7 +7,7 @@ use bitfun_core::service::snapshot::{ }; use log::{info, warn}; use serde::{Deserialize, Serialize}; -use std::{path::PathBuf, sync::Arc}; +use std::{path::PathBuf, sync::Arc, time::Duration}; use tauri::{AppHandle, Emitter}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -321,6 +321,22 @@ pub async fn rollback_to_turn( app_handle: AppHandle, request: RollbackTurnRequest, ) -> Result, String> { + { + use bitfun_core::agentic::coordination::get_global_coordinator; + + if let Some(coordinator) = get_global_coordinator() { + if let Err(e) = coordinator + .cancel_active_turn_for_session(&request.session_id, Duration::from_secs(2)) + .await + { + warn!( + "Failed to cancel active turn before rollback: session_id={}, turn_index={}, error={}", + request.session_id, request.turn_index, e + ); + } + } + } + let manager = ensure_snapshot_manager_ready(&request.workspace_path).await?; let restored_files = manager diff --git a/src/crates/core/src/agentic/coordination/coordinator.rs b/src/crates/core/src/agentic/coordination/coordinator.rs index e73eabd5..4a8dcf20 100644 --- a/src/crates/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/core/src/agentic/coordination/coordinator.rs @@ -23,6 +23,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::OnceLock; use tokio::sync::mpsc; +use tokio::time::{sleep, Duration, Instant}; use tokio_util::sync::CancellationToken; /// Subagent execution result @@ -1342,6 +1343,41 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet Ok(()) } + pub async fn cancel_active_turn_for_session( + &self, + session_id: &str, + wait_timeout: Duration, + ) -> BitFunResult> { + let Some(session) = self.session_manager.get_session(session_id) else { + return Ok(None); + }; + + let SessionState::Processing { + current_turn_id, .. + } = session.state + else { + return Ok(None); + }; + + self.cancel_dialog_turn(session_id, ¤t_turn_id).await?; + + let deadline = Instant::now() + wait_timeout; + while self.execution_engine.has_active_turn(¤t_turn_id) { + if Instant::now() >= deadline { + warn!( + "Timed out waiting for active turn cancellation: session_id={}, dialog_turn_id={}, timeout_ms={}", + session_id, + current_turn_id, + wait_timeout.as_millis() + ); + break; + } + sleep(Duration::from_millis(50)).await; + } + + Ok(Some(current_turn_id)) + } + /// Delete session pub async fn delete_session( &self, diff --git a/src/crates/core/src/service/snapshot/manager.rs b/src/crates/core/src/service/snapshot/manager.rs index b0502e95..615c699d 100644 --- a/src/crates/core/src/service/snapshot/manager.rs +++ b/src/crates/core/src/service/snapshot/manager.rs @@ -367,8 +367,8 @@ impl Tool for WrappedTool { self.original_tool.is_concurrency_safe(input) } - fn needs_permissions(&self, _input: Option<&Value>) -> bool { - false + fn needs_permissions(&self, input: Option<&Value>) -> bool { + self.original_tool.needs_permissions(input) } async fn validate_input(