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:
Douwe Osinga
2026-03-19 16:30:24 -04:00
committed by GitHub
parent 53236ff327
commit 3f8eec3237
8 changed files with 794 additions and 11 deletions
@@ -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;
}));
+47 -1
View File
@@ -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::*;
+17 -5
View File
@@ -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))
}
@@ -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
+1 -1
View File
@@ -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()
+52
View File
@@ -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)]