From 0240d579fbdc3d6210e42fd260ea9cbf87299cb7 Mon Sep 17 00:00:00 2001 From: Blushyes Date: Mon, 15 Sep 2025 12:06:37 +0800 Subject: [PATCH 1/2] fixed issue where execution could not be canceled --- cli/src/interactive/task_executor.rs | 21 +++++--- core/src/agent/abort.rs | 5 +- core/src/agent/core.rs | 73 ++++++++++++++++++++++++---- 3 files changed, 80 insertions(+), 19 deletions(-) diff --git a/cli/src/interactive/task_executor.rs b/cli/src/interactive/task_executor.rs index 3af138d..684b959 100644 --- a/cli/src/interactive/task_executor.rs +++ b/cli/src/interactive/task_executor.rs @@ -98,6 +98,9 @@ pub async fn execute_agent_task_with_context( } }); + // Create abort controller for this task execution (outside of agent lock) + let (abort_controller, _) = coro_core::agent::AbortController::new(); + // Lock the agent for the duration of this task let mut agent_guard = agent.lock().await; @@ -129,17 +132,22 @@ pub async fn execute_agent_task_with_context( ui_sender.clone(), ))); - // Create new agent + // Create new agent with abort controller let new_agent = coro_core::agent::AgentCore::new_with_output_and_registry( agent_config, llm_config, token_tracking_output, tool_registry, - None, + Some(abort_controller.clone()), ) .await?; *agent_guard = Some(new_agent); + } else { + // Agent exists, update its abort controller for this task + if let Some(existing_agent) = agent_guard.as_mut() { + existing_agent.set_abort_controller(abort_controller.clone()); + } } // Get mutable reference to the agent @@ -148,16 +156,13 @@ pub async fn execute_agent_task_with_context( // Execute task with conversation continuation let task_future = agent_ref.execute_task_with_context(&task, &project_path); - // Listen for interruption signals - cancel the persistent agent when triggered - let agent_for_cancel = agent.clone(); + // Listen for interruption signals - cancel via external abort controller + let abort_controller_for_cancel = abort_controller.clone(); let interrupt_future = async move { loop { match interrupt_receiver.recv().await { Ok(AppMessage::AgentExecutionInterrupted { .. }) => { - if let Some(a) = agent_for_cancel.lock().await.as_ref() { - a.cancel(); - } - tracing::warn!("Task interrupted by user"); + abort_controller_for_cancel.cancel(); return Err(anyhow::anyhow!("Task interrupted by user")); } Ok(_) => continue, // Ignore other messages diff --git a/core/src/agent/abort.rs b/core/src/agent/abort.rs index f3d4d0f..cfe8944 100644 --- a/core/src/agent/abort.rs +++ b/core/src/agent/abort.rs @@ -26,7 +26,10 @@ impl AbortController { /// Trigger cancellation (idempotent) pub fn cancel(&self) { - let _ = self.tx.send(true); + let result = self.tx.send(true); + // We can't use output handler here as AbortController doesn't have access to it + // This is fine as this is the lowest level and should work silently + let _ = result; // Suppress unused variable warning } } diff --git a/core/src/agent/core.rs b/core/src/agent/core.rs index a0c2c4b..ad27b1f 100644 --- a/core/src/agent/core.rs +++ b/core/src/agent/core.rs @@ -29,8 +29,10 @@ pub struct AgentCore { current_task_displayed: bool, execution_context: Option, conversation_manager: ConversationManager, - // Global cancellation controller injected at creation + // Global cancellation controller for external cancel calls abort_controller: crate::agent::AbortController, + // Registration derived from the abort controller for checking cancellation state + abort_registration: crate::agent::AbortRegistration, } impl AgentCore { @@ -69,8 +71,14 @@ impl AgentCore { let max_tokens = llm_config.params.max_tokens.unwrap_or(8192); let conversation_manager = ConversationManager::new(max_tokens, llm_client.clone()); - // Configure cancellation controller - let ac = abort_controller.unwrap_or_else(|| crate::agent::AbortController::new().0); + // Configure cancellation controller and registration + let (abort_controller, abort_registration) = if let Some(controller) = abort_controller { + let registration = controller.subscribe(); + (controller, registration) + } else { + // Create a new controller/registration pair for standalone operation + crate::agent::AbortController::new() + }; Ok(Self { config: agent_config, @@ -82,7 +90,8 @@ impl AgentCore { current_task_displayed: false, execution_context: None, conversation_manager, - abort_controller: ac, + abort_controller, + abort_registration, }) } @@ -96,6 +105,12 @@ impl AgentCore { self.abort_controller.cancel(); } + /// Set a new abort controller for this agent (used for task-specific cancellation) + pub fn set_abort_controller(&mut self, abort_controller: crate::agent::AbortController) { + self.abort_registration = abort_controller.subscribe(); + self.abort_controller = abort_controller; + } + /// Create a new TraeAgent with custom tool registry and output handler pub async fn new_with_output_and_registry( agent_config: AgentConfig, @@ -131,8 +146,14 @@ impl AgentCore { let max_tokens = llm_config.params.max_tokens.unwrap_or(8192); let conversation_manager = ConversationManager::new(max_tokens, llm_client.clone()); - // Configure cancellation controller - let ac = abort_controller.unwrap_or_else(|| crate::agent::AbortController::new().0); + // Configure cancellation controller and registration + let (abort_controller, abort_registration) = if let Some(controller) = abort_controller { + let registration = controller.subscribe(); + (controller, registration) + } else { + // Create a new controller/registration pair for standalone operation + crate::agent::AbortController::new() + }; Ok(Self { config: agent_config, @@ -144,7 +165,8 @@ impl AgentCore { current_task_displayed: false, execution_context: None, conversation_manager, - abort_controller: ac, + abort_controller, + abort_registration, }) } @@ -194,6 +216,24 @@ impl AgentCore { /// Execute a single step of the agent async fn execute_step(&mut self, step: usize, project_path: &Path) -> Result { + // Clone the stored registration for this step + let mut cancel_reg = self.abort_registration.clone(); + + // Race the entire step execution with cancellation + tokio::select! { + _ = cancel_reg.cancelled() => { + // Step was cancelled + let _ = self.output.normal("⏹ Task interrupted by user").await; + return Err("Task interrupted by user".into()); + } + result = self.execute_step_inner(step, project_path) => { + result + } + } + } + + /// Execute the actual step logic + async fn execute_step_inner(&mut self, step: usize, project_path: &Path) -> Result { // Prepare messages - only add system prompt if conversation history doesn't start with one let mut messages = Vec::new(); let needs_system_prompt = self.conversation_history.is_empty() @@ -679,16 +719,28 @@ impl AgentCore { let mut task_completed = false; let mut interrupted = false; - // Subscribe to global cancellation - let mut cancel_reg = self.abort_controller.subscribe(); + // Clone the stored registration for global cancellation + let mut cancel_reg = self.abort_registration.clone(); // Execute steps until completion or max steps reached while step < self.config.max_steps && !task_completed { step += 1; + // Check for cancellation before each step + if cancel_reg.is_cancelled() { + interrupted = true; + break; + } + // Apply intelligent compression before each step to manage token usage self.apply_intelligent_compression().await?; + // Check again after compression + if cancel_reg.is_cancelled() { + interrupted = true; + break; + } + // Race step execution with cancellation tokio::select! { _ = cancel_reg.cancelled() => { @@ -940,7 +992,7 @@ mod tests { std::sync::Arc::new(MockLlmClient::new()), ); - let (ac, _reg) = crate::agent::AbortController::new(); + let (ac, reg) = crate::agent::AbortController::new(); let agent = AgentCore { config: agent_config, @@ -953,6 +1005,7 @@ mod tests { execution_context: None, conversation_manager, abort_controller: ac, + abort_registration: reg, }; let project_path = PathBuf::from("/some/project/path"); From 575db46dabcd3226fb0c3ededa642cdfefdd4201 Mon Sep 17 00:00:00 2001 From: Blushyes Date: Mon, 15 Sep 2025 12:53:59 +0800 Subject: [PATCH 2/2] v0.0.5 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 66f9770..baf0e46 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ members = ["core", "cli"] resolver = "2" [workspace.package] -version = "0.0.4" +version = "0.0.5" edition = "2021" authors = ["Blushyes"] license = "MIT OR Apache-2.0"