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

fix: restore subagent tool call notifications after summon refactor (#7243)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-02-17 20:10:56 +05:30
committed by GitHub
parent 8533e8d73b
commit f2939483df
3 changed files with 225 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/subagent_handler.rs
+37 -198
View File
@@ -34,78 +34,20 @@ pub struct SubagentPromptContext {
type AgentMessagesFuture =
Pin<Box<dyn Future<Output = Result<(Conversation, Option<String>)>> + Send>>;
pub async fn run_complete_subagent_task(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
return_last_only: bool,
session_id: String,
cancellation_token: Option<CancellationToken>,
) -> Result<String, anyhow::Error> {
run_complete_subagent_task_with_notifications(
config,
recipe,
task_config,
return_last_only,
session_id,
cancellation_token,
None,
)
.await
pub struct SubagentRunParams {
pub config: AgentConfig,
pub recipe: Recipe,
pub task_config: TaskConfig,
pub return_last_only: bool,
pub session_id: String,
pub cancellation_token: Option<CancellationToken>,
pub on_message: Option<OnMessageCallback>,
pub notification_tx: Option<tokio::sync::mpsc::UnboundedSender<ServerNotification>>,
}
pub async fn run_subagent_task_with_callback(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
return_last_only: bool,
session_id: String,
cancellation_token: Option<CancellationToken>,
on_message: Option<OnMessageCallback>,
) -> Result<String, anyhow::Error> {
let (messages, final_output) = get_agent_messages_with_callback(
config,
recipe,
task_config,
session_id,
cancellation_token,
on_message,
)
.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);
}
Ok(extract_response_text(&messages, return_last_only))
}
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| {
pub async fn run_subagent_task(params: SubagentRunParams) -> Result<String, anyhow::Error> {
let return_last_only = params.return_last_only;
let (messages, final_output) = get_agent_messages(params).await.map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to execute task: {}", e),
@@ -177,15 +119,19 @@ fn extract_response_text(messages: &Conversation, return_last_only: bool) -> Str
pub const SUBAGENT_TOOL_REQUEST_TYPE: &str = "subagent_tool_request";
fn get_agent_messages_with_callback(
config: AgentConfig,
recipe: Recipe,
task_config: TaskConfig,
session_id: String,
cancellation_token: Option<CancellationToken>,
on_message: Option<OnMessageCallback>,
) -> AgentMessagesFuture {
fn get_agent_messages(params: SubagentRunParams) -> AgentMessagesFuture {
Box::pin(async move {
let SubagentRunParams {
config,
recipe,
task_config,
session_id,
cancellation_token,
on_message,
notification_tx,
..
} = params;
let system_instructions = recipe.instructions.clone().unwrap_or_default();
let user_task = recipe
.prompt
@@ -248,6 +194,18 @@ fn get_agent_messages_with_callback(
if let Some(ref callback) = on_message {
callback(&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 { .. }) => {}
@@ -267,79 +225,6 @@ fn get_agent_messages_with_callback(
})
}
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();
let user_task = recipe
.prompt
.clone()
.unwrap_or_else(|| "Begin.".to_string());
let agent = Arc::new(Agent::with_config(config));
agent
.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 {
if let Err(e) = agent.add_extension(extension.clone(), &session_id).await {
debug!(
"Failed to add extension '{}' to subagent: {}",
extension.name(),
e
);
}
}
let has_response_schema = recipe.response.is_some();
agent
.apply_recipe_components(recipe.response.clone(), true)
.await;
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);
let mut conversation = Conversation::new_unvalidated(vec![user_message.clone()]);
if let Some(activities) = recipe.activities {
for activity in activities {
info!("Recipe activity: {}", activity);
}
}
let session_config = SessionConfig {
id: session_id.clone(),
schedule_id: None,
max_turns: task_config.max_turns.map(|v| v as u32),
retry_config: recipe.retry,
};
conversation = run_subagent_stream(
agent.clone(),
user_message,
session_config,
cancellation_token,
&session_id,
&notification_tx,
conversation,
)
.await?;
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,
@@ -366,52 +251,6 @@ async fn build_subagent_prompt(
.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
@@ -425,7 +264,7 @@ async fn get_final_output(agent: &Agent, has_response_schema: bool) -> Option<St
}
}
fn create_tool_notification(
pub fn create_tool_notification(
content: &MessageContent,
subagent_id: &str,
) -> Option<ServerNotification> {