fix: update peer info on duplicate initialize (#862)
This commit is contained in:
parent
2d3d1879ad
commit
2536a05992
2 changed files with 69 additions and 12 deletions
|
|
@ -184,9 +184,7 @@ macro_rules! server_handler_methods {
|
|||
request: InitializeRequestParams,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> 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()))
|
||||
}
|
||||
fn complete(
|
||||
|
|
|
|||
|
|
@ -384,7 +384,7 @@ pub struct Peer<R: ServiceRole> {
|
|||
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
||||
request_id_provider: Arc<dyn RequestIdProvider>,
|
||||
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> {
|
||||
|
|
@ -423,7 +423,7 @@ impl<R: ServiceRole> Peer<R> {
|
|||
tx,
|
||||
request_id_provider,
|
||||
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,
|
||||
)
|
||||
|
|
@ -484,16 +484,14 @@ impl<R: ServiceRole> Peer<R> {
|
|||
peer: self.clone(),
|
||||
})
|
||||
}
|
||||
pub fn peer_info(&self) -> Option<&R::PeerInfo> {
|
||||
self.info.get()
|
||||
/// Snapshot of the peer's handshake info.
|
||||
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) {
|
||||
if self.info.initialized() {
|
||||
tracing::warn!("trying to set peer info, which is already initialized");
|
||||
} else {
|
||||
let _ = self.info.set(info);
|
||||
}
|
||||
*self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info));
|
||||
}
|
||||
|
||||
pub fn is_transport_closed(&self) -> bool {
|
||||
|
|
|
|||
Loading…
Reference in a new issue