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

tunacode/services/mcp.py

@@ -5,9 +5,11 @@ 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, List, Optional, Tuple
+from typing import TYPE_CHECKING, AsyncIterator, Dict, List, Optional, Tuple
 
 from pydantic_ai.mcp import MCPServerStdio
 
@@ -16,9 +18,17 @@ 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):
     """A version of ``MCPServerStdio`` that suppresses *all* output coming from the
@@ -54,28 +64,136 @@ 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.
+    """Load MCP servers from configuration with caching.
 
     Args:
         state_manager: The state manager containing user configuration
 
     Returns:
-        List of MCP server instances
+        List of MCP server instances (cached when possible)
 
     Raises:
         MCPError: If a server configuration is invalid
     """
+    global _MCP_CONFIG_HASH
+
     mcp_servers: MCPServers = state_manager.session.user_config.get("mcpServers", {})
+
+    # Calculate hash of current config
+    current_hash = hash(str(mcp_servers))
+
+    # Check if config has changed
+    if _MCP_CONFIG_HASH == current_hash and _MCP_SERVER_CACHE:
+        # Return cached servers
+        return list(_MCP_SERVER_CACHE.values())
+
+    # 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
+
     loaded_servers: List[MCPServerStdio] = []
     MCPServerStdio.log_level = "critical"
 
     for server_name, conf in mcp_servers.items():
         try:
-            # loaded_servers.append(QuietMCPServer(**conf))
-            mcp_instance = MCPServerStdio(**conf)
-            # mcp_instance.log_level = "critical"
-            loaded_servers.append(mcp_instance)
+            # Check if this server is already cached
+            if server_name not in _MCP_SERVER_CACHE:
+                # 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:
             raise MCPError(
                 server_name=server_name,

tunacode/setup.py

@@ -5,20 +5,18 @@ 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,
 )
 from tunacode.core.state import StateManager
 
 
-async def setup(run_setup: bool, state_manager: StateManager, cli_config: dict = None) -> None:
+async def setup(
+    run_setup: bool, state_manager: StateManager, cli_config: dict = None, wizard_mode: bool = False
+) -> None:
     """
     Setup TunaCode on startup using the new setup coordinator.
 
@@ -26,6 +24,7 @@ async def setup(run_setup: bool, state_manager: StateManager, cli_config: dict =
         run_setup (bool): If True, force run the setup process, resetting current config.
         state_manager (StateManager): The state manager instance.
         cli_config (dict): Optional CLI configuration with baseurl, model, and key.
+        wizard_mode (bool): If True, run interactive setup wizard.
     """
     coordinator = SetupCoordinator(state_manager)
 
@@ -36,25 +35,6 @@ async def setup(run_setup: bool, state_manager: StateManager, cli_config: dict =
     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)
-
-
-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")
+    await coordinator.run_setup(force_setup=run_setup, wizard_mode=wizard_mode)

tunacode/tools/base.py

@@ -4,7 +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, List, Optional
 
 from pydantic_ai.exceptions import ModelRetry
 
@@ -24,6 +26,9 @@ class BaseTool(ABC):
         """
         self.ui = ui_logger
         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.
@@ -32,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
@@ -59,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
@@ -81,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.
 
@@ -138,6 +194,113 @@ class BaseTool(ABC):
         """
         return f"in {self.tool_name}"
 
+    def prompt(self, context: Optional[Dict[str, Any]] = None) -> str:
+        """Generate the prompt for this tool.
+
+        Args:
+            context: Optional context including model, permissions, environment
+
+        Returns:
+            str: The generated prompt for this tool
+        """
+        # Update context if provided
+        if context:
+            self._context.update(context)
+
+        # Check cache if context hasn't changed
+        cache_key = str(sorted(self._context.items()))
+        if self._prompt_cache and cache_key == getattr(self, "_cache_key", None):
+            return self._prompt_cache
+
+        # Generate new prompt
+        prompt = self._generate_prompt()
+
+        # Cache the result
+        self._prompt_cache = prompt
+        self._cache_key = cache_key
+
+        return prompt
+
+    def _generate_prompt(self) -> str:
+        """Generate the actual prompt based on current context.
+
+        Override this method in subclasses to provide tool-specific prompts.
+
+        Returns:
+            str: The generated prompt
+        """
+        # Default prompt generation
+        base_prompt = self._get_base_prompt()
+
+        # Apply model-specific adjustments
+        if "model" in self._context:
+            base_prompt = self._adjust_for_model(base_prompt, self._context["model"])
+
+        # Apply permission-specific adjustments
+        if "permissions" in self._context:
+            base_prompt = self._adjust_for_permissions(base_prompt, self._context["permissions"])
+
+        return base_prompt
+
+    def _get_base_prompt(self) -> str:
+        """Get the base prompt for this tool.
+
+        Override this in subclasses to provide tool-specific base prompts.
+
+        Returns:
+            str: The base prompt template
+        """
+        return f"Execute the {self.tool_name} tool to perform its designated operation."
+
+    def _adjust_for_model(self, prompt: str, model: str) -> str:
+        """Adjust prompt based on the model being used.
+
+        Args:
+            prompt: The base prompt
+            model: The model identifier
+
+        Returns:
+            str: Adjusted prompt
+        """
+        # Default implementation - override in subclasses for specific adjustments
+        return prompt
+
+    def _adjust_for_permissions(self, prompt: str, permissions: Dict[str, Any]) -> str:
+        """Adjust prompt based on permissions.
+
+        Args:
+            prompt: The base prompt
+            permissions: Permission settings
+
+        Returns:
+            str: Adjusted prompt
+        """
+        # Default implementation - override in subclasses for specific adjustments
+        return prompt
+
+    def get_tool_schema(self) -> Dict[str, Any]:
+        """Generate the tool schema for API integration.
+
+        Returns:
+            Dict containing the tool schema in OpenAI function format
+        """
+        return {
+            "name": self.tool_name,
+            "description": self.prompt(),
+            "parameters": self._get_parameters_schema(),
+        }
+
+    @abstractmethod
+    def _get_parameters_schema(self) -> Dict[str, Any]:
+        """Get the parameters schema for this tool.
+
+        Must be implemented by subclasses.
+
+        Returns:
+            Dict containing the JSON schema for tool parameters
+        """
+        pass
+
 
 class FileBasedTool(BaseTool):
     """Base class for tools that work with files.

tunacode/tools/bash.py

@@ -7,10 +7,14 @@ environment variables, timeouts, and improved output handling.
 """
 
 import asyncio
+import logging
 import os
 import subprocess
-from typing import Dict, Optional
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 
+import defusedxml.ElementTree as ET
 from pydantic_ai.exceptions import ModelRetry
 
 from tunacode.constants import MAX_COMMAND_OUTPUT
@@ -18,6 +22,8 @@ from tunacode.exceptions import ToolExecutionError
 from tunacode.tools.base import BaseTool
 from tunacode.types import ToolResult
 
+logger = logging.getLogger(__name__)
+
 
 class BashTool(BaseTool):
     """Enhanced shell command execution tool with advanced features."""
@@ -26,6 +32,95 @@ class BashTool(BaseTool):
     def tool_name(self) -> str:
         return "Bash"
 
+    @lru_cache(maxsize=1)
+    def _get_base_prompt(self) -> str:
+        """Load and return the base prompt from XML file.
+
+        Returns:
+            str: The loaded prompt from XML or a default prompt
+        """
+        try:
+            # Load prompt from XML file
+            prompt_file = Path(__file__).parent / "prompts" / "bash_prompt.xml"
+            if prompt_file.exists():
+                tree = ET.parse(prompt_file)
+                root = tree.getroot()
+                description = root.find("description")
+                if description is not None:
+                    return description.text.strip()
+        except Exception as e:
+            logger.warning(f"Failed to load XML prompt for bash: {e}")
+
+        # Fallback to default prompt
+        return (
+            """Executes a given bash command in a persistent shell session with optional timeout"""
+        )
+
+    @lru_cache(maxsize=1)
+    def _get_parameters_schema(self) -> Dict[str, Any]:
+        """Get the parameters schema for bash tool.
+
+        Returns:
+            Dict containing the JSON schema for tool parameters
+        """
+        # Try to load from XML first
+        try:
+            prompt_file = Path(__file__).parent / "prompts" / "bash_prompt.xml"
+            if prompt_file.exists():
+                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": []}
+                    required_fields: List[str] = []
+
+                    for param in parameters.findall("parameter"):
+                        name = param.get("name")
+                        required = param.get("required", "false").lower() == "true"
+                        param_type = param.find("type")
+                        description = param.find("description")
+
+                        if name and param_type is not None:
+                            prop = {
+                                "type": param_type.text.strip(),
+                                "description": description.text.strip()
+                                if description is not None
+                                else "",
+                            }
+
+                            schema["properties"][name] = prop
+                            if required:
+                                required_fields.append(name)
+
+                    schema["required"] = required_fields
+                    return schema
+        except Exception as e:
+            logger.warning(f"Failed to load parameters from XML for bash: {e}")
+
+        # Fallback to hardcoded schema
+        return {
+            "type": "object",
+            "properties": {
+                "command": {
+                    "type": "string",
+                    "description": "The command to execute",
+                },
+                "description": {
+                    "type": "string",
+                    "description": "Clear, concise description of what this command does",
+                },
+                "timeout": {
+                    "type": "number",
+                    "description": "Optional timeout in milliseconds",
+                },
+                "run_in_background": {
+                    "type": "boolean",
+                    "description": "Set to true to run this command in the background",
+                },
+            },
+            "required": ["command"],
+        }
+
     async def _execute(
         self,
         command: str,
@@ -83,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(
@@ -143,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,