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

feat: migrate JsonRpcMessage/Request/Response/Error/Notification from internal mcp crates to rmcp versions (#3564)

Browse Source
This commit is contained in:
Alex Hancock
2025-07-22 13:45:00 -04:00
committed by GitHub
parent 246ba19256
commit 3e5c9b8505
27 changed files with 256 additions and 246 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/mcp-client/src/client.rs
+40 -32
View File
@@ -1,8 +1,12 @@
use mcp_core::protocol::{
CallToolResult, GetPromptResult, Implementation, InitializeResult, JsonRpcError,
JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, ListPromptsResult,
CallToolResult, GetPromptResult, Implementation, InitializeResult, ListPromptsResult,
ListResourcesResult, ListToolsResult, ReadResourceResult, ServerCapabilities, METHOD_NOT_FOUND,
};
use rmcp::model::{
JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse,
JsonRpcVersion2_0, Notification, NumberOrString, Request, RequestId,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::sync::{
@@ -112,7 +116,7 @@ where
T: TransportHandle + Send + Sync + 'static,
{
service: Mutex<tower::timeout::Timeout<McpService<T>>>,
next_id: AtomicU64,
next_id_counter: AtomicU64, // Added for atomic ID generation
server_capabilities: Option<ServerCapabilities>,
server_info: Option<Implementation>,
notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<JsonRpcMessage>>>>,
@@ -135,8 +139,14 @@ where
Ok(message) => {
tracing::info!("Received message: {:?}", message);
match message {
JsonRpcMessage::Response(JsonRpcResponse { id: Some(id), .. })
| JsonRpcMessage::Error(JsonRpcError { id: Some(id), .. }) => {
JsonRpcMessage::Response(JsonRpcResponse {
id: NumberOrString::Number(id),
..
})
| JsonRpcMessage::Error(JsonRpcError {
id: NumberOrString::Number(id),
..
}) => {
service_ptr.respond(&id.to_string(), Ok(message)).await;
}
_ => {
@@ -158,7 +168,7 @@ where
Ok(Self {
service: Mutex::new(middleware.layer(service)),
next_id: AtomicU64::new(1),
next_id_counter: AtomicU64::new(1),
server_capabilities: None,
server_info: None,
notification_subscribers,
@@ -172,7 +182,8 @@ where
{
let mut service = self.service.lock().await;
service.ready().await.map_err(|_| Error::NotReady)?;
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let id_num = self.next_id_counter.fetch_add(1, Ordering::SeqCst);
let id = RequestId::Number(id_num as u32);
let mut params = params.clone();
params["_meta"] = json!({
@@ -180,10 +191,13 @@ where
});
let request = JsonRpcMessage::Request(JsonRpcRequest {
jsonrpc: "2.0".to_string(),
id: Some(id),
method: method.to_string(),
params: Some(params),
jsonrpc: JsonRpcVersion2_0,
id,
request: Request {
method: method.to_string(),
params: params.as_object().unwrap().clone(),
extensions: Default::default(),
},
});
let response_msg = service
@@ -201,35 +215,26 @@ where
})?;
match response_msg {
JsonRpcMessage::Response(JsonRpcResponse {
id, result, error, ..
}) => {
// Verify id matches
if id != Some(self.next_id.load(Ordering::SeqCst) - 1) {
JsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => {
// Verify id matches - convert current id to match expected format
let expected_id = RequestId::Number((id_num) as u32);
if id != expected_id {
return Err(Error::UnexpectedResponse(
"id mismatch for JsonRpcResponse".to_string(),
));
}
if let Some(err) = error {
Err(Error::RpcError {
code: err.code,
message: err.message,
})
} else if let Some(r) = result {
Ok(serde_json::from_value(r)?)
} else {
Err(Error::UnexpectedResponse("missing result".to_string()))
}
Ok(serde_json::from_value(serde_json::to_value(result)?)?)
}
JsonRpcMessage::Error(JsonRpcError { id, error, .. }) => {
if id != Some(self.next_id.load(Ordering::SeqCst) - 1) {
let expected_id = RequestId::Number((id_num) as u32);
if id != expected_id {
return Err(Error::UnexpectedResponse(
"id mismatch for JsonRpcError".to_string(),
));
}
Err(Error::RpcError {
code: error.code,
message: error.message,
code: error.code.0, // Extract the i32 from ErrorCode
message: error.message.to_string(), // Convert Cow to String
})
}
_ => {
@@ -247,9 +252,12 @@ where
service.ready().await.map_err(|_| Error::NotReady)?;
let notification = JsonRpcMessage::Notification(JsonRpcNotification {
jsonrpc: "2.0".to_string(),
method: method.to_string(),
params: Some(params.clone()),
jsonrpc: JsonRpcVersion2_0,
notification: Notification {
method: method.to_string(),
params: params.as_object().unwrap().clone(),
extensions: Default::default(),
},
});
service
crates/mcp-client/src/service.rs
+11 -9
View File
@@ -1,5 +1,5 @@
use futures::future::BoxFuture;
use mcp_core::protocol::{JsonRpcMessage, JsonRpcRequest};
use rmcp::model::{JsonRpcMessage, JsonRpcRequest};
use std::collections::HashMap;
use std::sync::Arc;
use std::task::{Context, Poll};
@@ -50,8 +50,8 @@ where
let pending_requests = self.pending_requests.clone();
Box::pin(async move {
match request {
JsonRpcMessage::Request(JsonRpcRequest { id: Some(id), .. }) => {
match &request {
JsonRpcMessage::Request(JsonRpcRequest { id, .. }) => {
// Create a channel to receive the response
let (sender, receiver) = oneshot::channel();
pending_requests.insert(id.to_string(), sender).await;
@@ -59,15 +59,17 @@ where
transport.send(request).await?;
receiver.await.map_err(|_| Error::ChannelClosed)?
}
JsonRpcMessage::Request(_) => {
// Handle notifications without waiting for a response
transport.send(request).await?;
Ok(JsonRpcMessage::Nil)
}
JsonRpcMessage::Notification(_) => {
// Handle notifications without waiting for a response
transport.send(request).await?;
Ok(JsonRpcMessage::Nil)
// Return a dummy response for notifications
let dummy_response: JsonRpcMessage =
JsonRpcMessage::Response(rmcp::model::JsonRpcResponse {
jsonrpc: rmcp::model::JsonRpcVersion2_0,
id: rmcp::model::RequestId::Number(0),
result: serde_json::Map::new(),
});
Ok(dummy_response)
}
_ => Err(Error::UnsupportedMessage),
}
crates/mcp-client/src/transport/mod.rs
+1 -1
View File
@@ -1,5 +1,5 @@
use async_trait::async_trait;
use mcp_core::protocol::JsonRpcMessage;
use rmcp::model::JsonRpcMessage;
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};
crates/mcp-client/src/transport/sse.rs
+1 -1
View File
@@ -2,8 +2,8 @@ use crate::transport::Error;
use async_trait::async_trait;
use eventsource_client::{Client, SSE};
use futures::TryStreamExt;
use mcp_core::protocol::JsonRpcMessage;
use reqwest::Client as HttpClient;
use rmcp::model::JsonRpcMessage;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex, RwLock};
crates/mcp-client/src/transport/stdio.rs
+1 -1
View File
@@ -4,7 +4,7 @@ use std::sync::Arc;
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
use async_trait::async_trait;
use mcp_core::protocol::JsonRpcMessage;
use rmcp::model::JsonRpcMessage;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::sync::{mpsc, Mutex};
crates/mcp-client/src/transport/streamable_http.rs
+2 -2
View File
@@ -3,8 +3,8 @@ use crate::transport::Error;
use async_trait::async_trait;
use eventsource_client::{Client, SSE};
use futures::TryStreamExt;
use mcp_core::protocol::{JsonRpcMessage, JsonRpcRequest};
use reqwest::Client as HttpClient;
use rmcp::model::{JsonRpcMessage, JsonRpcRequest, NumberOrString::Number};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex, RwLock};
@@ -89,7 +89,7 @@ impl StreamableHttpActor {
let expects_response = matches!(
parsed_message,
JsonRpcMessage::Request(JsonRpcRequest { id: Some(_), .. })
JsonRpcMessage::Request(JsonRpcRequest { id: Number(_), .. })
);
// Try to send the request