fix: clean up subagent (#3565)
This commit is contained in:
Wendy Tang
@@ -10,8 +10,6 @@ use crate::{
|
||||
use anyhow::anyhow;
|
||||
use chrono::{DateTime, Utc};
|
||||
use mcp_core::{handler::ToolError, tool::Tool};
|
||||
use rmcp::model::{JsonRpcMessage, JsonRpcNotification, JsonRpcVersion2_0, Notification};
|
||||
use rmcp::object;
|
||||
use serde::{Deserialize, Serialize};
|
||||
// use serde_json::{self};
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
@@ -52,9 +50,7 @@ pub struct SubAgent {
|
||||
impl SubAgent {
|
||||
/// Create a new subagent with the given configuration and provider
|
||||
#[instrument(skip(task_config))]
|
||||
pub async fn new(
|
||||
task_config: TaskConfig,
|
||||
) -> Result<(Arc<Self>, tokio::task::JoinHandle<()>), anyhow::Error> {
|
||||
pub async fn new(task_config: TaskConfig) -> Result<Arc<Self>, anyhow::Error> {
|
||||
debug!("Creating new subagent with id: {}", task_config.id);
|
||||
|
||||
// Create a new extension manager for this subagent
|
||||
@@ -90,21 +86,8 @@ impl SubAgent {
|
||||
extension_manager: Arc::new(RwLock::new(extension_manager)),
|
||||
});
|
||||
|
||||
// Send initial MCP notification
|
||||
let subagent_clone = Arc::clone(&subagent);
|
||||
subagent_clone
|
||||
.send_mcp_notification("subagent_created", "Subagent created and ready")
|
||||
.await;
|
||||
|
||||
// Create a background task handle (for future use with streaming/monitoring)
|
||||
let subagent_clone = Arc::clone(&subagent);
|
||||
let handle = tokio::spawn(async move {
|
||||
// This could be used for background monitoring, cleanup, etc.
|
||||
debug!("Subagent {} background task started", subagent_clone.id);
|
||||
});
|
||||
|
||||
debug!("Subagent {} created successfully", subagent.id);
|
||||
Ok((subagent, handle))
|
||||
Ok(subagent)
|
||||
}
|
||||
|
||||
/// Get the current status of the subagent
|
||||
@@ -119,51 +102,6 @@ impl SubAgent {
|
||||
let mut current_status = self.status.write().await;
|
||||
*current_status = status.clone();
|
||||
} // Write lock is released here!
|
||||
|
||||
// Send MCP notifications based on status
|
||||
match &status {
|
||||
SubAgentStatus::Processing => {
|
||||
self.send_mcp_notification("status_changed", "Processing request")
|
||||
.await;
|
||||
}
|
||||
SubAgentStatus::Completed(msg) => {
|
||||
self.send_mcp_notification("completed", &format!("Completed: {}", msg))
|
||||
.await;
|
||||
}
|
||||
SubAgentStatus::Terminated => {
|
||||
self.send_mcp_notification("terminated", "Subagent terminated")
|
||||
.await;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send an MCP notification about the subagent's activity
|
||||
pub async fn send_mcp_notification(&self, notification_type: &str, message: &str) {
|
||||
let notification = JsonRpcMessage::Notification(JsonRpcNotification {
|
||||
jsonrpc: JsonRpcVersion2_0,
|
||||
notification: Notification {
|
||||
method: "notifications/message".to_string(),
|
||||
params: object!({
|
||||
"level": "info",
|
||||
"logger": format!("subagent_{}", self.id),
|
||||
"data": {
|
||||
"subagent_id": self.id,
|
||||
"type": notification_type,
|
||||
"message": message,
|
||||
"timestamp": Utc::now().to_rfc3339()
|
||||
}
|
||||
}),
|
||||
extensions: Default::default(),
|
||||
},
|
||||
});
|
||||
|
||||
if let Err(e) = self.config.mcp_tx.send(notification).await {
|
||||
error!(
|
||||
"Failed to send MCP notification from subagent {}: {}",
|
||||
self.id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current progress information
|
||||
@@ -192,10 +130,8 @@ impl SubAgent {
|
||||
&self,
|
||||
message: String,
|
||||
task_config: TaskConfig,
|
||||
) -> Result<Message, anyhow::Error> {
|
||||
) -> Result<Vec<Message>, anyhow::Error> {
|
||||
debug!("Processing message for subagent {}", self.id);
|
||||
self.send_mcp_notification("message_processing", &format!("Processing: {}", message))
|
||||
.await;
|
||||
|
||||
// Get provider from task config
|
||||
let provider = self
|
||||
@@ -234,6 +170,7 @@ impl SubAgent {
|
||||
// Generate response from provider with loop for tool processing (max_turns iterations)
|
||||
let mut loop_count = 0;
|
||||
let max_turns = self.config.max_turns.unwrap_or(DEFAULT_SUBAGENT_MAX_TURNS);
|
||||
let mut last_error: Option<anyhow::Error> = None;
|
||||
|
||||
// Generate response from provider
|
||||
loop {
|
||||
@@ -265,18 +202,12 @@ impl SubAgent {
|
||||
// If there are no tool requests, we're done
|
||||
if tool_requests.is_empty() || loop_count >= max_turns {
|
||||
self.add_message(response.clone()).await;
|
||||
messages.push(response.clone());
|
||||
|
||||
// Send notification about response
|
||||
self.send_mcp_notification(
|
||||
"response_generated",
|
||||
&format!("Responded: {}", response.as_concat_text()),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Set status back to ready and return the final response
|
||||
// Set status back to ready
|
||||
self.set_status(SubAgentStatus::Completed("Completed!".to_string()))
|
||||
.await;
|
||||
break Ok(response);
|
||||
break;
|
||||
}
|
||||
|
||||
// Add the assistant message with tool calls to the conversation
|
||||
@@ -285,13 +216,6 @@ impl SubAgent {
|
||||
// Process each tool request and create user response messages
|
||||
for request in &tool_requests {
|
||||
if let Ok(tool_call) = &request.tool_call {
|
||||
// Send notification about tool usage
|
||||
self.send_mcp_notification(
|
||||
"tool_usage",
|
||||
&format!("Using tool: {}", tool_call.name),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Handle platform tools or dispatch to extension manager
|
||||
let tool_result = match self
|
||||
.extension_manager
|
||||
@@ -310,13 +234,6 @@ impl SubAgent {
|
||||
let tool_response_message = Message::user()
|
||||
.with_tool_response(request.id.clone(), Ok(result.clone()));
|
||||
messages.push(tool_response_message);
|
||||
|
||||
// Send notification about tool completion
|
||||
self.send_mcp_notification(
|
||||
"tool_completed",
|
||||
&format!("Tool {} completed successfully", tool_call.name),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
// Create a user message with the tool error
|
||||
@@ -325,13 +242,6 @@ impl SubAgent {
|
||||
Err(ToolError::ExecutionError(e.to_string())),
|
||||
);
|
||||
messages.push(tool_error_message);
|
||||
|
||||
// Send notification about tool error
|
||||
self.send_mcp_notification(
|
||||
"tool_error",
|
||||
&format!("Tool {} error: {}", tool_call.name, e),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -344,24 +254,31 @@ impl SubAgent {
|
||||
"Context length exceeded".to_string(),
|
||||
))
|
||||
.await;
|
||||
break Ok(Message::assistant().with_context_length_exceeded(
|
||||
"The context length of the model has been exceeded. Please start a new session and try again.",
|
||||
));
|
||||
last_error = Some(anyhow::anyhow!("Context length exceeded"));
|
||||
break;
|
||||
}
|
||||
Err(ProviderError::RateLimitExceeded(_)) => {
|
||||
self.set_status(SubAgentStatus::Completed("Rate limit exceeded".to_string()))
|
||||
.await;
|
||||
break Ok(Message::assistant()
|
||||
.with_text("Rate limit exceeded. Please try again later."));
|
||||
last_error = Some(anyhow::anyhow!("Rate limit exceeded"));
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
self.set_status(SubAgentStatus::Completed(format!("Error: {}", e)))
|
||||
.await;
|
||||
error!("Error: {}", e);
|
||||
break Ok(Message::assistant().with_text(format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")));
|
||||
last_error = Some(anyhow::anyhow!("Provider error: {}", e));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Handle error cases or return the last message
|
||||
if let Some(error) = last_error {
|
||||
Err(error)
|
||||
} else {
|
||||
Ok(messages)
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a message to the conversation (for tracking agent responses)
|
||||
|
||||
Reference in New Issue
Block a user