fix: update peer info on duplicate initialize (#862)
This commit is contained in:
parent
2d3d1879ad
commit
2536a05992
3 changed files with 69 additions and 12 deletions
|
|
@ -184,9 +184,7 @@ macro_rules! server_handler_methods {
|
||||||
request: InitializeRequestParams,
|
request: InitializeRequestParams,
|
||||||
context: RequestContext<RoleServer>,
|
context: RequestContext<RoleServer>,
|
||||||
) -> impl Future<Output = Result<InitializeResult, McpError>> + MaybeSendFuture + '_ {
|
) -> impl Future<Output = Result<InitializeResult, McpError>> + MaybeSendFuture + '_ {
|
||||||
if context.peer.peer_info().is_none() {
|
context.peer.set_peer_info(request);
|
||||||
context.peer.set_peer_info(request);
|
|
||||||
}
|
|
||||||
std::future::ready(Ok(self.get_info()))
|
std::future::ready(Ok(self.get_info()))
|
||||||
}
|
}
|
||||||
fn complete(
|
fn complete(
|
||||||
|
|
|
||||||
|
|
@ -384,7 +384,7 @@ pub struct Peer<R: ServiceRole> {
|
||||||
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
||||||
request_id_provider: Arc<dyn RequestIdProvider>,
|
request_id_provider: Arc<dyn RequestIdProvider>,
|
||||||
progress_token_provider: Arc<dyn ProgressTokenProvider>,
|
progress_token_provider: Arc<dyn ProgressTokenProvider>,
|
||||||
info: Arc<tokio::sync::OnceCell<R::PeerInfo>>,
|
info: Arc<std::sync::RwLock<Option<Arc<R::PeerInfo>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
|
impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
|
||||||
|
|
@ -423,7 +423,7 @@ impl<R: ServiceRole> Peer<R> {
|
||||||
tx,
|
tx,
|
||||||
request_id_provider,
|
request_id_provider,
|
||||||
progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
|
progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
|
||||||
info: Arc::new(tokio::sync::OnceCell::new_with(peer_info)),
|
info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))),
|
||||||
},
|
},
|
||||||
rx,
|
rx,
|
||||||
)
|
)
|
||||||
|
|
@ -484,16 +484,14 @@ impl<R: ServiceRole> Peer<R> {
|
||||||
peer: self.clone(),
|
peer: self.clone(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
pub fn peer_info(&self) -> Option<&R::PeerInfo> {
|
/// Snapshot of the peer's handshake info.
|
||||||
self.info.get()
|
pub fn peer_info(&self) -> Option<Arc<R::PeerInfo>> {
|
||||||
|
self.info.read().expect("peer info lock poisoned").clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Stores the peer's handshake info, overwriting any previous value.
|
||||||
pub fn set_peer_info(&self, info: R::PeerInfo) {
|
pub fn set_peer_info(&self, info: R::PeerInfo) {
|
||||||
if self.info.initialized() {
|
*self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info));
|
||||||
tracing::warn!("trying to set peer info, which is already initialized");
|
|
||||||
} else {
|
|
||||||
let _ = self.info.set(info);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn is_transport_closed(&self) -> bool {
|
pub fn is_transport_closed(&self) -> bool {
|
||||||
|
|
|
||||||
|
|
@ -299,6 +299,67 @@ async fn server_pinned_version_used_as_fallback_for_unknown_client_request() {
|
||||||
assert_eq!(negotiated, ProtocolVersion::V_2025_06_18);
|
assert_eq!(negotiated, ProtocolVersion::V_2025_06_18);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn duplicate_init_request(id: u64, version: &str) -> ClientJsonRpcMessage {
|
||||||
|
msg(&format!(
|
||||||
|
r#"{{
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": {id},
|
||||||
|
"method": "initialize",
|
||||||
|
"params": {{
|
||||||
|
"protocolVersion": "{version}",
|
||||||
|
"capabilities": {{ "sampling": {{}} }},
|
||||||
|
"clientInfo": {{ "name": "renegotiated-client", "version": "9.9.9" }}
|
||||||
|
}}
|
||||||
|
}}"#
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn server_accepts_duplicate_initialize() {
|
||||||
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||||
|
let _server = tokio::spawn(async move { TestServer::new().serve(server_transport).await });
|
||||||
|
let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
|
||||||
|
|
||||||
|
do_initialize(&mut client).await;
|
||||||
|
client.send(initialized_notification()).await.unwrap();
|
||||||
|
|
||||||
|
client
|
||||||
|
.send(duplicate_init_request(2, "2025-11-25"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let response = client.receive().await.unwrap();
|
||||||
|
assert!(
|
||||||
|
matches!(response, ServerJsonRpcMessage::Response(_)),
|
||||||
|
"expected successful InitializeResult, got: {response:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn server_session_remains_usable_after_renegotiation() {
|
||||||
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||||
|
let _server = tokio::spawn(async move { TestServer::new().serve(server_transport).await });
|
||||||
|
let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
|
||||||
|
|
||||||
|
do_initialize(&mut client).await;
|
||||||
|
client.send(initialized_notification()).await.unwrap();
|
||||||
|
client
|
||||||
|
.send(duplicate_init_request(2, "2025-11-25"))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let _renegotiated = client.receive().await.unwrap();
|
||||||
|
|
||||||
|
client.send(ping_request(3)).await.unwrap();
|
||||||
|
let pong = client.receive().await.unwrap();
|
||||||
|
assert!(
|
||||||
|
matches!(
|
||||||
|
pong,
|
||||||
|
ServerJsonRpcMessage::Response(ref r)
|
||||||
|
if matches!(r.result, ServerResult::EmptyResult(_))
|
||||||
|
),
|
||||||
|
"expected EmptyResult ping after renegotiation, got: {pong:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Server buffers multiple requests before initialized and processes them in order.
|
// Server buffers multiple requests before initialized and processes them in order.
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn server_init_buffers_multiple_requests_before_initialized() {
|
async fn server_init_buffers_multiple_requests_before_initialized() {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue