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

chore: use typed notifications from rmcp (#3653)

Browse Source
This commit is contained in:
Jack Amadeo
2025-07-25 14:04:18 -04:00
committed by GitHub
parent 31a5f9cbbc
commit 0ef38c6658
17 changed files with 198 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+98 -108
View File
@@ -36,7 +36,8 @@ use goose::providers::pricing::initialize_pricing_cache;
use goose::session; use goose::session;
use input::InputResult; use input::InputResult;
use mcp_core::handler::ToolError; use mcp_core::handler::ToolError;
use rmcp::model::{JsonRpcMessage, JsonRpcNotification, Notification, PromptMessage}; use rmcp::model::PromptMessage;
use rmcp::model::ServerNotification;
use rand::{distributions::Alphanumeric, Rng}; use rand::{distributions::Alphanumeric, Rng};
use rustyline::EditMode; use rustyline::EditMode;
@@ -1023,126 +1024,115 @@ impl Session {
} }
} }
Some(Ok(AgentEvent::McpNotification((_id, message)))) => { Some(Ok(AgentEvent::McpNotification((_id, message)))) => {
if let JsonRpcMessage::Notification( JsonRpcNotification { match &message {
notification: Notification { ServerNotification::LoggingMessageNotification(notification) => {
method, let data = &notification.params.data;
params: o,.. let (formatted_message, subagent_id, message_notification_type) = match data {
},.. Value::String(s) => (s.clone(), None, None),
}) = message { Value::Object(o) => {
match method.as_str() { // Check for subagent notification structure first
"notifications/message" => { if let Some(Value::String(msg)) = o.get("message") {
let data = o.get("data").unwrap_or(&Value::Null); // Extract subagent info for better display
let (formatted_message, subagent_id, message_notification_type) = match data { let subagent_id = o.get("subagent_id")
Value::String(s) => (s.clone(), None, None), .and_then(|v| v.as_str())
Value::Object(o) => { .unwrap_or("unknown");
// Check for subagent notification structure first let notification_type = o.get("type")
if let Some(Value::String(msg)) = o.get("message") { .and_then(|v| v.as_str())
// Extract subagent info for better display .unwrap_or("");
let subagent_id = o.get("subagent_id")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let notification_type = o.get("type")
.and_then(|v| v.as_str())
.unwrap_or("");
let formatted = match notification_type { let formatted = match notification_type {
"subagent_created" | "completed" | "terminated" => { "subagent_created" | "completed" | "terminated" => {
format!("🤖 {}", msg) format!("🤖 {}", msg)
} }
"tool_usage" | "tool_completed" | "tool_error" => { "tool_usage" | "tool_completed" | "tool_error" => {
format!("🔧 {}", msg) format!("🔧 {}", msg)
} }
"message_processing" | "turn_progress" => { "message_processing" | "turn_progress" => {
format!("💭 {}", msg) format!("💭 {}", msg)
} }
"response_generated" => { "response_generated" => {
// Check verbosity setting for subagent response content // Check verbosity setting for subagent response content
let config = Config::global(); let config = Config::global();
let min_priority = config let min_priority = config
.get_param::<f32>("GOOSE_CLI_MIN_PRIORITY") .get_param::<f32>("GOOSE_CLI_MIN_PRIORITY")
.ok() .ok()
.unwrap_or(0.5); .unwrap_or(0.5);
if min_priority > 0.1 && !self.debug { if min_priority > 0.1 && !self.debug {
// High/Medium verbosity: show truncated response // High/Medium verbosity: show truncated response
if let Some(response_content) = msg.strip_prefix("Responded: ") { if let Some(response_content) = msg.strip_prefix("Responded: ") {
format!("🤖 Responded: {}", safe_truncate(response_content, 100)) format!("🤖 Responded: {}", safe_truncate(response_content, 100))
} else {
format!("🤖 {}", msg)
}
} else { } else {
// All verbosity or debug: show full response
format!("🤖 {}", msg) format!("🤖 {}", msg)
} }
} else {
// All verbosity or debug: show full response
format!("🤖 {}", msg)
} }
_ => { }
msg.to_string() _ => {
} msg.to_string()
}; }
(formatted, Some(subagent_id.to_string()), Some(notification_type.to_string())) };
} else if let Some(Value::String(output)) = o.get("output") { (formatted, Some(subagent_id.to_string()), Some(notification_type.to_string()))
// Fallback for other MCP notification types } else if let Some(Value::String(output)) = o.get("output") {
(output.to_owned(), None, None) // Fallback for other MCP notification types
} else if let Some(result) = format_task_execution_notification(data) { (output.to_owned(), None, None)
result } else if let Some(result) = format_task_execution_notification(data) {
} else { result
(data.to_string(), None, None) } else {
} (data.to_string(), None, None)
}, }
v => { },
(v.to_string(), None, None) v => {
}, (v.to_string(), None, None)
}; },
};
// Handle subagent notifications - show immediately // Handle subagent notifications - show immediately
if let Some(_id) = subagent_id { if let Some(_id) = subagent_id {
// TODO: proper display for subagent notifications // TODO: proper display for subagent notifications
if interactive {
let _ = progress_bars.hide();
println!("{}", console::style(&formatted_message).green().dim());
} else {
progress_bars.log(&formatted_message);
}
} else if let Some(ref notification_type) = message_notification_type {
if notification_type == TASK_EXECUTION_NOTIFICATION_TYPE {
if interactive { if interactive {
let _ = progress_bars.hide(); let _ = progress_bars.hide();
println!("{}", console::style(&formatted_message).green().dim()); print!("{}", formatted_message);
std::io::stdout().flush().unwrap();
} else { } else {
progress_bars.log(&formatted_message); print!("{}", formatted_message);
} std::io::stdout().flush().unwrap();
} else if let Some(ref notification_type) = message_notification_type {
if notification_type == TASK_EXECUTION_NOTIFICATION_TYPE {
if interactive {
let _ = progress_bars.hide();
print!("{}", formatted_message);
std::io::stdout().flush().unwrap();
} else {
print!("{}", formatted_message);
std::io::stdout().flush().unwrap();
}
} }
} }
else { }
// Non-subagent notification, display immediately with compact spacing else {
if interactive { // Non-subagent notification, display immediately with compact spacing
let _ = progress_bars.hide(); if interactive {
println!("{}", console::style(&formatted_message).green().dim()); let _ = progress_bars.hide();
} else { println!("{}", console::style(&formatted_message).green().dim());
progress_bars.log(&formatted_message); } else {
} progress_bars.log(&formatted_message);
} }
}, }
"notifications/progress" => { },
let progress = o.get("progress").and_then(|v| v.as_f64()); ServerNotification::ProgressNotification(notification) => {
let token = o.get("progressToken").map(|v| v.to_string()); let progress = notification.params.progress;
let message = o.get("message").and_then(|v| v.as_str()); let text = notification.params.message.as_deref();
let total = o let total = notification.params.total;
.get("total") let token = &notification.params.progress_token;
.and_then(|v| v.as_f64()); progress_bars.update(
if let (Some(progress), Some(token)) = (progress, token) { &token.0.to_string(),
progress_bars.update( progress,
token.as_str(), total,
progress, text,
total, );
message, },
); _ => (),
}
},
_ => (),
}
} }
} }
Some(Ok(AgentEvent::ModelChange { model, mode })) => { Some(Ok(AgentEvent::ModelChange { model, mode })) => {
crates/goose-cli/src/session/output.rs
+3 -3
View File
@@ -799,11 +799,11 @@ impl McpSpinners {
spinner.set_message(message.to_string()); spinner.set_message(message.to_string());
} }
pub fn update(&mut self, token: &str, value: f64, total: Option<f64>, message: Option<&str>) { pub fn update(&mut self, token: &str, value: u32, total: Option<u32>, message: Option<&str>) {
let bar = self.bars.entry(token.to_string()).or_insert_with(|| { let bar = self.bars.entry(token.to_string()).or_insert_with(|| {
if let Some(total) = total { if let Some(total) = total {
self.multi_bar.add( self.multi_bar.add(
ProgressBar::new((total * 100.0) as u64).with_style( ProgressBar::new((total * 100) as u64).with_style(
ProgressStyle::with_template("[{elapsed}] {bar:40} {pos:>3}/{len:3} {msg}") ProgressStyle::with_template("[{elapsed}] {bar:40} {pos:>3}/{len:3} {msg}")
.unwrap(), .unwrap(),
), ),
@@ -812,7 +812,7 @@ impl McpSpinners {
self.multi_bar.add(ProgressBar::new_spinner()) self.multi_bar.add(ProgressBar::new_spinner())
} }
}); });
bar.set_position((value * 100.0) as u64); bar.set_position((value * 100) as u64);
if let Some(msg) = message { if let Some(msg) = message {
bar.set_message(msg.to_string()); bar.set_message(msg.to_string());
} }
crates/goose-mcp/src/developer/mod.rs
+2
View File
@@ -672,6 +672,7 @@ impl DeveloperRouter {
notification: Notification { notification: Notification {
method: "notifications/message".to_string(), method: "notifications/message".to_string(),
params: object!({ params: object!({
"level": "info",
"data": { "data": {
"type": "shell", "type": "shell",
"stream": "stdout", "stream": "stdout",
@@ -698,6 +699,7 @@ impl DeveloperRouter {
notification: Notification { notification: Notification {
method: "notifications/message".to_string(), method: "notifications/message".to_string(),
params: object!({ params: object!({
"level": "info",
"data": { "data": {
"type": "shell", "type": "shell",
"stream": "stderr", "stream": "stderr",
crates/goose-server/src/routes/reply.rs
+2 -2
View File
@@ -19,7 +19,7 @@ use goose::{
session, session,
}; };
use mcp_core::ToolResult; use mcp_core::ToolResult;
use rmcp::model::{Content, JsonRpcMessage}; use rmcp::model::{Content, ServerNotification};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::json; use serde_json::json;
use serde_json::Value; use serde_json::Value;
@@ -97,7 +97,7 @@ enum MessageEvent {
}, },
Notification { Notification {
request_id: String, request_id: String,
message: JsonRpcMessage, message: ServerNotification,
}, },
} }
crates/goose/src/agents/agent.rs
+5 -6
View File
@@ -45,8 +45,7 @@ use crate::tool_monitor::{ToolCall, ToolMonitor};
use crate::utils::is_token_cancelled; use crate::utils::is_token_cancelled;
use mcp_core::{ToolError, ToolResult}; use mcp_core::{ToolError, ToolResult};
use regex::Regex; use regex::Regex;
use rmcp::model::Tool; use rmcp::model::{Content, GetPromptResult, Prompt, ServerNotification, Tool};
use rmcp::model::{Content, GetPromptResult, JsonRpcMessage, Prompt};
use serde_json::Value; use serde_json::Value;
use tokio::sync::{mpsc, Mutex, RwLock}; use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
@@ -83,7 +82,7 @@ pub struct Agent {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum AgentEvent { pub enum AgentEvent {
Message(Message), Message(Message),
McpNotification((String, JsonRpcMessage)), McpNotification((String, ServerNotification)),
ModelChange { model: String, mode: String }, ModelChange { model: String, mode: String },
} }
@@ -94,19 +93,19 @@ impl Default for Agent {
} }
pub enum ToolStreamItem<T> { pub enum ToolStreamItem<T> {
Message(JsonRpcMessage), Message(ServerNotification),
Result(T), Result(T),
} }
pub type ToolStream = Pin<Box<dyn Stream<Item = ToolStreamItem<ToolResult<Vec<Content>>>> + Send>>; 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 // tool_stream combines a stream of ServerNotifications with a future representing the
// final result of the tool call. MCP notifications are not request-scoped, but // 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 // this lets us capture all notifications emitted during the tool call for
// simpler consumption // simpler consumption
pub fn tool_stream<S, F>(rx: S, done: F) -> ToolStream pub fn tool_stream<S, F>(rx: S, done: F) -> ToolStream
where where
S: Stream<Item = JsonRpcMessage> + Send + Unpin + 'static, S: Stream<Item = ServerNotification> + Send + Unpin + 'static,
F: Future<Output = ToolResult<Vec<Content>>> + Send + 'static, F: Future<Output = ToolResult<Vec<Content>>> + Send + 'static,
{ {
Box::pin(async_stream::stream! { Box::pin(async_stream::stream! {
crates/goose/src/agents/extension_manager.rs
+2 -2
View File
@@ -835,7 +835,7 @@ mod tests {
CallToolResult, InitializeResult, ListPromptsResult, ListResourcesResult, ListToolsResult, CallToolResult, InitializeResult, ListPromptsResult, ListResourcesResult, ListToolsResult,
ReadResourceResult, ReadResourceResult,
}; };
use rmcp::model::{GetPromptResult, JsonRpcMessage}; use rmcp::model::{GetPromptResult, ServerNotification};
use serde_json::json; use serde_json::json;
use tokio::sync::mpsc; use tokio::sync::mpsc;
@@ -891,7 +891,7 @@ mod tests {
Err(Error::NotInitialized) Err(Error::NotInitialized)
} }
async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage> { async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
mpsc::channel(1).1 mpsc::channel(1).1
} }
} }
crates/goose/src/agents/subagent_execution_tool/executor/mod.rs
+3 -3
View File
@@ -7,7 +7,7 @@ use crate::agents::subagent_execution_tool::task_execution_tracker::{
use crate::agents::subagent_execution_tool::tasks::process_task; use crate::agents::subagent_execution_tool::tasks::process_task;
use crate::agents::subagent_execution_tool::workers::spawn_worker; use crate::agents::subagent_execution_tool::workers::spawn_worker;
use crate::agents::subagent_task_config::TaskConfig; use crate::agents::subagent_task_config::TaskConfig;
use rmcp::model::JsonRpcMessage; use rmcp::model::ServerNotification;
use std::sync::atomic::AtomicUsize; use std::sync::atomic::AtomicUsize;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::mpsc; use tokio::sync::mpsc;
@@ -20,7 +20,7 @@ const DEFAULT_MAX_WORKERS: usize = 10;
pub async fn execute_single_task( pub async fn execute_single_task(
task: &Task, task: &Task,
notifier: mpsc::Sender<JsonRpcMessage>, notifier: mpsc::Sender<ServerNotification>,
task_config: TaskConfig, task_config: TaskConfig,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
) -> ExecutionResponse { ) -> ExecutionResponse {
@@ -56,7 +56,7 @@ pub async fn execute_single_task(
pub async fn execute_tasks_in_parallel( pub async fn execute_tasks_in_parallel(
tasks: Vec<Task>, tasks: Vec<Task>,
notifier: Sender<JsonRpcMessage>, notifier: Sender<ServerNotification>,
task_config: TaskConfig, task_config: TaskConfig,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
) -> ExecutionResponse { ) -> ExecutionResponse {
crates/goose/src/agents/subagent_execution_tool/lib/mod.rs
+2 -2
View File
@@ -6,7 +6,7 @@ use crate::agents::subagent_execution_tool::{
tasks_manager::TasksManager, tasks_manager::TasksManager,
}; };
use crate::agents::subagent_task_config::TaskConfig; use crate::agents::subagent_task_config::TaskConfig;
use rmcp::model::JsonRpcMessage; use rmcp::model::ServerNotification;
use serde_json::{json, Value}; use serde_json::{json, Value};
use tokio::sync::mpsc::Sender; use tokio::sync::mpsc::Sender;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
@@ -14,7 +14,7 @@ use tokio_util::sync::CancellationToken;
pub async fn execute_tasks( pub async fn execute_tasks(
input: Value, input: Value,
execution_mode: ExecutionMode, execution_mode: ExecutionMode,
notifier: Sender<JsonRpcMessage>, notifier: Sender<ServerNotification>,
task_config: TaskConfig, task_config: TaskConfig,
tasks_manager: &TasksManager, tasks_manager: &TasksManager,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs
+2 -3
View File
@@ -1,5 +1,5 @@
use mcp_core::ToolError; use mcp_core::ToolError;
use rmcp::model::{Content, Tool, ToolAnnotations}; use rmcp::model::{Content, ServerNotification, Tool, ToolAnnotations};
use serde_json::Value; use serde_json::Value;
use crate::agents::subagent_task_config::TaskConfig; use crate::agents::subagent_task_config::TaskConfig;
@@ -8,7 +8,6 @@ use crate::agents::{
subagent_execution_tool::task_types::ExecutionMode, subagent_execution_tool::task_types::ExecutionMode,
subagent_execution_tool::tasks_manager::TasksManager, tool_execution::ToolCallResult, subagent_execution_tool::tasks_manager::TasksManager, tool_execution::ToolCallResult,
}; };
use rmcp::model::JsonRpcMessage;
use rmcp::object; use rmcp::object;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tokio_stream; use tokio_stream;
@@ -67,7 +66,7 @@ pub async fn run_tasks(
tasks_manager: &TasksManager, tasks_manager: &TasksManager,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
) -> ToolCallResult { ) -> ToolCallResult {
let (notification_tx, notification_rx) = mpsc::channel::<JsonRpcMessage>(100); let (notification_tx, notification_rx) = mpsc::channel::<ServerNotification>(100);
let tasks_manager_clone = tasks_manager.clone(); let tasks_manager_clone = tasks_manager.clone();
let result_future = async move { let result_future = async move {
crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs
+16 -13
View File
@@ -1,5 +1,7 @@
use rmcp::model::{JsonRpcMessage, JsonRpcNotification, JsonRpcVersion2_0, Notification}; use rmcp::model::{
use rmcp::object; LoggingLevel, LoggingMessageNotification, LoggingMessageNotificationMethod,
LoggingMessageNotificationParam, ServerNotification,
};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::{mpsc, RwLock}; use tokio::sync::{mpsc, RwLock};
@@ -52,7 +54,7 @@ fn format_task_metadata(task_info: &TaskInfo) -> String {
pub struct TaskExecutionTracker { pub struct TaskExecutionTracker {
tasks: Arc<RwLock<HashMap<String, TaskInfo>>>, tasks: Arc<RwLock<HashMap<String, TaskInfo>>>,
last_refresh: Arc<RwLock<Instant>>, last_refresh: Arc<RwLock<Instant>>,
notifier: mpsc::Sender<JsonRpcMessage>, notifier: mpsc::Sender<ServerNotification>,
display_mode: DisplayMode, display_mode: DisplayMode,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
} }
@@ -61,7 +63,7 @@ impl TaskExecutionTracker {
pub fn new( pub fn new(
tasks: Vec<Task>, tasks: Vec<Task>,
display_mode: DisplayMode, display_mode: DisplayMode,
notifier: Sender<JsonRpcMessage>, notifier: Sender<ServerNotification>,
cancellation_token: Option<CancellationToken>, cancellation_token: Option<CancellationToken>,
) -> Self { ) -> Self {
let task_map = tasks let task_map = tasks
@@ -97,7 +99,7 @@ impl TaskExecutionTracker {
fn log_notification_error( fn log_notification_error(
&self, &self,
error: &mpsc::error::TrySendError<JsonRpcMessage>, error: &mpsc::error::TrySendError<ServerNotification>,
context: &str, context: &str,
) { ) {
if !self.is_cancelled() { if !self.is_cancelled() {
@@ -108,16 +110,17 @@ impl TaskExecutionTracker {
fn try_send_notification(&self, event: TaskExecutionNotificationEvent, context: &str) { fn try_send_notification(&self, event: TaskExecutionNotificationEvent, context: &str) {
if let Err(e) = self if let Err(e) = self
.notifier .notifier
.try_send(JsonRpcMessage::Notification(JsonRpcNotification { .try_send(ServerNotification::LoggingMessageNotification(
jsonrpc: JsonRpcVersion2_0, LoggingMessageNotification {
notification: Notification { method: LoggingMessageNotificationMethod,
method: "notifications/message".to_string(), params: LoggingMessageNotificationParam {
params: object!({ data: event.to_notification_data(),
"data": event.to_notification_data() level: LoggingLevel::Info,
}), logger: None,
},
extensions: Default::default(), extensions: Default::default(),
}, },
})) ))
{ {
self.log_notification_error(&e, context); self.log_notification_error(&e, context);
} }
crates/goose/src/agents/tool_execution.rs
+2 -2
View File
@@ -4,7 +4,7 @@ use std::sync::Arc;
use async_stream::try_stream; use async_stream::try_stream;
use futures::stream::{self, BoxStream}; use futures::stream::{self, BoxStream};
use futures::{Stream, StreamExt}; use futures::{Stream, StreamExt};
use rmcp::model::JsonRpcMessage; use rmcp::model::ServerNotification;
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
@@ -19,7 +19,7 @@ use rmcp::model::Content;
// can be used to receive notifications from the tool. // can be used to receive notifications from the tool.
pub struct ToolCallResult { pub struct ToolCallResult {
pub result: Box<dyn Future<Output = ToolResult<Vec<Content>>> + Send + Unpin>, pub result: Box<dyn Future<Output = ToolResult<Vec<Content>>> + Send + Unpin>,
pub notification_stream: Option<Box<dyn Stream<Item = JsonRpcMessage> + Send + Unpin>>, pub notification_stream: Option<Box<dyn Stream<Item = ServerNotification> + Send + Unpin>>,
} }
impl From<ToolResult<Vec<Content>>> for ToolCallResult { impl From<ToolResult<Vec<Content>>> for ToolCallResult {
crates/mcp-client/src/client.rs
+20 -6
View File
@@ -5,6 +5,7 @@ use mcp_core::protocol::{
use rmcp::model::{ use rmcp::model::{
GetPromptResult, JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, GetPromptResult, JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
JsonRpcResponse, JsonRpcVersion2_0, Notification, NumberOrString, Request, RequestId, JsonRpcResponse, JsonRpcVersion2_0, Notification, NumberOrString, Request, RequestId,
ServerNotification,
}; };
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::{json, Value}; use serde_json::{json, Value};
@@ -106,7 +107,7 @@ pub trait McpClientTrait: Send + Sync {
async fn get_prompt(&self, name: &str, arguments: Value) -> Result<GetPromptResult, Error>; async fn get_prompt(&self, name: &str, arguments: Value) -> Result<GetPromptResult, Error>;
async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage>; async fn subscribe(&self) -> mpsc::Receiver<ServerNotification>;
} }
/// The MCP client is the interface for MCP operations. /// The MCP client is the interface for MCP operations.
@@ -118,7 +119,7 @@ where
next_id_counter: AtomicU64, // Added for atomic ID generation next_id_counter: AtomicU64, // Added for atomic ID generation
server_capabilities: Option<ServerCapabilities>, server_capabilities: Option<ServerCapabilities>,
server_info: Option<Implementation>, server_info: Option<Implementation>,
notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<JsonRpcMessage>>>>, notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<ServerNotification>>>>,
} }
impl<T> McpClient<T> impl<T> McpClient<T>
@@ -129,7 +130,7 @@ where
let service = McpService::new(transport.clone()); let service = McpService::new(transport.clone());
let service_ptr = service.clone(); let service_ptr = service.clone();
let notification_subscribers = let notification_subscribers =
Arc::new(Mutex::new(Vec::<mpsc::Sender<JsonRpcMessage>>::new())); Arc::new(Mutex::new(Vec::<mpsc::Sender<ServerNotification>>::new()));
let subscribers_ptr = notification_subscribers.clone(); let subscribers_ptr = notification_subscribers.clone();
tokio::spawn(async move { tokio::spawn(async move {
@@ -148,9 +149,22 @@ where
}) => { }) => {
service_ptr.respond(&id.to_string(), Ok(message)).await; service_ptr.respond(&id.to_string(), Ok(message)).await;
} }
_ => { JsonRpcMessage::Notification(JsonRpcNotification {
notification,
..
}) => {
let mut subs = subscribers_ptr.lock().await; let mut subs = subscribers_ptr.lock().await;
subs.retain(|sub| sub.try_send(message.clone()).is_ok()); if let Some(server_notification) = notification.into() {
subs.retain(|sub| {
sub.try_send(server_notification.clone()).is_ok()
});
}
}
_ => {
tracing::warn!(
"Received unexpected received message type: {:?}",
message
);
} }
} }
} }
@@ -437,7 +451,7 @@ where
self.send_request("prompts/get", params).await self.send_request("prompts/get", params).await
} }
async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage> { async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
let (tx, rx) = mpsc::channel(16); let (tx, rx) = mpsc::channel(16);
self.notification_subscribers.lock().await.push(tx); self.notification_subscribers.lock().await.push(tx);
rx rx
crates/mcp-client/src/service.rs
+11 -7
View File
@@ -6,7 +6,7 @@ use std::task::{Context, Poll};
use tokio::sync::{oneshot, RwLock}; use tokio::sync::{oneshot, RwLock};
use tower::{timeout::Timeout, Service, ServiceBuilder}; use tower::{timeout::Timeout, Service, ServiceBuilder};
use crate::transport::{Error, TransportHandle}; use crate::transport::{Error, TransportHandle, TransportMessageRecv};
/// A wrapper service that implements Tower's Service trait for MCP transport /// A wrapper service that implements Tower's Service trait for MCP transport
#[derive(Clone)] #[derive(Clone)]
@@ -23,7 +23,7 @@ impl<T: TransportHandle> McpService<T> {
} }
} }
pub async fn respond(&self, id: &str, response: Result<JsonRpcMessage, Error>) { pub async fn respond(&self, id: &str, response: Result<TransportMessageRecv, Error>) {
self.pending_requests.respond(id, response).await self.pending_requests.respond(id, response).await
} }
@@ -36,7 +36,7 @@ impl<T> Service<JsonRpcMessage> for McpService<T>
where where
T: TransportHandle + Send + Sync + 'static, T: TransportHandle + Send + Sync + 'static,
{ {
type Response = JsonRpcMessage; type Response = TransportMessageRecv;
type Error = Error; type Error = Error;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>; type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
@@ -63,7 +63,7 @@ where
// Handle notifications without waiting for a response // Handle notifications without waiting for a response
transport.send(request).await?; transport.send(request).await?;
// Return a dummy response for notifications // Return a dummy response for notifications
let dummy_response: JsonRpcMessage = let dummy_response: Self::Response =
JsonRpcMessage::Response(rmcp::model::JsonRpcResponse { JsonRpcMessage::Response(rmcp::model::JsonRpcResponse {
jsonrpc: rmcp::model::JsonRpcVersion2_0, jsonrpc: rmcp::model::JsonRpcVersion2_0,
id: rmcp::model::RequestId::Number(0), id: rmcp::model::RequestId::Number(0),
@@ -91,7 +91,7 @@ where
// A data structure to store pending requests and their response channels // A data structure to store pending requests and their response channels
pub struct PendingRequests { pub struct PendingRequests {
requests: RwLock<HashMap<String, oneshot::Sender<Result<JsonRpcMessage, Error>>>>, requests: RwLock<HashMap<String, oneshot::Sender<Result<TransportMessageRecv, Error>>>>,
} }
impl Default for PendingRequests { impl Default for PendingRequests {
@@ -107,11 +107,15 @@ impl PendingRequests {
} }
} }
pub async fn insert(&self, id: String, sender: oneshot::Sender<Result<JsonRpcMessage, Error>>) { pub async fn insert(
&self,
id: String,
sender: oneshot::Sender<Result<TransportMessageRecv, Error>>,
) {
self.requests.write().await.insert(id, sender); self.requests.write().await.insert(id, sender);
} }
pub async fn respond(&self, id: &str, response: Result<JsonRpcMessage, Error>) { pub async fn respond(&self, id: &str, response: Result<TransportMessageRecv, Error>) {
if let Some(tx) = self.requests.write().await.remove(id) { if let Some(tx) = self.requests.write().await.remove(id) {
let _ = tx.send(response); let _ = tx.send(response);
} }
crates/mcp-client/src/transport/mod.rs
+5 -12
View File
@@ -1,7 +1,7 @@
use async_trait::async_trait; use async_trait::async_trait;
use rmcp::model::JsonRpcMessage; use rmcp::model::{JsonObject, JsonRpcMessage, Request, ServerNotification};
use thiserror::Error; use thiserror::Error;
use tokio::sync::{mpsc, oneshot}; use tokio::sync::mpsc;
pub type BoxError = Box<dyn std::error::Error + Sync + Send>; pub type BoxError = Box<dyn std::error::Error + Sync + Send>;
/// A generic error type for transport operations. /// A generic error type for transport operations.
@@ -38,15 +38,6 @@ pub enum Error {
SessionError(String), SessionError(String),
} }
/// A message that can be sent through the transport
#[derive(Debug)]
pub struct TransportMessage {
/// The JSON-RPC message to send
pub message: JsonRpcMessage,
/// Channel to receive the response on (None for notifications)
pub response_tx: Option<oneshot::Sender<Result<JsonRpcMessage, Error>>>,
}
/// A generic asynchronous transport trait with channel-based communication /// A generic asynchronous transport trait with channel-based communication
#[async_trait] #[async_trait]
pub trait Transport { pub trait Transport {
@@ -60,10 +51,12 @@ pub trait Transport {
async fn close(&self) -> Result<(), Error>; async fn close(&self) -> Result<(), Error>;
} }
pub type TransportMessageRecv = JsonRpcMessage<Request, JsonObject, ServerNotification>;
#[async_trait] #[async_trait]
pub trait TransportHandle: Send + Sync + Clone + 'static { pub trait TransportHandle: Send + Sync + Clone + 'static {
async fn send(&self, message: JsonRpcMessage) -> Result<(), Error>; async fn send(&self, message: JsonRpcMessage) -> Result<(), Error>;
async fn receive(&self) -> Result<JsonRpcMessage, Error>; async fn receive(&self) -> Result<TransportMessageRecv, Error>;
} }
pub async fn serialize_and_send( pub async fn serialize_and_send(
crates/mcp-client/src/transport/sse.rs
+7 -7
View File
@@ -1,4 +1,4 @@
use crate::transport::Error; use crate::transport::{Error, TransportMessageRecv};
use async_trait::async_trait; use async_trait::async_trait;
use eventsource_client::{Client, SSE}; use eventsource_client::{Client, SSE};
use futures::TryStreamExt; use futures::TryStreamExt;
@@ -23,7 +23,7 @@ pub struct SseActor {
/// Receives messages (requests/notifications) from the handle /// Receives messages (requests/notifications) from the handle
receiver: mpsc::Receiver<String>, receiver: mpsc::Receiver<String>,
/// Sends messages (responses) back to the handle /// Sends messages (responses) back to the handle
sender: mpsc::Sender<JsonRpcMessage>, sender: mpsc::Sender<TransportMessageRecv>,
/// Base SSE URL /// Base SSE URL
sse_url: String, sse_url: String,
/// For sending HTTP POST requests /// For sending HTTP POST requests
@@ -35,7 +35,7 @@ pub struct SseActor {
impl SseActor { impl SseActor {
pub fn new( pub fn new(
receiver: mpsc::Receiver<String>, receiver: mpsc::Receiver<String>,
sender: mpsc::Sender<JsonRpcMessage>, sender: mpsc::Sender<TransportMessageRecv>,
sse_url: String, sse_url: String,
post_endpoint: Arc<RwLock<Option<String>>>, post_endpoint: Arc<RwLock<Option<String>>>,
) -> Self { ) -> Self {
@@ -71,7 +71,7 @@ impl SseActor {
/// - If a `message` event is received, parse it as `JsonRpcMessage` /// - If a `message` event is received, parse it as `JsonRpcMessage`
/// and respond to pending requests if it's a `Response`. /// and respond to pending requests if it's a `Response`.
async fn handle_incoming_messages( async fn handle_incoming_messages(
sender: mpsc::Sender<JsonRpcMessage>, sender: mpsc::Sender<TransportMessageRecv>,
sse_url: String, sse_url: String,
post_endpoint: Arc<RwLock<Option<String>>>, post_endpoint: Arc<RwLock<Option<String>>>,
) { ) {
@@ -109,7 +109,7 @@ impl SseActor {
match event { match event {
SSE::Event(e) if e.event_type == "message" => { SSE::Event(e) if e.event_type == "message" => {
// Attempt to parse the SSE data as a JsonRpcMessage // Attempt to parse the SSE data as a JsonRpcMessage
match serde_json::from_str::<JsonRpcMessage>(&e.data) { match serde_json::from_str::<TransportMessageRecv>(&e.data) {
Ok(message) => { Ok(message) => {
let _ = sender.send(message).await; let _ = sender.send(message).await;
} }
@@ -184,7 +184,7 @@ impl SseActor {
#[derive(Clone)] #[derive(Clone)]
pub struct SseTransportHandle { pub struct SseTransportHandle {
sender: mpsc::Sender<String>, sender: mpsc::Sender<String>,
receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>, receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>,
} }
#[async_trait::async_trait] #[async_trait::async_trait]
@@ -193,7 +193,7 @@ impl TransportHandle for SseTransportHandle {
serialize_and_send(&self.sender, message).await serialize_and_send(&self.sender, message).await
} }
async fn receive(&self) -> Result<JsonRpcMessage, Error> { async fn receive(&self) -> Result<TransportMessageRecv, Error> {
let mut receiver = self.receiver.lock().await; let mut receiver = self.receiver.lock().await;
receiver.recv().await.ok_or(Error::ChannelClosed) receiver.recv().await.ok_or(Error::ChannelClosed)
} }
crates/mcp-client/src/transport/stdio.rs
+8 -6
View File
@@ -14,6 +14,8 @@ use nix::sys::signal::{kill, Signal};
#[cfg(unix)] #[cfg(unix)]
use nix::unistd::{getpgid, Pid}; use nix::unistd::{getpgid, Pid};
use crate::transport::TransportMessageRecv;
use super::{serialize_and_send, Error, Transport, TransportHandle}; use super::{serialize_and_send, Error, Transport, TransportHandle};
// Global to track process groups we've created // Global to track process groups we've created
@@ -24,7 +26,7 @@ static PROCESS_GROUP: AtomicI32 = AtomicI32::new(-1);
/// It uses channels for message passing and handles responses asynchronously through a background task. /// It uses channels for message passing and handles responses asynchronously through a background task.
pub struct StdioActor { pub struct StdioActor {
receiver: Option<mpsc::Receiver<String>>, receiver: Option<mpsc::Receiver<String>>,
sender: Option<mpsc::Sender<JsonRpcMessage>>, sender: Option<mpsc::Sender<TransportMessageRecv>>,
process: Child, // we store the process to keep it alive process: Child, // we store the process to keep it alive
error_sender: mpsc::Sender<Error>, error_sender: mpsc::Sender<Error>,
stdin: Option<ChildStdin>, stdin: Option<ChildStdin>,
@@ -98,7 +100,7 @@ impl StdioActor {
} }
} }
async fn handle_proc_output(stdout: ChildStdout, sender: mpsc::Sender<JsonRpcMessage>) { async fn handle_proc_output(stdout: ChildStdout, sender: mpsc::Sender<TransportMessageRecv>) {
let mut reader = BufReader::new(stdout); let mut reader = BufReader::new(stdout);
let mut line = String::new(); let mut line = String::new();
loop { loop {
@@ -108,7 +110,7 @@ impl StdioActor {
break; break;
} // EOF } // EOF
Ok(_) => { Ok(_) => {
if let Ok(message) = serde_json::from_str::<JsonRpcMessage>(&line) { if let Ok(message) = serde_json::from_str::<TransportMessageRecv>(&line) {
tracing::debug!( tracing::debug!(
message = ?message, message = ?message,
"Received incoming message" "Received incoming message"
@@ -149,8 +151,8 @@ impl StdioActor {
#[derive(Clone)] #[derive(Clone)]
pub struct StdioTransportHandle { pub struct StdioTransportHandle {
sender: mpsc::Sender<String>, // to process sender: mpsc::Sender<String>, // to process
receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>, // from process receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>, // from process
error_receiver: Arc<Mutex<mpsc::Receiver<Error>>>, error_receiver: Arc<Mutex<mpsc::Receiver<Error>>>,
} }
@@ -163,7 +165,7 @@ impl TransportHandle for StdioTransportHandle {
result result
} }
async fn receive(&self) -> Result<JsonRpcMessage, Error> { async fn receive(&self) -> Result<TransportMessageRecv, Error> {
let mut receiver = self.receiver.lock().await; let mut receiver = self.receiver.lock().await;
match receiver.recv().await { match receiver.recv().await {
Some(message) => Ok(message), Some(message) => Ok(message),
crates/mcp-client/src/transport/streamable_http.rs
+10 -10
View File
@@ -1,5 +1,5 @@
use crate::oauth::{authenticate_service, ServiceConfig}; use crate::oauth::{authenticate_service, ServiceConfig};
use crate::transport::Error; use crate::transport::{Error, TransportMessageRecv};
use async_trait::async_trait; use async_trait::async_trait;
use eventsource_client::{Client, SSE}; use eventsource_client::{Client, SSE};
use futures::TryStreamExt; use futures::TryStreamExt;
@@ -25,7 +25,7 @@ pub struct StreamableHttpActor {
/// Receives messages (requests/notifications) from the handle /// Receives messages (requests/notifications) from the handle
receiver: mpsc::Receiver<String>, receiver: mpsc::Receiver<String>,
/// Sends messages (responses) back to the handle /// Sends messages (responses) back to the handle
sender: mpsc::Sender<JsonRpcMessage>, sender: mpsc::Sender<TransportMessageRecv>,
/// MCP endpoint URL /// MCP endpoint URL
mcp_endpoint: String, mcp_endpoint: String,
/// HTTP client for sending requests /// HTTP client for sending requests
@@ -41,7 +41,7 @@ pub struct StreamableHttpActor {
impl StreamableHttpActor { impl StreamableHttpActor {
pub fn new( pub fn new(
receiver: mpsc::Receiver<String>, receiver: mpsc::Receiver<String>,
sender: mpsc::Sender<JsonRpcMessage>, sender: mpsc::Sender<TransportMessageRecv>,
mcp_endpoint: String, mcp_endpoint: String,
session_id: Arc<RwLock<Option<String>>>, session_id: Arc<RwLock<Option<String>>>,
env: HashMap<String, String>, env: HashMap<String, String>,
@@ -84,8 +84,8 @@ impl StreamableHttpActor {
debug!("Sending message to MCP endpoint: {}", message_str); debug!("Sending message to MCP endpoint: {}", message_str);
// Parse the message to determine if it's a request that expects a response // Parse the message to determine if it's a request that expects a response
let parsed_message: JsonRpcMessage = let parsed_message = serde_json::from_str::<TransportMessageRecv>(&message_str)
serde_json::from_str(&message_str).map_err(Error::Serialization)?; .map_err(Error::Serialization)?;
let expects_response = matches!( let expects_response = matches!(
parsed_message, parsed_message,
@@ -196,8 +196,8 @@ impl StreamableHttpActor {
})?; })?;
if !response_text.is_empty() { if !response_text.is_empty() {
let json_message: JsonRpcMessage = let json_message = serde_json::from_str::<TransportMessageRecv>(&response_text)
serde_json::from_str(&response_text).map_err(Error::Serialization)?; .map_err(Error::Serialization)?;
let _ = self.sender.send(json_message).await; let _ = self.sender.send(json_message).await;
} }
@@ -267,7 +267,7 @@ impl StreamableHttpActor {
// Empty line indicates end of event // Empty line indicates end of event
if !event_data.is_empty() { if !event_data.is_empty() {
// Parse the streamed data as JSON-RPC message // Parse the streamed data as JSON-RPC message
match serde_json::from_str::<JsonRpcMessage>(&event_data) { match serde_json::from_str::<TransportMessageRecv>(&event_data) {
Ok(message) => { Ok(message) => {
debug!("Received streaming HTTP response message: {:?}", message); debug!("Received streaming HTTP response message: {:?}", message);
let _ = self.sender.send(message).await; let _ = self.sender.send(message).await;
@@ -301,7 +301,7 @@ impl StreamableHttpActor {
#[derive(Clone)] #[derive(Clone)]
pub struct StreamableHttpTransportHandle { pub struct StreamableHttpTransportHandle {
sender: mpsc::Sender<String>, sender: mpsc::Sender<String>,
receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>, receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>,
session_id: Arc<RwLock<Option<String>>>, session_id: Arc<RwLock<Option<String>>>,
mcp_endpoint: String, mcp_endpoint: String,
http_client: HttpClient, http_client: HttpClient,
@@ -314,7 +314,7 @@ impl TransportHandle for StreamableHttpTransportHandle {
serialize_and_send(&self.sender, message).await serialize_and_send(&self.sender, message).await
} }
async fn receive(&self) -> Result<JsonRpcMessage, Error> { async fn receive(&self) -> Result<TransportMessageRecv, Error> {
let mut receiver = self.receiver.lock().await; let mut receiver = self.receiver.lock().await;
receiver.recv().await.ok_or(Error::ChannelClosed) receiver.recv().await.ok_or(Error::ChannelClosed)
} }