From 7079ad2e9bae17f7dc342ae249f32b8c46226a6a Mon Sep 17 00:00:00 2001 From: WuJunkai2004 Date: Wed, 4 Mar 2026 19:21:56 +0800 Subject: [PATCH 1/3] Add workspace-safe I/O, shell helpers and some toolkits --- py/utils.py | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/py/utils.py b/py/utils.py index c6e946d..bb79dee 100644 --- a/py/utils.py +++ b/py/utils.py @@ -10,10 +10,65 @@ import traceback import configparser import base64 +from pathlib import Path utils_py_imported = True DEFAULT_ROLE_NAME = 'default' +WORKDIR = Path.cwd() + +def safe_path(subpath: str) -> Path: + path = (WORKDIR / subpath).resolve() + if not path.is_relative_to(WORKDIR): + raise ValueError(f"Path escapes workspace: {subpath}") + return path + +def run_bash(command: str) -> str: + dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"] + if any(d in command for d in dangerous): + return "Error: Dangerous command blocked" + try: + r = subprocess.run( + command, + shell=True, + cwd=WORKDIR, + capture_output=True, + text=True, + timeout=120, + ) + out = (r.stdout + r.stderr).strip() + return out[:50000] if out else "(no output)" + except subprocess.TimeoutExpired: + return "Error: Timeout (120s)" + +def run_read(path: str, max_line: int = -1) -> str: + try: + lines = safe_path(path).read_text().splitlines() + if max_line != -1 and max_line < len(lines): + lines = lines[:max_line] + [f"... ({len(lines) - max_line} more)"] + return "\n".join(lines)[:50000] + except Exception as e: + return f"Error: {e}" + +def run_write(path: str, content: str) -> str: + try: + fp = safe_path(path) + fp.parent.mkdir(parents=True, exist_ok=True) + fp.write_text(content) + return f"Wrote {len(content)} bytes to {path}" + except Exception as e: + return f"Error: {e}" + +def run_edit(path: str, old_text: str, new_text: str) -> str: + try: + fp = safe_path(path) + c = fp.read_text() + if old_text not in c: + return f"Error: Text not found in {path}" + fp.write_text(c.replace(old_text, new_text, 1)) + return f"Edited {path}" + except Exception as e: + return f"Error: {e}" _vimai_thread_is_debug_active = vim.eval("g:vim_ai_debug") == "1" _vimai_thread_log_file_path = vim.eval("g:vim_ai_debug_log_file") From 96c99232543436a1b76f8832088d589e09552128 Mon Sep 17 00:00:00 2001 From: WuJunkai2004 Date: Wed, 4 Mar 2026 19:22:44 +0800 Subject: [PATCH 2/3] Add MCP server and manager utilities --- py/extensions.py | 126 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 py/extensions.py diff --git a/py/extensions.py b/py/extensions.py new file mode 100644 index 0000000..4d3540c --- /dev/null +++ b/py/extensions.py @@ -0,0 +1,126 @@ +import json +import subprocess + + +class Pipe: + def __init__(self, command: list[str]): + self.process = subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + + def write(self, data: str): + if self.process.stdin: + self.process.stdin.write(data) + self.process.stdin.flush() + + def readline(self) -> str: + if self.process.stdout: + return self.process.stdout.readline().strip() + return "" + + def terminate(self): + self.process.terminate() + + +class MCPServer: + def __init__(self, name: str, command: str, args: list[str]): + self.name = name + self.command = [command] + args + self._id_counter = 0 + self.process: Pipe = Pipe(self.command) + + # 1. Initialize handshake + self._send_request( + "initialize", + { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "fittencode-vim-agent", "version": "1.0.0"}, + }, + ) + # 2. Send initialized notification + notification = {"jsonrpc": "2.0", "method": "notifications/initialized"} + self.process.write(json.dumps(notification) + "\n") + # 3. Cache available tools + self.tools_cache = self._send_request("tools/list", {}).get("tools", []) + + def _send_request(self, method: str, params: dict) -> dict: + self._id_counter += 1 + req_id = self._id_counter + request = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params} + + try: + self.process.write(json.dumps(request) + "\n") + + line = self.process.readline() + if not line: + return {"error": "Server disconnected"} + + response = json.loads(line) + if response.get("id") == req_id: + if "error" in response: + return {"error": response["error"]} + return response.get("result", {}) + except Exception as e: + return {"error": str(e)} + return {} + + def call_tool(self, tool_name: str, arguments: dict) -> str: + res = self._send_request( + "tools/call", {"name": tool_name, "arguments": arguments} + ) + + if "error" in res: + return f"MCP Error: {res['error']}" + + content_list = res.get("content", []) + texts = [c["text"] for c in content_list if c.get("type") == "text"] + return "\n".join(texts) if texts else "Success (no output)" + + def stop(self): + self.process.terminate() + + +class MCPManager: + """ + Manages multiple MCP servers and provides a unified tool interface for the Agent. + """ + + def __init__(self): + self.servers: dict[str, MCPServer] = {} + self.schema_cache: list[dict] = [] + + def add_server(self, name: str, command: str, args: list[str] | None = None): + """Register and start an MCP server.""" + try: + self.servers[name] = MCPServer(name, command, args or []) + except Exception: + # Silent fail for robustness in CLI + pass + finally: + for tool in self.servers[name].tools_cache: + self.schema_cache.append( + { + "tool": tool["name"], + "description": tool.get("description", ""), + "arguments": tool.get("inputSchema", {}).get("properties", {}), + } + ) + + def get_tools_schema(self) -> list: + return self.schema_cache + + def execute(self, tool_name: str, **kwargs) -> str: + for server in self.servers.values(): + if any(t["name"] == tool_name for t in server.tools_cache): + return server.call_tool(tool_name, kwargs) + return f"Unknown tool: {tool_name}" + + def shutdown(self): + for server in self.servers.values(): + server.stop() From e1dfd119a127ec7802bd716852a7de7435c94eb7 Mon Sep 17 00:00:00 2001 From: WuJunkai2004 Date: Wed, 4 Mar 2026 19:34:18 +0800 Subject: [PATCH 3/3] Add tool execution loop and agent tests Inject tools description into system messages and parse tool calls from assistant responses using ```tool```/```json``` blocks. Add MCPManager integration for external tools and a built-in tool_map for bash/read_file/write_file/edit_file. Execute tools iteratively, append results as user messages, limit loop iterations, and call MCPManager shutdown. Add tests covering write/read/edit/bash and safety checks. --- py/chat.py | 127 +++++++++++++++++++++++++++++++++++++------- tests/agent_test.py | 48 +++++++++++++++++ 2 files changed, 157 insertions(+), 18 deletions(-) create mode 100644 tests/agent_test.py diff --git a/py/chat.py b/py/chat.py index 781b065..d0b8f7d 100644 --- a/py/chat.py +++ b/py/chat.py @@ -4,9 +4,29 @@ import copy import json import traceback +import re +from extensions import MCPManager chat_py_imported = True +TOOLS_DESCRIPTION = """ +To perform actions, you MUST use the following format: +```tool +{ + "tool": "tool_name", + "arguments": { "arg1": "val1" } +} +``` + +Available tools: +- bash(command: str): Run a shell command. +- read_file(path: str, max_line: int = -1): Read file contents. +- write_file(path: str, content: str): Write content to file. +- edit_file(path: str, old_text: str, new_text: str): Replace exact text in file. + +If you are providing a JSON example for the user, use standard ```json blocks. Only use ```tool for actual execution. +""" + def _populate_options(provider, options, default_options, show_default = False): vim.command("normal! O[chat]") vim.command("normal! o") @@ -103,6 +123,14 @@ def initialize_chat_window(): initial_prompt = '\n'.join(options.get('initial_prompt', [])) initial_messages = parse_chat_messages(initial_prompt) + # Inject tools description if it's a chat + if command_type == 'chat': + tools_msg = TOOLS_DESCRIPTION + if initial_messages and initial_messages[0]['role'] == 'system': + initial_messages[0]['content'][0]['text'] += "\n\n" + tools_msg + else: + initial_messages.insert(0, {'role': 'system', 'content': [{'type': 'text', 'text': tools_msg}]}) + chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))') print_debug(f"[{command_type}] text:\n" + chat_content) chat_messages = parse_chat_messages(chat_content) @@ -164,29 +192,91 @@ def __init__(self, context, messages, provider): self.provider = provider self.done = False self.lock = threading.RLock() + self.mcp_manager = MCPManager() def run(self): print_debug("AI_chat_job thread STARTED") + tool_map = { + "bash": run_bash, + "read_file": run_read, + "write_file": run_write, + "edit_file": run_edit, + } try: - for chunk in self.provider.request(self.messages): + loop_count = 0 + while loop_count < 10 and not self.cancelled: + loop_count += 1 + full_response = "" + for chunk in self.provider.request(self.messages): + with self.lock: + if self.previous_type != chunk["type"] or "newsegment" in chunk: + if self.previous_type != "": + self.buffer += "\n" + self.buffer += "\n<<< " + chunk["type"] + "\n\n" + self.previous_type = chunk["type"] + self.buffer += chunk["content"] + full_response += chunk["content"] + if self.cancelled: + self.buffer += "\n\nCANCELLED by user" + break + if "\n" in self.buffer: + parts = self.buffer.split("\n") + self.lines.extend(parts[:-1]) + self.buffer = parts[-1] + + if self.cancelled: + break + + self.messages.append({"role": "assistant", "content": [{"type": "text", "text": full_response}]}) + + # call the tools if there are any calls in the response + tool_calls = re.findall(r"```(?:tool|json)\n(.*?)\n```", full_response, re.DOTALL) + if not tool_calls: + break + + executed_any = False + results = [] + for call_str in tool_calls: + try: + call = json.loads(call_str) + if not isinstance(call, dict) or "tool" not in call: + continue + + tool_name = call.get("tool") + args = call.get("arguments", {}) + + with self.lock: + self.lines.append(self.buffer) + self.buffer = f"\n>>> exec\n\n{tool_name}: {json.dumps(args)}" + self.previous_type = "exec" + + if tool_name in tool_map: + output = tool_map[tool_name](**args) + else: + output = self.mcp_manager.execute(tool_name, **args) + + results.append(f"Tool '{tool_name}' output:\n{output}") + executed_any = True + + with self.lock: + self.lines.append(self.buffer) + self.buffer = f"\n{output}\n" + + except Exception as e: + print_debug(f"Tool execution failed: {e}") + continue + + if not executed_any: + break + + tool_results_content = "\n\n".join(results) + self.messages.append({"role": "user", "content": [{"type": "text", "text": tool_results_content}]}) + with self.lock: - # For now, we only append whole lines to the buffer - print_debug(f"Received chunk: '{chunk['type']}' => '{chunk['content']}'") - if self.previous_type != chunk["type"] or "newsegment" in chunk: - if self.previous_type != "": - self.buffer += "\n" - self.buffer += "\n<<< " + chunk["type"] + "\n\n" - self.previous_type = chunk["type"] - self.buffer += chunk["content"] - if self.cancelled: - self.buffer += "\n\nCANCELLED by user" - print_debug("AI_chat_job cancelled during provider request") - if "\n" in self.buffer: - parts = self.buffer.split("\n") - self.lines.extend(parts[:-1]) - self.buffer = parts[-1] - if self.cancelled: - break # Exit the loop + self.lines.append(self.buffer) + self.buffer = "\n>>> user\n\n(Agent feedback received, thinking...)\n" + self.previous_type = "user" + except Exception as e: with self.lock: self.lines.append("") @@ -206,6 +296,7 @@ def run(self): if self.previous_type == "assistant": self.lines.extend("\n>>> user\n\n".split("\n")) self.done = True + self.mcp_manager.shutdown() print_debug("AI_chat_job thread DONE") def pickup_lines(self): diff --git a/tests/agent_test.py b/tests/agent_test.py new file mode 100644 index 0000000..fc4221d --- /dev/null +++ b/tests/agent_test.py @@ -0,0 +1,48 @@ +import os +import pytest +from pathlib import Path +from utils import run_write, run_read, run_edit, run_bash, WORKDIR + +def test_run_write_and_read(): + test_file = "test_write.txt" + test_content = "Hello, Agent!" + + # Write + result_write = run_write(test_file, test_content) + assert "Wrote" in result_write + + # Read + result_read = run_read(test_file) + assert result_read == test_content + + # Cleanup + (WORKDIR / test_file).unlink() + +def test_run_edit(): + test_file = "test_edit.txt" + test_content = "Original text." + run_write(test_file, test_content) + + # Edit + result_edit = run_edit(test_file, "Original", "Modified") + assert "Edited" in result_edit + + # Verify + result_read = run_read(test_file) + assert result_read == "Modified text." + + # Cleanup + (WORKDIR / test_file).unlink() + +def test_run_bash(): + result = run_bash("echo 'Hello World'") + assert result == "Hello World" + +def test_safe_path_violation(): + from utils import safe_path + with pytest.raises(ValueError): + safe_path("../outside.txt") + +def test_dangerous_command_blocked(): + result = run_bash("sudo rm -rf /") + assert "Dangerous command blocked" in result