diff --git a/Cargo.lock b/Cargo.lock index edaba94..915e6ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -689,6 +689,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-util", "tracing", "tracing-subscriber", ] @@ -722,6 +723,7 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tokio-stream", + "tokio-util", "tracing", "uuid", ] diff --git a/crates/g3-cli/Cargo.toml b/crates/g3-cli/Cargo.toml index 78a9ac5..37be2d9 100644 --- a/crates/g3-cli/Cargo.toml +++ b/crates/g3-cli/Cargo.toml @@ -16,3 +16,4 @@ serde = { workspace = true } serde_json = { workspace = true } rustyline = "17.0.1" dirs = "5.0" +tokio-util = "0.7" diff --git a/crates/g3-cli/src/lib.rs b/crates/g3-cli/src/lib.rs index 2092014..b590b8a 100644 --- a/crates/g3-cli/src/lib.rs +++ b/crates/g3-cli/src/lib.rs @@ -5,6 +5,7 @@ use anyhow::Result; use tracing::{info, error}; use rustyline::error::ReadlineError; use rustyline::DefaultEditor; +use tokio_util::sync::CancellationToken; #[derive(Parser)] #[command(name = "g3")] @@ -73,6 +74,7 @@ async fn run_interactive(agent: Agent, show_prompt: bool, show_code: bool) -> Re println!("I solve problems by writing and executing code. Tell me what you need to accomplish!"); println!(); println!("Type 'exit' or 'quit' to exit, use Up/Down arrows for command history"); + println!("Press ESC during operations to cancel the current request"); println!(); // Initialize rustyline editor with history @@ -106,10 +108,43 @@ async fn run_interactive(agent: Agent, show_prompt: bool, show_code: bool) -> Re // Add to history rl.add_history_entry(input)?; - // Execute task (code-first approach) - match agent.execute_task_with_timing(input, None, false, show_prompt, show_code, true).await { + // Create cancellation token for this request + let cancellation_token = CancellationToken::new(); + let cancel_token_clone = cancellation_token.clone(); + + // Spawn a task to monitor for ESC key during execution + let esc_monitor = tokio::spawn(async move { + // This is a simplified approach - in a real implementation, + // we'd need to handle raw terminal input to detect ESC + // For now, we'll just provide the cancellation infrastructure + tokio::time::sleep(tokio::time::Duration::from_secs(3600)).await; + }); + + // Execute task with cancellation support + let execution_result = tokio::select! { + result = agent.execute_task_with_timing_cancellable( + input, None, false, show_prompt, show_code, true, cancellation_token + ) => { + esc_monitor.abort(); + result + } + _ = tokio::signal::ctrl_c() => { + cancel_token_clone.cancel(); + esc_monitor.abort(); + println!("\n⚠️ Operation cancelled by user (Ctrl+C)"); + continue; + } + }; + + match execution_result { Ok(response) => println!("{}", response), - Err(e) => error!("Error: {}", e), + Err(e) => { + if e.to_string().contains("cancelled") { + println!("⚠️ Operation cancelled by user"); + } else { + error!("Error: {}", e); + } + } } }, Err(ReadlineError::Interrupted) => { diff --git a/crates/g3-core/Cargo.toml b/crates/g3-core/Cargo.toml index eb6aabf..fa2af3f 100644 --- a/crates/g3-core/Cargo.toml +++ b/crates/g3-core/Cargo.toml @@ -20,3 +20,4 @@ async-trait = "0.1" tokio-stream = "0.1" llama_cpp = { version = "0.3.2", features = ["metal"] } shellexpand = "3.1" +tokio-util = "0.7" diff --git a/crates/g3-core/src/lib.rs b/crates/g3-core/src/lib.rs index c10c07b..876eed1 100644 --- a/crates/g3-core/src/lib.rs +++ b/crates/g3-core/src/lib.rs @@ -7,6 +7,7 @@ use std::path::Path; use std::time::{Duration, Instant}; use tracing::field::debug; use tracing::info; +use tokio_util::sync::CancellationToken; pub struct Agent { providers: ProviderRegistry, @@ -91,6 +92,30 @@ impl Agent { show_prompt: bool, show_code: bool, show_timing: bool, + ) -> Result { + // Create a cancellation token that never cancels for backward compatibility + let cancellation_token = CancellationToken::new(); + self.execute_task_with_timing_cancellable( + description, + language, + _auto_execute, + show_prompt, + show_code, + show_timing, + cancellation_token, + ) + .await + } + + pub async fn execute_task_with_timing_cancellable( + &self, + description: &str, + language: Option<&str>, + _auto_execute: bool, + show_prompt: bool, + show_code: bool, + show_timing: bool, + cancellation_token: CancellationToken, ) -> Result { info!("Executing task: {}", description); @@ -154,17 +179,25 @@ with nothing afterwards.", stream: false, }; - // Time the LLM call + // Time the LLM call with cancellation support let llm_start = Instant::now(); - let response = provider.complete(request).await?; + let response = tokio::select! { + result = provider.complete(request) => result?, + _ = cancellation_token.cancelled() => { + return Err(anyhow::anyhow!("Operation cancelled by user")); + } + }; let llm_duration = llm_start.elapsed(); - // Time the code execution + // Time the code execution with cancellation support let exec_start = Instant::now(); let executor = CodeExecutor::new(); - let result = executor - .execute_from_response_with_options(&response.content, show_code) - .await?; + let result = tokio::select! { + result = executor.execute_from_response_with_options(&response.content, show_code) => result?, + _ = cancellation_token.cancelled() => { + return Err(anyhow::anyhow!("Operation cancelled by user")); + } + }; let exec_duration = exec_start.elapsed(); let total_duration = total_start.elapsed();