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::{Config, ExtensionConfig, GooseMode}; use crate::context_mgmt::format_message_for_compacting; use crate::conversation::message::Message; use crate::execution::manager::AgentManager; use crate::providers; use crate::providers::base::Provider; use crate::session::extension_data::EnabledExtensionsState; 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"; struct CancelTokenGuard { manager: Arc, session_id: String, disarmed: bool, } impl CancelTokenGuard { fn new(manager: Arc, session_id: String) -> Self { Self { manager, session_id, disarmed: false, } } fn disarm(&mut self) { self.disarmed = true; } } impl Drop for CancelTokenGuard { fn drop(&mut self) { if !self.disarmed { let manager = self.manager.clone(); let session_id = self.session_id.clone(); tokio::spawn(async move { manager.unregister_cancel_token(&session_id).await; }); } } } 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()) } fn parent_extensions(&self) -> Vec { let extension_data = self.context.session.as_ref().map(|s| &s.extension_data); EnabledExtensionsState::extensions_or_default(extension_data, Config::global()) } 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 raw_path = PathBuf::from(&working_dir); let path = if raw_path.is_absolute() { raw_path } else { let base = self .context .session .as_ref() .map(|s| s.working_dir.clone()) .unwrap_or_else(|| PathBuf::from(".")); base.join(&raw_path) }; let path = path .canonicalize() .map_err(|e| format!("Invalid working directory '{}': {}", working_dir, e))?; if !path.is_dir() { return Err(format!("'{}' 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))?; let parent_provider = self.get_provider().await?; let extensions = self.parent_extensions(); let provider = providers::create( parent_provider.get_name(), parent_provider.get_model_config(), extensions, ) .await .map_err(|e| format!("Failed to create provider for new agent: {}", e))?; 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, parent_session_id: &str, parent_cancel: &CancellationToken, 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")?; if session_id == parent_session_id { return Err("Cannot send a message to the orchestrator's own session".into()); } let manager = self.get_agent_manager().await?; 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 agent.provider().await.is_err() { if let Ok(parent_provider) = self.get_provider().await { let extensions = self.parent_extensions(); if let Ok(provider) = providers::create( parent_provider.get_name(), parent_provider.get_model_config(), extensions, ) .await { agent .update_provider(provider, &session_id) .await .map_err(|e| format!("Failed to set provider: {}", e))?; } } } let cancel_token = CancellationToken::new(); manager .try_register_cancel_token(&session_id, cancel_token.clone()) .await .map_err(|_| { format!( "Session '{}' is currently busy. Use interrupt_agent first, or wait.", session_id ) })?; let mut guard = CancelTokenGuard::new(manager.clone(), session_id.clone()); 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 mut stream = agent .reply(user_message, session_config, Some(cancel_token.clone())) .await .map_err(|e| format!("Failed to start reply: {}", e))?; let mut response_parts: Vec = Vec::new(); let mut cancelled = false; loop { tokio::select! { _ = parent_cancel.cancelled() => { cancel_token.cancel(); cancelled = true; break; } event = stream.next() => { match event { Some(Ok(AgentEvent::Message(msg))) => { let text = msg.as_concat_text(); if !text.is_empty() { response_parts.push(text); } } Some(Ok(_)) => {} Some(Err(e)) => { response_parts.push(format!("Error during agent processing: {}", e)); break; } None => break, } } } } drop(stream); guard.disarm(); manager.unregister_cancel_token(&session_id).await; if cancelled { return Err("Cancelled by parent session".into()); } if response_parts.is_empty() { Ok(CallToolResult::success(vec![Content::text( "Agent completed without producing text output.", )])) } else { Ok(CallToolResult::success(vec![Content::text(format!( "## Response from session {}\n\n{}", session_id, response_parts.join("\n\n") ))])) } } 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(&ctx.session_id, &cancel_token, 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)) }