Add the mcp support and a simple agent loop

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):

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()

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")