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

feat: Handle MCP server notification messages (#2613)

Browse Source 
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Jack Amadeo
2025-05-30 11:50:14 -04:00
committed by GitHub
parent eeb61ace22
commit 03e5549b54
40 changed files with 1186 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+127 -46
View File
@@ -1,9 +1,14 @@
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use anyhow::{anyhow, Result};
use futures::stream::BoxStream;
use futures::TryStreamExt;
use futures::{FutureExt, Stream, TryStreamExt};
use futures_util::stream;
use futures_util::stream::StreamExt;
use mcp_core::protocol::JsonRpcMessage;
use crate::config::{Config, ExtensionConfigManager, PermissionManager};
use crate::message::Message;
@@ -39,7 +44,7 @@ use mcp_core::{
use super::platform_tools;
use super::router_tools;
use super::tool_execution::{ToolFuture, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
/// The main goose Agent
pub struct Agent {
@@ -56,6 +61,12 @@ pub struct Agent {
pub(super) router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
}
#[derive(Clone, Debug)]
pub enum AgentEvent {
Message(Message),
McpNotification((String, JsonRpcMessage)),
}
impl Agent {
pub fn new() -> Self {
// Create channels with buffer size 32 (adjust if needed)
@@ -100,6 +111,40 @@ impl Default for Agent {
}
}
pub enum ToolStreamItem<T> {
Message(JsonRpcMessage),
Result(T),
}
pub type ToolStream = Pin<Box<dyn Stream<Item = ToolStreamItem<ToolResult<Vec<Content>>>> + Send>>;
// tool_stream combines a stream of JsonRpcMessages with a future representing the
// final result of the tool call. MCP notifications are not request-scoped, but
// this lets us capture all notifications emitted during the tool call for
// simpler consumption
pub fn tool_stream<S, F>(rx: S, done: F) -> ToolStream
where
S: Stream<Item = JsonRpcMessage> + Send + Unpin + 'static,
F: Future<Output = ToolResult<Vec<Content>>> + Send + 'static,
{
Box::pin(async_stream::stream! {
tokio::pin!(done);
let mut rx = rx;
loop {
tokio::select! {
Some(msg) = rx.next() => {
yield ToolStreamItem::Message(msg);
}
r = &mut done => {
yield ToolStreamItem::Result(r);
break;
}
}
}
})
}
impl Agent {
/// Get a reference count clone to the provider
pub async fn provider(&self) -> Result<Arc<dyn Provider>, anyhow::Error> {
@@ -143,7 +188,7 @@ impl Agent {
&self,
tool_call: mcp_core::tool::ToolCall,
request_id: String,
) -> (String, Result<Vec<Content>, ToolError>) {
) -> (String, Result<ToolCallResult, ToolError>) {
// Check if this tool call should be allowed based on repetition monitoring
if let Some(monitor) = self.tool_monitor.lock().await.as_mut() {
let tool_call_info = ToolCall::new(tool_call.name.clone(), tool_call.arguments.clone());
@@ -171,52 +216,65 @@ impl Agent {
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
return self
let (request_id, result) = self
.manage_extensions(action, extension_name, request_id)
.await;
return (request_id, Ok(ToolCallResult::from(result)));
}
let extension_manager = self.extension_manager.lock().await;
let result = if tool_call.name == PLATFORM_READ_RESOURCE_TOOL_NAME {
let result: ToolCallResult = if tool_call.name == PLATFORM_READ_RESOURCE_TOOL_NAME {
// Check if the tool is read_resource and handle it separately
extension_manager
.read_resource(tool_call.arguments.clone())
.await
ToolCallResult::from(
extension_manager
.read_resource(tool_call.arguments.clone())
.await,
)
} else if tool_call.name == PLATFORM_LIST_RESOURCES_TOOL_NAME {
extension_manager
.list_resources(tool_call.arguments.clone())
.await
ToolCallResult::from(
extension_manager
.list_resources(tool_call.arguments.clone())
.await,
)
} else if tool_call.name == PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME {
extension_manager.search_available_extensions().await
ToolCallResult::from(extension_manager.search_available_extensions().await)
} else if self.is_frontend_tool(&tool_call.name).await {
// For frontend tools, return an error indicating we need frontend execution
Err(ToolError::ExecutionError(
ToolCallResult::from(Err(ToolError::ExecutionError(
"Frontend tool execution required".to_string(),
))
)))
} else if tool_call.name == ROUTER_VECTOR_SEARCH_TOOL_NAME {
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
ToolCallResult::from(if let Some(selector) = selector {
selector.select_tools(tool_call.arguments.clone()).await
} else {
Err(ToolError::ExecutionError(
"Encountered vector search error.".to_string(),
))
}
})
} else {
extension_manager
// Clone the result to ensure no references to extension_manager are returned
let result = extension_manager
.dispatch_tool_call(tool_call.clone())
.await
.await;
match result {
Ok(call_result) => call_result,
Err(e) => ToolCallResult::from(Err(ToolError::ExecutionError(e.to_string()))),
}
};
debug!(
"input" = serde_json::to_string(&tool_call).unwrap(),
"output" = serde_json::to_string(&result).unwrap(),
);
// Process the response to handle large text content
let processed_result = super::large_response_handler::process_tool_response(result);
(request_id, processed_result)
(
request_id,
Ok(ToolCallResult {
notification_stream: result.notification_stream,
result: Box::new(
result
.result
.map(super::large_response_handler::process_tool_response),
),
}),
)
}
pub(super) async fn manage_extensions(
@@ -466,7 +524,7 @@ impl Agent {
&self,
messages: &[Message],
session: Option<SessionConfig>,
) -> anyhow::Result<BoxStream<'_, anyhow::Result<Message>>> {
) -> anyhow::Result<BoxStream<'_, anyhow::Result<AgentEvent>>> {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
@@ -532,9 +590,8 @@ impl Agent {
}
}
}
// Yield the assistant's response with frontend tool requests filtered out
yield filtered_response.clone();
yield AgentEvent::Message(filtered_response.clone());
tokio::task::yield_now().await;
@@ -556,7 +613,7 @@ impl Agent {
// execution is yeield back to this reply loop, and is of the same Message
// type, so we can yield that back up to be handled
while let Some(msg) = frontend_tool_stream.try_next().await? {
yield msg;
yield AgentEvent::Message(msg);
}
// Clone goose_mode once before the match to avoid move issues
@@ -584,13 +641,23 @@ impl Agent {
self.provider().await?).await;
// Handle pre-approved and read-only tools in parallel
let mut tool_futures: Vec<ToolFuture> = Vec::new();
let mut tool_futures: Vec<(String, ToolStream)> = Vec::new();
// Skip the confirmation for approved tools
for request in &permission_check_result.approved {
if let Ok(tool_call) = request.tool_call.clone() {
let tool_future = self.dispatch_tool_call(tool_call, request.id.clone());
tool_futures.push(Box::pin(tool_future));
let (req_id, tool_result) = self.dispatch_tool_call(tool_call, request.id.clone()).await;
tool_futures.push((req_id, match tool_result {
Ok(result) => tool_stream(
result.notification_stream.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => tool_stream(
Box::new(stream::empty()),
futures::future::ready(Err(e)),
),
}));
}
}
@@ -618,7 +685,7 @@ impl Agent {
// type, so we can yield the Message back up to be handled and grab any
// confirmations or denials
while let Some(msg) = tool_approval_stream.try_next().await? {
yield msg;
yield AgentEvent::Message(msg);
}
tool_futures = {
@@ -628,16 +695,30 @@ impl Agent {
futures_lock.drain(..).collect::<Vec<_>>()
};
// Wait for all tool calls to complete
let results = futures::future::join_all(tool_futures).await;
let with_id = tool_futures
.into_iter()
.map(|(request_id, stream)| {
stream.map(move |item| (request_id.clone(), item))
})
.collect::<Vec<_>>();
let mut combined = stream::select_all(with_id);
let mut all_install_successful = true;
for (request_id, output) in results.into_iter() {
if enable_extension_request_ids.contains(&request_id) && output.is_err(){
all_install_successful = false;
while let Some((request_id, item)) = combined.next().await {
match item {
ToolStreamItem::Result(output) => {
if enable_extension_request_ids.contains(&request_id) && output.is_err(){
all_install_successful = false;
}
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(request_id, output);
},
ToolStreamItem::Message(msg) => {
yield AgentEvent::McpNotification((request_id, msg))
}
}
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(request_id, output);
}
// Update system prompt and tools if installations were successful
@@ -647,7 +728,7 @@ impl Agent {
}
let final_message_tool_resp = message_tool_response.lock().await.clone();
yield final_message_tool_resp.clone();
yield AgentEvent::Message(final_message_tool_resp.clone());
messages.push(response);
messages.push(final_message_tool_resp);
@@ -656,15 +737,15 @@ impl Agent {
// At this point, the last message should be a user message
// because call to provider led to context length exceeded error
// Immediately yield a special message and break
yield Message::assistant().with_context_length_exceeded(
yield AgentEvent::Message(Message::assistant().with_context_length_exceeded(
"The context length of the model has been exceeded. Please start a new session and try again.",
);
));
break;
},
Err(e) => {
// Create an error message & terminate the stream
error!("Error: {}", e);
yield Message::assistant().with_text(format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error."));
yield AgentEvent::Message(Message::assistant().with_text(format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")));
break;
}
}
crates/goose/src/agents/extension_manager.rs
+68 -43
View File
@@ -1,8 +1,7 @@
use anyhow::Result;
use chrono::{DateTime, TimeZone, Utc};
use futures::future;
use futures::stream::{FuturesUnordered, StreamExt};
use mcp_client::McpService;
use futures::{future, FutureExt};
use mcp_core::protocol::GetPromptResult;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
@@ -10,15 +9,17 @@ use std::sync::LazyLock;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::task;
use tracing::{debug, error, warn};
use tokio_stream::wrappers::ReceiverStream;
use tracing::{error, warn};
use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, ToolInfo};
use super::tool_execution::ToolCallResult;
use crate::agents::extension::Envs;
use crate::config::{Config, ExtensionConfigManager};
use crate::prompt_template;
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport};
use mcp_core::{prompt::Prompt, Content, Tool, ToolCall, ToolError, ToolResult};
use mcp_core::{prompt::Prompt, Content, Tool, ToolCall, ToolError};
use serde_json::Value;
// By default, we set it to Jan 1, 2020 if the resource does not have a timestamp
@@ -113,7 +114,8 @@ impl ExtensionManager {
/// Add a new MCP extension based on the provided client type
// TODO IMPORTANT need to ensure this times out if the extension command is broken!
pub async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()> {
let sanitized_name = normalize(config.key().to_string());
let config_name = config.key().to_string();
let sanitized_name = normalize(config_name.clone());
/// Helper function to merge environment variables from direct envs and keychain-stored env_keys
async fn merge_environments(
@@ -183,13 +185,15 @@ impl ExtensionManager {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = SseTransport::new(uri, all_envs);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
);
Box::new(McpClient::new(service))
Box::new(
McpClient::connect(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
)
.await?,
)
}
ExtensionConfig::Stdio {
cmd,
@@ -202,13 +206,15 @@ impl ExtensionManager {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = StdioTransport::new(cmd, args.to_vec(), all_envs);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
);
Box::new(McpClient::new(service))
Box::new(
McpClient::connect(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
)
.await?,
)
}
ExtensionConfig::Builtin {
name,
@@ -227,13 +233,15 @@ impl ExtensionManager {
HashMap::new(),
);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
);
Box::new(McpClient::new(service))
Box::new(
McpClient::connect(
handle,
Duration::from_secs(
timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
),
)
.await?,
)
}
_ => unreachable!(),
};
@@ -609,7 +617,7 @@ impl ExtensionManager {
}
}
pub async fn dispatch_tool_call(&self, tool_call: ToolCall) -> ToolResult<Vec<Content>> {
pub async fn dispatch_tool_call(&self, tool_call: ToolCall) -> Result<ToolCallResult> {
// Dispatch tool call based on the prefix naming convention
let (client_name, client) = self
.get_client_for_tool(&tool_call.name)
@@ -620,22 +628,26 @@ impl ExtensionManager {
.name
.strip_prefix(client_name)
.and_then(|s| s.strip_prefix("__"))
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?
.to_string();
let client_guard = client.lock().await;
let arguments = tool_call.arguments.clone();
let client = client.clone();
let notifications_receiver = client.lock().await.subscribe().await;
let result = client_guard
.call_tool(tool_name, tool_call.clone().arguments)
.await
.map(|result| result.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()));
let fut = async move {
let client_guard = client.lock().await;
client_guard
.call_tool(&tool_name, arguments)
.await
.map(|call| call.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()))
};
debug!(
"input" = serde_json::to_string(&tool_call).unwrap(),
"output" = serde_json::to_string(&result).unwrap(),
);
result
Ok(ToolCallResult {
result: Box::new(fut.boxed()),
notification_stream: Some(Box::new(ReceiverStream::new(notifications_receiver))),
})
}
pub async fn list_prompts_from_extension(
@@ -793,10 +805,11 @@ mod tests {
use mcp_client::client::Error;
use mcp_client::client::McpClientTrait;
use mcp_core::protocol::{
CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult, ListResourcesResult,
ListToolsResult, ReadResourceResult,
CallToolResult, GetPromptResult, InitializeResult, JsonRpcMessage, ListPromptsResult,
ListResourcesResult, ListToolsResult, ReadResourceResult,
};
use serde_json::json;
use tokio::sync::mpsc;
struct MockClient {}
@@ -849,6 +862,10 @@ mod tests {
) -> Result<GetPromptResult, Error> {
Err(Error::NotInitialized)
}
async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage> {
mpsc::channel(1).1
}
}
#[test]
@@ -970,6 +987,9 @@ mod tests {
let result = extension_manager
.dispatch_tool_call(invalid_tool_call)
.await
.unwrap()
.result
.await;
assert!(matches!(
result.err().unwrap(),
@@ -986,6 +1006,11 @@ mod tests {
let result = extension_manager
.dispatch_tool_call(invalid_tool_call)
.await;
assert!(matches!(result.err().unwrap(), ToolError::NotFound(_)));
if let Err(err) = result {
let tool_err = err.downcast_ref::<ToolError>().expect("Expected ToolError");
assert!(matches!(tool_err, ToolError::NotFound(_)));
} else {
panic!("Expected ToolError::NotFound");
}
}
}
crates/goose/src/agents/mod.rs
+1 -1
View File
@@ -13,7 +13,7 @@ mod tool_router_index_manager;
pub(crate) mod tool_vectordb;
mod types;
pub use agent::Agent;
pub use agent::{Agent, AgentEvent};
pub use extension::ExtensionConfig;
pub use extension_manager::ExtensionManager;
pub use prompt_manager::PromptManager;
crates/goose/src/agents/tool_execution.rs
+33 -11
View File
@@ -1,23 +1,35 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use async_stream::try_stream;
use futures::stream::BoxStream;
use futures::StreamExt;
use futures::stream::{self, BoxStream};
use futures::{Stream, StreamExt};
use mcp_core::protocol::JsonRpcMessage;
use tokio::sync::Mutex;
use crate::config::permission::PermissionLevel;
use crate::config::PermissionManager;
use crate::message::{Message, ToolRequest};
use crate::permission::Permission;
use mcp_core::{Content, ToolError};
use mcp_core::{Content, ToolResult};
// Type alias for ToolFutures - used in the agent loop to join all futures together
pub(crate) type ToolFuture<'a> =
Pin<Box<dyn Future<Output = (String, Result<Vec<Content>, ToolError>)> + Send + 'a>>;
pub(crate) type ToolFuturesVec<'a> = Arc<Mutex<Vec<ToolFuture<'a>>>>;
// ToolCallResult combines the result of a tool call with an optional notification stream that
// can be used to receive notifications from the tool.
pub struct ToolCallResult {
pub result: Box<dyn Future<Output = ToolResult<Vec<Content>>> + Send + Unpin>,
pub notification_stream: Option<Box<dyn Stream<Item = JsonRpcMessage> + Send + Unpin>>,
}
impl From<ToolResult<Vec<Content>>> for ToolCallResult {
fn from(result: ToolResult<Vec<Content>>) -> Self {
Self {
result: Box::new(futures::future::ready(result)),
notification_stream: None,
}
}
}
use super::agent::{tool_stream, ToolStream};
use crate::agents::Agent;
pub const DECLINED_RESPONSE: &str = "The user has declined to run this tool. \
@@ -37,7 +49,7 @@ impl Agent {
pub(crate) fn handle_approval_tool_requests<'a>(
&'a self,
tool_requests: &'a [ToolRequest],
tool_futures: ToolFuturesVec<'a>,
tool_futures: Arc<Mutex<Vec<(String, ToolStream)>>>,
permission_manager: &'a mut PermissionManager,
message_tool_response: Arc<Mutex<Message>>,
) -> BoxStream<'a, anyhow::Result<Message>> {
@@ -56,9 +68,19 @@ impl Agent {
while let Some((req_id, confirmation)) = rx.recv().await {
if req_id == request.id {
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
let tool_future = self.dispatch_tool_call(tool_call.clone(), request.id.clone());
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone()).await;
let mut futures = tool_futures.lock().await;
futures.push(Box::pin(tool_future));
futures.push((req_id, match tool_result {
Ok(result) => tool_stream(
result.notification_stream.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => tool_stream(
Box::new(stream::empty()),
futures::future::ready(Err(e)),
),
}));
if confirmation.permission == Permission::AlwaysAllow {
permission_manager.update_user_permission(&tool_call.name, PermissionLevel::AlwaysAllow);
crates/goose/src/scheduler.rs
+5 -1
View File
@@ -11,6 +11,7 @@ use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use tokio_cron_scheduler::{job::JobId, Job, JobScheduler as TokioJobScheduler};
use crate::agents::AgentEvent;
use crate::agents::{Agent, SessionConfig};
use crate::config::{self, Config};
use crate::message::Message;
@@ -1102,12 +1103,15 @@ async fn run_scheduled_job_internal(
tokio::task::yield_now().await;
match message_result {
Ok(msg) => {
Ok(AgentEvent::Message(msg)) => {
if msg.role == mcp_core::role::Role::Assistant {
tracing::info!("[Job {}] Assistant: {:?}", job.id, msg.content);
}
all_session_messages.push(msg);
}
Ok(AgentEvent::McpNotification(_)) => {
// Handle notifications if needed
}
Err(e) => {
tracing::error!(
"[Job {}] Error receiving message from agent: {}",