kubrick-cli 0.1.4__py3-none-any.whl

kubrick_cli/safety.py

@@ -0,0 +1,204 @@
+"""Safety manager for dangerous command detection and validation."""
+
+import re
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Confirm
+
+console = Console()
+
+
+@dataclass
+class SafetyConfig:
+    """Safety configuration settings."""
+
+    max_iterations: int = 15
+    max_tools_per_turn: int = 5
+    total_timeout_seconds: int = 600
+    tool_timeout_seconds: int = 30
+    max_file_size_mb: int = 10
+    require_dangerous_command_confirmation: bool = True
+
+    @classmethod
+    def from_config(cls, config: Dict) -> "SafetyConfig":
+        """
+        Create SafetyConfig from configuration dict.
+
+        Args:
+            config: Configuration dictionary
+
+        Returns:
+            SafetyConfig instance
+        """
+        return cls(
+            max_iterations=config.get("max_iterations", 15),
+            max_tools_per_turn=config.get("max_tools_per_turn", 5),
+            total_timeout_seconds=config.get("total_timeout_seconds", 600),
+            tool_timeout_seconds=config.get("tool_timeout_seconds", 30),
+            max_file_size_mb=config.get("max_file_size_mb", 10),
+            require_dangerous_command_confirmation=config.get(
+                "require_dangerous_command_confirmation", True
+            ),
+        )
+
+
+class SafetyManager:
+    """
+    Manages safety checks for tool execution.
+
+    Features:
+    - Dangerous command detection
+    - User confirmation prompts
+    - File size limits
+    - Timeout enforcement
+    """
+
+    # Dangerous bash command patterns
+    DANGEROUS_PATTERNS = [
+        (r"\brm\s+-rf\s+/", "Recursive delete from root"),
+        (r"\brm\s+-rf\s+~", "Recursive delete from home"),
+        (r"\brm\s+-rf\s+\*", "Recursive delete all files"),
+        (r"\brm\s+-rf", "Recursive force delete"),
+        (r"\bsudo\b", "Elevated privileges"),
+        (r"\bchmod\s+777", "Overly permissive permissions"),
+        (r"\bchmod\s+-R\s+777", "Recursive overly permissive permissions"),
+        (r">\s*/dev/", "Writing to device files"),
+        (r"\bgit\s+push\s+--force", "Force push to git"),
+        (r"\bgit\s+push\s+-f", "Force push to git (short form)"),
+        (r"\bmkfs\b", "Format filesystem"),
+        (r"\bdd\s+.*of=/dev", "Writing to block device"),
+        (r"\b:(\()\|:\1&", "Fork bomb"),
+        (r"\bcurl\b.*\|\s*bash", "Pipe curl to bash"),
+        (r"\bwget\b.*\|\s*bash", "Pipe wget to bash"),
+        (r"\beval\b.*\$\(", "Eval with command substitution"),
+    ]
+
+    def __init__(self, config: SafetyConfig):
+        """
+        Initialize safety manager.
+
+        Args:
+            config: SafetyConfig instance
+        """
+        self.config = config
+
+    def validate_bash_command(self, command: str) -> Tuple[bool, Optional[str]]:
+        """
+        Validate a bash command for dangerous patterns.
+
+        Args:
+            command: The bash command to validate
+
+        Returns:
+            Tuple of (is_safe, warning_message)
+            - is_safe: True if safe, False if dangerous
+            - warning_message: Description of the danger (if any)
+        """
+        command_lower = command.lower()
+
+        for pattern, description in self.DANGEROUS_PATTERNS:
+            if re.search(pattern, command_lower, re.IGNORECASE):
+                return False, f"Dangerous command detected: {description}"
+
+        return True, None
+
+    def get_user_confirmation(self, warning: str, command: str) -> bool:
+        """
+        Get user confirmation for a dangerous command.
+
+        Args:
+            warning: Warning message
+            command: The dangerous command
+
+        Returns:
+            True if user confirms, False otherwise
+        """
+        if not self.config.require_dangerous_command_confirmation:
+            return True
+
+        console.print(
+            Panel(
+                f"[bold red]⚠️  WARNING[/bold red]\n\n"
+                f"{warning}\n\n"
+                f"[bold]Command:[/bold]\n"
+                f"[yellow]{command}[/yellow]\n\n"
+                f"This command could be dangerous.",
+                border_style="red",
+                title="Safety Check",
+            )
+        )
+
+        confirmed = Confirm.ask(
+            "[bold red]Do you want to execute this command?[/bold red]",
+            default=False,
+        )
+
+        if confirmed:
+            console.print("[yellow]⚠️  Proceeding with caution...[/yellow]")
+        else:
+            console.print("[green]✓ Command cancelled[/green]")
+
+        return confirmed
+
+    def validate_file_size(self, file_path: str, size_bytes: int) -> bool:
+        """
+        Validate file size against limits.
+
+        Args:
+            file_path: Path to the file
+            size_bytes: File size in bytes
+
+        Returns:
+            True if within limits, False otherwise
+        """
+        max_bytes = self.config.max_file_size_mb * 1024 * 1024
+
+        if size_bytes > max_bytes:
+            console.print(
+                f"[yellow]⚠️  File {file_path} exceeds size limit "
+                f"({size_bytes / 1024 / 1024:.1f}MB > {self.config.max_file_size_mb}MB)[/yellow]"
+            )
+            return False
+
+        return True
+
+    def check_iteration_limit(self, current: int, max_iterations: int) -> bool:
+        """
+        Check if iteration limit has been reached.
+
+        Args:
+            current: Current iteration number
+            max_iterations: Maximum allowed iterations
+
+        Returns:
+            True if within limit, False otherwise
+        """
+        if current >= max_iterations:
+            console.print(
+                f"[yellow]⚠️  Max iteration limit reached ({max_iterations})[/yellow]"
+            )
+            return False
+
+        return True
+
+    def check_tool_limit(self, current: int, max_tools: int) -> bool:
+        """
+        Check if tool call limit has been reached.
+
+        Args:
+            current: Current number of tool calls
+            max_tools: Maximum allowed tools
+
+        Returns:
+            True if within limit, False otherwise
+        """
+        if current > max_tools:
+            console.print(
+                f"[yellow]⚠️  Too many tool calls ({current} > {max_tools})[/yellow]"
+            )
+            return False
+
+        return True

kubrick_cli/scheduler.py

@@ -0,0 +1,183 @@
+"""Parallel tool execution scheduler for improved performance."""
+
+import concurrent.futures
+from typing import Dict, List, Tuple
+
+from rich.console import Console
+
+console = Console()
+
+
+# Read-only tools that can be executed in parallel
+READ_ONLY_TOOLS = {
+    "read_file",
+    "list_files",
+    "search_files",
+}
+
+# Write tools that must be executed sequentially
+WRITE_TOOLS = {
+    "write_file",
+    "edit_file",
+    "create_directory",
+    "run_bash",
+}
+
+
+class ToolScheduler:
+    """
+    Schedules and executes tools with intelligent parallelization.
+
+    Strategy: Conservative parallelization
+    - Read-only tools (read_file, list_files, search_files) run in parallel
+    - Write tools (write_file, edit_file, run_bash, etc.) run sequentially
+    - Graceful error handling per tool
+    """
+
+    def __init__(
+        self, tool_executor, max_workers: int = 3, enable_parallel: bool = True
+    ):
+        """
+        Initialize tool scheduler.
+
+        Args:
+            tool_executor: ToolExecutor instance
+            max_workers: Maximum parallel workers (default: 3)
+            enable_parallel: Whether to enable parallel execution
+        """
+        self.tool_executor = tool_executor
+        self.max_workers = max_workers
+        self.enable_parallel = enable_parallel
+
+    def execute_tools(self, tool_calls: List[Tuple[str, Dict]]) -> List[Dict]:
+        """
+        Execute a list of tool calls with intelligent scheduling.
+
+        Args:
+            tool_calls: List of (tool_name, parameters) tuples
+
+        Returns:
+            List of result dictionaries in the same order as input
+        """
+        if not self.enable_parallel or len(tool_calls) <= 1:
+            return self._execute_sequential(tool_calls)
+
+        read_only_calls = []
+        write_calls = []
+        call_order = []
+
+        for i, (tool_name, params) in enumerate(tool_calls):
+            if tool_name in READ_ONLY_TOOLS:
+                read_only_calls.append((i, tool_name, params))
+                call_order.append(("read", len(read_only_calls) - 1))
+            else:
+                write_calls.append((i, tool_name, params))
+                call_order.append(("write", len(write_calls) - 1))
+
+        read_results = {}
+        if read_only_calls:
+            console.print(
+                f"[dim]→ Executing {len(read_only_calls)} read-only tool(s) in parallel[/dim]"
+            )
+            read_results = self._execute_parallel(read_only_calls)
+
+        write_results = {}
+        if write_calls:
+            write_results = self._execute_sequential_indexed(write_calls)
+
+        results = []
+        for call_type, index in call_order:
+            if call_type == "read":
+                results.append(read_results[index])
+            else:
+                results.append(write_results[index])
+
+        return results
+
+    def _execute_parallel(
+        self, indexed_calls: List[Tuple[int, str, Dict]]
+    ) -> Dict[int, Dict]:
+        """
+        Execute tools in parallel using ThreadPoolExecutor.
+
+        Args:
+            indexed_calls: List of (index, tool_name, parameters) tuples
+
+        Returns:
+            Dict mapping index to result
+        """
+        results = {}
+
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=self.max_workers
+        ) as executor:
+            future_to_index = {}
+            for index, tool_name, params in indexed_calls:
+                future = executor.submit(self._execute_single, tool_name, params)
+                future_to_index[future] = index
+
+            for future in concurrent.futures.as_completed(future_to_index):
+                index = future_to_index[future]
+                try:
+                    result = future.result()
+                    results[index] = result
+                except Exception as e:
+                    results[index] = {
+                        "success": False,
+                        "error": f"Parallel execution error: {str(e)}",
+                    }
+
+        return results
+
+    def _execute_sequential(self, tool_calls: List[Tuple[str, Dict]]) -> List[Dict]:
+        """
+        Execute tools sequentially.
+
+        Args:
+            tool_calls: List of (tool_name, parameters) tuples
+
+        Returns:
+            List of result dictionaries
+        """
+        results = []
+        for tool_name, params in tool_calls:
+            result = self._execute_single(tool_name, params)
+            results.append(result)
+        return results
+
+    def _execute_sequential_indexed(
+        self, indexed_calls: List[Tuple[int, str, Dict]]
+    ) -> Dict[int, Dict]:
+        """
+        Execute tools sequentially with index mapping.
+
+        Args:
+            indexed_calls: List of (index, tool_name, parameters) tuples
+
+        Returns:
+            Dict mapping index to result
+        """
+        results = {}
+        for index, tool_name, params in indexed_calls:
+            result = self._execute_single(tool_name, params)
+            results[index] = result
+        return results
+
+    def _execute_single(self, tool_name: str, params: Dict) -> Dict:
+        """
+        Execute a single tool.
+
+        Args:
+            tool_name: Name of tool
+            params: Tool parameters
+
+        Returns:
+            Result dictionary
+        """
+        try:
+            return self.tool_executor.execute(tool_name, params)
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"Tool execution error: {str(e)}",
+            }

kubrick_cli/setup_wizard.py

@@ -0,0 +1,161 @@
+"""Setup wizard for first-time Kubrick configuration."""
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.prompt import Prompt
+from rich.table import Table
+
+from .providers.factory import ProviderFactory
+
+console = Console()
+
+
+class SetupWizard:
+    """
+    Interactive setup wizard for Kubrick configuration.
+
+    This wizard automatically discovers all available providers and
+    generates the configuration UI based on provider metadata.
+    """
+
+    @staticmethod
+    def run() -> dict:
+        """
+        Run the setup wizard.
+
+        Returns:
+            Dictionary with provider configuration
+        """
+        console.print(
+            Panel.fit(
+                "[bold cyan]Welcome to Kubrick![/bold cyan]\n\n"
+                "Let's set up your AI provider.\n"
+                "You can change these settings later in ~/.kubrick/config.json",
+                border_style="cyan",
+            )
+        )
+
+        provider_metadata = SetupWizard._select_provider()
+
+        config = SetupWizard._configure_provider(provider_metadata)
+        config["provider"] = provider_metadata.name
+
+        SetupWizard._show_summary(provider_metadata, config)
+
+        console.print("\n[green]✓ Setup complete![/green]")
+        console.print(
+            "[dim]Your configuration has been saved to ~/.kubrick/config.json[/dim]\n"
+        )
+
+        return config
+
+    @staticmethod
+    def _select_provider():
+        """
+        Prompt user to select a provider from all discovered providers.
+
+        Returns:
+            ProviderMetadata object for selected provider
+        """
+        console.print("\n[bold]Step 1: Select Your AI Provider[/bold]\n")
+
+        providers = ProviderFactory.list_available_providers()
+
+        if not providers:
+            raise RuntimeError(
+                "No providers found! Please ensure provider files are in kubrick_cli/providers/"
+            )
+
+        table = Table(show_header=True, header_style="bold cyan")
+        table.add_column("Option", style="cyan")
+        table.add_column("Provider", style="green")
+        table.add_column("Description")
+
+        choice_map = {}
+        for idx, provider in enumerate(providers, start=1):
+            choice_map[str(idx)] = provider
+            table.add_row(
+                str(idx),
+                provider.display_name,
+                provider.description,
+            )
+
+        console.print(table)
+        console.print()
+
+        choices = list(choice_map.keys())
+        choice = Prompt.ask(
+            "[bold yellow]Choose your provider[/bold yellow]",
+            choices=choices,
+            default="1",
+        )
+
+        return choice_map[choice]
+
+    @staticmethod
+    def _configure_provider(metadata) -> dict:
+        """
+        Get provider-specific configuration based on metadata.
+
+        Args:
+            metadata: ProviderMetadata object
+
+        Returns:
+            Configuration dictionary
+        """
+        console.print(f"\n[bold]Step 2: Configure {metadata.display_name}[/bold]\n")
+
+        config = {}
+
+        for field in metadata.config_fields:
+            key = field["key"]
+            label = field["label"]
+            field_type = field.get("type", "text")
+            default = field.get("default")
+            help_text = field.get("help_text")
+
+            if help_text:
+                console.print(f"[dim]{help_text}[/dim]\n")
+
+            if field_type == "password":
+                value = Prompt.ask(
+                    f"[cyan]{label}[/cyan]",
+                    password=True,
+                )
+            elif default is not None:
+                value = Prompt.ask(
+                    f"[cyan]{label}[/cyan]",
+                    default=str(default),
+                )
+            else:
+                value = Prompt.ask(f"[cyan]{label}[/cyan]")
+
+            config[key] = value
+
+        return config
+
+    @staticmethod
+    def _show_summary(metadata, config: dict):
+        """
+        Show configuration summary.
+
+        Args:
+            metadata: ProviderMetadata object
+            config: Configuration dict
+        """
+        console.print("\n[bold]Configuration Summary[/bold]\n")
+
+        console.print(f"[cyan]Provider:[/cyan] {metadata.display_name}")
+
+        for field in metadata.config_fields:
+            key = field["key"]
+            label = field["label"]
+            field_type = field.get("type", "text")
+            value = config.get(key)
+
+            if field_type == "password" and value:
+                display_value = f"{'*' * 20}{value[-4:]}"
+            else:
+                display_value = value
+
+            console.print(f"[cyan]{label}:[/cyan] {display_value}")