perf(nvim-core): eliminate JSON serialization overhead and potential Lua injection across all tools by passing arguments directly to nvim_exec_lua
This commit is contained in:
Riz Ashraf committed
1 parent
78e94c6aea
commit
05672cc884
1 file changed
+83
-84
+83
-84
@@ -460,12 +460,49 @@ fn rmpv_to_json(val: &rmpv::Value) -> serde_json::Value {
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_nvim_lua(code: &str) -> Result<String, String> {
|
||||
fn json_to_rmpv(val: &serde_json::Value) -> rmpv::Value {
|
||||
match val {
|
||||
serde_json::Value::Null => rmpv::Value::Nil,
|
||||
serde_json::Value::Bool(b) => rmpv::Value::Boolean(*b),
|
||||
serde_json::Value::Number(n) => {
|
||||
if let Some(i) = n.as_i64() {
|
||||
rmpv::Value::Integer(i.into())
|
||||
} else if let Some(u) = n.as_u64() {
|
||||
rmpv::Value::Integer(u.into())
|
||||
} else if let Some(f) = n.as_f64() {
|
||||
rmpv::Value::F64(f)
|
||||
} else {
|
||||
rmpv::Value::Nil
|
||||
}
|
||||
}
|
||||
serde_json::Value::String(s) => rmpv::Value::String(s.as_str().into()),
|
||||
serde_json::Value::Array(arr) => {
|
||||
rmpv::Value::Array(arr.iter().map(json_to_rmpv).collect())
|
||||
}
|
||||
serde_json::Value::Object(obj) => {
|
||||
rmpv::Value::Map(obj.iter().map(|(k, v)| {
|
||||
(rmpv::Value::String(k.as_str().into()), json_to_rmpv(v))
|
||||
}).collect())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_nvim_lua_with_args(code: &str, args: Vec<rmpv::Value>) -> Result<String, String> {
|
||||
let result = call_nvim_method("nvim_exec_lua", vec![
|
||||
rmpv::Value::String(code.into()),
|
||||
rmpv::Value::Array(vec![]),
|
||||
rmpv::Value::Array(args),
|
||||
]).await?;
|
||||
Ok(serde_json::to_string_pretty(&rmpv_to_json(&result)).unwrap_or_default())
|
||||
let json_res = rmpv_to_json(&result);
|
||||
// If it's a simple string, return it unquoted, otherwise format it pretty
|
||||
if let serde_json::Value::String(s) = json_res {
|
||||
Ok(s)
|
||||
} else {
|
||||
Ok(serde_json::to_string_pretty(&json_res).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_nvim_lua(code: &str) -> Result<String, String> {
|
||||
execute_nvim_lua_with_args(code, vec![]).await
|
||||
}
|
||||
|
||||
macro_rules! send_text_result {
|
||||
@@ -744,33 +781,24 @@ 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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
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 {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local buf = vim.api.nvim_create_buf(true, true)
|
||||
if args.name and args.name ~= '' then
|
||||
pcall(vim.api.nvim_buf_set_name, buf, args.name)
|
||||
@@ -784,40 +812,30 @@ 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 {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
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 }})
|
||||
vim.api.nvim_buf_delete(buf, { force = force })
|
||||
return 'Closed buffer ' .. tostring(buf)
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local cmd = args.direction == 'horizontal' and 'split' or 'vsplit'
|
||||
vim.cmd(cmd)
|
||||
if args.file and args.file ~= '' then
|
||||
@@ -826,21 +844,16 @@ 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 {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||
vim.api.nvim_buf_call(buf, function()
|
||||
if args.force then
|
||||
@@ -850,9 +863,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 {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
@@ -868,14 +880,10 @@ 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!(
|
||||
"
|
||||
local args = vim.json.decode('{json_str}')
|
||||
local items = args.items or {{}}
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local items = args.items or {}
|
||||
local action = ' '
|
||||
if args.action == 'append' then action = 'a' end
|
||||
if args.action == 'prepend' then action = 'p' end
|
||||
@@ -883,20 +891,16 @@ 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 {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).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 code = format!("
|
||||
local args = vim.json.decode('{json_str}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local buf = args.buf_id or vim.api.nvim_get_current_buf()
|
||||
local group = args.group or 'IncSearch'
|
||||
local ns = vim.api.nvim_create_namespace('antigravity_highlight')
|
||||
@@ -918,33 +922,28 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
|
||||
end, duration)
|
||||
end
|
||||
return 'Highlighted lines ' .. tostring(args.start_line) .. ' to ' .. tostring(args.end_line)
|
||||
");
|
||||
match execute_nvim_lua(&code).await {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
}
|
||||
"nvim_get_messages" => {
|
||||
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}')
|
||||
let args_rmp = json_to_rmpv(args);
|
||||
let code = "
|
||||
local args = ...
|
||||
local msg = vim.fn.execute('messages')
|
||||
local lines = vim.split(msg, '\\n')
|
||||
if args.tail and args.tail > 0 and #lines > args.tail then
|
||||
local tail_lines = {{}}
|
||||
local tail_lines = {}
|
||||
for i = #lines - args.tail + 1, #lines do
|
||||
table.insert(tail_lines, lines[i])
|
||||
end
|
||||
return table.concat(tail_lines, '\\n')
|
||||
end
|
||||
return msg
|
||||
"
|
||||
);
|
||||
match execute_nvim_lua(&code).await {
|
||||
";
|
||||
match execute_nvim_lua_with_args(code, vec![args_rmp]).await {
|
||||
Ok(res) => send_text_result!(id.clone(), res),
|
||||
Err(e) => send_error(id, -32603, &e).await,
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user