mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
Orchestration (#7999)
Signed-off-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Tyler Longwell <tlongwell@squareup.com>
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
use crate::routes::errors::ErrorResponse;
|
||||
use crate::routes::reply::{get_token_state, track_tool_telemetry, MessageEvent};
|
||||
use crate::session_event_bus::RequestGuard;
|
||||
use crate::state::AppState;
|
||||
use axum::{
|
||||
extract::{DefaultBodyLimit, Path, State},
|
||||
@@ -320,7 +321,13 @@ pub async fn session_reply(
|
||||
}
|
||||
|
||||
let bus = state.get_or_create_event_bus(&session_id).await;
|
||||
let cancel_token = bus.register_request(request_id.clone()).await;
|
||||
|
||||
let cancel_token = bus
|
||||
.try_register_request(request_id.clone())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
ErrorResponse::bad_request("Session already has an active request. Cancel it first.")
|
||||
})?;
|
||||
|
||||
let user_message = request.user_message;
|
||||
let override_conversation = request.override_conversation;
|
||||
@@ -332,6 +339,8 @@ pub async fn session_reply(
|
||||
let task_bus = bus.clone();
|
||||
|
||||
drop(tokio::spawn(async move {
|
||||
let mut _guard = RequestGuard::new(task_bus.clone(), task_request_id.clone());
|
||||
|
||||
let publish = |rid: Option<String>, event: MessageEvent| {
|
||||
let bus = task_bus.clone();
|
||||
async move {
|
||||
@@ -350,7 +359,6 @@ pub async fn session_reply(
|
||||
},
|
||||
)
|
||||
.await;
|
||||
task_bus.cleanup_request(&task_request_id).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
@@ -370,7 +378,6 @@ pub async fn session_reply(
|
||||
},
|
||||
)
|
||||
.await;
|
||||
task_bus.cleanup_request(&task_request_id).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
@@ -420,7 +427,6 @@ pub async fn session_reply(
|
||||
},
|
||||
)
|
||||
.await;
|
||||
task_bus.cleanup_request(&task_request_id).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
@@ -562,6 +568,7 @@ pub async fn session_reply(
|
||||
)
|
||||
.await;
|
||||
|
||||
_guard.disarm();
|
||||
task_bus.cleanup_request(&task_request_id).await;
|
||||
}));
|
||||
|
||||
|
||||
@@ -116,7 +116,7 @@ impl SessionEventBus {
|
||||
requests.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Register a new request and return its cancellation token.
|
||||
#[cfg(test)]
|
||||
pub async fn register_request(&self, request_id: String) -> CancellationToken {
|
||||
let token = CancellationToken::new();
|
||||
let mut requests = self.active_requests.lock().await;
|
||||
@@ -124,6 +124,20 @@ impl SessionEventBus {
|
||||
token
|
||||
}
|
||||
|
||||
/// Atomically check no requests are active and register one. Returns Err if busy.
|
||||
pub async fn try_register_request(
|
||||
&self,
|
||||
request_id: String,
|
||||
) -> Result<CancellationToken, String> {
|
||||
let mut requests = self.active_requests.lock().await;
|
||||
if !requests.is_empty() {
|
||||
return Err("Session already has an active request".into());
|
||||
}
|
||||
let token = CancellationToken::new();
|
||||
requests.insert(request_id, token.clone());
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Cancel a specific request by request_id.
|
||||
pub async fn cancel_request(&self, request_id: &str) -> bool {
|
||||
let requests = self.active_requests.lock().await;
|
||||
@@ -156,6 +170,38 @@ impl Default for SessionEventBus {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RequestGuard {
|
||||
bus: std::sync::Arc<SessionEventBus>,
|
||||
request_id: String,
|
||||
disarmed: bool,
|
||||
}
|
||||
|
||||
impl RequestGuard {
|
||||
pub fn new(bus: std::sync::Arc<SessionEventBus>, request_id: String) -> Self {
|
||||
Self {
|
||||
bus,
|
||||
request_id,
|
||||
disarmed: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn disarm(&mut self) {
|
||||
self.disarmed = true;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RequestGuard {
|
||||
fn drop(&mut self) {
|
||||
if !self.disarmed {
|
||||
let bus = self.bus.clone();
|
||||
let request_id = self.request_id.clone();
|
||||
tokio::spawn(async move {
|
||||
bus.cleanup_request(&request_id).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -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<Vec<Content>, 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<String> = 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<String> =
|
||||
self.extensions.lock().await.keys().cloned().collect();
|
||||
// Get currently enabled extensions that can be disabled (skip hidden ones)
|
||||
let enabled_extensions: Vec<String> = self
|
||||
.extensions
|
||||
.lock()
|
||||
.await
|
||||
.keys()
|
||||
.filter(|name| !is_hidden_extension(name))
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
// Build output string
|
||||
if !disabled_extensions.is_empty() {
|
||||
|
||||
@@ -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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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<HashMap<&'static str, PlatformExtensionDef>
|
||||
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<HashMap<&'static str, PlatformExtensionDef>
|
||||
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<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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(
|
||||
@@ -147,10 +156,25 @@ pub static PLATFORM_EXTENSIONS: Lazy<HashMap<&'static str, PlatformExtensionDef>
|
||||
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 {
|
||||
@@ -160,6 +184,7 @@ pub static PLATFORM_EXTENSIONS: Lazy<HashMap<&'static str, PlatformExtensionDef>
|
||||
"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()),
|
||||
},
|
||||
);
|
||||
@@ -217,5 +242,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<dyn McpClientTrait>,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,635 @@
|
||||
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";
|
||||
|
||||
struct CancelTokenGuard {
|
||||
manager: Arc<AgentManager>,
|
||||
session_id: String,
|
||||
disarmed: bool,
|
||||
}
|
||||
|
||||
impl CancelTokenGuard {
|
||||
fn new(manager: Arc<AgentManager>, 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<String>,
|
||||
/// Maximum number of sessions to return (most recent first). Defaults to 10.
|
||||
last_n: Option<usize>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
// 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<Self> {
|
||||
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<Arc<AgentManager>, String> {
|
||||
AgentManager::instance()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to get agent manager: {}", e))
|
||||
}
|
||||
|
||||
async fn get_provider(&self) -> Result<Arc<dyn Provider>, 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<JsonObject>,
|
||||
) -> Result<CallToolResult, String> {
|
||||
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<JsonObject>,
|
||||
) -> Result<CallToolResult, String> {
|
||||
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<String, String> {
|
||||
let provider = self.get_provider().await?;
|
||||
|
||||
let conversation_text = messages
|
||||
.iter()
|
||||
.filter(|m| m.is_agent_visible())
|
||||
.map(format_message_for_compacting)
|
||||
.collect::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.join("\n"))
|
||||
}
|
||||
|
||||
async fn handle_start_agent(
|
||||
&self,
|
||||
arguments: Option<JsonObject>,
|
||||
) -> Result<CallToolResult, String> {
|
||||
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))?;
|
||||
|
||||
// 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,
|
||||
parent_session_id: &str,
|
||||
parent_cancel: &CancellationToken,
|
||||
arguments: Option<JsonObject>,
|
||||
) -> Result<CallToolResult, String> {
|
||||
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(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
|
||||
.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<String> = 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<JsonObject>,
|
||||
) -> Result<CallToolResult, String> {
|
||||
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<String>,
|
||||
_cancel_token: CancellationToken,
|
||||
) -> Result<ListToolsResult, Error> {
|
||||
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::<ListSessionsParams>(),
|
||||
),
|
||||
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::<ViewSessionParams>(),
|
||||
),
|
||||
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::<StartAgentParams>(),
|
||||
),
|
||||
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::<SendMessageParams>(),
|
||||
),
|
||||
Tool::new(
|
||||
"interrupt_agent".to_string(),
|
||||
"Interrupt a busy agent by cancelling its current operation."
|
||||
.to_string(),
|
||||
schema::<InterruptAgentParams>(),
|
||||
),
|
||||
];
|
||||
|
||||
Ok(ListToolsResult {
|
||||
tools,
|
||||
next_cursor: None,
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
ctx: &ToolCallContext,
|
||||
name: &str,
|
||||
arguments: Option<JsonObject>,
|
||||
cancel_token: CancellationToken,
|
||||
) -> Result<CallToolResult, Error> {
|
||||
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<T: JsonSchema>() -> 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<String, String> {
|
||||
args.get(key)
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| format!("Missing or invalid '{}'", key))
|
||||
}
|
||||
+4
@@ -89,6 +89,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
|
||||
|
||||
|
||||
|
||||
@@ -342,7 +342,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<String> = msg
|
||||
.content
|
||||
.iter()
|
||||
|
||||
@@ -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<SessionManager>,
|
||||
default_provider: Arc<RwLock<Option<Arc<dyn crate::providers::base::Provider>>>>,
|
||||
default_mode: GooseMode,
|
||||
cancel_tokens: Arc<RwLock<HashMap<String, CancellationToken>>>,
|
||||
}
|
||||
|
||||
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,9 @@ impl AgentManager {
|
||||
}
|
||||
|
||||
pub async fn remove_session(&self, session_id: &str) -> Result<()> {
|
||||
if let Some(token) = self.cancel_tokens.write().await.remove(session_id) {
|
||||
token.cancel();
|
||||
}
|
||||
let mut sessions = self.sessions.write().await;
|
||||
sessions
|
||||
.pop(session_id)
|
||||
@@ -166,6 +173,51 @@ impl AgentManager {
|
||||
pub async fn session_count(&self) -> usize {
|
||||
self.sessions.read().await.len()
|
||||
}
|
||||
|
||||
/// Atomically check if busy and register a cancel token. Returns Err if already busy.
|
||||
pub async fn try_register_cancel_token(
|
||||
&self,
|
||||
session_id: &str,
|
||||
token: CancellationToken,
|
||||
) -> Result<()> {
|
||||
let mut tokens = self.cancel_tokens.write().await;
|
||||
if tokens.contains_key(session_id) {
|
||||
anyhow::bail!("Session '{}' is currently busy", session_id);
|
||||
}
|
||||
tokens.insert(session_id.to_string(), token);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 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<String> {
|
||||
self.sessions
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.map(|(id, _)| id.clone())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
Reference in New Issue
Block a user