fix(server): safely drop detached tokio join handles and refactor timestamp boilerplate
This commit is contained in:
Riz Ashraf committed
1 parent
13c7eca4c2
commit
478655698e
15 files changed
+457
-312
No files matched your search
+94
-37
@@ -199,8 +199,10 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
|
||||
let msgid = &arr[1];
|
||||
let msgid_str = format!("{msgid:?}");
|
||||
|
||||
if let Some(reply_sender) =
|
||||
pending_clone2.lock().unwrap_or_else(|e| e.into_inner()).remove(&msgid_str)
|
||||
if let Some(reply_sender) = pending_clone2
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.remove(&msgid_str)
|
||||
{
|
||||
let _ = reply_sender.send(Ok(val));
|
||||
}
|
||||
@@ -555,7 +557,9 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
tracing::info!("{} MCP server started", app_name);
|
||||
let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
|
||||
loop {
|
||||
let raw_msg = if let Some(m) = mcp_stdio::read_mcp_message(&mut stdin).await { m } else {
|
||||
let raw_msg = if let Some(m) = mcp_stdio::read_mcp_message(&mut stdin).await {
|
||||
m
|
||||
} else {
|
||||
tracing::info!("Stdin closed, exiting loop");
|
||||
break;
|
||||
};
|
||||
@@ -595,7 +599,6 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
};
|
||||
|
||||
match msg.method.as_str() {
|
||||
|
||||
"initialize" => {
|
||||
let init = rmcp::model::InitializeResult::new(
|
||||
rmcp::model::ServerCapabilities::builder()
|
||||
@@ -719,7 +722,13 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
let cmd = format!("e {escaped_file} | {line} | normal! zz");
|
||||
match send_nvim_command(&cmd).await {
|
||||
Ok(()) => {
|
||||
send_text_result!(id.clone(), format!("Successfully jumped to {} line {}", file, line));
|
||||
send_text_result!(
|
||||
id.clone(),
|
||||
format!(
|
||||
"Successfully jumped to {} line {}",
|
||||
file, line
|
||||
)
|
||||
);
|
||||
}
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
@@ -752,7 +761,10 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
) {
|
||||
match set_nvim_diagnostics(line, message).await {
|
||||
Ok(()) => {
|
||||
send_text_result!(id.clone(), format!("Successfully set diagnostic on line {}", line));
|
||||
send_text_result!(
|
||||
id.clone(),
|
||||
format!("Successfully set diagnostic on line {}", line)
|
||||
);
|
||||
}
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
@@ -804,23 +816,32 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
}
|
||||
|
||||
"nvim_open_file" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
vim.cmd('edit ' .. vim.fn.fnameescape(args.file))
|
||||
if args.filetype and args.filetype ~= '' then
|
||||
vim.bo.filetype = args.filetype
|
||||
end
|
||||
return 'Opened file ' .. args.file
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_open_buffer" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local buf = vim.api.nvim_create_buf(true, true)
|
||||
if args.name and args.name ~= '' then
|
||||
@@ -835,29 +856,39 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
end
|
||||
vim.api.nvim_win_set_buf(0, buf)
|
||||
return 'Opened buffer ' .. tostring(buf)
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_close_buffer" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||
local force = args.force or false
|
||||
vim.api.nvim_buf_delete(buf, {{ force = force }})
|
||||
return 'Closed buffer ' .. tostring(buf)
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_split_window" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local cmd = args.direction == 'horizontal' and 'split' or 'vsplit'
|
||||
vim.cmd(cmd)
|
||||
@@ -867,15 +898,20 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
vim.api.nvim_win_set_buf(0, args.buf_id)
|
||||
end
|
||||
return 'Split window created'
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_reload_buffer" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||
vim.api.nvim_buf_call(buf, function()
|
||||
@@ -886,7 +922,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
end
|
||||
end)
|
||||
return 'Reloaded buffer ' .. tostring(buf)
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
@@ -903,8 +940,12 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
}
|
||||
}
|
||||
"nvim_set_quickfix" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local items = args.items or {{}}
|
||||
local action = ' '
|
||||
@@ -914,14 +955,18 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
vim.fn.setqflist(items, action)
|
||||
vim.cmd('copen')
|
||||
return 'Populated quickfix with ' .. tostring(#items) .. ' items'
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_highlight_lines" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!("
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||
@@ -952,8 +997,12 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
}
|
||||
}
|
||||
"nvim_get_messages" => {
|
||||
let json_str = serde_json::to_string(args).unwrap_or_else(|_| "{}".to_string()).replace('\\', "\\\\").replace('\'', "\\'");
|
||||
let code = format!("
|
||||
let json_str = serde_json::to_string(args)
|
||||
.unwrap_or_else(|_| "{}".to_string())
|
||||
.replace('\\', "\\\\")
|
||||
.replace('\'', "\\'");
|
||||
let code = format!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local msg = vim.fn.execute('messages')
|
||||
local lines = vim.split(msg, '\\n')
|
||||
@@ -965,7 +1014,8 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
return table.concat(tail_lines, '\\n')
|
||||
end
|
||||
return msg
|
||||
");
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
@@ -992,14 +1042,26 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
if let Some(code) = args.get("code").and_then(|v| v.as_str()) {
|
||||
// BAKE IN: Block interactive prompts that cause server deadlocks
|
||||
let lower_code = code.to_lowercase();
|
||||
if lower_code.contains("vim.fn.input") || lower_code.contains("vim.ui.select") || lower_code.contains("vim.fn.confirm") || lower_code.contains("vim.ui.input") {
|
||||
if lower_code.contains("vim.fn.input")
|
||||
|| lower_code.contains("vim.ui.select")
|
||||
|| lower_code.contains("vim.fn.confirm")
|
||||
|| lower_code.contains("vim.ui.input")
|
||||
{
|
||||
send_error(id, -32600, "CRITICAL ERROR: Interactive Lua functions (input, select, confirm) are strictly forbidden via MCP as they will hang the headless server.").await;
|
||||
return;
|
||||
}
|
||||
// Block the 'c' confirm flag in vim.cmd substitutions
|
||||
if (code.contains("vim.cmd") || code.contains("vim.api.nvim_command")) && code.contains("%s") && (code.contains("gc'") || code.contains("gc\"") || code.contains("gc\n") || code.contains("c'") || code.contains("c\"")) {
|
||||
send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/g' or '/ge' instead.").await;
|
||||
return;
|
||||
if (code.contains("vim.cmd")
|
||||
|| code.contains("vim.api.nvim_command"))
|
||||
&& code.contains("%s")
|
||||
&& (code.contains("gc'")
|
||||
|| code.contains("gc\"")
|
||||
|| code.contains("gc\n")
|
||||
|| code.contains("c'")
|
||||
|| code.contains("c\""))
|
||||
{
|
||||
send_error(id, -32600, "CRITICAL ERROR: The 'c' (confirm) flag in Neovim substitutions is strictly forbidden via MCP as it triggers an interactive prompt that hangs the headless server. Use '/g' or '/ge' instead.").await;
|
||||
return;
|
||||
}
|
||||
|
||||
match execute_nvim_lua(code).await {
|
||||
@@ -1116,8 +1178,3 @@ mod tests {
|
||||
assert!(req.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user