From 36d23b1659b045920212492d359f5b0541fac4c8 Mon Sep 17 00:00:00 2001 From: Douwe Osinga Date: Wed, 18 Mar 2026 19:26:13 -0400 Subject: [PATCH 1/2] Orchestration --- crates/goose/src/agents/extension_manager.rs | 22 +- .../src/agents/platform_extensions/mod.rs | 27 + .../platform_extensions/orchestrator.rs | 572 ++++++++++++++++++ ...nager__tests__all_platform_extensions.snap | 4 + crates/goose/src/context_mgmt/mod.rs | 2 +- crates/goose/src/execution/manager.rs | 44 ++ 6 files changed, 665 insertions(+), 6 deletions(-) create mode 100644 crates/goose/src/agents/platform_extensions/orchestrator.rs diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 30476210f2..6dd055de52 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -222,6 +222,12 @@ pub fn is_first_class_extension(name: &str) -> bool { .is_some_and(|def| def.unprefixed_tools) } +pub fn is_hidden_extension(name: &str) -> bool { + PLATFORM_EXTENSIONS + .get(name_to_key(name).as_str()) + .is_some_and(|def| def.hidden) +} + /// Result of resolving a tool call to its owning extension struct ResolvedTool { extension_name: String, @@ -1543,10 +1549,10 @@ impl ExtensionManager { pub async fn search_available_extensions(&self) -> Result, ErrorData> { let mut output_parts = vec![]; - // First get disabled extensions from current config + // First get disabled extensions from current config (skip hidden ones) let mut disabled_extensions: Vec = vec![]; for extension in get_all_extensions() { - if !extension.enabled { + if !extension.enabled && !is_hidden_extension(&extension.config.name()) { let config = extension.config.clone(); let description = match &config { ExtensionConfig::Builtin { @@ -1571,9 +1577,15 @@ impl ExtensionManager { } } - // Get currently enabled extensions that can be disabled - let enabled_extensions: Vec = - self.extensions.lock().await.keys().cloned().collect(); + // Get currently enabled extensions that can be disabled (skip hidden ones) + let enabled_extensions: Vec = self + .extensions + .lock() + .await + .keys() + .filter(|name| !is_hidden_extension(name)) + .cloned() + .collect(); // Build output string if !disabled_extensions.is_empty() { diff --git a/crates/goose/src/agents/platform_extensions/mod.rs b/crates/goose/src/agents/platform_extensions/mod.rs index cd577d71e8..b54751d6a1 100644 --- a/crates/goose/src/agents/platform_extensions/mod.rs +++ b/crates/goose/src/agents/platform_extensions/mod.rs @@ -5,6 +5,7 @@ pub mod chatrecall; pub mod code_execution; pub mod developer; pub mod ext_manager; +pub mod orchestrator; pub mod summarize; pub mod summon; pub mod todo; @@ -37,6 +38,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Analyze code structure with tree-sitter: directory overviews, file details, symbol call graphs", default_enabled: true, unprefixed_tools: true, + hidden: false, client_factory: |ctx| Box::new(analyze::AnalyzeClient::new(ctx).unwrap()), }, ); @@ -50,6 +52,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Enable a todo list for goose so it can keep track of what it is doing", default_enabled: true, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(todo::TodoClient::new(ctx).unwrap()), }, ); @@ -63,6 +66,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Create and manage custom Goose apps through chat. Apps are HTML/CSS/JavaScript and run in sandboxed windows.", default_enabled: true, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(apps::AppsManagerClient::new(ctx).unwrap()), }, ); @@ -76,6 +80,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Search past conversations and load session summaries for contextual memory", default_enabled: false, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(chatrecall::ChatRecallClient::new(ctx).unwrap()), }, ); @@ -89,6 +94,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Enable extension management tools for discovering, enabling, and disabling extensions", default_enabled: true, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(ext_manager::ExtensionManagerClient::new(ctx).unwrap()), }, ); @@ -101,6 +107,7 @@ pub static PLATFORM_EXTENSIONS: Lazy description: "Load knowledge and delegate tasks to subagents", default_enabled: true, unprefixed_tools: true, + hidden: false, client_factory: |ctx| Box::new(summon::SummonClient::new(ctx).unwrap()), }, ); @@ -113,6 +120,7 @@ pub static PLATFORM_EXTENSIONS: Lazy description: "Load files/directories and get an LLM summary in a single call", default_enabled: false, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(summarize::SummarizeClient::new(ctx).unwrap()), }, ); @@ -127,6 +135,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Goose will make extension calls through code execution, saving tokens", default_enabled: false, unprefixed_tools: true, + hidden: false, client_factory: |ctx| { Box::new(code_execution::CodeExecutionClient::new(ctx).unwrap()) }, @@ -141,10 +150,25 @@ pub static PLATFORM_EXTENSIONS: Lazy description: "Write and edit files, and execute shell commands", default_enabled: true, unprefixed_tools: true, + hidden: false, client_factory: |ctx| Box::new(developer::DeveloperClient::new(ctx).unwrap()), }, ); + map.insert( + orchestrator::EXTENSION_NAME, + PlatformExtensionDef { + name: orchestrator::EXTENSION_NAME, + display_name: "Orchestrator", + description: + "Manage agent sessions: list, view, start, send messages, interrupt, and stop agents", + default_enabled: false, + unprefixed_tools: false, + hidden: true, + client_factory: |ctx| Box::new(orchestrator::OrchestratorClient::new(ctx).unwrap()), + }, + ); + map.insert( tom::EXTENSION_NAME, PlatformExtensionDef { @@ -154,6 +178,7 @@ pub static PLATFORM_EXTENSIONS: Lazy "Inject custom context into every turn via GOOSE_MOIM_MESSAGE_TEXT and GOOSE_MOIM_MESSAGE_FILE environment variables", default_enabled: true, unprefixed_tools: false, + hidden: false, client_factory: |ctx| Box::new(tom::TomClient::new(ctx).unwrap()), }, ); @@ -211,5 +236,7 @@ pub struct PlatformExtensionDef { pub default_enabled: bool, /// If true, tools are exposed without extension prefix for intuitive first-class use. pub unprefixed_tools: bool, + /// If true, the extension is not shown in the UI or discoverable via search_available_extensions. + pub hidden: bool, pub client_factory: fn(PlatformExtensionContext) -> Box, } diff --git a/crates/goose/src/agents/platform_extensions/orchestrator.rs b/crates/goose/src/agents/platform_extensions/orchestrator.rs new file mode 100644 index 0000000000..dfdcfe9970 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/orchestrator.rs @@ -0,0 +1,572 @@ +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::agents::tool_execution::ToolCallContext; +use crate::agents::{AgentEvent, SessionConfig}; +use crate::config::GooseMode; +use crate::context_mgmt::format_message_for_compacting; +use crate::conversation::message::Message; +use crate::execution::manager::AgentManager; +use crate::providers::base::Provider; +use crate::session::session_manager::SessionType; +use anyhow::Result; +use async_trait::async_trait; +use futures::StreamExt; +use rmcp::model::{ + CallToolResult, Content, Implementation, InitializeResult, JsonObject, ListToolsResult, + ServerCapabilities, Tool, +}; +use schemars::{schema_for, JsonSchema}; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "orchestrator"; + +const DEFAULT_LIST_LIMIT: usize = 10; + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct ListSessionsParams { + /// Filter by session type: "user", "sub_agent", "scheduled", "hidden", "terminal", "gateway". + /// If omitted, returns all session types. + session_type: Option, + /// Maximum number of sessions to return (most recent first). Defaults to 10. + last_n: Option, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct ViewSessionParams { + /// The session ID to inspect + session_id: String, + /// How to view the conversation: "first_last" returns the first and last message, + /// "summarize" calls the LLM to produce a summary. If omitted, returns first and last. + mode: Option, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct StartAgentParams { + /// Working directory for the new agent session + working_dir: String, + /// Human-readable name for the session + name: Option, + // TODO: add a "model_tier" parameter (e.g. "fast" vs "normal") to let the orchestrator + // choose between a fast/cheap model and the default one. For now we inherit the + // orchestrator's own provider and model. +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct SendMessageParams { + /// The session ID of the agent to send a message to + session_id: String, + /// The message text to send + message: String, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct InterruptAgentParams { + /// The session ID of the agent to interrupt + session_id: String, +} + +pub struct OrchestratorClient { + info: InitializeResult, + context: PlatformExtensionContext, +} + +impl OrchestratorClient { + pub fn new(context: PlatformExtensionContext) -> Result { + let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info( + Implementation::new(EXTENSION_NAME, "1.0.0").with_title("Orchestrator"), + ) + .with_instructions( + "Manage agent sessions: list, view, start, send messages, and interrupt agents.", + ); + + Ok(Self { info, context }) + } + + async fn get_agent_manager(&self) -> Result, String> { + AgentManager::instance() + .await + .map_err(|e| format!("Failed to get agent manager: {}", e)) + } + + async fn get_provider(&self) -> Result, String> { + let extension_manager = self + .context + .extension_manager + .as_ref() + .and_then(|weak| weak.upgrade()) + .ok_or("Extension manager not available")?; + + let provider_guard = extension_manager.get_provider().lock().await; + provider_guard + .as_ref() + .cloned() + .ok_or_else(|| "Provider not available".to_string()) + } + + async fn handle_list_sessions( + &self, + arguments: Option, + ) -> Result { + let type_filter = arguments + .as_ref() + .and_then(|args| args.get("session_type")) + .and_then(|v| v.as_str()); + + let limit = arguments + .as_ref() + .and_then(|args| args.get("last_n")) + .and_then(|v| v.as_u64()) + .map(|v| v as usize) + .unwrap_or(DEFAULT_LIST_LIMIT); + + let manager = self.get_agent_manager().await?; + + let mut sessions = if let Some(type_str) = type_filter { + let session_type: SessionType = type_str + .parse() + .map_err(|e| format!("Invalid session type '{}': {}", type_str, e))?; + self.context + .session_manager + .list_sessions_by_types(&[session_type]) + .await + .map_err(|e| format!("Failed to list sessions: {}", e))? + } else { + self.context + .session_manager + .list_sessions() + .await + .map_err(|e| format!("Failed to list sessions: {}", e))? + }; + + // Most recent first + sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); + let total = sessions.len(); + sessions.truncate(limit); + + if sessions.is_empty() { + return Ok(CallToolResult::success(vec![Content::text( + "No sessions found.", + )])); + } + + let active_ids = manager.list_active_session_ids().await; + + let mut lines = vec![format!( + "Showing {} of {} session(s):\n", + sessions.len(), + total + )]; + for session in &sessions { + let is_loaded = active_ids.contains(&session.id); + let is_busy = if is_loaded { + manager.is_session_busy(&session.id).await + } else { + false + }; + + let status = if is_busy { + "🔄 busy" + } else if is_loaded { + "✓ loaded" + } else { + "○ idle" + }; + + lines.push(format!( + "- **{}** ({})\n Type: {} | Status: {} | Messages: {} | Updated: {}", + session.name, + session.id, + session.session_type, + status, + session.message_count, + session.updated_at.format("%Y-%m-%d %H:%M"), + )); + } + + Ok(CallToolResult::success(vec![Content::text( + lines.join("\n"), + )])) + } + + async fn handle_view_session( + &self, + session_id_for_llm: &str, + arguments: Option, + ) -> Result { + let args = arguments.ok_or("Missing arguments")?; + let session_id = extract_string(&args, "session_id")?; + let mode = args + .get("mode") + .and_then(|v| v.as_str()) + .unwrap_or("first_last"); + + let session = self + .context + .session_manager + .get_session(&session_id, true) + .await + .map_err(|e| format!("Session '{}' not found: {}", session_id, e))?; + + let manager = self.get_agent_manager().await?; + let is_busy = manager.is_session_busy(&session_id).await; + + let mut output = vec![format!( + "# Session: {} ({})\n\nType: {} | Status: {} | Working dir: {}\nMessages: {} | Updated: {}\n", + session.name, + session.id, + session.session_type, + if is_busy { "🔄 busy" } else { "idle" }, + session.working_dir.display(), + session.message_count, + session.updated_at.format("%Y-%m-%d %H:%M"), + )]; + + match mode { + "first_last" => { + if let Some(conversation) = &session.conversation { + let messages = conversation.messages(); + if messages.is_empty() { + output.push("No messages in this session.".to_string()); + } else { + output.push("## First message\n".to_string()); + output.push(format_message_for_compacting(&messages[0])); + + if messages.len() > 1 { + output.push(format!("\n*({} messages omitted)*\n", messages.len() - 2)); + output.push("## Last message\n".to_string()); + output + .push(format_message_for_compacting(&messages[messages.len() - 1])); + } + } + } else { + output.push("No messages in this session.".to_string()); + } + } + "summarize" => { + if let Some(conversation) = &session.conversation { + let messages = conversation.messages(); + if messages.is_empty() { + output.push("No messages to summarize.".to_string()); + } else { + let summary = self + .summarize_conversation(session_id_for_llm, messages) + .await?; + output.push(format!("## Summary\n\n{}", summary)); + } + } else { + output.push("No messages to summarize.".to_string()); + } + } + other => { + return Err(format!( + "Unknown mode '{}'. Use 'first_last' or 'summarize'.", + other + )); + } + } + + Ok(CallToolResult::success(vec![Content::text( + output.join("\n"), + )])) + } + + async fn summarize_conversation( + &self, + session_id: &str, + messages: &[Message], + ) -> Result { + let provider = self.get_provider().await?; + + let conversation_text = messages + .iter() + .filter(|m| m.is_agent_visible()) + .map(format_message_for_compacting) + .collect::>() + .join("\n"); + + let system = + "You are a helpful assistant. Summarize the following conversation concisely, \ + capturing the key topics, decisions, and current state. Be brief."; + + let user_message = Message::user().with_text(format!( + "Summarize this conversation ({} messages):\n\n{}", + messages.len(), + conversation_text + )); + + let (response, _usage) = provider + .complete_fast(session_id, system, &[user_message], &[]) + .await + .map_err(|e| format!("LLM summarization failed: {}", e))?; + + Ok(response + .content + .iter() + .filter_map(|c| { + if let crate::conversation::message::MessageContent::Text(t) = c { + Some(t.text.clone()) + } else { + None + } + }) + .collect::>() + .join("\n")) + } + + async fn handle_start_agent( + &self, + arguments: Option, + ) -> Result { + let args = arguments.ok_or("Missing arguments")?; + let working_dir = extract_string(&args, "working_dir")?; + let name = args + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("Orchestrated Agent") + .to_string(); + + let path = PathBuf::from(&working_dir); + if !path.exists() || !path.is_dir() { + return Err(format!( + "Working directory '{}' does not exist or is not a directory", + working_dir + )); + } + + let mode = GooseMode::default(); + + let session = self + .context + .session_manager + .create_session(path, name.clone(), SessionType::User, mode) + .await + .map_err(|e| format!("Failed to create session: {}", e))?; + + let manager = self.get_agent_manager().await?; + let agent = manager + .get_or_create_agent(session.id.clone()) + .await + .map_err(|e| format!("Failed to create agent: {}", e))?; + + // Inherit the orchestrator's provider and model + let provider = self.get_provider().await?; + agent + .update_provider(provider, &session.id) + .await + .map_err(|e| format!("Failed to set provider on new agent: {}", e))?; + + Ok(CallToolResult::success(vec![Content::text(format!( + "Started agent session '{}' with ID: {}\n\nUse send_message with this session_id to interact with it.", + name, session.id + ))])) + } + + async fn handle_send_message( + &self, + arguments: Option, + ) -> Result { + let args = arguments.ok_or("Missing arguments")?; + let session_id = extract_string(&args, "session_id")?; + let message_text = extract_string(&args, "message")?; + + let manager = self.get_agent_manager().await?; + + if manager.is_session_busy(&session_id).await { + return Err(format!( + "Session '{}' is currently busy processing a request. Use interrupt_agent to cancel it first, or wait for it to finish.", + session_id + )); + // TODO: post the message so it gets picked up on the next turn instead of erroring + } + + let agent = manager + .get_or_create_agent(session_id.clone()) + .await + .map_err(|e| format!("Failed to get agent for session '{}': {}", session_id, e))?; + + // If the agent has no provider yet (e.g. resumed from DB), inherit ours + if agent.provider().await.is_err() { + if let Ok(provider) = self.get_provider().await { + agent + .update_provider(provider, &session_id) + .await + .map_err(|e| format!("Failed to set provider: {}", e))?; + } + } + + let cancel_token = CancellationToken::new(); + manager + .register_cancel_token(&session_id, cancel_token.clone()) + .await; + + let user_message = Message::user().with_text(&message_text); + + let session_config = SessionConfig { + id: session_id.clone(), + schedule_id: None, + max_turns: None, + retry_config: None, + }; + + let result: Result = async { + let mut stream = agent + .reply(user_message, session_config, Some(cancel_token)) + .await + .map_err(|e| format!("Failed to start reply: {}", e))?; + + let mut response_parts: Vec = Vec::new(); + + while let Some(event_result) = stream.next().await { + match event_result { + Ok(AgentEvent::Message(msg)) => { + let text = msg.as_concat_text(); + if !text.is_empty() { + response_parts.push(text); + } + } + Ok(AgentEvent::HistoryReplaced(_)) + | Ok(AgentEvent::McpNotification(_)) + | Ok(AgentEvent::ModelChange { .. }) => {} + Err(e) => { + response_parts.push(format!("Error during agent processing: {}", e)); + break; + } + } + } + + if response_parts.is_empty() { + Ok("Agent completed without producing text output.".to_string()) + } else { + Ok(response_parts.join("\n\n")) + } + } + .await; + + manager.unregister_cancel_token(&session_id).await; + + let response_text = result?; + + Ok(CallToolResult::success(vec![Content::text(format!( + "## Response from session {}\n\n{}", + session_id, response_text + ))])) + } + + async fn handle_interrupt_agent( + &self, + arguments: Option, + ) -> Result { + let args = arguments.ok_or("Missing arguments")?; + let session_id = extract_string(&args, "session_id")?; + + let manager = self.get_agent_manager().await?; + + manager + .cancel_session(&session_id) + .await + .map_err(|e| format!("Failed to interrupt session '{}': {}", session_id, e))?; + + Ok(CallToolResult::success(vec![Content::text(format!( + "Interrupted agent session '{}'.", + session_id + ))])) + } +} + +#[async_trait] +impl McpClientTrait for OrchestratorClient { + async fn list_tools( + &self, + _session_id: &str, + _next_cursor: Option, + _cancel_token: CancellationToken, + ) -> Result { + let tools = vec![ + Tool::new( + "list_sessions".to_string(), + "List agent sessions with their status (loaded, busy, idle). Returns the most recent 10 by default. Optionally filter by session type." + .to_string(), + schema::(), + ), + Tool::new( + "view_session".to_string(), + "View a session's details and conversation. Mode 'first_last' (default) returns the first and last message. Mode 'summarize' calls the LLM to produce a conversation summary." + .to_string(), + schema::(), + ), + Tool::new( + "start_agent".to_string(), + "Start a new agent session with its own working directory. Inherits the current provider and model. Returns a session_id for future interaction." + .to_string(), + schema::(), + ), + Tool::new( + "send_message".to_string(), + "Send a message to an existing agent session and get the response. Returns an error if the agent is currently busy." + .to_string(), + schema::(), + ), + Tool::new( + "interrupt_agent".to_string(), + "Interrupt a busy agent by cancelling its current operation." + .to_string(), + schema::(), + ), + ]; + + Ok(ListToolsResult { + tools, + next_cursor: None, + meta: None, + }) + } + + async fn call_tool( + &self, + ctx: &ToolCallContext, + name: &str, + arguments: Option, + _cancel_token: CancellationToken, + ) -> Result { + let result = match name { + "list_sessions" => self.handle_list_sessions(arguments).await, + "view_session" => self.handle_view_session(&ctx.session_id, arguments).await, + "start_agent" => self.handle_start_agent(arguments).await, + "send_message" => self.handle_send_message(arguments).await, + "interrupt_agent" => self.handle_interrupt_agent(arguments).await, + _ => Err(format!("Unknown tool: {}", name)), + }; + + match result { + Ok(result) => Ok(result), + Err(error) => Ok(CallToolResult::error(vec![Content::text(format!( + "Error: {}", + error + ))])), + } + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } +} + +fn schema() -> JsonObject { + let mut obj = serde_json::to_value(schema_for!(T)) + .map(|v| v.as_object().unwrap().clone()) + .expect("valid schema"); + obj.entry("properties") + .or_insert_with(|| serde_json::json!({})); + obj +} + +fn extract_string(args: &JsonObject, key: &str) -> Result { + args.get(key) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| format!("Missing or invalid '{}'", key)) +} diff --git a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap index 1b00f5374e..798d43e26a 100644 --- a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap +++ b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap @@ -94,6 +94,10 @@ and file sizes. When you need to search, prefer rg which correctly respects giti content. Then use cat or sed to gather the context you need, always reading before editing. Use write and edit to efficiently make changes. Test and verify as appropriate. +## orchestrator + +### Instructions +Manage agent sessions: list, view, start, send messages, and interrupt agents. ## summarize diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index f6847e460c..71085dfd50 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -339,7 +339,7 @@ async fn do_compact( )) } -fn format_message_for_compacting(msg: &Message) -> String { +pub fn format_message_for_compacting(msg: &Message) -> String { let content_parts: Vec = msg .content .iter() diff --git a/crates/goose/src/execution/manager.rs b/crates/goose/src/execution/manager.rs index 49311086d1..04de97be60 100644 --- a/crates/goose/src/execution/manager.rs +++ b/crates/goose/src/execution/manager.rs @@ -7,9 +7,11 @@ use crate::scheduler_trait::SchedulerTrait; use crate::session::SessionManager; use anyhow::Result; use lru::LruCache; +use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::Arc; use tokio::sync::{OnceCell, RwLock}; +use tokio_util::sync::CancellationToken; use tracing::{debug, info}; const DEFAULT_MAX_SESSION: usize = 100; @@ -22,6 +24,7 @@ pub struct AgentManager { session_manager: Arc, default_provider: Arc>>>, default_mode: GooseMode, + cancel_tokens: Arc>>, } impl AgentManager { @@ -42,6 +45,7 @@ impl AgentManager { session_manager, default_provider: Arc::new(RwLock::new(None)), default_mode, + cancel_tokens: Arc::new(RwLock::new(HashMap::new())), }; Ok(manager) @@ -151,6 +155,7 @@ impl AgentManager { } pub async fn remove_session(&self, session_id: &str) -> Result<()> { + self.cancel_tokens.write().await.remove(session_id); let mut sessions = self.sessions.write().await; sessions .pop(session_id) @@ -166,6 +171,45 @@ impl AgentManager { pub async fn session_count(&self) -> usize { self.sessions.read().await.len() } + + /// Register a cancellation token for a session (called when reply starts) + pub async fn register_cancel_token(&self, session_id: &str, token: CancellationToken) { + self.cancel_tokens + .write() + .await + .insert(session_id.to_string(), token); + } + + /// Remove the cancellation token for a session (called when reply finishes) + pub async fn unregister_cancel_token(&self, session_id: &str) { + self.cancel_tokens.write().await.remove(session_id); + } + + /// Cancel a running agent by triggering its cancellation token + pub async fn cancel_session(&self, session_id: &str) -> Result<()> { + let tokens = self.cancel_tokens.read().await; + let token = tokens + .get(session_id) + .ok_or_else(|| anyhow::anyhow!("No active operation for session {}", session_id))?; + token.cancel(); + Ok(()) + } + + /// Check if a session has an active reply in progress + pub async fn is_session_busy(&self, session_id: &str) -> bool { + let tokens = self.cancel_tokens.read().await; + tokens.contains_key(session_id) + } + + /// List session IDs that currently have active agents loaded + pub async fn list_active_session_ids(&self) -> Vec { + self.sessions + .read() + .await + .iter() + .map(|(id, _)| id.clone()) + .collect() + } } #[cfg(test)] From 1d559a7eebc91e8f2369f64844c70c24eb7c52a9 Mon Sep 17 00:00:00 2001 From: Douwe Osinga Date: Thu, 19 Mar 2026 08:31:29 -0400 Subject: [PATCH 2/2] fix: remove ModelChange variant after merge with main Signed-off-by: Douwe Osinga --- crates/goose/src/agents/platform_extensions/orchestrator.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/crates/goose/src/agents/platform_extensions/orchestrator.rs b/crates/goose/src/agents/platform_extensions/orchestrator.rs index dfdcfe9970..439363ee52 100644 --- a/crates/goose/src/agents/platform_extensions/orchestrator.rs +++ b/crates/goose/src/agents/platform_extensions/orchestrator.rs @@ -428,9 +428,7 @@ impl OrchestratorClient { response_parts.push(text); } } - Ok(AgentEvent::HistoryReplaced(_)) - | Ok(AgentEvent::McpNotification(_)) - | Ok(AgentEvent::ModelChange { .. }) => {} + Ok(AgentEvent::HistoryReplaced(_)) | Ok(AgentEvent::McpNotification(_)) => {} Err(e) => { response_parts.push(format!("Error during agent processing: {}", e)); break;