Skip to content

Commit ffbfbb8

Browse files
dingfelidingfeli
committed
adds pid field to messenger message
1 parent decbd4d commit ffbfbb8

File tree

4 files changed

+91
-9
lines changed

4 files changed

+91
-9
lines changed

‎crates/chat-cli/src/cli/chat/server_messenger.rs‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,24 +19,30 @@ pub enum UpdateEventMessage {
1919
ToolsListResult {
2020
server_name: String,
2121
result: eyre::Result<ToolsListResult>,
22+
pid: Option<u32>,
2223
},
2324
PromptsListResult {
2425
server_name: String,
2526
result: eyre::Result<PromptsListResult>,
27+
pid: Option<u32>,
2628
},
2729
ResourcesListResult {
2830
server_name: String,
2931
result: eyre::Result<ResourcesListResult>,
32+
pid: Option<u32>,
3033
},
3134
ResourceTemplatesListResult {
3235
server_name: String,
3336
result: eyre::Result<ResourceTemplatesListResult>,
37+
pid: Option<u32>,
3438
},
3539
InitStart {
3640
server_name: String,
41+
pid: Option<u32>,
3742
},
3843
Deinit {
3944
server_name: String,
45+
pid: Option<u32>,
4046
},
4147
}
4248

@@ -58,6 +64,7 @@ impl ServerMessengerBuilder {
5864
ServerMessenger {
5965
server_name,
6066
update_event_sender: self.update_event_sender.clone(),
67+
pid: None,
6168
}
6269
}
6370
}
@@ -66,6 +73,7 @@ impl ServerMessengerBuilder {
6673
pub struct ServerMessenger {
6774
pub server_name: String,
6875
pub update_event_sender: Sender<UpdateEventMessage>,
76+
pub pid: Option<u32>,
6977
}
7078

7179
#[async_trait::async_trait]
@@ -76,6 +84,7 @@ impl Messenger for ServerMessenger {
7684
.send(UpdateEventMessage::ToolsListResult {
7785
server_name: self.server_name.clone(),
7886
result,
87+
pid: self.pid,
7988
})
8089
.await
8190
.map_err(|e| MessengerError::Custom(e.to_string()))?)
@@ -87,6 +96,7 @@ impl Messenger for ServerMessenger {
8796
.send(UpdateEventMessage::PromptsListResult {
8897
server_name: self.server_name.clone(),
8998
result,
99+
pid: self.pid,
90100
})
91101
.await
92102
.map_err(|e| MessengerError::Custom(e.to_string()))?)
@@ -101,6 +111,7 @@ impl Messenger for ServerMessenger {
101111
.send(UpdateEventMessage::ResourcesListResult {
102112
server_name: self.server_name.clone(),
103113
result,
114+
pid: self.pid,
104115
})
105116
.await
106117
.map_err(|e| MessengerError::Custom(e.to_string()))?)
@@ -115,6 +126,7 @@ impl Messenger for ServerMessenger {
115126
.send(UpdateEventMessage::ResourceTemplatesListResult {
116127
server_name: self.server_name.clone(),
117128
result,
129+
pid: self.pid,
118130
})
119131
.await
120132
.map_err(|e| MessengerError::Custom(e.to_string()))?)
@@ -125,6 +137,7 @@ impl Messenger for ServerMessenger {
125137
.update_event_sender
126138
.send(UpdateEventMessage::InitStart {
127139
server_name: self.server_name.clone(),
140+
pid: self.pid,
128141
})
129142
.await
130143
.map_err(|e| MessengerError::Custom(e.to_string()))?)
@@ -133,8 +146,9 @@ impl Messenger for ServerMessenger {
133146
fn send_deinit_msg(&self) {
134147
let sender = self.update_event_sender.clone();
135148
let server_name = self.server_name.clone();
149+
let pid = self.pid;
136150
tokio::spawn(async move {
137-
let _ = sender.send(UpdateEventMessage::Deinit { server_name }).await;
151+
let _ = sender.send(UpdateEventMessage::Deinit { server_name, pid }).await;
138152
});
139153
}
140154

‎crates/chat-cli/src/cli/chat/tool_manager.rs‎

Lines changed: 69 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ use tokio::sync::{
4747
use tokio::task::JoinHandle;
4848
use tracing::{
4949
error,
50+
info,
5051
warn,
5152
};
5253

@@ -527,7 +528,18 @@ impl ToolManagerBuilder {
527528
// request method on the mcp client no longer buffers all the pages from
528529
// list calls.
529530
match msg {
530-
UpdateEventMessage::ToolsListResult { server_name, result } => {
531+
UpdateEventMessage::ToolsListResult {
532+
server_name,
533+
result,
534+
pid,
535+
} => {
536+
let pid = pid.unwrap();
537+
if !is_process_running(pid) {
538+
info!(
539+
"Received tool list result from {server_name} but its associated process {pid} is no longer running. Ignoring."
540+
);
541+
return;
542+
}
531543
let time_taken =
532544
loading_servers
533545
.remove(&server_name)
@@ -708,9 +720,19 @@ impl ToolManagerBuilder {
708720
}
709721
}
710722
},
711-
UpdateEventMessage::PromptsListResult { server_name, result } => match result {
712-
Ok(prompt_list_result) => {
713-
tracing::error!("## swap: received prompts from {server_name}");
723+
UpdateEventMessage::PromptsListResult {
724+
server_name,
725+
result,
726+
pid,
727+
} => match result {
728+
Ok(prompt_list_result) if pid.is_some() => {
729+
let pid = pid.unwrap();
730+
if !is_process_running(pid) {
731+
info!(
732+
"Received prompt list result from {server_name} but its associated process {pid} is no longer running. Ignoring."
733+
);
734+
return;
735+
}
714736
// We first need to clear all the PromptGets that are associated with
715737
// this server because PromptsListResult is declaring what is available
716738
// (and not the diff)
@@ -740,6 +762,9 @@ impl ToolManagerBuilder {
740762
});
741763
}
742764
},
765+
Ok(_) => {
766+
error!("Received prompt list result without pid from {server_name}. Ignoring.");
767+
},
743768
Err(e) => {
744769
error!("Error fetching prompts from server {server_name}: {:?}", e);
745770
let mut buf_writer = BufWriter::new(&mut *record_temp_buf);
@@ -761,16 +786,18 @@ impl ToolManagerBuilder {
761786
UpdateEventMessage::ResourcesListResult {
762787
server_name: _,
763788
result: _,
789+
pid: _,
764790
} => {},
765791
UpdateEventMessage::ResourceTemplatesListResult {
766792
server_name: _,
767793
result: _,
794+
pid: _,
768795
} => {},
769-
UpdateEventMessage::InitStart { server_name } => {
796+
UpdateEventMessage::InitStart { server_name, .. } => {
770797
pending_clone.write().await.insert(server_name.clone());
771798
loading_servers.insert(server_name, std::time::Instant::now());
772799
},
773-
UpdateEventMessage::Deinit { server_name } => {
800+
UpdateEventMessage::Deinit { server_name, .. } => {
774801
// Only prompts are stored here so we'll just be clearing that
775802
// In the future if we are also storing tools, we need to make sure that
776803
// the tools are also pruned.
@@ -821,9 +848,11 @@ impl ToolManagerBuilder {
821848
debug_assert!(messenger_builder.is_some());
822849
let messenger_builder = messenger_builder.unwrap();
823850
for (mut name, init_res) in pre_initialized {
824-
let messenger = messenger_builder.build_with_name(name.clone());
851+
let mut messenger = messenger_builder.build_with_name(name.clone());
825852
match init_res {
826853
Ok(mut client) => {
854+
let pid = client.get_pid();
855+
messenger.pid = pid;
827856
client.assign_messenger(Box::new(messenger));
828857
let mut client = Arc::new(client);
829858
while let Some(collided_client) = clients.insert(name.clone(), client) {
@@ -1611,6 +1640,39 @@ fn sanitize_name(orig: String, regex: &regex::Regex, hasher: &mut impl Hasher) -
16111640
}
16121641
}
16131642

1643+
// Add this function to check if a process is still running
1644+
fn is_process_running(pid: u32) -> bool {
1645+
#[cfg(unix)]
1646+
{
1647+
// On Unix systems, we can use kill with signal 0 to check if process exists
1648+
std::process::Command::new("ps")
1649+
.arg("-p")
1650+
.arg(pid.to_string())
1651+
.output()
1652+
.map(|output| output.status.success())
1653+
.unwrap_or(false)
1654+
}
1655+
#[cfg(windows)]
1656+
{
1657+
// On Windows, try to open the process handle
1658+
use std::ptr;
1659+
1660+
use winapi::um::handleapi::CloseHandle;
1661+
use winapi::um::processthreadsapi::OpenProcess;
1662+
use winapi::um::winnt::PROCESS_QUERY_INFORMATION;
1663+
1664+
unsafe {
1665+
let handle = OpenProcess(PROCESS_QUERY_INFORMATION, 0, pid);
1666+
if handle != ptr::null_mut() {
1667+
CloseHandle(handle);
1668+
true
1669+
} else {
1670+
false
1671+
}
1672+
}
1673+
}
1674+
}
1675+
16141676
fn queue_success_message(name: &str, time_taken: &str, output: &mut impl Write) -> eyre::Result<()> {
16151677
Ok(queue!(
16161678
output,

‎crates/chat-cli/src/cli/chat/tools/custom_tool.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,12 @@ impl CustomToolClient {
170170
}
171171
}
172172

173+
pub fn get_pid(&self) -> Option<u32> {
174+
match self {
175+
CustomToolClient::Stdio { client, .. } => client.server_process_id.as_ref().map(|pid| pid.as_u32()),
176+
}
177+
}
178+
173179
#[allow(dead_code)]
174180
pub async fn notify(&self, method: &str, params: Option<serde_json::Value>) -> Result<()> {
175181
match self {

‎crates/chat-cli/src/mcp_client/client.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ pub struct Client<T: Transport> {
124124
server_name: String,
125125
transport: Arc<T>,
126126
timeout: u64,
127-
server_process_id: Option<Pid>,
127+
pub server_process_id: Option<Pid>,
128128
client_info: serde_json::Value,
129129
current_id: Arc<AtomicU64>,
130130
pub messenger: Option<Box<dyn Messenger>>,

0 commit comments

Comments
 (0)