diff --git a/crates/goose-cli/src/commands/web.rs b/crates/goose-cli/src/commands/web.rs index ba5206b4ba2e..bac587e16b76 100644 --- a/crates/goose-cli/src/commands/web.rs +++ b/crates/goose-cli/src/commands/web.rs @@ -10,6 +10,7 @@ use axum::{ }; use futures::{sink::SinkExt, stream::StreamExt}; use goose::agents::{Agent, AgentEvent}; +use goose::context_mgmt::auto_compact::check_and_compact_messages; use goose::message::Message as GooseMessage; use goose::session; use serde::{Deserialize, Serialize}; @@ -455,6 +456,45 @@ async fn process_message_streaming( session_msgs.clone() }; + // Check and compact messages if needed before calling reply + let compact_result = check_and_compact_messages(agent, &messages, None).await?; + if compact_result.compacted { + messages = compact_result.messages.clone(); + + // Update session messages + { + let mut session_msgs = session_messages.lock().await; + *session_msgs = compact_result.messages; + } + + // Notify client of compaction + let msg = if let (Some(before), Some(after)) = + (compact_result.tokens_before, compact_result.tokens_after) + { + format!( + "Auto-compacted context: {} → {} tokens ({:.0}% reduction)", + before, + after, + (1.0 - (after as f64 / before as f64)) * 100.0 + ) + } else { + "Auto-compacted context to prevent overflow".to_string() + }; + + let mut sender_lock = sender.lock().await; + let _ = sender_lock + .send(Message::Text( + serde_json::to_string(&WebSocketMessage::Response { + content: msg, + role: "system".to_string(), + timestamp: chrono::Utc::now().timestamp_millis(), + }) + .unwrap() + .into(), + )) + .await; + } + // Persist messages to JSONL file with provider for automatic description generation let provider = agent.provider().await; if provider.is_err() { diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 6ebc79547008..fb6f8a318290 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -16,6 +16,7 @@ pub use self::export::message_to_markdown; pub use builder::{build_session, SessionBuilderConfig, SessionSettings}; use console::Color; use goose::agents::AgentEvent; +use goose::context_mgmt::auto_compact::check_and_compact_messages; use goose::message::push_message; use goose::permission::permission_confirmation::PrincipalType; use goose::permission::Permission; @@ -840,6 +841,41 @@ impl Session { } async fn process_agent_response(&mut self, interactive: bool) -> Result<()> { + // Check and compact messages if needed before calling reply + let compact_result = check_and_compact_messages(&self.agent, &self.messages, None).await?; + if compact_result.compacted { + self.messages = compact_result.messages; + + // Notify user of compaction + let msg = if let (Some(before), Some(after)) = + (compact_result.tokens_before, compact_result.tokens_after) + { + format!( + "Auto-compacted context: {} → {} tokens ({:.0}% reduction)", + before, + after, + (1.0 - (after as f64 / before as f64)) * 100.0 + ) + } else { + "Auto-compacted context to prevent overflow".to_string() + }; + output::render_text(&msg, Some(Color::Yellow), true); + + // Persist the compacted messages + if let Some(session_file) = &self.session_file { + let working_dir = std::env::current_dir().ok(); + let provider = self.agent.provider().await.ok(); + session::persist_messages_with_schedule_id( + session_file, + &self.messages, + provider, + self.scheduled_job_id.clone(), + working_dir, + ) + .await?; + } + } + let cancel_token = CancellationToken::new(); let cancel_token_clone = cancel_token.clone(); diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 24a1f7eb104d..63140ac5656e 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -11,6 +11,7 @@ use bytes::Bytes; use futures::{stream::StreamExt, Stream}; use goose::{ agents::{AgentEvent, SessionConfig}, + context_mgmt::auto_compact::check_and_compact_messages, message::{push_message, Message}, permission::permission_confirmation::PrincipalType, }; @@ -159,8 +160,43 @@ async fn reply_handler( retry_config: None, }; + // Check and compact messages if needed before calling reply + let mut messages_to_process = messages.clone(); + let compact_result = check_and_compact_messages(&agent, &messages_to_process, None).await; + if let Ok(result) = compact_result { + if result.compacted { + messages_to_process = result.messages; + + // Notify client of compaction + let msg = if let (Some(before), Some(after)) = + (result.tokens_before, result.tokens_after) + { + format!( + "Auto-compacted context: {} → {} tokens ({:.0}% reduction)", + before, + after, + (1.0 - (after as f64 / before as f64)) * 100.0 + ) + } else { + "Auto-compacted context to prevent overflow".to_string() + }; + + let _ = stream_event( + MessageEvent::Message { + message: Message::assistant().with_text(&msg), + }, + &task_tx, + ) + .await; + } + } + let mut stream = match agent - .reply(&messages, Some(session_config), Some(task_cancel.clone())) + .reply( + &messages_to_process, + Some(session_config), + Some(task_cancel.clone()), + ) .await { Ok(stream) => stream, diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 40fe5563eee0..c0e83fb8063d 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -273,6 +273,7 @@ impl Agent { &self, tool_call: mcp_core::tool::ToolCall, request_id: String, + cancellation_token: Option, ) -> (String, Result) { // Check if this tool call should be allowed based on repetition monitoring if let Some(monitor) = self.tool_monitor.lock().await.as_mut() { @@ -345,10 +346,12 @@ impl Agent { let task_config = TaskConfig::new(provider, Some(Arc::clone(&self.extension_manager)), mcp_tx); + subagent_execute_task_tool::run_tasks( tool_call.arguments.clone(), task_config, &self.tasks_manager, + cancellation_token, ) .await } else if tool_call.name == DYNAMIC_TASK_TOOL_NAME_PREFIX { @@ -914,7 +917,7 @@ impl Agent { for request in &permission_check_result.approved { if let Ok(tool_call) = request.tool_call.clone() { let (req_id, tool_result) = self - .dispatch_tool_call(tool_call, request.id.clone()) + .dispatch_tool_call(tool_call, request.id.clone(), cancel_token.clone()) .await; tool_futures.push(( @@ -951,6 +954,7 @@ impl Agent { tool_futures_arc.clone(), &mut permission_manager, message_tool_response.clone(), + cancel_token.clone(), ); while let Some(msg) = tool_approval_stream.try_next().await? { diff --git a/crates/goose/src/agents/sub_recipe_execution_tool/mod.rs b/crates/goose/src/agents/sub_recipe_execution_tool/mod.rs deleted file mode 100644 index 49fcc194c56a..000000000000 --- a/crates/goose/src/agents/sub_recipe_execution_tool/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -mod executor; -pub mod lib; -pub mod notification_events; -pub mod sub_recipe_execute_task_tool; -mod task_execution_tracker; -mod task_types; -mod tasks; -pub mod tasks_manager; -pub mod utils; -mod workers; - diff --git a/crates/goose/src/agents/sub_recipe_execution_tool/tasks.rs b/crates/goose/src/agents/sub_recipe_execution_tool/tasks.rs deleted file mode 100644 index 66f67729e69e..000000000000 --- a/crates/goose/src/agents/sub_recipe_execution_tool/tasks.rs +++ /dev/null @@ -1,186 +0,0 @@ -use serde_json::Value; -use std::process::Stdio; -use std::sync::Arc; -use tokio::io::{AsyncBufReadExt, BufReader}; -use tokio::process::Command; - -use crate::agents::sub_recipe_execution_tool::task_execution_tracker::TaskExecutionTracker; -use crate::agents::sub_recipe_execution_tool::task_types::{Task, TaskResult, TaskStatus}; - -pub async fn process_task( - task: &Task, - task_execution_tracker: Arc, -) -> TaskResult { - match get_task_result(task.clone(), task_execution_tracker).await { - Ok(data) => TaskResult { - task_id: task.id.clone(), - status: TaskStatus::Completed, - data: Some(data), - error: None, - }, - Err(error) => TaskResult { - task_id: task.id.clone(), - status: TaskStatus::Failed, - data: None, - error: Some(error), - }, - } -} - -async fn get_task_result( - task: Task, - task_execution_tracker: Arc, -) -> Result { - let (command, output_identifier) = build_command(&task)?; - let (stdout_output, stderr_output, success) = run_command( - command, - &output_identifier, - &task.id, - task_execution_tracker, - ) - .await?; - - if success { - process_output(stdout_output) - } else { - Err(format!("Command failed:\n{}", stderr_output)) - } -} - -fn build_command(task: &Task) -> Result<(Command, String), String> { - let task_error = |field: &str| format!("Task {}: Missing {}", task.id, field); - - let mut output_identifier = task.id.clone(); - let mut command = if task.task_type == "sub_recipe" { - let sub_recipe_name = task - .get_sub_recipe_name() - .ok_or_else(|| task_error("sub_recipe name"))?; - let path = task - .get_sub_recipe_path() - .ok_or_else(|| task_error("sub_recipe path"))?; - let command_parameters = task - .get_command_parameters() - .ok_or_else(|| task_error("command_parameters"))?; - - output_identifier = format!("sub-recipe {}", sub_recipe_name); - let mut cmd = Command::new("goose"); - cmd.arg("run").arg("--recipe").arg(path).arg("--no-session"); - - for (key, value) in command_parameters { - let key_str = key.to_string(); - let value_str = value.as_str().unwrap_or(&value.to_string()).to_string(); - cmd.arg("--params") - .arg(format!("{}={}", key_str, value_str)); - } - cmd - } else { - let text = task - .get_text_instruction() - .ok_or_else(|| task_error("text_instruction"))?; - let mut cmd = Command::new("goose"); - cmd.arg("run").arg("--text").arg(text); - cmd - }; - - command.stdout(Stdio::piped()); - command.stderr(Stdio::piped()); - Ok((command, output_identifier)) -} - -async fn run_command( - mut command: Command, - output_identifier: &str, - task_id: &str, - task_execution_tracker: Arc, -) -> Result<(String, String, bool), String> { - let mut child = command - .spawn() - .map_err(|e| format!("Failed to spawn goose: {}", e))?; - - let stdout = child.stdout.take().expect("Failed to capture stdout"); - let stderr = child.stderr.take().expect("Failed to capture stderr"); - - let stdout_task = spawn_output_reader( - stdout, - output_identifier, - false, - task_id, - task_execution_tracker.clone(), - ); - let stderr_task = spawn_output_reader( - stderr, - output_identifier, - true, - task_id, - task_execution_tracker.clone(), - ); - - let status = child - .wait() - .await - .map_err(|e| format!("Failed to wait for process: {}", e))?; - - let stdout_output = stdout_task.await.unwrap(); - let stderr_output = stderr_task.await.unwrap(); - - Ok((stdout_output, stderr_output, status.success())) -} - -fn spawn_output_reader( - reader: impl tokio::io::AsyncRead + Unpin + Send + 'static, - output_identifier: &str, - is_stderr: bool, - task_id: &str, - task_execution_tracker: Arc, -) -> tokio::task::JoinHandle { - let output_identifier = output_identifier.to_string(); - let task_id = task_id.to_string(); - tokio::spawn(async move { - let mut buffer = String::new(); - let mut lines = BufReader::new(reader).lines(); - while let Ok(Some(line)) = lines.next_line().await { - buffer.push_str(&line); - buffer.push('\n'); - - if !is_stderr { - task_execution_tracker - .send_live_output(&task_id, &line) - .await; - } else { - tracing::warn!("Task stderr [{}]: {}", output_identifier, line); - } - } - buffer - }) -} - -fn extract_json_from_line(line: &str) -> Option { - let start = line.find('{')?; - let end = line.rfind('}')?; - - if start >= end { - return None; - } - - let potential_json = &line[start..=end]; - if serde_json::from_str::(potential_json).is_ok() { - Some(potential_json.to_string()) - } else { - None - } -} - -fn process_output(stdout_output: String) -> Result { - let last_line = stdout_output - .lines() - .filter(|line| !line.trim().is_empty()) - .next_back() - .unwrap_or(""); - - if let Some(json_string) = extract_json_from_line(last_line) { - Ok(Value::String(json_string)) - } else { - Ok(Value::String(stdout_output)) - } -} - diff --git a/crates/goose/src/agents/sub_recipe_execution_tool/workers.rs b/crates/goose/src/agents/sub_recipe_execution_tool/workers.rs deleted file mode 100644 index 89473f7c6a65..000000000000 --- a/crates/goose/src/agents/sub_recipe_execution_tool/workers.rs +++ /dev/null @@ -1,31 +0,0 @@ -use crate::agents::sub_recipe_execution_tool::task_types::{SharedState, Task}; -use crate::agents::sub_recipe_execution_tool::tasks::process_task; -use std::sync::Arc; - -async fn receive_task(state: &SharedState) -> Option { - let mut receiver = state.task_receiver.lock().await; - receiver.recv().await -} - -pub fn spawn_worker(state: Arc, worker_id: usize) -> tokio::task::JoinHandle<()> { - state.increment_active_workers(); - - tokio::spawn(async move { - worker_loop(state, worker_id).await; - }) -} - -async fn worker_loop(state: Arc, _worker_id: usize) { - while let Some(task) = receive_task(&state).await { - state.task_execution_tracker.start_task(&task.id).await; - let result = process_task(&task, state.task_execution_tracker.clone()).await; - - if let Err(e) = state.result_sender.send(result).await { - tracing::error!("Worker failed to send result: {}", e); - break; - } - } - - state.decrement_active_workers(); -} - diff --git a/crates/goose/src/agents/subagent_execution_tool/executor/mod.rs b/crates/goose/src/agents/subagent_execution_tool/executor/mod.rs index 14726a75470c..665bf9e14b64 100644 --- a/crates/goose/src/agents/subagent_execution_tool/executor/mod.rs +++ b/crates/goose/src/agents/subagent_execution_tool/executor/mod.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; use tokio::time::Instant; +use tokio_util::sync::CancellationToken; const EXECUTION_STATUS_COMPLETED: &str = "completed"; const DEFAULT_MAX_WORKERS: usize = 10; @@ -21,14 +22,22 @@ pub async fn execute_single_task( task: &Task, notifier: mpsc::Sender, task_config: TaskConfig, + cancellation_token: Option, ) -> ExecutionResponse { let start_time = Instant::now(); let task_execution_tracker = Arc::new(TaskExecutionTracker::new( vec![task.clone()], DisplayMode::SingleTaskOutput, notifier, + cancellation_token.clone(), )); - let result = process_task(task, task_execution_tracker.clone(), task_config).await; + let result = process_task( + task, + task_execution_tracker.clone(), + task_config, + cancellation_token.unwrap_or_default(), + ) + .await; // Complete the task in the tracker task_execution_tracker @@ -49,11 +58,13 @@ pub async fn execute_tasks_in_parallel( tasks: Vec, notifier: Sender, task_config: TaskConfig, + cancellation_token: Option, ) -> ExecutionResponse { let task_execution_tracker = Arc::new(TaskExecutionTracker::new( tasks.clone(), DisplayMode::MultipleTasksOutput, notifier, + cancellation_token.clone(), )); let start_time = Instant::now(); let task_count = tasks.len(); @@ -71,7 +82,12 @@ pub async fn execute_tasks_in_parallel( return create_error_response(e); } - let shared_state = create_shared_state(task_rx, result_tx, task_execution_tracker.clone()); + let shared_state = create_shared_state( + task_rx, + result_tx, + task_execution_tracker.clone(), + cancellation_token.unwrap_or_default(), + ); let worker_count = std::cmp::min(task_count, DEFAULT_MAX_WORKERS); let mut worker_handles = Vec::new(); @@ -135,12 +151,14 @@ fn create_shared_state( task_rx: mpsc::Receiver, result_tx: mpsc::Sender, task_execution_tracker: Arc, + cancellation_token: CancellationToken, ) -> Arc { Arc::new(SharedState { task_receiver: Arc::new(tokio::sync::Mutex::new(task_rx)), result_sender: result_tx, active_workers: Arc::new(AtomicUsize::new(0)), task_execution_tracker, + cancellation_token, }) } diff --git a/crates/goose/src/agents/subagent_execution_tool/lib/mod.rs b/crates/goose/src/agents/subagent_execution_tool/lib/mod.rs index d6f431ede629..172ad03c6146 100644 --- a/crates/goose/src/agents/subagent_execution_tool/lib/mod.rs +++ b/crates/goose/src/agents/subagent_execution_tool/lib/mod.rs @@ -9,6 +9,7 @@ use crate::agents::subagent_task_config::TaskConfig; use rmcp::model::JsonRpcMessage; use serde_json::{json, Value}; use tokio::sync::mpsc::Sender; +use tokio_util::sync::CancellationToken; pub async fn execute_tasks( input: Value, @@ -16,6 +17,7 @@ pub async fn execute_tasks( notifier: Sender, task_config: TaskConfig, tasks_manager: &TasksManager, + cancellation_token: Option, ) -> Result { let task_ids: Vec = serde_json::from_value( input @@ -31,7 +33,8 @@ pub async fn execute_tasks( match execution_mode { ExecutionMode::Sequential => { if task_count == 1 { - let response = execute_single_task(&tasks[0], notifier, task_config).await; + let response = + execute_single_task(&tasks[0], notifier, task_config, cancellation_token).await; handle_response(response) } else { Err("Sequential execution mode requires exactly one task".to_string()) @@ -47,8 +50,13 @@ pub async fn execute_tasks( } )) } else { - let response: ExecutionResponse = - execute_tasks_in_parallel(tasks, notifier.clone(), task_config).await; + let response: ExecutionResponse = execute_tasks_in_parallel( + tasks, + notifier.clone(), + task_config, + cancellation_token, + ) + .await; handle_response(response) } } diff --git a/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs b/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs index fc400dad326b..e06da4061566 100644 --- a/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs +++ b/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs @@ -11,6 +11,7 @@ use crate::agents::{ use rmcp::model::JsonRpcMessage; use tokio::sync::mpsc; use tokio_stream; +use tokio_util::sync::CancellationToken; pub const SUBAGENT_EXECUTE_TASK_TOOL_NAME: &str = "subagent__execute_task"; pub fn create_subagent_execute_task_tool() -> Tool { @@ -64,6 +65,7 @@ pub async fn run_tasks( execute_data: Value, task_config: TaskConfig, tasks_manager: &TasksManager, + cancellation_token: Option, ) -> ToolCallResult { let (notification_tx, notification_rx) = mpsc::channel::(100); @@ -81,6 +83,7 @@ pub async fn run_tasks( notification_tx, task_config, &tasks_manager_clone, + cancellation_token, ) .await { diff --git a/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs b/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs index dab102392799..7d6854b3e229 100644 --- a/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs +++ b/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs @@ -3,7 +3,8 @@ use rmcp::object; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{mpsc, RwLock}; -use tokio::time::{sleep, Duration, Instant}; +use tokio::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; use crate::agents::subagent_execution_tool::notification_events::{ FailedTaskInfo, TaskCompletionStats, TaskExecutionNotificationEvent, TaskExecutionStats, @@ -21,7 +22,6 @@ pub enum DisplayMode { } const THROTTLE_INTERVAL_MS: u64 = 250; -const COMPLETION_NOTIFICATION_DELAY_MS: u64 = 500; fn format_task_metadata(task_info: &TaskInfo) -> String { if let Some(params) = task_info.task.get_command_parameters() { @@ -62,6 +62,7 @@ pub struct TaskExecutionTracker { last_refresh: Arc>, notifier: mpsc::Sender, display_mode: DisplayMode, + cancellation_token: Option, } impl TaskExecutionTracker { @@ -69,6 +70,7 @@ impl TaskExecutionTracker { tasks: Vec, display_mode: DisplayMode, notifier: Sender, + cancellation_token: Option, ) -> Self { let task_map = tasks .into_iter() @@ -93,6 +95,41 @@ impl TaskExecutionTracker { last_refresh: Arc::new(RwLock::new(Instant::now())), notifier, display_mode, + cancellation_token, + } + } + + fn is_cancelled(&self) -> bool { + self.cancellation_token + .as_ref() + .is_some_and(|t| t.is_cancelled()) + } + + fn log_notification_error( + &self, + error: &mpsc::error::TrySendError, + context: &str, + ) { + if !self.is_cancelled() { + tracing::warn!("Failed to send {} notification: {}", context, error); + } + } + + fn try_send_notification(&self, event: TaskExecutionNotificationEvent, context: &str) { + if let Err(e) = self + .notifier + .try_send(JsonRpcMessage::Notification(JsonRpcNotification { + jsonrpc: JsonRpcVersion2_0, + notification: Notification { + method: "notifications/message".to_string(), + params: object!({ + "data": event.to_notification_data() + }), + extensions: Default::default(), + }, + })) + { + self.log_notification_error(&e, context); } } @@ -153,21 +190,7 @@ impl TaskExecutionTracker { formatted_line, ); - if let Err(e) = - self.notifier - .try_send(JsonRpcMessage::Notification(JsonRpcNotification { - jsonrpc: JsonRpcVersion2_0, - notification: Notification { - method: "notifications/message".to_string(), - params: object!({ - "data": event.to_notification_data() - }), - extensions: Default::default(), - }, - })) - { - tracing::warn!("Failed to send live output notification: {}", e); - } + self.try_send_notification(event, "live output"); } DisplayMode::MultipleTasksOutput => { let mut tasks = self.tasks.write().await; @@ -197,6 +220,10 @@ impl TaskExecutionTracker { } async fn send_tasks_update(&self) { + if self.is_cancelled() { + return; + } + let tasks = self.tasks.read().await; let task_list: Vec<_> = tasks.values().collect(); let (total, pending, running, completed, failed) = count_by_status(&tasks); @@ -229,21 +256,7 @@ impl TaskExecutionTracker { let event = TaskExecutionNotificationEvent::tasks_update(stats, event_tasks); - if let Err(e) = self - .notifier - .try_send(JsonRpcMessage::Notification(JsonRpcNotification { - jsonrpc: JsonRpcVersion2_0, - notification: Notification { - method: "notifications/message".to_string(), - params: object!({ - "data": event.to_notification_data() - }), - extensions: Default::default(), - }, - })) - { - tracing::warn!("Failed to send tasks update notification: {}", e); - } + self.try_send_notification(event, "tasks update"); } pub async fn refresh_display(&self) { @@ -276,6 +289,10 @@ impl TaskExecutionTracker { } pub async fn send_tasks_complete(&self) { + if self.is_cancelled() { + return; + } + let tasks = self.tasks.read().await; let (total, _, _, completed, failed) = count_by_status(&tasks); @@ -293,23 +310,6 @@ impl TaskExecutionTracker { let event = TaskExecutionNotificationEvent::tasks_complete(stats, failed_tasks); - if let Err(e) = self - .notifier - .try_send(JsonRpcMessage::Notification(JsonRpcNotification { - jsonrpc: JsonRpcVersion2_0, - notification: Notification { - method: "notifications/message".to_string(), - params: object!({ - "data": event.to_notification_data() - }), - extensions: Default::default(), - }, - })) - { - tracing::warn!("Failed to send tasks complete notification: {}", e); - } - - // Brief delay to ensure completion notification is processed - sleep(Duration::from_millis(COMPLETION_NOTIFICATION_DELAY_MS)).await; + self.try_send_notification(event, "tasks complete"); } } diff --git a/crates/goose/src/agents/subagent_execution_tool/task_types.rs b/crates/goose/src/agents/subagent_execution_tool/task_types.rs index 796491f624f2..6bdcce33a7f9 100644 --- a/crates/goose/src/agents/subagent_execution_tool/task_types.rs +++ b/crates/goose/src/agents/subagent_execution_tool/task_types.rs @@ -3,6 +3,7 @@ use serde_json::{Map, Value}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; use crate::agents::subagent_execution_tool::task_execution_tracker::TaskExecutionTracker; @@ -117,6 +118,7 @@ pub struct SharedState { pub result_sender: mpsc::Sender, pub active_workers: Arc, pub task_execution_tracker: Arc, + pub cancellation_token: CancellationToken, } impl SharedState { diff --git a/crates/goose/src/agents/subagent_execution_tool/tasks.rs b/crates/goose/src/agents/subagent_execution_tool/tasks.rs index a330711e0a0a..4ecd5b628ffa 100644 --- a/crates/goose/src/agents/subagent_execution_tool/tasks.rs +++ b/crates/goose/src/agents/subagent_execution_tool/tasks.rs @@ -4,6 +4,7 @@ use std::process::Stdio; use std::sync::Arc; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::Command; +use tokio_util::sync::CancellationToken; use crate::agents::subagent_execution_tool::task_execution_tracker::TaskExecutionTracker; use crate::agents::subagent_execution_tool::task_types::{Task, TaskResult, TaskStatus}; @@ -14,8 +15,16 @@ pub async fn process_task( task: &Task, task_execution_tracker: Arc, task_config: TaskConfig, + cancellation_token: CancellationToken, ) -> TaskResult { - match get_task_result(task.clone(), task_execution_tracker, task_config).await { + match get_task_result( + task.clone(), + task_execution_tracker, + task_config, + cancellation_token, + ) + .await + { Ok(data) => TaskResult { task_id: task.id.clone(), status: TaskStatus::Completed, @@ -35,10 +44,17 @@ async fn get_task_result( task: Task, task_execution_tracker: Arc, task_config: TaskConfig, + cancellation_token: CancellationToken, ) -> Result { if task.task_type == "text_instruction" { // Handle text_instruction tasks using subagent system - handle_text_instruction_task(task, task_execution_tracker, task_config).await + handle_text_instruction_task( + task, + task_execution_tracker, + task_config, + cancellation_token, + ) + .await } else { // Handle sub_recipe tasks using command execution let (command, output_identifier) = build_command(&task)?; @@ -47,6 +63,7 @@ async fn get_task_result( &output_identifier, &task.id, task_execution_tracker, + cancellation_token, ) .await?; @@ -62,6 +79,7 @@ async fn handle_text_instruction_task( task: Task, task_execution_tracker: Arc, task_config: TaskConfig, + cancellation_token: CancellationToken, ) -> Result { let text_instruction = task .get_text_instruction() @@ -76,7 +94,14 @@ async fn handle_text_instruction_task( // "instructions": "You are a helpful assistant. Execute the given task and provide a clear, concise response.", }); - match run_complete_subagent_task(task_arguments, task_config).await { + let result = tokio::select! { + result = run_complete_subagent_task(task_arguments, task_config) => result, + _ = cancellation_token.cancelled() => { + return Err("Task cancelled".to_string()); + } + }; + + match result { Ok(contents) => { // Extract the text content from the result let result_text = contents @@ -141,6 +166,7 @@ async fn run_command( output_identifier: &str, task_id: &str, task_execution_tracker: Arc, + cancellation_token: CancellationToken, ) -> Result<(String, String, bool), String> { let mut child = command .spawn() @@ -164,15 +190,25 @@ async fn run_command( task_execution_tracker.clone(), ); - let status = child - .wait() - .await - .map_err(|e| format!("Failed to wait for process: {}", e))?; + let result = tokio::select! { + _ = cancellation_token.cancelled() => { + if let Err(e) = child.kill().await { + tracing::warn!("Failed to kill child process: {}", e); + } + // Abort the output reading tasks + stdout_task.abort(); + stderr_task.abort(); + return Err("Command cancelled".to_string()); + } + status_result = child.wait() => { + status_result.map_err(|e| format!("Failed to wait for process: {}", e))? + } + }; let stdout_output = stdout_task.await.unwrap(); let stderr_output = stderr_task.await.unwrap(); - Ok((stdout_output, stderr_output, status.success())) + Ok((stdout_output, stderr_output, result.success())) } fn spawn_output_reader( diff --git a/crates/goose/src/agents/subagent_execution_tool/workers.rs b/crates/goose/src/agents/subagent_execution_tool/workers.rs index 4ae0ab250737..d28808cf5d59 100644 --- a/crates/goose/src/agents/subagent_execution_tool/workers.rs +++ b/crates/goose/src/agents/subagent_execution_tool/workers.rs @@ -21,18 +21,35 @@ pub fn spawn_worker( } async fn worker_loop(state: Arc, _worker_id: usize, task_config: TaskConfig) { - while let Some(task) = receive_task(&state).await { - state.task_execution_tracker.start_task(&task.id).await; - let result = process_task( - &task, - state.task_execution_tracker.clone(), - task_config.clone(), - ) - .await; + loop { + tokio::select! { + task_option = receive_task(&state) => { + match task_option { + Some(task) => { + state.task_execution_tracker.start_task(&task.id).await; + let result = process_task( + &task, + state.task_execution_tracker.clone(), + task_config.clone(), + state.cancellation_token.clone(), + ) + .await; - if let Err(e) = state.result_sender.send(result).await { - tracing::error!("Worker failed to send result: {}", e); - break; + if let Err(e) = state.result_sender.send(result).await { + // Only log error if not cancelled (channel close is expected during cancellation) + if !state.cancellation_token.is_cancelled() { + tracing::error!("Worker failed to send result: {}", e); + } + break; + } + } + None => break, // No more tasks + } + } + _ = state.cancellation_token.cancelled() => { + tracing::debug!("Worker cancelled"); + break; + } } } diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index 9af001fe7666..bc9f4292f72f 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -6,6 +6,7 @@ use futures::stream::{self, BoxStream}; use futures::{Stream, StreamExt}; use rmcp::model::JsonRpcMessage; use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; use crate::config::permission::PermissionLevel; use crate::config::PermissionManager; @@ -53,6 +54,7 @@ impl Agent { tool_futures: Arc>>, permission_manager: &'a mut PermissionManager, message_tool_response: Arc>, + cancellation_token: Option, ) -> BoxStream<'a, anyhow::Result> { try_stream! { for request in tool_requests { @@ -69,7 +71,7 @@ impl Agent { while let Some((req_id, confirmation)) = rx.recv().await { if req_id == request.id { if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow { - let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone()).await; + let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone()).await; let mut futures = tool_futures.lock().await; futures.push((req_id, match tool_result { diff --git a/crates/goose/src/context_mgmt/auto_compact.rs b/crates/goose/src/context_mgmt/auto_compact.rs new file mode 100644 index 000000000000..6c45aa9ac1c4 --- /dev/null +++ b/crates/goose/src/context_mgmt/auto_compact.rs @@ -0,0 +1,336 @@ +use crate::{ + agents::Agent, + config::Config, + context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async}, + message::Message, + token_counter::create_async_token_counter, +}; +use anyhow::Result; +use tracing::{debug, info}; + +/// Result of auto-compaction check +#[derive(Debug)] +pub struct AutoCompactResult { + /// Whether compaction was performed + pub compacted: bool, + /// The messages after potential compaction + pub messages: Vec, + /// Token count before compaction (if compaction occurred) + pub tokens_before: Option, + /// Token count after compaction (if compaction occurred) + pub tokens_after: Option, +} + +/// Check if messages need compaction and compact them if necessary +/// +/// This function checks the current token usage against a configurable threshold +/// and automatically compacts the messages using the summarization algorithm if needed. +/// +/// # Arguments +/// * `agent` - The agent to use for context management +/// * `messages` - The current message history +/// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config) +/// +/// # Returns +/// * `AutoCompactResult` containing the potentially compacted messages and metadata +pub async fn check_and_compact_messages( + agent: &Agent, + messages: &[Message], + threshold_override: Option, +) -> Result { + // Get threshold from config or use override + let config = Config::global(); + let threshold = threshold_override.unwrap_or_else(|| { + config + .get_param::("GOOSE_AUTO_COMPACT_THRESHOLD") + .unwrap_or(0.3) // Default to 30% + }); + + // Check if auto-compaction is disabled + if threshold <= 0.0 || threshold >= 1.0 { + debug!("Auto-compaction disabled (threshold: {})", threshold); + return Ok(AutoCompactResult { + compacted: false, + messages: messages.to_vec(), + tokens_before: None, + tokens_after: None, + }); + } + + // Get provider and token counter + let provider = agent.provider().await?; + let token_counter = create_async_token_counter() + .await + .map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?; + + // Calculate current token usage + let token_counts = get_messages_token_counts_async(&token_counter, messages); + let total_tokens: usize = token_counts.iter().sum(); + let context_limit = estimate_target_context_limit(provider); + + // Calculate usage ratio + let usage_ratio = total_tokens as f64 / context_limit as f64; + + debug!( + "Context usage: {} / {} ({:.1}%)", + total_tokens, + context_limit, + usage_ratio * 100.0 + ); + + // Check if compaction is needed + if usage_ratio <= threshold { + debug!( + "No compaction needed (usage: {:.1}% <= threshold: {:.1}%)", + usage_ratio * 100.0, + threshold * 100.0 + ); + return Ok(AutoCompactResult { + compacted: false, + messages: messages.to_vec(), + tokens_before: None, + tokens_after: None, + }); + } + + info!( + "Auto-compacting messages (usage: {:.1}% > threshold: {:.1}%)", + usage_ratio * 100.0, + threshold * 100.0 + ); + + // Perform compaction + let (compacted_messages, compacted_token_counts) = agent.summarize_context(messages).await?; + let tokens_after: usize = compacted_token_counts.iter().sum(); + + info!( + "Compaction complete: {} tokens -> {} tokens ({:.1}% reduction)", + total_tokens, + tokens_after, + (1.0 - (tokens_after as f64 / total_tokens as f64)) * 100.0 + ); + + Ok(AutoCompactResult { + compacted: true, + messages: compacted_messages, + tokens_before: Some(total_tokens), + tokens_after: Some(tokens_after), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + agents::Agent, + message::{Message, MessageContent}, + model::ModelConfig, + providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage}, + providers::errors::ProviderError, + }; + use chrono::Utc; + use mcp_core::tool::Tool; + use rmcp::model::{AnnotateAble, RawTextContent, Role}; + use std::sync::Arc; + + #[derive(Clone)] + struct MockProvider { + model_config: ModelConfig, + } + + #[async_trait::async_trait] + impl Provider for MockProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::empty() + } + + fn get_model_config(&self) -> ModelConfig { + self.model_config.clone() + } + + async fn complete( + &self, + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + // Return a short summary message + Ok(( + Message::new( + Role::Assistant, + Utc::now().timestamp(), + vec![MessageContent::Text( + RawTextContent { + text: "Summary of conversation".to_string(), + } + .no_annotation(), + )], + ), + ProviderUsage::new("mock".to_string(), Usage::default()), + )) + } + } + + fn create_test_message(text: &str) -> Message { + Message::new( + Role::User, + Utc::now().timestamp(), + vec![MessageContent::text(text.to_string())], + ) + } + + #[tokio::test] + async fn test_auto_compact_disabled() { + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model".to_string()) + .with_context_limit(10_000.into()), + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + let messages = vec![create_test_message("Hello"), create_test_message("World")]; + + // Test with threshold 0 (disabled) + let result = check_and_compact_messages(&agent, &messages, Some(0.0)) + .await + .unwrap(); + + assert!(!result.compacted); + assert_eq!(result.messages.len(), messages.len()); + assert!(result.tokens_before.is_none()); + assert!(result.tokens_after.is_none()); + + // Test with threshold 1.0 (disabled) + let result = check_and_compact_messages(&agent, &messages, Some(1.0)) + .await + .unwrap(); + + assert!(!result.compacted); + } + + #[tokio::test] + async fn test_auto_compact_below_threshold() { + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model".to_string()) + .with_context_limit(100_000.into()), // Increased to ensure overhead doesn't dominate + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + // Create small messages that won't trigger compaction + let messages = vec![create_test_message("Hello"), create_test_message("World")]; + + let result = check_and_compact_messages(&agent, &messages, Some(0.3)) + .await + .unwrap(); + + assert!(!result.compacted); + assert_eq!(result.messages.len(), messages.len()); + } + + #[tokio::test] + async fn test_auto_compact_above_threshold() { + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model".to_string()) + .with_context_limit(50_000.into()), // Realistic context limit that won't underflow + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + // Create messages that will exceed 30% of the context limit + // With 50k context limit, after overhead we have ~27k usable tokens + // 30% of that is ~8.1k tokens, so we need messages that exceed that + let mut messages = Vec::new(); + + // Create longer messages with more content to reach the threshold + for i in 0..200 { + messages.push(create_test_message(&format!( + "This is message number {} with significantly more content to increase token count. \ + We need to ensure that our total token usage exceeds 30% of the available context \ + limit after accounting for system prompt and tools overhead. This message contains \ + multiple sentences to increase the token count substantially.", + i + ))); + } + + let result = check_and_compact_messages(&agent, &messages, Some(0.3)) + .await + .unwrap(); + + assert!(result.compacted); + assert!(result.tokens_before.is_some()); + assert!(result.tokens_after.is_some()); + + // Should have fewer tokens after compaction + if let (Some(before), Some(after)) = (result.tokens_before, result.tokens_after) { + assert!( + after < before, + "Token count should decrease after compaction" + ); + } + + // Should have fewer messages (summarized) + assert!(result.messages.len() <= messages.len()); + } + + #[tokio::test] + async fn test_auto_compact_respects_config() { + let mock_provider = Arc::new(MockProvider { + model_config: ModelConfig::new("test-model".to_string()) + .with_context_limit(50_000.into()), // Realistic context limit that won't underflow + }); + + let agent = Agent::new(); + let _ = agent.update_provider(mock_provider).await; + + // Create enough messages to trigger compaction with low threshold + let mut messages = Vec::new(); + // Need to create more messages since we have a 27k usable token limit + // 10% of 27k = 2.7k tokens + for i in 0..150 { + messages.push(create_test_message(&format!( + "Message {} with enough content to ensure we exceed 10% of the context limit. Adding more content.", + i + ))); + } + + // Set config value + let config = Config::global(); + config + .set_param("GOOSE_AUTO_COMPACT_THRESHOLD", serde_json::Value::from(0.1)) + .unwrap(); + + // Should use config value when no override provided + let result = check_and_compact_messages(&agent, &messages, None) + .await + .unwrap(); + + // Debug info if not compacted + if !result.compacted { + let provider = agent.provider().await.unwrap(); + let token_counter = create_async_token_counter().await.unwrap(); + let token_counts = get_messages_token_counts_async(&token_counter, &messages); + let total_tokens: usize = token_counts.iter().sum(); + let context_limit = estimate_target_context_limit(provider); + let usage_ratio = total_tokens as f64 / context_limit as f64; + + eprintln!( + "Config test not compacted - tokens: {} / {} ({:.1}%)", + total_tokens, + context_limit, + usage_ratio * 100.0 + ); + } + + // With such a low threshold (10%), it should compact + assert!(result.compacted); + + // Clean up config + config + .set_param("GOOSE_AUTO_COMPACT_THRESHOLD", serde_json::Value::from(0.3)) + .unwrap(); + } +} diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 838e27fece54..00d11d6b871b 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -1,3 +1,4 @@ +pub mod auto_compact; mod common; pub mod summarize; pub mod truncate; diff --git a/crates/goose/src/context_mgmt/summarize.rs b/crates/goose/src/context_mgmt/summarize.rs index 2f24ca94542c..27f74fbbbd6f 100644 --- a/crates/goose/src/context_mgmt/summarize.rs +++ b/crates/goose/src/context_mgmt/summarize.rs @@ -116,19 +116,13 @@ pub async fn summarize_messages_oneshot( token_counter: &TokenCounter, _context_limit: usize, ) -> Result<(Vec, Vec), anyhow::Error> { - // Preprocess messages to handle tool response edge case. - let (preprocessed_messages, removed_messages) = preprocess_messages(messages); - - if preprocessed_messages.is_empty() { - // If no messages to summarize, just return the removed messages - return Ok(( - removed_messages.clone(), - get_messages_token_counts(token_counter, &removed_messages), - )); + if messages.is_empty() { + // If no messages to summarize, return empty + return Ok((vec![], vec![])); } // Format all messages as a single string for the summarization prompt - let messages_text = preprocessed_messages + let messages_text = messages .iter() .map(|msg| format!("{:?}", msg)) .collect::>() @@ -155,8 +149,8 @@ pub async fn summarize_messages_oneshot( // Set role to user as it will be used in following conversation as user content. response.role = Role::User; - // Add back removed messages. - let final_summary = reintegrate_removed_messages(&[response], &removed_messages); + // Return just the summary without any tool response preservation + let final_summary = vec![response]; Ok(( final_summary.clone(), @@ -180,17 +174,14 @@ pub async fn summarize_messages_chunked( let summary_prompt_tokens = token_counter.count_tokens(SUMMARY_PROMPT); let mut accumulated_summary = Vec::new(); - // Preprocess messages to handle tool response edge case. - let (preprocessed_messages, removed_messages) = preprocess_messages(messages); - // Get token counts for each message. - let token_counts = get_messages_token_counts(token_counter, &preprocessed_messages); + let token_counts = get_messages_token_counts(token_counter, messages); // Tokenize and break messages into chunks. let mut current_chunk: Vec = Vec::new(); let mut current_chunk_tokens = 0; - for (message, message_tokens) in preprocessed_messages.iter().zip(token_counts.iter()) { + for (message, message_tokens) in messages.iter().zip(token_counts.iter()) { if current_chunk_tokens + message_tokens > chunk_size - summary_prompt_tokens { // Summarize the current chunk with the accumulated summary. accumulated_summary = @@ -213,12 +204,10 @@ pub async fn summarize_messages_chunked( summarize_combined_messages(&provider, &accumulated_summary, ¤t_chunk).await?; } - // Add back removed messages. - let final_summary = reintegrate_removed_messages(&accumulated_summary, &removed_messages); - + // Return just the summary without any tool response preservation Ok(( - final_summary.clone(), - get_messages_token_counts(token_counter, &final_summary), + accumulated_summary.clone(), + get_messages_token_counts(token_counter, &accumulated_summary), )) } @@ -281,17 +270,14 @@ pub async fn summarize_messages_async( let summary_prompt_tokens = token_counter.count_tokens(SUMMARY_PROMPT); let mut accumulated_summary = Vec::new(); - // Preprocess messages to handle tool response edge case. - let (preprocessed_messages, removed_messages) = preprocess_messages(messages); - // Get token counts for each message. - let token_counts = get_messages_token_counts_async(token_counter, &preprocessed_messages); + let token_counts = get_messages_token_counts_async(token_counter, messages); // Tokenize and break messages into chunks. let mut current_chunk: Vec = Vec::new(); let mut current_chunk_tokens = 0; - for (message, message_tokens) in preprocessed_messages.iter().zip(token_counts.iter()) { + for (message, message_tokens) in messages.iter().zip(token_counts.iter()) { if current_chunk_tokens + message_tokens > chunk_size - summary_prompt_tokens { // Summarize the current chunk with the accumulated summary. accumulated_summary = @@ -314,12 +300,10 @@ pub async fn summarize_messages_async( summarize_combined_messages(&provider, &accumulated_summary, ¤t_chunk).await?; } - // Add back removed messages. - let final_summary = reintegrate_removed_messages(&accumulated_summary, &removed_messages); - + // Return just the summary without any tool response preservation Ok(( - final_summary.clone(), - get_messages_token_counts_async(token_counter, &final_summary), + accumulated_summary.clone(), + get_messages_token_counts_async(token_counter, &accumulated_summary), )) } @@ -418,7 +402,7 @@ mod tests { async fn test_summarize_messages_single_chunk() { let provider = create_mock_provider(); let token_counter = TokenCounter::new(); - let context_limit = 100; // Set a high enough limit to avoid chunking. + let context_limit = 10_000; // Higher limit to avoid underflow let messages = create_test_messages(); let result = summarize_messages( @@ -454,7 +438,7 @@ mod tests { async fn test_summarize_messages_multiple_chunks() { let provider = create_mock_provider(); let token_counter = TokenCounter::new(); - let context_limit = 30; + let context_limit = 10_000; // Higher limit to avoid underflow let messages = create_test_messages(); let result = summarize_messages( @@ -490,7 +474,7 @@ mod tests { async fn test_summarize_messages_empty_input() { let provider = create_mock_provider(); let token_counter = TokenCounter::new(); - let context_limit = 100; + let context_limit = 10_000; // Higher limit to avoid underflow let messages: Vec = Vec::new(); let result = summarize_messages( @@ -616,7 +600,7 @@ mod tests { async fn test_summarize_messages_uses_chunked_for_large_context() { let provider = create_mock_provider(); let token_counter = TokenCounter::new(); - let context_limit = 100; // Small context limit but not too small to cause overflow + let context_limit = 10_000; // Higher limit to avoid underflow let messages = create_test_messages(); let result = summarize_messages( @@ -772,7 +756,7 @@ mod tests { async fn test_summarize_messages_chunked_direct_call() { let provider = create_mock_provider(); let token_counter = TokenCounter::new(); - let context_limit = 30; // Small to force chunking + let context_limit = 10_000; // Higher limit to avoid underflow let messages = create_test_messages(); let result = summarize_messages_chunked( diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index ab8b8cb155d8..497ebcaab715 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -603,7 +603,7 @@ mod final_output_tool_tests { }), ); let (_, result) = agent - .dispatch_tool_call(tool_call, "request_id".to_string()) + .dispatch_tool_call(tool_call, "request_id".to_string(), None) .await; assert!(result.is_ok(), "Tool call should succeed"); diff --git a/crates/goose/tests/private_tests.rs b/crates/goose/tests/private_tests.rs index d2ec7a06e8ae..e23d0c09e319 100644 --- a/crates/goose/tests/private_tests.rs +++ b/crates/goose/tests/private_tests.rs @@ -885,7 +885,7 @@ async fn test_schedule_tool_dispatch() { }; let (request_id, result) = agent - .dispatch_tool_call(tool_call, "test_dispatch".to_string()) + .dispatch_tool_call(tool_call, "test_dispatch".to_string(), None) .await; assert_eq!(request_id, "test_dispatch"); assert!(result.is_ok()); diff --git a/documentation/docs/getting-started/using-extensions.md b/documentation/docs/getting-started/using-extensions.md index 71fec2537df6..068efbf42340 100644 --- a/documentation/docs/getting-started/using-extensions.md +++ b/documentation/docs/getting-started/using-extensions.md @@ -5,7 +5,7 @@ title: Using Extensions import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -import { PanelLeft } from 'lucide-react'; +import { PanelLeft, Settings } from 'lucide-react'; Extensions are add-ons that provide a way to extend the functionality of Goose by connecting with applications and tools you already use in your workflow. These extensions can be used to add new features, access data and resources, or integrate with other systems. @@ -22,10 +22,10 @@ Out of the box, Goose is installed with a few extensions but with only the `Deve Here are the built-in extensions: -1. **Developer**: provides a set of general development tools that are useful for software development. -2. **Computer Controller**: provides general computer control tools for webscraping, file caching, and automations. -3. **Memory**: teaches goose to remember your preferences as you use it -4. **JetBrains**: provides an integration for working with JetBrains IDEs. +- [Developer](/docs/mcp/developer-mcp): Provides a set of general development tools that are useful for software development. +- [Computer Controller](/docs/mcp/computer-controller-mcp): Provides general computer control tools for webscraping, file caching, and automations. +- [Memory](/docs/mcp/memory-mcp): Teaches Goose to remember your preferences as you use it. +- [Tutorial](/docs/mcp/tutorial-mcp): Provides interactive tutorials for learning about Goose. #### Toggling Built-in Extensions @@ -500,6 +500,29 @@ extension_name: postgresql +## Updating Extension Properties + +Goose relies on extension properties to determine how to handle an extension. You can edit these properties if you want to change the extension's display settings and behavior, such as the name, timeout, or environment variables. + + + + + 1. Click the button in the top-left to open the sidebar. + 2. Click the `Extensions` button on the sidebar. + 3. Under `Extensions`, click the button on the extension you'd like to edit. + 4. In the dialog that appears, edit the extension's properties as needed. + 5. Click `Save Changes`. + + + + + + 1. Navigate to the Goose [configuration file](/docs/guides/config-file). For example, navigate to `~/.config/goose/config.yaml` on macOS. + 2. Edit the extension properties as needed and save your changes. + + + + ## Removing Extensions You can remove installed extensions. @@ -509,7 +532,7 @@ You can remove installed extensions. 1. Click the button in the top-left to open the sidebar. 2. Click the `Extensions` button on the sidebar. - 3. Under `Extensions`, find the extension you'd like to remove and click on the settings icon beside it. + 3. Under `Extensions`, click the button on the extension you'd like to remove. 4. In the dialog that appears, click `Remove Extension`. diff --git a/documentation/docs/guides/recipes/recipe-reference.md b/documentation/docs/guides/recipes/recipe-reference.md index 27ea188812bf..89366be2b355 100644 --- a/documentation/docs/guides/recipes/recipe-reference.md +++ b/documentation/docs/guides/recipes/recipe-reference.md @@ -18,6 +18,49 @@ Files should be named either: After creating recipe files, you can use [`goose` CLI commands](/docs/guides/goose-cli-commands) to run or validate the files and to manage recipe sharing. +### CLI and Desktop Formats + +The Goose CLI supports CLI and Desktop recipe formats: + +- **CLI Format**: Recipe fields (like `title`, `description`, `instructions`) are at the root level of the YAML/JSON file +- **Desktop Format**: Recipe fields are nested inside a `recipe` object, with additional metadata fields at the root level + +The CLI automatically detects and handles both formats when running `goose run --recipe ` and `goose recipe` commands. + +
+Format Examples + +**CLI Format:** +```yaml +version: "1.0.0" +title: "Code Review Assistant" +description: "Automated code review with best practices" +instructions: "You are a code reviewer..." +prompt: "Review the code in this repository" +extensions: [] +``` + +**Desktop Format:** +```yaml +name: "Code Review Assistant" +recipe: + version: "1.0.0" + title: "Code Review Assistant" + description: "Automated code review with best practices" + instructions: "You are a code reviewer..." + prompt: "Review the code in this repository" + extensions: [] +isGlobal: true +lastModified: 2025-07-02T03:46:46.778Z +isArchived: false +``` + +:::note +Goose automatically adds metadata fields to recipes saved from the Desktop app. +::: + +
+ ## Recipe Structure ### Required Fields @@ -40,6 +83,17 @@ After creating recipe files, you can use [`goose` CLI commands](/docs/guides/goo | `response` | Object | Configuration for structured output validation | | `retry` | Object | Configuration for automated retry logic with success validation | +### Desktop Format Metadata Fields + +When recipes are saved from Goose Desktop, additional metadata fields are included at the top level (outside the `recipe` key). These fields are used by the Desktop app for organization and management but are ignored by CLI operations. + +| Field | Type | Description | +|-------|------|-------------| +| `name` | String | Display name used in Desktop Recipe Library | +| `isGlobal` | Boolean | Whether the recipe is available globally or locally to a project | +| `lastModified` | String | ISO timestamp of when the recipe was last modified | +| `isArchived` | Boolean | Whether the recipe is archived in the Desktop interface | + ## Parameters Each parameter in the `parameters` array has the following structure: diff --git a/documentation/docs/guides/recipes/session-recipes.md b/documentation/docs/guides/recipes/session-recipes.md index ff82d46d2eb2..f218ffa1c283 100644 --- a/documentation/docs/guides/recipes/session-recipes.md +++ b/documentation/docs/guides/recipes/session-recipes.md @@ -588,5 +588,14 @@ To protect your privacy and system integrity, Goose excludes: This means others may need to supply their own credentials or memory context if the recipe depends on those elements. +## CLI and Desktop Formats + +The Goose CLI supports both CLI and Desktop recipe formats: + +- **CLI Format**: Recipe fields are at the root level. This format is used when recipes are created via the CLI `/recipe` command and Recipe Generator YAML option. +- **Desktop Format**: Recipe fields are nested under a `recipe` key. This format is used when recipes are saved in Goose Desktop. + +Both formats work seamlessly with `goose run --recipe ` and `goose recipe` CLI commands - you don't need to convert between them. For more details, see [CLI and Desktop Formats](/docs/guides/recipes/recipe-reference#cli-and-desktop-formats). + ## Learn More Check out the [Goose Recipes](/docs/guides/recipes) guide for more docs, tools, and resources to help you master Goose recipes. \ No newline at end of file diff --git a/documentation/docs/guides/recipes/storing-recipes.md b/documentation/docs/guides/recipes/storing-recipes.md index 99c1f862ea9c..b9075c697f40 100644 --- a/documentation/docs/guides/recipes/storing-recipes.md +++ b/documentation/docs/guides/recipes/storing-recipes.md @@ -124,5 +124,9 @@ Set up [custom recipe paths](/docs/guides/recipes/session-recipes#configure-reci Once you've located your recipe file, [run the recipe](/docs/guides/recipes/session-recipes#run-a-recipe). +:::tip Format Compatibility +The CLI can run recipes saved from Goose Desktop without any conversion. Both CLI-created and Desktop-saved recipes work with all recipe commands. +::: + diff --git a/documentation/docs/tutorials/mongodb-mcp.md b/documentation/docs/mcp/mongodb-mcp.md similarity index 90% rename from documentation/docs/tutorials/mongodb-mcp.md rename to documentation/docs/mcp/mongodb-mcp.md index 03c78cd1143c..2c9c430fa3c9 100644 --- a/documentation/docs/tutorials/mongodb-mcp.md +++ b/documentation/docs/mcp/mongodb-mcp.md @@ -5,6 +5,7 @@ description: Add MongoDB MCP Server as a Goose Extension import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; +import GooseDesktopInstaller from '@site/src/components/GooseDesktopInstaller'; The MongoDB MCP Server extension allows Goose to interact directly with your MongoDB databases, enabling comprehensive database operations including querying, document manipulation, collection management, and database administration. This makes it easy to work with your MongoDB databases through natural language interactions. @@ -55,11 +56,18 @@ Note that you'll need [Node.js](https://nodejs.org/) installed on your system to - 1. [Launch the installer](goose://extension?cmd=npx&arg=-y&arg=mongodb-mcp-server&arg=--connection-string&arg=mongodb://localhost:27017&id=mongodb&name=MongoDB&description=MongoDB%20database%20integration) - 2. Press `Yes` to confirm the installation - 3. Enter your MongoDB connection string in the format: `mongodb://username:password@hostname:27017/database` - 4. Click `Save Configuration` - 5. Scroll to the top and click `Exit` from the upper left corner + + + :::info Configure Your Connection String + If needed, [update the extension](/docs/getting-started/using-extensions#updating-extension-properties) to match to your [MongoDB environment](#customizing-your-connection). For example, change the connection string in the `command` property to use the `mongodb://username:password@hostname:27017/database` format. + ::: + 1. Run the `configure` command: @@ -100,7 +108,7 @@ Note that you'll need [Node.js](https://nodejs.org/) installed on your system to └ ``` - 4. Enter the command with your database connection string + 4. Enter the command with the database connection string that matches your [MongoDB environment](#customizing-your-connection) ```sh ┌ goose-configure │ diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index 44fd7e7216ad..e51e74744069 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -157,6 +157,7 @@ function BaseChatContent({ updateMessageStreamBody, sessionMetadata, isUserMessage, + clearError, } = useChatEngine({ chat, setChat, @@ -434,21 +435,51 @@ function BaseChatContent({ {error.message || 'Honk! Goose experienced an error while responding'} - {/* Regular retry button for non-token-limit errors */} -
{ - // Find the last user message - const lastUserMessage = messages.reduceRight( - (found, m) => found || (m.role === 'user' ? m : null), - null as Message | null - ); - if (lastUserMessage) { - append(lastUserMessage); - } - }} - > - Retry Last Message + {/* Action buttons for non-token-limit errors */} +
+
{ + // Create a contextLengthExceeded message similar to token limit errors + const contextMessage: Message = { + id: `context-${Date.now()}`, + role: 'assistant', + created: Math.floor(Date.now() / 1000), + content: [ + { + type: 'contextLengthExceeded', + msg: 'Summarization requested due to error. Creating summary to help resolve the issue.', + }, + ], + display: true, + sendToLLM: false, + }; + + // Add the context message to trigger ContextHandler + const updatedMessages = [...messages, contextMessage]; + setMessages(updatedMessages); + + // Clear the error state since we're handling it with summarization + clearError(); + }} + > + Summarize Conversation +
+
{ + // Find the last user message + const lastUserMessage = messages.reduceRight( + (found, m) => found || (m.role === 'user' ? m : null), + null as Message | null + ); + if (lastUserMessage) { + append(lastUserMessage); + } + }} + > + Retry Last Message +
@@ -472,7 +503,7 @@ function BaseChatContent({ {/* Fixed loading indicator at bottom left of chat container */} {isLoading && (
- { clearAlerts(); - // Only show token alerts if we have loaded the real token limit + // Always show token alerts if we have loaded the real token limit and have tokens if (isTokenLimitLoaded && tokenLimit && numTokens && numTokens > 0) { if (numTokens >= tokenLimit) { // Only show error alert when limit reached @@ -409,6 +409,16 @@ export default function ChatInput({ }, }); } + } else if (isTokenLimitLoaded && tokenLimit) { + // Always show context window info even when no tokens are present (start of conversation) + addAlert({ + type: AlertType.Info, + message: 'Context window', + progress: { + current: 0, + total: tokenLimit, + }, + }); } // Add tool count alert if we have the data diff --git a/ui/desktop/src/components/bottom_menu/BottomMenuAlertPopover.tsx b/ui/desktop/src/components/bottom_menu/BottomMenuAlertPopover.tsx index 3cf85f7b2cf2..bd34285f1bd1 100644 --- a/ui/desktop/src/components/bottom_menu/BottomMenuAlertPopover.tsx +++ b/ui/desktop/src/components/bottom_menu/BottomMenuAlertPopover.tsx @@ -12,7 +12,6 @@ interface AlertPopoverProps { export default function BottomMenuAlertPopover({ alerts }: AlertPopoverProps) { const [isOpen, setIsOpen] = useState(false); - const [hasShownInitial, setHasShownInitial] = useState(false); const [isHovered, setIsHovered] = useState(false); const [wasAutoShown, setWasAutoShown] = useState(false); const [popoverPosition, setPopoverPosition] = useState({ top: 0, left: 0 }); @@ -116,17 +115,14 @@ export default function BottomMenuAlertPopover({ alerts }: AlertPopoverProps) { // Only auto-show if any of the new/changed alerts have autoShow: true const hasNewAutoShowAlert = changedAlerts.some((alert) => alert.autoShow === true); - // Auto show the popover only if: - // 1. There are new alerts that should auto-show AND - // 2. We haven't shown this specific alert before (tracked by hasShownInitial) - if (hasNewAutoShowAlert && !hasShownInitial) { + // Auto show the popover for new auto-show alerts + if (hasNewAutoShowAlert) { setIsOpen(true); - setHasShownInitial(true); setWasAutoShown(true); // Start 3 second timer for auto-show startHideTimer(3000); } - }, [alerts, hasShownInitial, startHideTimer]); + }, [alerts, startHideTimer]); // Handle auto-hide based on hover state changes useEffect(() => { diff --git a/ui/desktop/src/components/bottom_menu/DirSwitcher.tsx b/ui/desktop/src/components/bottom_menu/DirSwitcher.tsx index 5dd5483eb0c8..c51bea91c9b7 100644 --- a/ui/desktop/src/components/bottom_menu/DirSwitcher.tsx +++ b/ui/desktop/src/components/bottom_menu/DirSwitcher.tsx @@ -35,7 +35,7 @@ export const DirSwitcher: React.FC = ({
- + {window.appConfig.get('GOOSE_WORKING_DIR') as string} diff --git a/ui/desktop/src/components/settings/models/ModelsSection.tsx b/ui/desktop/src/components/settings/models/ModelsSection.tsx index 9d665d5b27b6..9d87472004c5 100644 --- a/ui/desktop/src/components/settings/models/ModelsSection.tsx +++ b/ui/desktop/src/components/settings/models/ModelsSection.tsx @@ -1,4 +1,4 @@ -import { useEffect, useState, useCallback } from 'react'; +import { useEffect, useState, useCallback, useRef } from 'react'; import type { View } from '../../../App'; import ModelSettingsButtons from './subcomponents/ModelSettingsButtons'; import { useConfig } from '../../ConfigContext'; @@ -18,14 +18,16 @@ export default function ModelsSection({ setView }: ModelsSectionProps) { const [displayModelName, setDisplayModelName] = useState(''); const [isLoading, setIsLoading] = useState(true); const { read, getProviders } = useConfig(); - const { getCurrentModelDisplayName, getCurrentProviderDisplayName } = useModelAndProvider(); + const { + getCurrentModelDisplayName, + getCurrentProviderDisplayName, + currentModel, + currentProvider, + } = useModelAndProvider(); - // Function to load model data const loadModelData = useCallback(async () => { try { setIsLoading(true); - const gooseProvider = (await read('GOOSE_PROVIDER', false)) as string; - const providers = await getProviders(true); // Get display name (alias if available, otherwise model name) const modelDisplayName = await getCurrentModelDisplayName(); @@ -37,6 +39,8 @@ export default function ModelsSection({ setView }: ModelsSectionProps) { setProvider(providerDisplayName); } else { // Fallback to original provider lookup + const gooseProvider = (await read('GOOSE_PROVIDER', false)) as string; + const providers = await getProviders(true); const providerDetailsList = providers.filter((provider) => provider.name === gooseProvider); if (providerDetailsList.length != 1) { @@ -59,8 +63,23 @@ export default function ModelsSection({ setView }: ModelsSectionProps) { useEffect(() => { loadModelData(); - // eslint-disable-next-line react-hooks/exhaustive-deps - }, []); + }, [loadModelData]); + + // Update display when model or provider changes - but only if they actually changed + const prevModelRef = useRef(null); + const prevProviderRef = useRef(null); + + useEffect(() => { + if ( + currentModel && + currentProvider && + (currentModel !== prevModelRef.current || currentProvider !== prevProviderRef.current) + ) { + prevModelRef.current = currentModel; + prevProviderRef.current = currentProvider; + loadModelData(); + } + }, [currentModel, currentProvider, loadModelData]); return (
diff --git a/ui/desktop/src/components/settings/models/model_list/BaseModelsList.tsx b/ui/desktop/src/components/settings/models/model_list/BaseModelsList.tsx index 0e25b3e20c3f..cfa0f1f6dcaa 100644 --- a/ui/desktop/src/components/settings/models/model_list/BaseModelsList.tsx +++ b/ui/desktop/src/components/settings/models/model_list/BaseModelsList.tsx @@ -1,4 +1,4 @@ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useState, useRef } from 'react'; import Model from '../modelInterface'; import { useRecentModels } from './recentModels'; import { useModelAndProvider } from '../../../ModelAndProviderContext'; @@ -29,7 +29,8 @@ export function BaseModelsList({ } else { modelList = providedModelList; } - const { changeModel, getCurrentModelAndProvider } = useModelAndProvider(); + const { changeModel, getCurrentModelAndProvider, currentModel, currentProvider } = + useModelAndProvider(); const [selectedModel, setSelectedModel] = useState(null); const [isInitialized, setIsInitialized] = useState(false); @@ -119,6 +120,33 @@ export function BaseModelsList({ } }; + // Update selected model when context changes - but only if they actually changed + const prevModelRef = useRef(null); + const prevProviderRef = useRef(null); + + useEffect(() => { + if ( + currentModel && + currentProvider && + isInitialized && + (currentModel !== prevModelRef.current || currentProvider !== prevProviderRef.current) + ) { + prevModelRef.current = currentModel; + prevProviderRef.current = currentProvider; + + const match = modelList.find( + (model) => model.name === currentModel && model.provider === currentProvider + ); + + if (match) { + setSelectedModel(match); + } else { + // Create a model object if not found in list + setSelectedModel({ name: currentModel, provider: currentProvider }); + } + } + }, [currentModel, currentProvider, modelList, isInitialized]); + // Don't render until we've loaded the initial model/provider if (!isInitialized) { return
Loading models...
; diff --git a/ui/desktop/src/hooks/useChatEngine.ts b/ui/desktop/src/hooks/useChatEngine.ts index 1d965c4a4af9..c183688f833d 100644 --- a/ui/desktop/src/hooks/useChatEngine.ts +++ b/ui/desktop/src/hooks/useChatEngine.ts @@ -79,6 +79,7 @@ export const useChatEngine = ({ updateMessageStreamBody, notifications, sessionMetadata, + setError, } = useMessageStream({ api: getApiUrl('/reply'), id: chat.id, @@ -402,5 +403,8 @@ export const useChatEngine = ({ // Utilities isUserMessage, + + // Error management + clearError: () => setError(undefined), }; }; diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts index 83ecc9e9ff1d..4850e8b1a2c0 100644 --- a/ui/desktop/src/hooks/useMessageStream.ts +++ b/ui/desktop/src/hooks/useMessageStream.ts @@ -173,6 +173,9 @@ export interface UseMessageStreamHelpers { /** Session metadata including token counts */ sessionMetadata: SessionMetadata | null; + + /** Clear error state */ + setError: (error: Error | undefined) => void; } /** @@ -709,5 +712,6 @@ export function useMessageStream({ notifications, currentModelInfo, sessionMetadata, + setError, }; }