tunacode-cli 0.0.70__py3-none-any.whl → 0.0.78.6__py3-none-any.whl

tunacode/services/mcp.py

@@ -5,6 +5,8 @@ Provides Model Context Protocol (MCP) server management functionality.
 Handles MCP server initialization, configuration validation, and client connections.
 """
 
+import asyncio
+import logging
 import os
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, AsyncIterator, Dict, List, Optional, Tuple
@@ -16,12 +18,16 @@ from tunacode.types import MCPServers
 
 if TYPE_CHECKING:
     from mcp.client.stdio import ReadStream, WriteStream
+    from pydantic_ai import Agent
 
     from tunacode.core.state import StateManager
 
 # Module-level cache for MCP server instances
 _MCP_SERVER_CACHE: Dict[str, MCPServerStdio] = {}
 _MCP_CONFIG_HASH: Optional[int] = None
+_MCP_SERVER_AGENTS: Dict[str, "Agent"] = {}
+
+logger = logging.getLogger(__name__)
 
 
 class QuietMCPServer(MCPServerStdio):
@@ -58,6 +64,80 @@ class QuietMCPServer(MCPServerStdio):
                 yield read_stream, write_stream
 
 
+async def cleanup_mcp_servers(server_names: Optional[List[str]] = None) -> None:
+    """Clean up MCP server connections and clear cache.
+
+    Args:
+        server_names: Optional list of specific server names to clean up.
+                     If None, all servers will be cleaned up.
+    """
+    global _MCP_SERVER_CACHE, _MCP_SERVER_AGENTS
+
+    servers_to_cleanup = server_names or list(_MCP_SERVER_CACHE.keys())
+
+    cleanup_tasks = []
+    for server_name in servers_to_cleanup:
+        if server_name in _MCP_SERVER_CACHE:
+            logger.debug(f"Cleaning up MCP server: {server_name}")
+            cleanup_tasks.append(_cleanup_single_server(server_name))
+
+    if cleanup_tasks:
+        try:
+            await asyncio.gather(*cleanup_tasks, return_exceptions=True)
+        except Exception as e:
+            logger.warning(f"Error during MCP server cleanup: {e}")
+
+    # Clear caches for cleaned up servers
+    if server_names is None:
+        # Clear all
+        _MCP_SERVER_CACHE.clear()
+        _MCP_SERVER_AGENTS.clear()
+        logger.debug("Cleared all MCP server caches")
+    else:
+        # Clear specific servers
+        for server_name in server_names:
+            _MCP_SERVER_CACHE.pop(server_name, None)
+            _MCP_SERVER_AGENTS.pop(server_name, None)
+            logger.debug(f"Cleared cache for MCP server: {server_name}")
+
+
+async def _cleanup_single_server(server_name: str) -> None:
+    """Clean up a single MCP server instance.
+
+    Args:
+        server_name: Name of the server to clean up
+    """
+    if server_name not in _MCP_SERVER_CACHE:
+        return
+
+    server = _MCP_SERVER_CACHE[server_name]
+    agent = _MCP_SERVER_AGENTS.get(server_name)
+
+    try:
+        # Use agent's run_mcp_servers context manager if available
+        if agent and hasattr(agent, "run_mcp_servers"):
+            # The agent should handle proper cleanup via context manager exit
+            logger.debug(f"Agent cleanup for {server_name} handled by run_mcp_servers context")
+        else:
+            # Fallback: try to stop the server subprocess directly
+            if hasattr(server, "is_running") and server.is_running:
+                # MCPServerStdio doesn't expose direct cleanup, so we log this
+                logger.debug(f"MCP server {server_name} is running but no direct cleanup available")
+    except Exception as e:
+        logger.warning(f"Error cleaning up MCP server {server_name}: {e}")
+
+
+def register_mcp_agent(server_name: str, agent: "Agent") -> None:
+    """Register an agent that manages MCP servers for cleanup tracking.
+
+    Args:
+        server_name: Name of the server configuration
+        agent: Agent instance that manages the server
+    """
+    _MCP_SERVER_AGENTS[server_name] = agent
+    logger.debug(f"Registered agent for MCP server: {server_name}")
+
+
 def get_mcp_servers(state_manager: "StateManager") -> List[MCPServerStdio]:
     """Load MCP servers from configuration with caching.
 
@@ -82,7 +162,22 @@ def get_mcp_servers(state_manager: "StateManager") -> List[MCPServerStdio]:
         # Return cached servers
         return list(_MCP_SERVER_CACHE.values())
 
-    # Config changed or first load - clear cache and rebuild
+    # Config changed or first load - cleanup old servers and rebuild cache
+    if _MCP_CONFIG_HASH is not None and _MCP_SERVER_CACHE:
+        # Config changed - schedule cleanup of stale servers
+        removed_servers = [name for name in _MCP_SERVER_CACHE.keys() if name not in mcp_servers]
+        if removed_servers:
+            logger.debug(
+                f"Configuration changed, cleaning up removed MCP servers: {removed_servers}"
+            )
+            # Schedule async cleanup - use asyncio.create_task to avoid blocking
+            try:
+                loop = asyncio.get_event_loop()
+                loop.create_task(cleanup_mcp_servers(removed_servers))
+            except RuntimeError:
+                # No event loop running, schedule for later
+                logger.debug("No event loop available, deferring MCP server cleanup")
+
     _MCP_SERVER_CACHE.clear()
     _MCP_CONFIG_HASH = current_hash
 
@@ -96,6 +191,7 @@ def get_mcp_servers(state_manager: "StateManager") -> List[MCPServerStdio]:
                 # Create new instance
                 mcp_instance = MCPServerStdio(**conf)
                 _MCP_SERVER_CACHE[server_name] = mcp_instance
+                logger.debug(f"Created MCP server instance: {server_name}")
 
             loaded_servers.append(_MCP_SERVER_CACHE[server_name])
         except Exception as e:

tunacode/setup.py

@@ -5,13 +5,9 @@ Package setup and metadata configuration for the TunaCode CLI.
 Provides high-level setup functions for initializing the application and its agents.
 """
 
-from typing import Any, Optional
-
 from tunacode.core.setup import (
-    AgentSetup,
     ConfigSetup,
     EnvironmentSetup,
-    GitSafetySetup,
     SetupCoordinator,
     TemplateSetup,
 )
@@ -39,25 +35,6 @@ async def setup(
     coordinator.register_step(config_setup)
     coordinator.register_step(EnvironmentSetup(state_manager))
     coordinator.register_step(TemplateSetup(state_manager))
-    coordinator.register_step(GitSafetySetup(state_manager))
 
     # Run all setup steps
     await coordinator.run_setup(force_setup=run_setup, wizard_mode=wizard_mode)
-
-
-async def setup_agent(agent: Optional[Any], state_manager: StateManager) -> None:
-    """
-    Setup the agent separately.
-
-    This is called from other parts of the codebase when an agent needs to be initialized.
-
-    Args:
-        agent: The agent instance to initialize.
-        state_manager (StateManager): The state manager instance.
-    """
-    if agent is not None:
-        agent_setup = AgentSetup(state_manager, agent)
-        if await agent_setup.should_run():
-            await agent_setup.execute()
-            if not await agent_setup.validate():
-                raise RuntimeError("Agent setup failed validation")

tunacode/tools/base.py

@@ -4,8 +4,9 @@ This module provides a base class that implements common patterns
 for all tools including error handling, UI logging, and ModelRetry support.
 """
 
+import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from pydantic_ai.exceptions import ModelRetry
 
@@ -27,6 +28,7 @@ class BaseTool(ABC):
         self.logger = get_logger(self.__class__.__name__)
         self._prompt_cache: Optional[str] = None
         self._context: Dict[str, Any] = {}
+        self._resources: List[Any] = []  # Track resources for cleanup
 
     async def execute(self, *args, **kwargs) -> ToolResult:
         """Execute the tool with error handling and logging.
@@ -35,6 +37,7 @@ class BaseTool(ABC):
         - UI logging of the operation
         - Exception handling (except ModelRetry and ToolExecutionError)
         - Consistent error message formatting
+        - Resource cleanup on any exception
 
         Returns:
             str: Success message
@@ -62,6 +65,9 @@ class BaseTool(ABC):
         except Exception as e:
             # Handle any other exceptions
             await self._handle_error(e, *args, **kwargs)
+        finally:
+            # Ensure resource cleanup even on success or failure
+            await self.cleanup()
 
     @property
     @abstractmethod
@@ -84,6 +90,53 @@ class BaseTool(ABC):
         """
         pass
 
+    async def cleanup(self) -> None:
+        """Clean up any resources created during tool execution.
+
+        This method is called automatically in a finally block after tool execution.
+        Subclasses should override this to implement tool-specific cleanup logic.
+
+        The base implementation handles cleanup of any resources registered via
+        register_resource(), attempting to call close() or cleanup() methods.
+        """
+        for resource in self._resources:
+            try:
+                if hasattr(resource, "close"):
+                    if asyncio.iscoroutinefunction(resource.close):
+                        await resource.close()
+                    else:
+                        resource.close()
+                elif hasattr(resource, "cleanup"):
+                    if asyncio.iscoroutinefunction(resource.cleanup):
+                        await resource.cleanup()
+                    else:
+                        resource.cleanup()
+            except Exception as e:
+                self.logger.warning(f"Failed to clean up resource {resource}: {e}")
+
+        # Clear the resource list
+        self._resources.clear()
+
+    def register_resource(self, resource: Any) -> None:
+        """Register a resource for automatic cleanup.
+
+        Resources registered here will be automatically cleaned up in the finally
+        block of execute(), regardless of success or failure.
+
+        Args:
+            resource: Any object with a close() or cleanup() method
+        """
+        self._resources.append(resource)
+
+    async def __aenter__(self):
+        """Enter async context manager."""
+        return self
+
+    async def __aexit__(self, exc_type, _, __):
+        """Exit async context manager and cleanup resources."""
+        await self.cleanup()
+        return False
+
     async def _handle_error(self, error: Exception, *args, **kwargs) -> ToolResult:
         """Handle errors by logging and raising proper exceptions.
 

tunacode/tools/bash.py

@@ -178,6 +178,7 @@ class BashTool(BaseTool):
         # Set working directory
         exec_cwd = cwd or os.getcwd()
 
+        process = None
         try:
             # Execute command with timeout
             process = await asyncio.create_subprocess_shell(
@@ -238,6 +239,19 @@ class BashTool(BaseTool):
                 f"Shell not found. Cannot execute command: {command}\n"
                 "This typically indicates a system configuration issue."
             )
+        finally:
+            # Ensure process cleanup regardless of success or failure
+            if process is not None and process.returncode is None:
+                try:
+                    # Multi-stage escalation: graceful → terminate → kill
+                    try:
+                        process.terminate()
+                        await asyncio.wait_for(process.wait(), timeout=5.0)
+                    except asyncio.TimeoutError:
+                        process.kill()
+                        await asyncio.wait_for(process.wait(), timeout=1.0)
+                except Exception as cleanup_error:
+                    self.logger.warning(f"Failed to cleanup process: {cleanup_error}")
 
     def _format_output(
         self,

tunacode/tools/glob.py

@@ -553,9 +553,11 @@ async def glob(
         directory: Directory to search in (default: current directory)
         recursive: Whether to search recursively (default: True)
         include_hidden: Whether to include hidden files/directories (default: False)
-        exclude_dirs: Additional directories to exclude from search (default: common build/cache dirs)
+        exclude_dirs: Additional directories to exclude from search
+                         (default: common build/cache dirs)
         max_results: Maximum number of results to return (default: 5000)
-        sort_by: How to sort results - "modified", "size", "alphabetical", or "depth" (default: "modified")
+        sort_by: How to sort results - "modified", "size", "alphabetical", or "depth"
+                 (default: "modified")
         case_sensitive: Whether pattern matching is case-sensitive (default: False)
         use_gitignore: Whether to respect .gitignore patterns (default: True)
 

tunacode/tools/grep.py

@@ -242,7 +242,10 @@ Usage:
                 raise ToolExecutionError(f"Unknown search type: {search_type}")
 
             # 5️⃣ Format and return results with strategy info
-            strategy_info = f"Strategy: {search_type} (was {original_search_type}), Files: {len(candidates)}/{5000}"
+            strategy_info = (
+                f"Strategy: {search_type} (was {original_search_type}), "
+                f"Files: {len(candidates)}/{5000}"
+            )
             formatted_results = self._result_formatter.format_results(results, pattern, config)
 
             if return_format == "list":
@@ -281,7 +284,6 @@ Usage:
         def run_enhanced_ripgrep():
             """Execute ripgrep search using the new executor."""
             start_time = time.time()
-            first_match_time = None
             results = []
 
             # Configure timeout from settings
@@ -306,17 +308,8 @@ Usage:
                     context_after=config.context_lines,
                 )
 
-                # Track first match time for metrics
-                if search_results and first_match_time is None:
-                    first_match_time = time.time() - start_time
-
-                    # Check if we exceeded the first match deadline
-                    if first_match_time > config.first_match_deadline:
-                        if self._config.get("debug", False):
-                            logger.debug(
-                                f"Search exceeded first match deadline: {first_match_time:.2f}s"
-                            )
-                        raise TooBroadPatternError(pattern, config.first_match_deadline)
+                # Ripgrep doesn't provide timing info for first match, so we rely on
+                # the overall timeout mechanism instead of first_match_deadline
 
                 # Parse results
                 for result_line in search_results:
@@ -363,10 +356,7 @@ Usage:
                 )
 
                 if self._config.get("debug", False):
-                    logger.debug(
-                        f"Ripgrep search completed in {total_time:.2f}s "
-                        f"(first match: {first_match_time:.2f}s if found)"
-                    )
+                    logger.debug(f"Ripgrep search completed in {total_time:.2f}s")
 
             return results
 

tunacode/tools/prompts/glob_prompt.xml

@@ -6,7 +6,7 @@
 - Returns matching file paths sorted by modification time
 - Use this tool when you need to find files by name patterns
 - When you are doing an open ended search that may require multiple rounds of globbing and grepping, use the Agent tool instead
-- You have the capability to call multiple tools in a single response. It is always better to speculatively perform multiple searches as a batch that are potentially useful.
+- You have the capability to call multiple tools in a single response. When you need multiple glob patterns, list each call up front so the read-only scheduler can execute them together as one batch.
     </description>
 
     <parameters>

tunacode/tools/prompts/grep_prompt.xml

@@ -11,6 +11,7 @@ A powerful search tool built on ripgrep
   - Use Task tool for open-ended searches requiring multiple rounds
   - Pattern syntax: Uses ripgrep (not grep) - literal braces need escaping (use `interface\{\}` to find `interface{}` in Go code)
   - Multiline matching: By default patterns match within single lines only. For cross-line patterns like `struct \{[\s\S]*?field`, use `multiline: true`
+  - When investigating several patterns or directories at once, queue every `grep` call within the same response so they form a single batched execution.
     </description>
 
     <parameters>

tunacode/tools/prompts/list_dir_prompt.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <tool_prompt>
     <description>
-Lists files and directories in a given path. The path parameter must be an absolute path, not a relative path. You can optionally provide an array of glob patterns to ignore with the ignore parameter. You should generally prefer the Glob and Grep tools, if you know which directories to search.
+Lists files and directories in a given path. The path parameter must be an absolute path, not a relative path. You can optionally provide an array of glob patterns to ignore with the ignore parameter. You should generally prefer the Glob and Grep tools, if you know which directories to search. When inspecting multiple directories, enumerate every `list_dir` call you intend to run in the same response so they execute together as a parallel batch.
     </description>
 
     <parameters>

tunacode/tools/prompts/react_prompt.xml

@@ -0,0 +1,23 @@
+<tool>
+  <description>
+    Record a ReAct-style think/observe timeline, retrieve it, or clear it for the current session.
+  </description>
+  <parameters>
+    <parameter name="action" required="true">
+      <type>string</type>
+      <description>One of think, observe, get, clear.</description>
+    </parameter>
+    <parameter name="thoughts" required="false">
+      <type>string</type>
+      <description>Reasoning text for think entries.</description>
+    </parameter>
+    <parameter name="next_action" required="false">
+      <type>string</type>
+      <description>Planned action to pair with think entries.</description>
+    </parameter>
+    <parameter name="result" required="false">
+      <type>string</type>
+      <description>Observation details for observe entries.</description>
+    </parameter>
+  </parameters>
+</tool>

tunacode/tools/prompts/read_file_prompt.xml

@@ -13,7 +13,7 @@ Usage:
 - This tool allows Claude Code to read images (eg PNG, JPG, etc). When reading an image file the contents are presented visually as Claude Code is a multimodal LLM.
 - This tool can read PDF files (.pdf). PDFs are processed page by page, extracting both text and visual content for analysis.
 - This tool can read Jupyter notebooks (.ipynb files) and returns all cells with their outputs, combining code, text, and visualizations.
-- You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful.
+- You have the capability to call multiple tools in a single response. Enumerate every file you plan to inspect so multiple `read_file` calls can run in parallel rather than waiting for sequential turns.
 - You will regularly be asked to read screenshots. If the user provides a path to a screenshot ALWAYS use this tool to view the file at the path. This tool will work with all temporary file paths like /var/folders/123/abc/T/TemporaryItems/NSIRD_screencaptureui_ZfB1tD/Screenshot.png
 - If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents.
     </description>

tunacode/tools/react.py

@@ -0,0 +1,153 @@
+"""Lightweight ReAct-style scratchpad tool."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, Literal
+
+import defusedxml.ElementTree as ET
+from pydantic_ai.exceptions import ModelRetry
+
+from tunacode.core.state import StateManager
+from tunacode.types import ToolResult, UILogger
+
+from .base import BaseTool
+
+
+# CLAUDE_ANCHOR[react-tool]: Minimal ReAct scratchpad tool surface
+class ReactTool(BaseTool):
+    """Minimal ReAct scratchpad for tracking think/observe steps."""
+
+    def __init__(self, state_manager: StateManager, ui_logger: UILogger | None = None):
+        super().__init__(ui_logger)
+        self.state_manager = state_manager
+
+    @property
+    def tool_name(self) -> str:
+        return "react"
+
+    async def _execute(
+        self,
+        action: Literal["think", "observe", "get", "clear"],
+        thoughts: str | None = None,
+        next_action: str | None = None,
+        result: str | None = None,
+    ) -> ToolResult:
+        scratchpad = self._ensure_scratchpad()
+
+        if action == "think":
+            if not thoughts:
+                raise ModelRetry("Provide thoughts when using react think action")
+            if not next_action:
+                raise ModelRetry("Specify next_action when recording react thoughts")
+
+            entry = {
+                "type": "think",
+                "thoughts": thoughts,
+                "next_action": next_action,
+            }
+            self.state_manager.append_react_entry(entry)
+            return "Recorded think step"
+
+        if action == "observe":
+            if not result:
+                raise ModelRetry("Provide result when using react observe action")
+
+            entry = {
+                "type": "observe",
+                "result": result,
+            }
+            self.state_manager.append_react_entry(entry)
+            return "Recorded observation"
+
+        if action == "get":
+            timeline = scratchpad.get("timeline", [])
+            if not timeline:
+                return "React scratchpad is empty"
+
+            formatted = [
+                f"{index + 1}. {item['type']}: {self._format_entry(item)}"
+                for index, item in enumerate(timeline)
+            ]
+            return "\n".join(formatted)
+
+        if action == "clear":
+            self.state_manager.clear_react_scratchpad()
+            return "React scratchpad cleared"
+
+        raise ModelRetry("Invalid react action. Use one of: think, observe, get, clear")
+
+    def _format_entry(self, item: Dict[str, Any]) -> str:
+        if item["type"] == "think":
+            return f"thoughts='{item['thoughts']}', next_action='{item['next_action']}'"
+        if item["type"] == "observe":
+            return f"result='{item['result']}'"
+        return str(item)
+
+    def _ensure_scratchpad(self) -> dict[str, Any]:
+        scratchpad = self.state_manager.get_react_scratchpad()
+        scratchpad.setdefault("timeline", [])
+        return scratchpad
+
+    def _get_base_prompt(self) -> str:
+        prompt_file = Path(__file__).parent / "prompts" / "react_prompt.xml"
+        if prompt_file.exists():
+            try:
+                tree = ET.parse(prompt_file)
+                root = tree.getroot()
+                description = root.find("description")
+                if description is not None and description.text:
+                    return description.text.strip()
+            except Exception:
+                pass
+        return "Use this tool to record think/observe notes and manage the react scratchpad"
+
+    def _get_parameters_schema(self) -> Dict[str, Any]:
+        prompt_file = Path(__file__).parent / "prompts" / "react_prompt.xml"
+        if prompt_file.exists():
+            try:
+                tree = ET.parse(prompt_file)
+                root = tree.getroot()
+                parameters = root.find("parameters")
+                if parameters is not None:
+                    schema: Dict[str, Any] = {
+                        "type": "object",
+                        "properties": {},
+                        "required": ["action"],
+                    }
+                    for param in parameters.findall("parameter"):
+                        name = param.get("name")
+                        param_type = param.find("type")
+                        description = param.find("description")
+                        if name and param_type is not None:
+                            schema["properties"][name] = {
+                                "type": param_type.text.strip(),
+                                "description": description.text.strip()
+                                if description is not None and description.text
+                                else "",
+                            }
+                    return schema
+            except Exception:
+                pass
+        return {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "type": "string",
+                    "description": "react operation to perform",
+                },
+                "thoughts": {
+                    "type": "string",
+                    "description": "Thought content for think action",
+                },
+                "next_action": {
+                    "type": "string",
+                    "description": "Planned next action for think action",
+                },
+                "result": {
+                    "type": "string",
+                    "description": "Observation message for observe action",
+                },
+            },
+            "required": ["action"],
+        }

tunacode/tools/run_command.py

@@ -144,6 +144,7 @@ class RunCommandTool(BaseTool):
             CommandSecurityError: If command fails security validation
             Exception: Any command execution errors
         """
+        process = None
         try:
             # Use secure subprocess execution with validation
             process = safe_subprocess_popen(
@@ -158,6 +159,20 @@ class RunCommandTool(BaseTool):
         except CommandSecurityError as e:
             # Security validation failed - return error without execution
             return f"Security validation failed: {str(e)}"
+        finally:
+            # Ensure process cleanup regardless of success or failure
+            if process is not None and process.poll() is None:
+                try:
+                    # Multi-stage escalation: graceful → terminate → kill
+                    process.terminate()
+                    try:
+                        process.wait(timeout=5.0)
+                    except subprocess.TimeoutExpired:
+                        process.kill()
+                        process.wait(timeout=1.0)
+                except Exception as cleanup_error:
+                    self.logger.warning(f"Failed to cleanup process: {cleanup_error}")
+
         output = stdout.strip() or CMD_OUTPUT_NO_OUTPUT
         error = stderr.strip() or CMD_OUTPUT_NO_ERRORS
         resp = CMD_OUTPUT_FORMAT.format(output=output, error=error).strip()