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::errors::ErrorResponse;
|
||||||
use crate::routes::reply::{get_token_state, track_tool_telemetry, MessageEvent};
|
use crate::routes::reply::{get_token_state, track_tool_telemetry, MessageEvent};
|
||||||
|
use crate::session_event_bus::RequestGuard;
|
||||||
use crate::state::AppState;
|
use crate::state::AppState;
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{DefaultBodyLimit, Path, State},
|
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 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 user_message = request.user_message;
|
||||||
let override_conversation = request.override_conversation;
|
let override_conversation = request.override_conversation;
|
||||||
@@ -332,6 +339,8 @@ pub async fn session_reply(
|
|||||||
let task_bus = bus.clone();
|
let task_bus = bus.clone();
|
||||||
|
|
||||||
drop(tokio::spawn(async move {
|
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 publish = |rid: Option<String>, event: MessageEvent| {
|
||||||
let bus = task_bus.clone();
|
let bus = task_bus.clone();
|
||||||
async move {
|
async move {
|
||||||
@@ -350,7 +359,6 @@ pub async fn session_reply(
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
task_bus.cleanup_request(&task_request_id).await;
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -370,7 +378,6 @@ pub async fn session_reply(
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
task_bus.cleanup_request(&task_request_id).await;
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -420,7 +427,6 @@ pub async fn session_reply(
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
task_bus.cleanup_request(&task_request_id).await;
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -562,6 +568,7 @@ pub async fn session_reply(
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
|
_guard.disarm();
|
||||||
task_bus.cleanup_request(&task_request_id).await;
|
task_bus.cleanup_request(&task_request_id).await;
|
||||||
}));
|
}));
|
||||||
|
|
||||||
|
|||||||
@@ -116,7 +116,7 @@ impl SessionEventBus {
|
|||||||
requests.keys().cloned().collect()
|
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 {
|
pub async fn register_request(&self, request_id: String) -> CancellationToken {
|
||||||
let token = CancellationToken::new();
|
let token = CancellationToken::new();
|
||||||
let mut requests = self.active_requests.lock().await;
|
let mut requests = self.active_requests.lock().await;
|
||||||
@@ -124,6 +124,20 @@ impl SessionEventBus {
|
|||||||
token
|
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.
|
/// Cancel a specific request by request_id.
|
||||||
pub async fn cancel_request(&self, request_id: &str) -> bool {
|
pub async fn cancel_request(&self, request_id: &str) -> bool {
|
||||||
let requests = self.active_requests.lock().await;
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|||||||
@@ -222,6 +222,12 @@ pub fn is_first_class_extension(name: &str) -> bool {
|
|||||||
.is_some_and(|def| def.unprefixed_tools)
|
.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
|
/// Result of resolving a tool call to its owning extension
|
||||||
struct ResolvedTool {
|
struct ResolvedTool {
|
||||||
extension_name: String,
|
extension_name: String,
|
||||||
@@ -1543,10 +1549,10 @@ impl ExtensionManager {
|
|||||||
pub async fn search_available_extensions(&self) -> Result<Vec<Content>, ErrorData> {
|
pub async fn search_available_extensions(&self) -> Result<Vec<Content>, ErrorData> {
|
||||||
let mut output_parts = vec![];
|
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![];
|
let mut disabled_extensions: Vec<String> = vec![];
|
||||||
for extension in get_all_extensions() {
|
for extension in get_all_extensions() {
|
||||||
if !extension.enabled {
|
if !extension.enabled && !is_hidden_extension(&extension.config.name()) {
|
||||||
let config = extension.config.clone();
|
let config = extension.config.clone();
|
||||||
let description = match &config {
|
let description = match &config {
|
||||||
ExtensionConfig::Builtin {
|
ExtensionConfig::Builtin {
|
||||||
@@ -1571,9 +1577,15 @@ impl ExtensionManager {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get currently enabled extensions that can be disabled
|
// Get currently enabled extensions that can be disabled (skip hidden ones)
|
||||||
let enabled_extensions: Vec<String> =
|
let enabled_extensions: Vec<String> = self
|
||||||
self.extensions.lock().await.keys().cloned().collect();
|
.extensions
|
||||||
|
.lock()
|
||||||
|
.await
|
||||||
|
.keys()
|
||||||
|
.filter(|name| !is_hidden_extension(name))
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
|
||||||
// Build output string
|
// Build output string
|
||||||
if !disabled_extensions.is_empty() {
|
if !disabled_extensions.is_empty() {
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ pub mod chatrecall;
|
|||||||
pub mod code_execution;
|
pub mod code_execution;
|
||||||
pub mod developer;
|
pub mod developer;
|
||||||
pub mod ext_manager;
|
pub mod ext_manager;
|
||||||
|
pub mod orchestrator;
|
||||||
pub mod summarize;
|
pub mod summarize;
|
||||||
pub mod summon;
|
pub mod summon;
|
||||||
pub mod todo;
|
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",
|
"Analyze code structure with tree-sitter: directory overviews, file details, symbol call graphs",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: true,
|
unprefixed_tools: true,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(analyze::AnalyzeClient::new(ctx).unwrap()),
|
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",
|
"Enable a todo list for goose so it can keep track of what it is doing",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(todo::TodoClient::new(ctx).unwrap()),
|
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.",
|
"Create and manage custom Goose apps through chat. Apps are HTML/CSS/JavaScript and run in sandboxed windows.",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(apps::AppsManagerClient::new(ctx).unwrap()),
|
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",
|
"Search past conversations and load session summaries for contextual memory",
|
||||||
default_enabled: false,
|
default_enabled: false,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(chatrecall::ChatRecallClient::new(ctx).unwrap()),
|
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",
|
"Enable extension management tools for discovering, enabling, and disabling extensions",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(ext_manager::ExtensionManagerClient::new(ctx).unwrap()),
|
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",
|
description: "Load knowledge and delegate tasks to subagents",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: true,
|
unprefixed_tools: true,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(summon::SummonClient::new(ctx).unwrap()),
|
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",
|
description: "Load files/directories and get an LLM summary in a single call",
|
||||||
default_enabled: false,
|
default_enabled: false,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(summarize::SummarizeClient::new(ctx).unwrap()),
|
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",
|
"Goose will make extension calls through code execution, saving tokens",
|
||||||
default_enabled: false,
|
default_enabled: false,
|
||||||
unprefixed_tools: true,
|
unprefixed_tools: true,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| {
|
client_factory: |ctx| {
|
||||||
Box::new(
|
Box::new(
|
||||||
code_execution::CodeExecutionClient::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",
|
description: "Write and edit files, and execute shell commands",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: true,
|
unprefixed_tools: true,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(developer::DeveloperClient::new(ctx).unwrap()),
|
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(
|
map.insert(
|
||||||
tom::EXTENSION_NAME,
|
tom::EXTENSION_NAME,
|
||||||
PlatformExtensionDef {
|
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",
|
"Inject custom context into every turn via GOOSE_MOIM_MESSAGE_TEXT and GOOSE_MOIM_MESSAGE_FILE environment variables",
|
||||||
default_enabled: true,
|
default_enabled: true,
|
||||||
unprefixed_tools: false,
|
unprefixed_tools: false,
|
||||||
|
hidden: false,
|
||||||
client_factory: |ctx| Box::new(tom::TomClient::new(ctx).unwrap()),
|
client_factory: |ctx| Box::new(tom::TomClient::new(ctx).unwrap()),
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
@@ -217,5 +242,7 @@ pub struct PlatformExtensionDef {
|
|||||||
pub default_enabled: bool,
|
pub default_enabled: bool,
|
||||||
/// If true, tools are exposed without extension prefix for intuitive first-class use.
|
/// If true, tools are exposed without extension prefix for intuitive first-class use.
|
||||||
pub unprefixed_tools: bool,
|
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>,
|
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.
|
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.
|
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
|
## 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
|
let content_parts: Vec<String> = msg
|
||||||
.content
|
.content
|
||||||
.iter()
|
.iter()
|
||||||
|
|||||||
@@ -7,9 +7,11 @@ use crate::scheduler_trait::SchedulerTrait;
|
|||||||
use crate::session::SessionManager;
|
use crate::session::SessionManager;
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use lru::LruCache;
|
use lru::LruCache;
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::num::NonZeroUsize;
|
use std::num::NonZeroUsize;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::{OnceCell, RwLock};
|
use tokio::sync::{OnceCell, RwLock};
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
use tracing::{debug, info};
|
use tracing::{debug, info};
|
||||||
|
|
||||||
const DEFAULT_MAX_SESSION: usize = 100;
|
const DEFAULT_MAX_SESSION: usize = 100;
|
||||||
@@ -22,6 +24,7 @@ pub struct AgentManager {
|
|||||||
session_manager: Arc<SessionManager>,
|
session_manager: Arc<SessionManager>,
|
||||||
default_provider: Arc<RwLock<Option<Arc<dyn crate::providers::base::Provider>>>>,
|
default_provider: Arc<RwLock<Option<Arc<dyn crate::providers::base::Provider>>>>,
|
||||||
default_mode: GooseMode,
|
default_mode: GooseMode,
|
||||||
|
cancel_tokens: Arc<RwLock<HashMap<String, CancellationToken>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl AgentManager {
|
impl AgentManager {
|
||||||
@@ -42,6 +45,7 @@ impl AgentManager {
|
|||||||
session_manager,
|
session_manager,
|
||||||
default_provider: Arc::new(RwLock::new(None)),
|
default_provider: Arc::new(RwLock::new(None)),
|
||||||
default_mode,
|
default_mode,
|
||||||
|
cancel_tokens: Arc::new(RwLock::new(HashMap::new())),
|
||||||
};
|
};
|
||||||
|
|
||||||
Ok(manager)
|
Ok(manager)
|
||||||
@@ -151,6 +155,9 @@ impl AgentManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub async fn remove_session(&self, session_id: &str) -> Result<()> {
|
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;
|
let mut sessions = self.sessions.write().await;
|
||||||
sessions
|
sessions
|
||||||
.pop(session_id)
|
.pop(session_id)
|
||||||
@@ -166,6 +173,51 @@ impl AgentManager {
|
|||||||
pub async fn session_count(&self) -> usize {
|
pub async fn session_count(&self) -> usize {
|
||||||
self.sessions.read().await.len()
|
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)]
|
#[cfg(test)]
|
||||||
|
|||||||
Reference in New Issue
Block a user