tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: display subagent tool calls in CLI and UI (#6535)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-02-04 02:26:59 +05:30
committed by GitHub
parent 7ee634eaa7
commit 1e645c0212
5 changed files with 486 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/subagent_handler.rs
+261 -96
View File
@@ -1,12 +1,18 @@
use crate::{
agents::{subagent_task_config::TaskConfig, Agent, AgentConfig, AgentEvent, SessionConfig},
conversation::{message::Message, Conversation},
conversation::{
message::{Message, MessageContent},
Conversation,
},
prompt_template::render_template,
recipe::Recipe,
};
use anyhow::{anyhow, Result};
use futures::StreamExt;
use rmcp::model::{ErrorCode, ErrorData};
use rmcp::model::{
ErrorCode, ErrorData, LoggingLevel, LoggingMessageNotification,
LoggingMessageNotificationMethod, LoggingMessageNotificationParam, ServerNotification,
};
use serde::Serialize;
use std::future::Future;
use std::pin::Pin;
@@ -15,18 +21,17 @@ use tokio_util::sync::CancellationToken;
use tracing::{debug, info};
#[derive(Serialize)]
struct SubagentPromptContext {
max_turns: usize,
subagent_id: String,
task_instructions: String,
tool_count: usize,
available_tools: String,
pub struct SubagentPromptContext {
pub max_turns: usize,
pub subagent_id: String,
pub task_instructions: String,
pub tool_count: usize,
pub available_tools: String,
}
type AgentMessagesFuture =
Pin<Box<dyn Future<Output = Result<(Conversation, Option<String>)>> + Send>>;
/// Standalone function to run a complete subagent task with output options
pub async fn run_complete_subagent_task(
config: AgentConfig,
recipe: Recipe,
@@ -35,16 +40,43 @@ pub async fn run_complete_subagent_task(
session_id: String,
cancellation_token: Option<CancellationToken>,
) -> Result<String, anyhow::Error> {
let (messages, final_output) =
get_agent_messages(config, recipe, task_config, session_id, cancellation_token)
.await
.map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to execute task: {}", e),
None,
)
})?;
run_complete_subagent_task_with_notifications(
config,
recipe,
task_config,
return_last_only,
session_id,
cancellation_token,
None,
)
.await
}
pub async fn run_complete_subagent_task_with_notifications(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
return_last_only: bool,
session_id: String,
cancellation_token: Option<CancellationToken>,
notification_tx: Option<tokio::sync::mpsc::UnboundedSender<rmcp::model::ServerNotification>>,
) -> Result<String, anyhow::Error> {
let (messages, final_output) = get_agent_messages_with_notifications(
config,
recipe,
task_config,
session_id,
cancellation_token,
notification_tx,
)
.await
.map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to execute task: {}", e),
None,
)
})?;
if let Some(output) = final_output {
return Ok(output);
@@ -67,40 +99,35 @@ pub async fn run_complete_subagent_task(
let all_text_content: Vec<String> = messages
.iter()
.flat_map(|message| {
message.content.iter().filter_map(|content| {
match content {
crate::conversation::message::MessageContent::Text(text_content) => {
Some(text_content.text.clone())
}
crate::conversation::message::MessageContent::ToolResponse(
tool_response,
) => {
// Extract text from tool response
if let Ok(result) = &tool_response.tool_result {
let texts: Vec<String> = result
.content
.iter()
.filter_map(|content| {
if let rmcp::model::RawContent::Text(raw_text_content) =
&content.raw
{
Some(raw_text_content.text.clone())
} else {
None
}
})
.collect();
if !texts.is_empty() {
Some(format!("Tool result: {}", texts.join("\n")))
} else {
None
}
message.content.iter().filter_map(|content| match content {
crate::conversation::message::MessageContent::Text(text_content) => {
Some(text_content.text.clone())
}
crate::conversation::message::MessageContent::ToolResponse(tool_response) => {
if let Ok(result) = &tool_response.tool_result {
let texts: Vec<String> = result
.content
.iter()
.filter_map(|content| {
if let rmcp::model::RawContent::Text(raw_text_content) =
&content.raw
{
Some(raw_text_content.text.clone())
} else {
None
}
})
.collect();
if !texts.is_empty() {
Some(format!("Tool result: {}", texts.join("\n")))
} else {
None
}
} else {
None
}
_ => None,
}
_ => None,
})
})
.collect();
@@ -111,12 +138,15 @@ pub async fn run_complete_subagent_task(
Ok(response_text)
}
fn get_agent_messages(
pub const SUBAGENT_TOOL_REQUEST_TYPE: &str = "subagent_tool_request";
fn get_agent_messages_with_notifications(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
session_id: String,
cancellation_token: Option<CancellationToken>,
notification_tx: Option<tokio::sync::mpsc::UnboundedSender<rmcp::model::ServerNotification>>,
) -> AgentMessagesFuture {
Box::pin(async move {
let system_instructions = recipe.instructions.clone().unwrap_or_default();
@@ -128,11 +158,11 @@ fn get_agent_messages(
let agent = Arc::new(Agent::with_config(config));
agent
.update_provider(task_config.provider, &session_id)
.update_provider(task_config.provider.clone(), &session_id)
.await
.map_err(|e| anyhow!("Failed to set provider on sub agent: {}", e))?;
for extension in task_config.extensions {
for extension in &task_config.extensions {
if let Err(e) = agent.add_extension(extension.clone(), &session_id).await {
debug!(
"Failed to add extension '{}' to subagent: {}",
@@ -147,24 +177,8 @@ fn get_agent_messages(
.apply_recipe_components(recipe.sub_recipes.clone(), recipe.response.clone(), true)
.await;
let tools = agent.list_tools(&session_id, None).await;
let subagent_prompt = render_template(
"subagent_system.md",
&SubagentPromptContext {
max_turns: task_config
.max_turns
.expect("TaskConfig always sets max_turns"),
subagent_id: session_id.clone(),
task_instructions: system_instructions,
tool_count: tools.len(),
available_tools: tools
.iter()
.map(|t| t.name.to_string())
.collect::<Vec<_>>()
.join(", "),
},
)
.map_err(|e| anyhow!("Failed to render subagent system prompt: {}", e))?;
let subagent_prompt =
build_subagent_prompt(&agent, &task_config, &session_id, system_instructions).await?;
agent.override_system_prompt(subagent_prompt).await;
let user_message = Message::user().with_text(user_task);
@@ -182,35 +196,186 @@ fn get_agent_messages(
retry_config: recipe.retry,
};
let mut stream = agent
.reply(user_message, session_config, cancellation_token)
.await
.map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
while let Some(message_result) = stream.next().await {
match message_result {
Ok(AgentEvent::Message(msg)) => conversation.push(msg),
Ok(AgentEvent::McpNotification(_)) | Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(updated_conversation)) => {
conversation = updated_conversation;
}
Err(e) => {
tracing::error!("Error receiving message from subagent: {}", e);
break;
}
}
}
conversation = run_subagent_stream(
agent.clone(),
user_message,
session_config,
cancellation_token,
&session_id,
&notification_tx,
conversation,
)
.await?;
let final_output = if has_response_schema {
agent
.final_output_tool
.lock()
.await
.as_ref()
.and_then(|tool| tool.final_output.clone())
} else {
None
};
let final_output = get_final_output(&agent, has_response_schema).await;
Ok((conversation, final_output))
})
}
async fn build_subagent_prompt(
agent: &Agent,
task_config: &TaskConfig,
session_id: &str,
system_instructions: String,
) -> Result<String> {
let tools = agent.list_tools(session_id, None).await;
render_template(
"subagent_system.md",
&SubagentPromptContext {
max_turns: task_config
.max_turns
.expect("TaskConfig always sets max_turns"),
subagent_id: session_id.to_string(),
task_instructions: system_instructions,
tool_count: tools.len(),
available_tools: tools
.iter()
.map(|t| t.name.to_string())
.collect::<Vec<_>>()
.join(", "),
},
)
.map_err(|e| anyhow!("Failed to render subagent system prompt: {}", e))
}
async fn run_subagent_stream(
agent: Arc<Agent>,
user_message: Message,
session_config: SessionConfig,
cancellation_token: Option<CancellationToken>,
session_id: &str,
notification_tx: &Option<tokio::sync::mpsc::UnboundedSender<rmcp::model::ServerNotification>>,
mut conversation: Conversation,
) -> Result<Conversation> {
let mut stream = crate::session_context::with_session_id(Some(session_id.to_string()), async {
agent
.reply(user_message, session_config, cancellation_token)
.await
})
.await
.map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
while let Some(message_result) = stream.next().await {
match message_result {
Ok(AgentEvent::Message(msg)) => {
if let Some(ref tx) = notification_tx {
for content in &msg.content {
if let Some(notif) = create_tool_notification(content, session_id) {
if tx.send(notif).is_err() {
debug!("Notification receiver dropped for subagent {}", session_id);
}
}
}
}
conversation.push(msg);
}
Ok(AgentEvent::McpNotification(_)) => {}
Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(updated_conversation)) => {
conversation = updated_conversation;
}
Err(e) => {
tracing::error!("Error receiving message from subagent: {}", e);
break;
}
}
}
Ok(conversation)
}
async fn get_final_output(agent: &Agent, has_response_schema: bool) -> Option<String> {
if has_response_schema {
agent
.final_output_tool
.lock()
.await
.as_ref()
.and_then(|tool| tool.final_output.clone())
} else {
None
}
}
fn create_tool_notification(
content: &MessageContent,
subagent_id: &str,
) -> Option<ServerNotification> {
if let MessageContent::ToolRequest(req) = content {
let tool_call = req.tool_call.as_ref().ok()?;
Some(ServerNotification::LoggingMessageNotification(
LoggingMessageNotification {
method: LoggingMessageNotificationMethod,
params: LoggingMessageNotificationParam {
level: LoggingLevel::Info,
logger: Some(format!("subagent:{}", subagent_id)),
data: serde_json::json!({
"type": SUBAGENT_TOOL_REQUEST_TYPE,
"subagent_id": subagent_id,
"tool_call": {
"name": tool_call.name,
"arguments": tool_call.arguments
}
}),
},
extensions: Default::default(),
},
))
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::{create_tool_notification, SUBAGENT_TOOL_REQUEST_TYPE};
use crate::conversation::message::MessageContent;
use rmcp::model::{CallToolRequestParams, ServerNotification};
use serde_json::json;
#[test]
fn create_tool_notification_for_tool_request() {
let tool_call = CallToolRequestParams {
meta: None,
task: None,
name: "developer__shell".to_string().into(),
arguments: Some(json!({"command": "ls"}).as_object().unwrap().clone()),
};
let content = MessageContent::tool_request("req1", Ok(tool_call));
let notification =
create_tool_notification(&content, "session_1").expect("expected notification");
let ServerNotification::LoggingMessageNotification(log_notif) = notification else {
panic!("expected logging notification");
};
let data = log_notif
.params
.data
.as_object()
.expect("expected object data");
assert_eq!(
data.get("type").and_then(|v| v.as_str()),
Some(SUBAGENT_TOOL_REQUEST_TYPE)
);
assert_eq!(
data.get("subagent_id").and_then(|v| v.as_str()),
Some("session_1")
);
let tool_call = data
.get("tool_call")
.and_then(|v| v.as_object())
.expect("expected tool_call object");
assert_eq!(
tool_call.get("name").and_then(|v| v.as_str()),
Some("developer__shell")
);
}
#[test]
fn create_tool_notification_ignores_non_tool_request() {
let content = MessageContent::text("hello");
assert!(create_tool_notification(&content, "session_1").is_none());
}
}
crates/goose/src/agents/subagent_tool.rs
+15 -8
View File
@@ -4,12 +4,14 @@ use std::path::PathBuf;
use anyhow::{anyhow, Result};
use futures::FutureExt;
use rmcp::model::{Content, ErrorCode, ErrorData, Tool};
use rmcp::model::{Content, ErrorCode, ErrorData, ServerNotification, Tool};
use serde::Deserialize;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_util::sync::CancellationToken;
use crate::agents::subagent_handler::run_complete_subagent_task;
use crate::agents::subagent_handler::run_complete_subagent_task_with_notifications;
use crate::agents::subagent_task_config::TaskConfig;
use crate::agents::tool_execution::ToolCallResult;
use crate::agents::AgentConfig;
@@ -31,7 +33,7 @@ Make sure your last message provides a comprehensive summary of:
Be concise but complete.
"#;
#[derive(Debug, Deserialize)]
#[derive(Debug, Deserialize, Clone)]
pub struct SubagentParams {
pub instructions: Option<String>,
pub subrecipe: Option<String>,
@@ -46,7 +48,7 @@ fn default_summary() -> bool {
true
}
#[derive(Debug, Deserialize)]
#[derive(Debug, Deserialize, Clone)]
pub struct SubagentSettings {
pub provider: Option<String>,
pub model: Option<String>,
@@ -224,29 +226,33 @@ pub fn handle_subagent_tool(
};
let config = config.clone();
let (notification_tx, notification_rx) = mpsc::unbounded_channel();
ToolCallResult {
notification_stream: None,
notification_stream: Some(Box::new(UnboundedReceiverStream::new(notification_rx))),
result: Box::new(
execute_subagent(
execute_subagent_with_notifications(
config,
recipe,
task_config,
parsed_params,
working_dir,
cancellation_token,
notification_tx,
)
.boxed(),
),
}
}
async fn execute_subagent(
async fn execute_subagent_with_notifications(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
params: SubagentParams,
working_dir: PathBuf,
cancellation_token: Option<CancellationToken>,
notification_tx: mpsc::UnboundedSender<ServerNotification>,
) -> Result<rmcp::model::CallToolResult, ErrorData> {
let session = config
.session_manager
@@ -270,13 +276,14 @@ async fn execute_subagent(
data: None,
})?;
let result = run_complete_subagent_task(
let result = run_complete_subagent_task_with_notifications(
config,
recipe,
task_config,
params.summary,
session.id,
cancellation_token,
Some(notification_tx),
)
.await;