From 410b0c42cae8861071b2301b0c8d772bd08ea2a2 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Thu, 8 Oct 2026 23:41:57 +0100 Subject: [PATCH] feat(ollama): enforce JSON logits & add nvim_buffer edit action - Added format parameter to Ollama generate for JSON logit enforcement. - Implemented safe edit action in nvim_buffer MCP tool to replace raw Lua. --- nvim-core/src/lib.rs | 50 ++++++++++++++++++++++++++++--- server/src/handlers/graph.rs | 2 +- server/src/handlers/meta.rs | 4 +-- server/src/handlers/workspaces.rs | 2 +- server/src/lib.rs | 1 + server/src/ollama.rs | 5 ++++ 6 files changed, 56 insertions(+), 8 deletions(-) diff --git a/nvim-core/src/lib.rs b/nvim-core/src/lib.rs index b3687c7..137333e 100644 --- a/nvim-core/src/lib.rs +++ b/nvim-core/src/lib.rs @@ -1088,16 +1088,17 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) { "properties": { "action": { "type": "string", - "enum": ["get_active", "read", "open", "create_scratch", "save", "reload", "close", "list", "search"], + "enum": ["get_active", "read", "open", "create_scratch", "save", "reload", "close", "list", "search", "edit"], "description": "Action to perform on the buffer" }, - "file": { "type": "string", "description": "Target file path (for open, read, search)" }, + "file": { "type": "string", "description": "Target file path (for open, read, search, edit)" }, "buf_id": { "type": "integer", "description": "Buffer ID (for close, reload, or split)" }, "content": { "type": "string", "description": "Initial text content (for create_scratch)" }, "name": { "type": "string", "description": "Buffer display name (for create_scratch)" }, "filetype": { "type": "string", "description": "Filetype syntax (for open, create_scratch)" }, - "start_line": { "type": "integer", "description": "1-indexed start line (for read)" }, - "end_line": { "type": "integer", "description": "1-indexed end line (for read)" }, + "start_line": { "type": "integer", "description": "1-indexed start line (for read, edit)" }, + "end_line": { "type": "integer", "description": "1-indexed end line (for read, edit)" }, + "replacement_content": { "type": "string", "description": "New content for the specified line range (for edit)" }, "pattern": { "type": "string", "description": "Regex or substring pattern to search for (for search)" }, "force": { "type": "boolean", "description": "Force reload or close (for reload, close)" } }, @@ -1324,6 +1325,47 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) { execute_nvim_lua_with_args(code, vec![args_rmp]) ) } + "edit" => { + let args_rmp = json_to_rmpv(args); + let code = " + local args = ... + local file = args.file + if not file or file == '' or file == vim.NIL then + error(\"Missing 'file'\") + end + local start_line = tonumber(args.start_line) + if not start_line then + error(\"Missing 'start_line'\") + end + local end_line = tonumber(args.end_line) + if not end_line then + error(\"Missing 'end_line'\") + end + local content = args.replacement_content + if not content or content == vim.NIL then + error(\"Missing 'replacement_content'\") + end + + local bufnr = vim.fn.bufnr(file, true) + if vim.fn.bufloaded(bufnr) == 0 then + vim.fn.bufload(bufnr) + end + + local lines = vim.split(content, '\\n', { plain = true }) + vim.api.nvim_buf_set_lines(bufnr, start_line - 1, end_line, false, lines) + + local win = vim.fn.bufwinid(bufnr) + if win ~= -1 then + pcall(vim.api.nvim_win_set_cursor, win, {start_line, 0}) + end + + return 'Successfully edited lines ' .. tostring(start_line) .. ' to ' .. tostring(end_line) .. ' in ' .. file + "; + handle_lua_result!( + id, + execute_nvim_lua_with_args(code, vec![args_rmp]) + ) + } "open" => { let args_rmp = json_to_rmpv(args); let code = " diff --git a/server/src/handlers/graph.rs b/server/src/handlers/graph.rs index b5232d4..cb421c2 100644 --- a/server/src/handlers/graph.rs +++ b/server/src/handlers/graph.rs @@ -824,7 +824,7 @@ impl McpTool for CondenseEntityHandler { req.entity_name, unique_obs.join("\n- ") ); - if let Ok(summary) = state.ollama.generate(&prompt, None, None).await { + if let Ok(summary) = state.ollama.generate(&prompt, None, None, None).await { let lines: Vec = summary .lines() .map(|l| l.trim().trim_start_matches('-').trim().to_string()) diff --git a/server/src/handlers/meta.rs b/server/src/handlers/meta.rs index bc8523d..31a4b45 100644 --- a/server/src/handlers/meta.rs +++ b/server/src/handlers/meta.rs @@ -33,7 +33,7 @@ impl McpTool for LogErrorFixHandler { ); if let Ok(summary) = state .ollama - .generate(&prompt, Some(&state.ollama.reasoning_model), None) + .generate(&prompt, Some(&state.ollama.reasoning_model), None, None) .await { let clean = summary.trim(); @@ -207,7 +207,7 @@ impl McpTool for LogCodeChangeHandler { "Summarize in 1 concise sentence the architectural impact of changing file '{}': {}", req.file_path, description ); - if let Ok(summary) = state.ollama.generate(&prompt, None, None).await { + if let Ok(summary) = state.ollama.generate(&prompt, None, None, None).await { let clean = summary.trim(); if !clean.is_empty() { description = format!("{} (AI Summary: {})", description, clean); diff --git a/server/src/handlers/workspaces.rs b/server/src/handlers/workspaces.rs index 0ca803c..c2ad03d 100644 --- a/server/src/handlers/workspaces.rs +++ b/server/src/handlers/workspaces.rs @@ -359,7 +359,7 @@ impl McpTool for SemanticCodeSearchHandler { ); if let Ok(summary) = state .ollama - .generate(&prompt, None, Some("Respond clearly and concisely.")) + .generate(&prompt, None, Some("Respond clearly and concisely."), None) .await { out.push_str("\n\n--- Local GraphRAG Summary ---\n"); diff --git a/server/src/lib.rs b/server/src/lib.rs index b0b7188..daaab5d 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -308,6 +308,7 @@ pub async fn memory_consolidation_worker(state: Arc) { &prompt, None, Some("You are a helpful JSON-only data deduplication assistant. Output only JSON."), + Some("json"), ) .await { diff --git a/server/src/ollama.rs b/server/src/ollama.rs index 359308b..13496c3 100644 --- a/server/src/ollama.rs +++ b/server/src/ollama.rs @@ -27,6 +27,8 @@ struct GenerateRequest<'a> { options: Option, #[serde(skip_serializing_if = "Option::is_none")] keep_alive: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + format: Option<&'a str>, } #[derive(Deserialize)] @@ -113,6 +115,7 @@ impl OllamaClient { prompt: &str, model_override: Option<&str>, system: Option<&str>, + format: Option<&str>, ) -> Result { let model = model_override.unwrap_or(&self.coder_model); let url = format!("{}/api/generate", self.base_url.trim_end_matches('/')); @@ -128,6 +131,7 @@ impl OllamaClient { "num_predict": 4096 // Give reasoning models plenty of output room })), keep_alive: Some("1h"), + format, }; let res = self @@ -172,6 +176,7 @@ impl OllamaClient { "num_predict": 1024 })), keep_alive: Some("1h"), + format: None, }; let res = self