From 73ed32224385ff10876e2aaa931a23c47b6dd7f3 Mon Sep 17 00:00:00 2001 From: Blushyes Date: Tue, 26 Aug 2025 10:50:01 +0800 Subject: [PATCH 1/4] update README --- README.md | 18 +++++++++++------- README_zh.md | 16 ++++++++++------ 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 1e6e66c..eff3830 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,11 @@ Create a `coro.json` file: "base_url": "https://api.deepseek.com", "api_key": "your-api-key", "model": "deepseek-chat", - "max_token": 8192 + "params": { + "max_tokens": 131072, + "temperature": 0.7, + "top_p": 0.9 + } } ``` @@ -152,12 +156,12 @@ coro --config custom.json "Analyze this codebase"
🤖 Phase 3: Intelligence & Performance -| Priority | Status | Feature | Description | -| --------- | ------ | ---------------------------------- | ----------------------------------------------------------------------------- | -| 🟡 Medium | 📋 | **Multi-model & Auto Routing** | Auto model selection by task type, failure auto-downgrade & retry strategies | -| 🟡 Medium | 📋 | **Context Optimization & Caching** | File summary caching, duplicate reference deduplication, token budget control | -| 🟡 Medium | ✅ | **Token Compression** | Intelligent context compression, selective token reduction, adaptive context windows | -| 🔵 Low | 📋 | **MCP Extension Ecosystem** | Common provider presets & templates, one-click start/stop external tools | +| Priority | Status | Feature | Description | +| --------- | ------ | ---------------------------------- | ------------------------------------------------------------------------------------ | +| 🟡 Medium | 📋 | **Multi-model & Auto Routing** | Auto model selection by task type, failure auto-downgrade & retry strategies | +| 🟡 Medium | 📋 | **Context Optimization & Caching** | File summary caching, duplicate reference deduplication, token budget control | +| 🟡 Medium | ✅ | **Token Compression** | Intelligent context compression, selective token reduction, adaptive context windows | +| 🔵 Low | 📋 | **MCP Extension Ecosystem** | Common provider presets & templates, one-click start/stop external tools |
diff --git a/README_zh.md b/README_zh.md index 1dd8c8b..e30c4ea 100644 --- a/README_zh.md +++ b/README_zh.md @@ -76,7 +76,11 @@ export CORO_MODEL="custom-model" "base_url": "https://api.deepseek.com", "api_key": "your-api-key", "model": "deepseek-chat", - "max_token": 8192 + "params": { + "max_tokens": 131072, + "temperature": 0.7, + "top_p": 0.9 + } } ``` @@ -138,12 +142,12 @@ export CORO_MODEL="custom-model"
🤖 第三阶段：智能化与性能 -| 优先级 | 状态 | 功能特性 | 描述 | -| ------ | ---- | -------------------- | ---------------------------------------------- | -| 🟡 中 | 📋 | **多模型与自动路由** | 按任务类型自动选择模型，失败自动降级与重试策略 | -| 🟡 中 | 📋 | **上下文优化与缓存** | 文件摘要缓存、重复引用去重、Token 预算控制 | +| 优先级 | 状态 | 功能特性 | 描述 | +| ------ | ---- | -------------------- | --------------------------------------------------- | +| 🟡 中 | 📋 | **多模型与自动路由** | 按任务类型自动选择模型，失败自动降级与重试策略 | +| 🟡 中 | 📋 | **上下文优化与缓存** | 文件摘要缓存、重复引用去重、Token 预算控制 | | 🟡 中 | ✅ | **Token 压缩** | 智能上下文压缩、选择性 Token 减少、自适应上下文窗口 | -| 🔵 低 | 📋 | **MCP 扩展生态** | 常用 Provider 预设与模板，一键启停外部工具 | +| 🔵 低 | 📋 | **MCP 扩展生态** | 常用 Provider 预设与模板，一键启停外部工具 |
From aea0dacada5e3a1d4623619cfd75b9274e1fb4bc Mon Sep 17 00:00:00 2001 From: Blushyes Date: Tue, 26 Aug 2025 10:55:16 +0800 Subject: [PATCH 2/4] fix utf8 panic --- core/src/agent/tokens/conversation_manager.rs | 35 ++++++++----------- 1 file changed, 14 insertions(+), 21 deletions(-) diff --git a/core/src/agent/tokens/conversation_manager.rs b/core/src/agent/tokens/conversation_manager.rs index a3f2b39..ce30a41 100644 --- a/core/src/agent/tokens/conversation_manager.rs +++ b/core/src/agent/tokens/conversation_manager.rs @@ -128,14 +128,7 @@ impl ConversationManager { }; match compression_level { - Some((level, reason)) => { - tracing::info!( - "Applying {} compression: {:.1}% token usage ({})", - level.as_str(), - usage_ratio * 100.0, - reason - ); - + Some((level, _reason)) => { let target_tokens = (self.max_tokens as f64 * self.get_compression_target(level)) as u32; let messages_before_count = messages.len() as u32; @@ -169,8 +162,6 @@ impl ConversationManager { ), }; - tracing::info!("{}", summary.summary); - // Update current token count self.current_tokens = tokens_after; @@ -229,8 +220,6 @@ impl ConversationManager { } async fn light_compression(&self, mut messages: Vec) -> Result> { - tracing::debug!("Applying light compression"); - for message in &mut messages { if let MessageContent::MultiModal(blocks) = &mut message.content { for block in blocks { @@ -251,8 +240,6 @@ impl ConversationManager { messages: Vec, context: Option<&AgentExecutionContext>, ) -> Result> { - tracing::debug!("Applying medium compression"); - if messages.is_empty() { return Ok(messages); } @@ -297,8 +284,6 @@ impl ConversationManager { } async fn heavy_compression(&self, messages: Vec) -> Result> { - tracing::debug!("Applying heavy compression"); - if messages.is_empty() { return Ok(messages); } @@ -340,14 +325,24 @@ impl ConversationManager { let sample_size = (self.tool_output_budget as usize / 2).min(2000); let half_sample = sample_size / 2; + // Safe UTF-8 character boundary slicing let beginning = if output.len() > half_sample { - &output[..half_sample] + let mut boundary = half_sample; + while boundary < output.len() && !output.is_char_boundary(boundary) { + boundary += 1; + } + &output[..boundary] } else { output }; let ending = if output.len() > sample_size { - &output[output.len().saturating_sub(half_sample)..] + let start_pos = output.len().saturating_sub(half_sample); + let mut boundary = start_pos; + while boundary > 0 && !output.is_char_boundary(boundary) { + boundary -= 1; + } + &output[boundary..] } else { "" }; @@ -390,8 +385,7 @@ impl ConversationManager { )) } } - Err(e) => { - tracing::warn!("Failed to generate tool output summary: {}", e); + Err(_e) => { // Fallback to truncation Ok(format!( "{}...[truncated {} chars]", @@ -574,7 +568,6 @@ The structure MUST be as follows: } } Err(e) => { - tracing::warn!("Failed to generate conversation summary: {}", e); let fallback_goal = context .map(|c| c.original_goal.as_str()) .unwrap_or("Summary generation error"); From 1498621150fe0e1eac47a082cd5e08a418ebc56c Mon Sep 17 00:00:00 2001 From: Blushyes Date: Tue, 26 Aug 2025 12:07:59 +0800 Subject: [PATCH 3/4] optimize prompt --- cli/src/commands/test.rs | 2 +- cli/src/commands/tools.rs | 2 +- core/src/agent/mod.rs | 2 +- core/src/agent/prompt.rs | 8 ++++---- core/src/error.rs | 6 +++--- core/src/output/mod.rs | 2 +- 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/cli/src/commands/test.rs b/cli/src/commands/test.rs index 55e679d..91a1125 100644 --- a/cli/src/commands/test.rs +++ b/cli/src/commands/test.rs @@ -7,7 +7,7 @@ use tracing::info; pub async fn test_command() -> Result<()> { info!("Testing basic functionality"); - println!("🧪 Running Trae Agent Tests"); + println!("🧪 Running Coro Code Tests"); // Test 1: Configuration loading println!("📋 Test 1: Configuration system"); diff --git a/cli/src/commands/tools.rs b/cli/src/commands/tools.rs index 09b70cc..92c42fe 100644 --- a/cli/src/commands/tools.rs +++ b/cli/src/commands/tools.rs @@ -25,7 +25,7 @@ pub async fn tools_command() -> Result<()> { println!("💡 Use these tools in your tasks to accomplish complex workflows!"); println!( - "📋 All tools follow the exact same specifications as the Python version of Trae Agent." + "📋 All tools follow the exact same specifications as the Python version of Coro Code." ); Ok(()) diff --git a/core/src/agent/mod.rs b/core/src/agent/mod.rs index fbaa8d1..8d15b6d 100644 --- a/core/src/agent/mod.rs +++ b/core/src/agent/mod.rs @@ -11,7 +11,7 @@ pub use base::{Agent, AgentResult}; pub use config::{AgentBuilder, AgentConfig, OutputMode}; pub use core::AgentCore; pub use execution::AgentExecution; -pub use prompt::{build_system_prompt_with_context, build_user_message, TRAE_AGENT_SYSTEM_PROMPT}; +pub use prompt::{build_system_prompt_with_context, build_user_message, CORO_CODE_SYSTEM_PROMPT}; pub use tokens::{ CompressionLevel, CompressionSummary, ConversationManager, ConversationTokenStats, MaybeCompressedResult, TokenCalculator, diff --git a/core/src/agent/prompt.rs b/core/src/agent/prompt.rs index 59f857b..109e761 100644 --- a/core/src/agent/prompt.rs +++ b/core/src/agent/prompt.rs @@ -1,7 +1,7 @@ //! Agent system prompts -/// Trae Agent system prompt (consistent with Python version) -pub const TRAE_AGENT_SYSTEM_PROMPT: &str = r#"You are an expert AI software engineering agent. +/// Coro Code system prompt (consistent with Python version) +pub const CORO_CODE_SYSTEM_PROMPT: &str = r#"You are an expert AI software engineering agent. File Path Rule: All tools that take a `file_path` as an argument require an **absolute path**. You MUST construct the full, absolute path by combining the `[Project root path]` provided in the user's message with the file's path inside the project. @@ -49,7 +49,7 @@ Follow these steps methodically: - The sequential_thinking tool can help you break down complex problems, analyze issues step-by-step, and ensure a thorough approach to problem-solving. - Don't hesitate to use it multiple times throughout your thought process to enhance the depth and accuracy of your solutions. -If you are sure the issue has been solved, you should call the `task_done` to finish the task. +If you are sure the issue has been solved, or you have already finished answering the user’s question, you should call the `task_done` to finish the task. "#; /// Build system context information @@ -83,7 +83,7 @@ pub fn build_system_prompt_with_context(project_path: &std::path::Path) -> Strin Construct absolute paths by combining the project root path above with relative file paths.\n\ Example: If you want to edit 'src/main.rs', use '{}/src/main.rs'\n\n\ [System Context]:\n{}", - TRAE_AGENT_SYSTEM_PROMPT, + CORO_CODE_SYSTEM_PROMPT, project_path_str, project_path_str, system_context diff --git a/core/src/error.rs b/core/src/error.rs index 1ddadee..243b361 100644 --- a/core/src/error.rs +++ b/core/src/error.rs @@ -1,11 +1,11 @@ -//! Error types and handling for Trae Agent Core +//! Error types and handling for Coro Code Core use thiserror::Error; -/// Result type alias for Trae Agent operations +/// Result type alias for Coro Code operations pub type Result = std::result::Result; -/// Main error type for Trae Agent Core +/// Main error type for Coro Code Core #[derive(Error, Debug)] pub enum Error { /// Configuration-related errors diff --git a/core/src/output/mod.rs b/core/src/output/mod.rs index 91157a7..b7ec6dd 100644 --- a/core/src/output/mod.rs +++ b/core/src/output/mod.rs @@ -1,4 +1,4 @@ -//! Output abstraction layer for the Trae Agent core +//! Output abstraction layer for the Coro Code core //! //! This module provides an abstract interface for outputting agent execution information, //! allowing different implementations for CLI, API, logging, etc. From 0c2d8559868acc317f1343cf43934d2d397932eb Mon Sep 17 00:00:00 2001 From: Blushyes Date: Mon, 8 Sep 2025 11:16:36 +0800 Subject: [PATCH 4/4] add abort controller --- Cargo.toml | 2 +- cli/src/interactive/README.md | 20 ++++ cli/src/interactive/task_executor.rs | 19 +++- cli/src/output/cli_handler.rs | 9 ++ core/examples/custom_system_prompt.rs | 3 +- core/src/agent/abort.rs | 147 ++++++++++++++++++++++++++ core/src/agent/config.rs | 17 ++- core/src/agent/core.rs | 120 +++++++++++++++------ core/src/agent/mod.rs | 3 + core/src/output/mod.rs | 5 + 10 files changed, 307 insertions(+), 38 deletions(-) create mode 100644 core/src/agent/abort.rs diff --git a/Cargo.toml b/Cargo.toml index fe964d5..3b248ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ members = ["core", "cli"] resolver = "2" [workspace.package] -version = "0.0.3" +version = "0.0.4" edition = "2021" authors = ["Blushyes"] license = "MIT OR Apache-2.0" diff --git a/cli/src/interactive/README.md b/cli/src/interactive/README.md index e9e9bca..6f691a1 100644 --- a/cli/src/interactive/README.md +++ b/cli/src/interactive/README.md @@ -165,6 +165,26 @@ Manages agent task execution with UI integration and token tracking. ### `terminal_output.rs` - Terminal Output Abstraction +#### Cancellation (Interrupt) + +- UI sends `AppMessage::AgentExecutionInterrupted` (e.g., on Ctrl+C) +- Task executor forwards this to core by calling the agent's global AbortController +- Core emits `AgentEvent::ExecutionInterrupted { context, reason }` and stops gracefully + +Programmatic example (non-UI): + +```rust +let (abort_controller, _reg) = coro_core::agent::AbortController::new(); +let mut agent = coro_core::agent::AgentCore::new_with_llm_config( + agent_config, + llm_config, + Box::new(coro_core::output::events::NullOutput), + Some(abort_controller.clone()), +).await?; +// In another task/thread: +abort_controller.cancel(); +``` + Provides terminal output utilities and formatting functions that work with the AgentOutput system. **Key Traits:** diff --git a/cli/src/interactive/task_executor.rs b/cli/src/interactive/task_executor.rs index 4a98aa6..3af138d 100644 --- a/cli/src/interactive/task_executor.rs +++ b/cli/src/interactive/task_executor.rs @@ -135,6 +135,7 @@ pub async fn execute_agent_task_with_context( llm_config, token_tracking_output, tool_registry, + None, ) .await?; @@ -147,11 +148,15 @@ 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 - let interrupt_future = async { + // Listen for interruption signals - cancel the persistent agent when triggered + let agent_for_cancel = agent.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"); return Err(anyhow::anyhow!("Task interrupted by user")); } @@ -223,23 +228,29 @@ pub async fn execute_agent_task( ui_sender.clone(), ))); + // Create an AbortController for this single-run agent + let (abort_controller, _reg) = coro_core::agent::AbortController::new(); + // Create and execute agent task let mut agent = coro_core::agent::AgentCore::new_with_output_and_registry( agent_config, llm_config, token_tracking_output, tool_registry, + Some(abort_controller.clone()), ) .await?; // Execute task with interruption support let task_future = agent.execute_task_with_context(&task, &project_path); - // Listen for interruption signals - let interrupt_future = async { + // Listen for interruption signals - cancel via AbortController when triggered + let abort_controller_for_cancel = abort_controller.clone(); + let interrupt_future = async move { loop { match interrupt_receiver.recv().await { Ok(AppMessage::AgentExecutionInterrupted { .. }) => { + abort_controller_for_cancel.cancel(); tracing::warn!("Task interrupted by user"); return Err(anyhow::anyhow!("Task interrupted by user")); } diff --git a/cli/src/output/cli_handler.rs b/cli/src/output/cli_handler.rs index e3fdba0..782fdf0 100644 --- a/cli/src/output/cli_handler.rs +++ b/cli/src/output/cli_handler.rs @@ -103,6 +103,15 @@ impl AgentOutput for CliOutputHandler { } } + AgentEvent::ExecutionInterrupted { context, reason } => { + warn!("Task interrupted: {}", reason); + debug!( + "Interrupted after {} steps, duration: {:.2}s", + context.current_step, + context.execution_time.as_secs_f64() + ); + } + AgentEvent::StepStarted { step_info } => { debug!("Step {}: {}", step_info.step_number, step_info.task); } diff --git a/core/examples/custom_system_prompt.rs b/core/examples/custom_system_prompt.rs index fc6c77c..53f644c 100644 --- a/core/examples/custom_system_prompt.rs +++ b/core/examples/custom_system_prompt.rs @@ -36,7 +36,8 @@ async fn main() -> Result<(), Box> { // Create agent with custom system prompt let mut agent = - AgentCore::new_with_llm_config(agent_config, llm_config, Box::new(NullOutput)).await?; + AgentCore::new_with_llm_config(agent_config, llm_config, Box::new(NullOutput), None) + .await?; // Verify the system prompt is set if let Some(prompt) = agent.get_configured_system_prompt() { diff --git a/core/src/agent/abort.rs b/core/src/agent/abort.rs new file mode 100644 index 0000000..f3d4d0f --- /dev/null +++ b/core/src/agent/abort.rs @@ -0,0 +1,147 @@ +//! Abort (cancellation) controller for AgentCore + +#[derive(Clone)] +pub struct AbortController { + tx: tokio::sync::watch::Sender, +} + +#[derive(Clone)] +pub struct AbortRegistration { + rx: tokio::sync::watch::Receiver, +} + +impl AbortController { + /// Create a new controller and its registration + pub fn new() -> (Self, AbortRegistration) { + let (tx, rx) = tokio::sync::watch::channel(false); + (Self { tx: tx.clone() }, AbortRegistration { rx }) + } + + /// Subscribe to this controller to obtain a registration + pub fn subscribe(&self) -> AbortRegistration { + AbortRegistration { + rx: self.tx.subscribe(), + } + } + + /// Trigger cancellation (idempotent) + pub fn cancel(&self) { + let _ = self.tx.send(true); + } +} + +impl AbortRegistration { + /// Check current cancellation state without blocking + pub fn is_cancelled(&self) -> bool { + *self.rx.borrow() + } + + /// Wait until cancellation is triggered (returns immediately if already cancelled) + pub async fn cancelled(&mut self) { + if !*self.rx.borrow() { + let _ = self.rx.changed().await; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::time::{sleep, timeout}; + + #[tokio::test] + async fn test_abort_controller_basic() { + let (controller, mut registration) = AbortController::new(); + + // Initially not cancelled + assert!(!registration.is_cancelled()); + + // Cancel the controller + controller.cancel(); + + // Should be cancelled now + assert!(registration.is_cancelled()); + + // cancelled() should return immediately + let result = timeout(Duration::from_millis(100), registration.cancelled()).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_abort_controller_multiple_subscribers() { + let (controller, mut reg1) = AbortController::new(); + let mut reg2 = controller.subscribe(); + + // Both should not be cancelled initially + assert!(!reg1.is_cancelled()); + assert!(!reg2.is_cancelled()); + + // Cancel the controller + controller.cancel(); + + // Both should be cancelled + assert!(reg1.is_cancelled()); + assert!(reg2.is_cancelled()); + + // Both should return immediately from cancelled() + let result1 = timeout(Duration::from_millis(100), reg1.cancelled()).await; + let result2 = timeout(Duration::from_millis(100), reg2.cancelled()).await; + assert!(result1.is_ok()); + assert!(result2.is_ok()); + } + + #[tokio::test] + async fn test_abort_controller_task_interruption() { + // Case 1: No cancel → should complete other branch quickly + let (_controller, mut registration) = AbortController::new(); + let result = tokio::select! { + _ = registration.cancelled() => "interrupted", + _ = sleep(Duration::from_millis(20)) => "completed", + }; + assert_eq!(result, "completed"); + + // Case 2: Cancel after short delay → should interrupt + let (controller2, mut registration2) = AbortController::new(); + let controller2_clone = controller2.clone(); + tokio::spawn(async move { + sleep(Duration::from_millis(30)).await; + controller2_clone.cancel(); + }); + let result2 = tokio::select! { + _ = registration2.cancelled() => "interrupted", + _ = sleep(Duration::from_secs(2)) => "completed", + }; + assert_eq!(result2, "interrupted"); + } + + #[tokio::test] + async fn test_abort_controller_idempotent_cancel() { + let (controller, mut registration) = AbortController::new(); + + // Cancel multiple times + controller.cancel(); + controller.cancel(); + controller.cancel(); + + // Should still work correctly + assert!(registration.is_cancelled()); + + let result = timeout(Duration::from_millis(100), registration.cancelled()).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_abort_controller_clone() { + let (controller, mut registration) = AbortController::new(); + let controller_clone = controller.clone(); + + // Clone should work the same way + controller_clone.cancel(); + + assert!(registration.is_cancelled()); + + let result = timeout(Duration::from_millis(100), registration.cancelled()).await; + assert!(result.is_ok()); + } +} diff --git a/core/src/agent/config.rs b/core/src/agent/config.rs index 7af512a..0e1b5ac 100644 --- a/core/src/agent/config.rs +++ b/core/src/agent/config.rs @@ -60,6 +60,7 @@ impl Default for AgentConfig { pub struct AgentBuilder { llm_config: crate::config::ResolvedLlmConfig, agent_config: AgentConfig, + abort_controller: Option, } impl AgentBuilder { @@ -68,6 +69,7 @@ impl AgentBuilder { Self { llm_config, agent_config: AgentConfig::default(), + abort_controller: None, } } @@ -101,12 +103,24 @@ impl AgentBuilder { self } + /// Inject a global AbortController for cancellation support + pub fn with_cancellation(mut self, controller: super::AbortController) -> Self { + self.abort_controller = Some(controller); + self + } + /// Build the agent with the given output handler pub async fn build_with_output( self, output: Box, ) -> crate::error::Result { - super::AgentCore::new_with_llm_config(self.agent_config, self.llm_config, output).await + super::AgentCore::new_with_llm_config( + self.agent_config, + self.llm_config, + output, + self.abort_controller, + ) + .await } /// Build the agent with custom output handler and tool registry @@ -120,6 +134,7 @@ impl AgentBuilder { self.llm_config, output, tool_registry, + self.abort_controller, ) .await } diff --git a/core/src/agent/core.rs b/core/src/agent/core.rs index f43573b..ea15677 100644 --- a/core/src/agent/core.rs +++ b/core/src/agent/core.rs @@ -29,6 +29,8 @@ pub struct AgentCore { current_task_displayed: bool, execution_context: Option, conversation_manager: ConversationManager, + // Global cancellation controller injected at creation + abort_controller: crate::agent::AbortController, } impl AgentCore { @@ -37,6 +39,7 @@ impl AgentCore { agent_config: AgentConfig, llm_config: crate::config::ResolvedLlmConfig, output: Box, + abort_controller: Option, ) -> Result { // Create LLM client based on protocol let llm_client: Arc = match llm_config.protocol { @@ -66,6 +69,9 @@ 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); + Ok(Self { config: agent_config, llm_client, @@ -76,6 +82,7 @@ impl AgentCore { current_task_displayed: false, execution_context: None, conversation_manager, + abort_controller: ac, }) } @@ -84,12 +91,18 @@ impl AgentCore { &self.config } + /// Request cancellation on this agent + pub fn cancel(&self) { + self.abort_controller.cancel(); + } + /// Create a new TraeAgent with custom tool registry and output handler pub async fn new_with_output_and_registry( agent_config: AgentConfig, llm_config: crate::config::ResolvedLlmConfig, output: Box, tool_registry: ToolRegistry, + abort_controller: Option, ) -> Result { // Create LLM client based on protocol let llm_client: Arc = match llm_config.protocol { @@ -118,6 +131,9 @@ 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); + Ok(Self { config: agent_config, llm_client, @@ -128,6 +144,7 @@ impl AgentCore { current_task_displayed: false, execution_context: None, conversation_manager, + abort_controller: ac, }) } @@ -137,7 +154,7 @@ impl AgentCore { llm_config: crate::config::ResolvedLlmConfig, ) -> Result { use crate::output::events::NullOutput; - Self::new_with_llm_config(agent_config, llm_config, Box::new(NullOutput)).await + Self::new_with_llm_config(agent_config, llm_config, Box::new(NullOutput), None).await } /// Set a custom system prompt for the agent @@ -619,6 +636,10 @@ impl AgentCore { let mut step = 0; let mut task_completed = false; + let mut interrupted = false; + // Subscribe to global cancellation + let mut cancel_reg = self.abort_controller.subscribe(); + // Execute steps until completion or max steps reached while step < self.config.max_steps && !task_completed { step += 1; @@ -626,39 +647,49 @@ impl AgentCore { // Apply intelligent compression before each step to manage token usage self.apply_intelligent_compression().await?; - match self.execute_step(step, project_path).await { - Ok(completed) => { - task_completed = completed; - - // Record step completion - if let Some(recorder) = &self.trajectory_recorder { - recorder - .record(TrajectoryEntry::step_complete( - format!("Step {} completed", step), - true, - step, - )) - .await?; - } + // Race step execution with cancellation + tokio::select! { + _ = cancel_reg.cancelled() => { + interrupted = true; + break; } - Err(e) => { - // Record error - if let Some(recorder) = &self.trajectory_recorder { - recorder - .record(TrajectoryEntry::error( - e.to_string(), - Some(format!("Step {}", step)), + result = self.execute_step(step, project_path) => { + match result { + Ok(completed) => { + task_completed = completed; + + // Record step completion + if let Some(recorder) = &self.trajectory_recorder { + recorder + .record(TrajectoryEntry::step_complete( + format!("Step {} completed", step), + true, + step, + )) + .await?; + } + } + Err(e) => { + // Record error + if let Some(recorder) = &self.trajectory_recorder { + recorder + .record(TrajectoryEntry::error( + e.to_string(), + Some(format!("Step {}", step)), + step, + )) + .await?; + } + + let duration = start_time.elapsed().as_millis() as u64; + return Ok(AgentExecution::failure( + format!("Error in step {}: {}", step, e), step, - )) - .await?; - } + duration, + )); - let duration = start_time.elapsed().as_millis() as u64; - return Ok(AgentExecution::failure( - format!("Error in step {}: {}", step, e), - step, - duration, - )); + } + } } } } @@ -695,6 +726,30 @@ impl AgentCore { format!("Task incomplete after {} steps", step) }; + // If interrupted, emit event and return immediately + if interrupted { + if let Some(context) = &self.execution_context { + self.output + .emit_event(AgentEvent::ExecutionInterrupted { + context: context.clone(), + reason: "Execution interrupted by user".to_string(), + }) + .await + .unwrap_or_else(|e| { + let _ = futures::executor::block_on(self.output.debug(&format!( + "Failed to emit execution interrupted event: {}", + e + ))); + }); + } + let duration_ms = duration.as_millis() as u64; + return Ok(AgentExecution::failure( + "Execution interrupted".to_string(), + step, + duration_ms, + )); + } + self.output .emit_event(AgentEvent::ExecutionCompleted { context: context.clone(), @@ -843,6 +898,8 @@ mod tests { std::sync::Arc::new(MockLlmClient::new()), ); + let (ac, _reg) = crate::agent::AbortController::new(); + let agent = AgentCore { config: agent_config, llm_client: std::sync::Arc::new(MockLlmClient::new()), @@ -853,6 +910,7 @@ mod tests { current_task_displayed: false, execution_context: None, conversation_manager, + abort_controller: ac, }; let project_path = PathBuf::from("/some/project/path"); diff --git a/core/src/agent/mod.rs b/core/src/agent/mod.rs index 8d15b6d..ad5b4c0 100644 --- a/core/src/agent/mod.rs +++ b/core/src/agent/mod.rs @@ -16,3 +16,6 @@ pub use tokens::{ CompressionLevel, CompressionSummary, ConversationManager, ConversationTokenStats, MaybeCompressedResult, TokenCalculator, }; + +pub mod abort; +pub use abort::{AbortController, AbortRegistration}; diff --git a/core/src/output/mod.rs b/core/src/output/mod.rs index b7ec6dd..dbc04d2 100644 --- a/core/src/output/mod.rs +++ b/core/src/output/mod.rs @@ -116,6 +116,11 @@ pub enum AgentEvent { success: bool, summary: String, }, + /// Agent execution interrupted (cancelled) + ExecutionInterrupted { + context: AgentExecutionContext, + reason: String, + }, /// New step started StepStarted { step_info: AgentStepInfo }, /// Step completed