Enhance execute

[bishoy-at-pieces]

close #299

Added some tests as well

[Copilot]

Pull Request Overview

This PR enhances the execute command to support custom handlers per asset classification, refactors context and LTM integration, and centralizes streamed-identifier caching.

• Introduces ExecuteCommand extensibility with per-format command mapping and persistence.
• Adds chat-level LTM enable/disable methods and integrates LTM checks into commands.
• Replaces individual AssetSnapshot/ConversationsSnapshot client assignments with a unified StreamedIdentifiersCache.

Reviewed Changes

Copilot reviewed 26 out of 26 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
src/pieces/commands/execute_command.py	New extensible execution flow with handle_execute, command map, and SafeDict.
src/pieces/wrapper/long_term_memory.py	Added chat_enable_ltm, chat_disable_ltm, and is_chat_ltm_enabled.
src/pieces/wrapper/copilot.py	Updated Context construction to include Copilot reference.
src/pieces/wrapper/context.py	Overhauled context storage to use ValidatedContextList and manage assets/paths/messages.
src/pieces/wrapper/client.py	Swapped AssetSnapshot/ConversationsSnapshot for StreamedIdentifiersCache.
src/pieces/wrapper/basic_identifier/basic.py	Filtered out deleted indices in _from_indices.

Comments suppressed due to low confidence (3)

src/pieces/wrapper/basic_identifier/basic.py:19

• [nitpick] Using the variable name id shadows the built-in id() function. Consider renaming it to key or idx for clarity.

for id, v in indices.items() if v != -1

src/pieces/commands/execute_command.py:47

• [nitpick] Variable named map shadows the Python built-in map(). Rename it to something like command_map to avoid confusion.

map = cls.get_command_map()

src/pieces/wrapper/ask_command.py:38

• The BasicChat class is used here but not imported in this module. Add from pieces.wrapper.basic_identifier.chat import BasicChat at the top.

Settings.pieces_client.copilot.chat = BasicChat(response.conversation)

[Copilot]

Pull Request Overview

This PR enhances the execution workflow by refactoring the execute command logic and adding comprehensive tests for multiple asset classifications.

• Added extensive tests for command execution across various languages.
• Refactored command execution logic and introduced a new command extensions map in settings.
• Updated CLI argument parsing to support custom handlers for different asset classifications.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/assets/execute_command_test.py	Comprehensive tests covering execution flows and error handling for various asset types.
src/pieces/settings.py	Added a configuration path for the command extensions map.
src/pieces/commands/execute_command.py	Refactored the execution logic, added support for custom command mappings, and improved error handling.
src/pieces/commands/assets_command.py	Updated asset file creation logic to use a centralized helper method.
src/pieces/app.py	Modified CLI parser to inject custom execute command handlers for different classifications.

Comments suppressed due to low confidence (1)

src/pieces/commands/assets_command.py:103

• The function 'get_file_extension' is used without an import or definition in this file. Ensure it is properly imported or defined.

file_extension = get_file_extension(asset.classification)

[Copilot]

Pull Request Overview

This PR enhances the functionality of the execute command by refining command mapping, safe formatting, and error handling, while also adding comprehensive tests for multiple language executions.

• Added multiple unit tests for various command executions (Python, JS, Lua, etc.).
• Introduced a SafeDict for formatting and extended the command map configuration.
• Updated CLI parsing and asset file handling logic.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/assets/execute_command_test.py	Added tests covering various execution scenarios and error conditions.
src/pieces/settings.py	Added the execute_command_extensions_map configuration for dynamic command mapping.
src/pieces/commands/execute_command.py	Enhanced command execution logic with safe formatting and improved error messages.
src/pieces/commands/assets_command.py	Refactored asset file creation via a new create_asset_file method.
src/pieces/app.py	Updated CLI argument definitions to support custom handler commands.

Comments suppressed due to low confidence (2)

src/pieces/commands/execute_command.py:48

• The variable name 'map' shadows the built-in function; consider renaming it to 'command_map' to avoid potential issues.

if asset.classification.value not in map:

src/pieces/commands/assets_command.py:103

• The function 'get_file_extension' is used but not imported or defined in this file; please import or define it to ensure correct functionality.

file_extension = get_file_extension(asset.classification)

[bishoy-at-pieces]

Screen.Recording.2025-05-29.at.5.42.02.PM.mov

[Copilot]

Pull Request Overview

This PR enhances the command execution functionality by introducing improved handling of command maps and execution flows, as well as adding extensive tests for different material classifications and error conditions.

• Added comprehensive tests covering various language handlers and error scenarios.
• Enhanced the execute command flow with handler support and improved command map management.
• Refactored asset file creation logic in the assets command module.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/assets/execute_command_test.py	Added tests to cover various execution paths and error conditions.
src/pieces/settings.py	Introduced a new setting for the execute command extensions map.
src/pieces/commands/execute_command.py	Updated command execution flow and command mapping logic.
src/pieces/commands/assets_command.py	Refactored asset file creation by consolidating file writing logic.
src/pieces/app.py	Updated the CLI parser to utilize the new command handler functionality.

Comments suppressed due to low confidence (1)

src/pieces/commands/execute_command.py:47

• The variable 'map' shadows the built-in function 'map'. It is recommended to rename it to 'command_map' or a similar non-conflicting name.

map = cls.get_command_map()

[tsavo-at-pieces]

Detailed Code Review

Overall Architecture Observations

The codebase represents a CLI tool for managing code snippets (called "materials" or "assets") with execution capabilities. While functional, there are several critical security, reliability, and design issues that need addressing.

---

1. src/pieces/commands/execute_command.py

Critical Security Issues

1.1 Command Injection Vulnerability

commands = (
    map[asset.classification.value]
    .format_map(
        SafeDict({
            "content": asset.raw_content,
            "file": file,
            "file_no_extension": file_no_extension,
        })
    )
    .split("&&")
)

Issue: The code directly interpolates user content into shell commands, creating a severe command injection vulnerability.

Example Attack Vector:

# If asset.raw_content contains: '; rm -rf / #
# The bash command becomes: bash -c '; rm -rf / #'

Fix:

@classmethod
def execute_command(cls, asset: Optional["BasicAsset"], **kwargs):
    if not asset:
        return
    try:
        if not asset.raw_content:
            return Settings.show_error("Couldn't get the material content")
        if not asset.classification:
            return Settings.show_error("Couldn't extract the material classification")
        
        map = cls.get_command_map()
        if asset.classification.value not in map:
            return Settings.show_error(
                f"No matching command found for material type: '{asset.classification.value}'.",
                f"Tip: Use `pieces execute --{asset.classification.value}` to configure a handler for this material type.",
            )
        
        file = AssetsCommands.create_asset_file(asset)
        file_no_extension = Path(file).with_suffix("")
        
        # Parse command template to identify placeholders
        command_template = map[asset.classification.value]
        
        # Use a safer approach - don't allow {content} in shell commands
        if "{content}" in command_template and asset.classification.value not in ["py", "js", "ts", "rb", "dart", "go"]:
            # For interpreted languages that accept code via -e flag, use stdin instead
            commands = command_template.replace("{content}", "").split("&&")
            stdin_content = asset.raw_content
        else:
            # For file-based execution, only allow {file} and {file_no_extension}
            commands = command_template.format(
                file=shlex.quote(str(file)),
                file_no_extension=shlex.quote(str(file_no_extension))
            ).split("&&")
            stdin_content = None
        
        stderr = ""
        stdout = ""
        for command in commands:
            out = shlex.split(command)
            result = subprocess.run(
                out,
                capture_output=True,
                text=True,
                input=stdin_content if stdin_content else None,
                cwd=os.path.dirname(file),  # Set working directory
                timeout=30  # Add timeout to prevent hanging
            )
            stdout += result.stdout
            if result.stderr:
                stderr += result.stderr
                
        Settings.logger.print(f"Executing {asset.classification.value} command:")
        Settings.logger.print(stdout)
        if stderr:
            Settings.logger.print("Errors:")
            Settings.logger.print(stderr)
    except subprocess.TimeoutExpired:
        Settings.logger.print("Error: Command execution timed out after 30 seconds")
    except subprocess.CalledProcessError as e:
        Settings.logger.print(f"Error executing command: {e}")
    except Exception as e:
        Settings.logger.print(f"An error occurred: {e}")

1.2 Unsafe Default Commands

commands_map.setdefault("bash", "bash -c {content}")
commands_map.setdefault("sh", "sh -c {content}")

Issue: Using -c with shell interpreters and direct content interpolation is dangerous.

Fix: Use stdin or temporary files:

commands_map.setdefault("py", "python {file}")
commands_map.setdefault("bash", "bash {file}")
commands_map.setdefault("sh", "sh {file}")
commands_map.setdefault("js", "node {file}")
# Remove all -c and -e flags that directly execute content

Design Issues

1.3 Missing Input Validation

No validation of file paths or command templates before execution.

Fix:

@staticmethod
def validate_command_template(template: str) -> bool:
    """Validate command template for safety."""
    # Disallow dangerous patterns
    dangerous_patterns = [
        "rm ", "del ", "format ", "dd ", 
        ">", ">>",  # Redirection that could overwrite files
        "|",  # Pipes that could chain commands
        ";",  # Command chaining
        "`",  # Command substitution
        "$(",  # Command substitution
        "&&", "||"  # Conditional execution (except when used as separator)
    ]
    
    # Check for dangerous patterns in the base template
    for pattern in dangerous_patterns:
        if pattern in template and pattern != "&&":  # Allow && as command separator
            return False
    
    # Only allow specific placeholders
    allowed_placeholders = ["{file}", "{file_no_extension}"]
    placeholders = re.findall(r'\{[^}]+\}', template)
    for placeholder in placeholders:
        if placeholder not in allowed_placeholders:
            return False
    
    return True

---

2. src/pieces/commands/assets_command.py

Security and Reliability Issues

2.1 Unsafe File Creation

@classmethod
def create_asset_file(cls, asset: BasicAsset):
    code_content = asset.raw_content
    file_extension = get_file_extension(asset.classification)
    
    if not os.path.exists(Settings.open_snippet_dir):
        os.makedirs(Settings.open_snippet_dir)
    
    file_path = os.path.join(
        Settings.open_snippet_dir, f"{asset.id}{file_extension}"
    )

Issues:

• No sanitization of asset.id (potential path traversal)
• No file permissions set
• No cleanup of old files

Fix:

@classmethod
def create_asset_file(cls, asset: BasicAsset):
    import tempfile
    import stat
    
    code_content = asset.raw_content
    file_extension = get_file_extension(asset.classification)
    
    # Ensure directory exists with proper permissions
    if not os.path.exists(Settings.open_snippet_dir):
        os.makedirs(Settings.open_snippet_dir, mode=0o700)  # User-only access
    
    # Sanitize asset ID to prevent path traversal
    safe_id = re.sub(r'[^\w\-]', '_', asset.id)
    
    # Use more secure file creation
    file_path = os.path.join(
        Settings.open_snippet_dir, f"{safe_id}{file_extension}"
    )
    
    try:
        # Write file with restricted permissions
        if isinstance(code_content, str):
            with open(file_path, "w", encoding='utf-8') as file:
                file.write(code_content)
        else:
            with open(file_path, "wb") as file:
                file.write(bytes(code_content))
        
        # Set file permissions to user-only
        os.chmod(file_path, stat.S_IRUSR | stat.S_IWUSR)
        
        # Clean up old files (older than 24 hours)
        cls._cleanup_old_files()
        
        return file_path
    except Exception as e:
        Settings.logger.print(f"Error creating asset file: {e}")
        raise

@classmethod
def _cleanup_old_files(cls):
    """Remove temporary files older than 24 hours."""
    import time
    now = time.time()
    for filename in os.listdir(Settings.open_snippet_dir):
        file_path = os.path.join(Settings.open_snippet_dir, filename)
        if os.path.isfile(file_path):
            if os.stat(file_path).st_mtime < now - 86400:  # 24 hours
                try:
                    os.remove(file_path)
                except:
                    pass  # Ignore cleanup errors

2.2 Editor Command Injection

subprocess.run([editor_exe, file_path])

Issue: No validation of editor command, could be exploited.

Fix:

def open_in_editor(self, file_path: str, editor: str):
    """Safely open file in editor."""
    # Whitelist of allowed editors
    allowed_editors = [
        'vim', 'vi', 'nano', 'emacs', 'code', 'subl', 
        'atom', 'gedit', 'notepad', 'notepad++', 'nvim'
    ]
    
    editor_base = os.path.basename(editor).lower()
    if editor_base not in allowed_editors:
        Settings.show_error(
            f"Editor '{editor}' is not in the allowed list.",
            "Please use one of: " + ", ".join(allowed_editors)
        )
        return
    
    editor_exe = shutil.which(editor)
    if not editor_exe:
        Settings.show_error(
            "Editor executable not found",
            "Please make sure it is added to the PATH"
        )
        return
    
    try:
        # Use subprocess with explicit arguments (no shell=True)
        subprocess.run([editor_exe, file_path], check=False)
    except Exception as e:
        Settings.show_error("Error opening editor", str(e))

2.3 Clipboard Security

text = pyperclip.paste()

Issue: No validation of clipboard content before saving.

Fix:

@classmethod
def create_asset(cls, **kwargs):
    MAX_CONTENT_SIZE = 10 * 1024 * 1024  # 10MB limit
    
    text = None
    content_flag = kwargs.get("content", None)
    
    if content_flag:
        text = sys.stdin.read(MAX_CONTENT_SIZE)
        if len(text) == MAX_CONTENT_SIZE:
            Settings.logger.print(
                "Warning: Content truncated at 10MB limit"
            )
    else:
        try:
            text = pyperclip.paste()
            # Validate clipboard content
            if len(text) > MAX_CONTENT_SIZE:
                Settings.show_error(
                    "Clipboard content too large",
                    f"Maximum size is {MAX_CONTENT_SIZE} bytes"
                )
                return
        except pyperclip.PyperclipException as e:
            Settings.show_error("Error accessing clipboard:", str(e))
            return
    
    if not text:
        Settings.logger.print(
            "No content found in the clipboard to create a material."
        )
        return
    
    # Additional validation
    if '\x00' in text:  # Null bytes can cause issues
        Settings.show_error(
            "Invalid content",
            "Content contains null bytes"
        )
        return
    
    # Rest of the function...

---

3. src/pieces/app.py

Design Issues

3.1 Argument Parsing Security

The current implementation allows arbitrary command-line argument injection.

Fix: Add argument validation:

def run(self):
    config = ConfigCommands.load_config()
    Settings.logger = Logger(config.get("debug", False), Settings.pieces_data_dir)
    
    # Validate arguments early
    for arg in sys.argv[1:]:
        if len(arg) > 1000:  # Prevent extremely long arguments
            self.parser.error("Argument too long")
        if '\x00' in arg:  # Null bytes
            self.parser.error("Invalid character in argument")
    
    # Rest of the method...

---

4. src/pieces/settings.py

Security and Design Issues

4.1 Pickle Usage (Critical Security Risk)

with open(file, "rb") as f:
    cache = pickle.load(f)

Issue: Pickle can execute arbitrary code during deserialization.

Fix: Replace with JSON:

import json

@classmethod
def get_from_json(cls, file, key):
    """Safely load data from JSON file."""
    try:
        cache = cls.file_cache.get(str(file))
        if not cache:
            with open(file, "r") as f:
                cache = json.load(f)
                cls.file_cache[str(file)] = cache
        return cache.get(key)
    except (FileNotFoundError, json.JSONDecodeError):
        return None

@classmethod
def dump_json(cls, file):
    """Store data in a JSON file."""
    with open(file, "w") as f:
        json.dump(cls.file_cache, f, indent=2)

# Update model storage to use JSON
models_file = Path(
    pieces_data_dir,
    "model_data.json",  # Changed from .pkl
)

4.2 Missing Path Validation

pieces_data_dir = user_data_dir(
    appauthor="pieces", appname="cli-agent", ensure_exists=True
)

Fix: Validate and secure directory creation:

@classmethod
def initialize_data_directory(cls):
    """Safely initialize data directory with proper permissions."""
    import stat
    
    cls.pieces_data_dir = user_data_dir(
        appauthor="pieces", 
        appname="cli-agent", 
        ensure_exists=False  # Don't auto-create
    )
    
    # Validate path
    if '..' in cls.pieces_data_dir or cls.pieces_data_dir.startswith('/'):
        raise ValueError("Invalid data directory path")
    
    # Create with restricted permissions
    if not os.path.exists(cls.pieces_data_dir):
        os.makedirs(cls.pieces_data_dir, mode=0o700)
    else:
        # Verify permissions
        current_mode = os.stat(cls.pieces_data_dir).st_mode
        if current_mode & 0o077:  # Check if group/other have any permissions
            os.chmod(cls.pieces_data_dir, 0o700)

---

5. Environment and Dependency Considerations

5.1 Working Directory Requirements

The execute command should run in an isolated directory:

@classmethod
def get_execution_environment(cls, asset: BasicAsset) -> dict:
    """Create safe execution environment."""
    import tempfile
    
    # Create temporary directory for execution
    exec_dir = tempfile.mkdtemp(prefix="pieces_exec_")
    
    # Copy only the necessary file
    file_path = cls.create_asset_file(asset)
    exec_file = os.path.join(exec_dir, os.path.basename(file_path))
    shutil.copy2(file_path, exec_file)
    
    return {
        'cwd': exec_dir,
        'env': {
            **os.environ,
            'PIECES_EXEC': '1',  # Flag for scripts to detect execution context
            'HOME': exec_dir,  # Prevent access to user home
        },
        'cleanup': lambda: shutil.rmtree(exec_dir, ignore_errors=True)
    }

5.2 Dependency Checks

Add runtime checks for required executables:

@classmethod
def check_runtime_dependencies(cls, classification: str) -> Tuple[bool, str]:
    """Check if runtime dependencies are available."""
    requirements = {
        'py': ['python', 'python3'],
        'js': ['node'],
        'ts': ['ts-node'],
        'rb': ['ruby'],
        'go': ['go'],
        'rs': ['rustc'],
        'c': ['gcc'],
        'cpp': ['g++'],
        'java': ['javac', 'java'],
    }
    
    if classification not in requirements:
        return True, ""
    
    for cmd in requirements.get(classification, []):
        if shutil.which(cmd):
            return True, cmd
    
    return False, f"No runtime found for {classification}. Install one of: {', '.join(requirements[classification])}"

5.3 Container/VM Execution

For enhanced security, consider containerized execution:

@classmethod
def execute_in_container(cls, asset: BasicAsset, command: str):
    """Execute code in a Docker container for isolation."""
    import docker
    
    client = docker.from_env()
    
    # Language to image mapping
    images = {
        'py': 'python:3.9-slim',
        'js': 'node:16-slim',
        'go': 'golang:1.17-alpine',
        # ... more mappings
    }
    
    image = images.get(asset.classification.value)
    if not image:
        raise ValueError(f"No container image for {asset.classification.value}")
    
    # Run with restrictions
    container = client.containers.run(
        image,
        command,
        detach=True,
        mem_limit='512m',
        cpu_quota=50000,  # 50% CPU
        network_mode='none',  # No network access
        read_only=True,
        remove=True,
        timeout=30
    )

---

6. Overlooked Edge Cases

6.1 Race Conditions

Multiple instances could conflict when accessing files:

import fcntl

def acquire_file_lock(file_path: str):
    """Acquire exclusive lock on file."""
    lock_file = f"{file_path}.lock"
    lock_fd = open(lock_file, 'w')
    try:
        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        return lock_fd
    except IOError:
        lock_fd.close()
        raise RuntimeError(f"File {file_path} is locked by another process")

6.2 Signal Handling

Add proper cleanup on interruption:

import signal
import atexit

class CleanupManager:
    def __init__(self):
        self.cleanup_tasks = []
        atexit.register(self.cleanup)
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
    
    def register(self, func):
        self.cleanup_tasks.append(func)
    
    def cleanup(self):
        for task in self.cleanup_tasks:
            try:
                task()
            except:
                pass
    
    def _signal_handler(self, signum, frame):
        self.cleanup()
        sys.exit(1)

6.3 Resource Limits

Implement resource limits for execution:

import resource

def set_resource_limits():
    """Set resource limits for child processes."""
    # CPU time limit (30 seconds)
    resource.setrlimit(resource.RLIMIT_CPU, (30, 30))
    
    # Memory limit (512 MB)
    resource.setrlimit(resource.RLIMIT_AS, (512 * 1024 * 1024, 512 * 1024 * 1024))
    
    # File size limit (10 MB)
    resource.setrlimit(resource.RLIMIT_FSIZE, (10 * 1024 * 1024, 10 * 1024 * 1024))
    
    # Number of processes
    resource.setrlimit(resource.RLIMIT_NPROC, (10, 10))

---

Summary of Critical Issues

1. Command Injection - Direct interpolation of user content into shell commands
2. Pickle Deserialization - Can execute arbitrary code
3. Path Traversal - No sanitization of file paths
4. Missing Input Validation - No limits on content size or validation
5. Insufficient Error Handling - Many operations can fail silently
6. No Resource Limits - Executed code can consume unlimited resources
7. Race Conditions - File operations not thread-safe
8. Insecure Defaults - Shell execution with -c flag

Recommended Implementation Priority

1. Immediate: Fix command injection vulnerabilities
2. High: Replace pickle with JSON
3. High: Add input validation and sanitization
4. Medium: Implement resource limits
5. Medium: Add proper error handling
6. Low: Consider containerized execution for additional security

[bishoy-at-pieces]

@tsavo-at-pieces I was thinking about #1 but if the snippet contains rm -rf / , the user chose to execute that snippet so the user needs to be careful but he is the one that chose to execute it not us,

[tsavo-at-pieces]

Complete Fixes for Remaining Issues

1. src/pieces/commands/execute_command.py

Fix: SafeDict Class Security

Old Code:

class SafeDict(dict):
    def __missing__(self, key):
        return "{" + key + "}"

New Code:

class SafeDict(dict):
    """Dictionary that returns empty string for missing keys to prevent injection."""
    def __missing__(self, key):
        # Log missing keys for debugging
        Settings.logger.debug(f"Missing template key: {key}")
        # Return empty string instead of the placeholder to prevent injection
        return ""
    
    def __getitem__(self, key):
        # Validate key format
        if not isinstance(key, str) or not key.replace('_', '').isalnum():
            raise KeyError(f"Invalid template key: {key}")
        return super().__getitem__(key)

Fix: handle_execute Method Validation

Old Code:

@classmethod
def handle_execute(cls, **kwargs):
    # parse the args manually for better performance
    passed_args = set()
    for i, arg in enumerate(sys.argv[1:]):
        if arg.startswith("--"):
            passed_args.add(arg.lstrip("--").split("=")[0] + "_handler")

    if passed_args:
        cls.save_commands_map(passed_args, kwargs)
        return
    cls.execute_command(**kwargs)

New Code:

@classmethod
def handle_execute(cls, **kwargs):
    """Handle execute command with validation."""
    # Validate command line arguments
    MAX_ARG_LENGTH = 1000
    VALID_HANDLERS = {f"{lang.value}_handler" for lang in ClassificationSpecificEnum}
    
    passed_args = set()
    for i, arg in enumerate(sys.argv[1:]):
        # Validate argument length
        if len(arg) > MAX_ARG_LENGTH:
            Settings.show_error(
                "Invalid argument",
                f"Argument too long: {arg[:50]}..."
            )
            return
        
        if arg.startswith("--"):
            handler_name = arg.lstrip("--").split("=")[0] + "_handler"
            
            # Validate handler name
            if handler_name not in VALID_HANDLERS:
                Settings.show_error(
                    "Invalid handler",
                    f"Unknown language handler: {arg}"
                )
                return
            
            passed_args.add(handler_name)

    if passed_args:
        # Validate handler commands before saving
        for handler in passed_args:
            command = kwargs.get(handler, "")
            if not cls.validate_command_template(command):
                Settings.show_error(
                    "Invalid command template",
                    f"Command template for {handler} contains dangerous patterns"
                )
                return
        
        cls.save_commands_map(passed_args, kwargs)
        return
    
    cls.execute_command(**kwargs)

@staticmethod
def validate_command_template(template: str) -> bool:
    """Validate command template for safety."""
    if not template:
        return False
    
    # Length check
    if len(template) > 500:
        return False
    
    # Disallow dangerous patterns
    dangerous_patterns = [
        "rm -rf", "del /f", "format", "dd if=",
        ":(){ :|:& };:",  # Fork bomb
        ">", ">>",  # Redirection
        "|",  # Pipes
        ";",  # Command chaining
        "`",  # Command substitution
        "$(",  # Command substitution
        "${",  # Variable expansion
        "eval",  # Eval commands
        "exec",  # Exec commands
    ]
    
    template_lower = template.lower()
    for pattern in dangerous_patterns:
        if pattern in template_lower:
            return False
    
    # Only allow specific placeholders
    allowed_placeholders = ["{file}", "{file_no_extension}", "{content}"]
    import re
    placeholders = re.findall(r'\{[^}]+\}', template)
    for placeholder in placeholders:
        if placeholder not in allowed_placeholders:
            return False
    
    return True

Fix: save_commands_map Input Validation

Old Code:

@classmethod
def save_commands_map(cls, commands: Iterable[str], default: Dict[str, str]):
    data = cls.get_command_map()
    for command in commands:
        data[command.removesuffix("_handler")] = default[command]
    with open(Settings.execute_command_extensions_map, "w") as f:
        json.dump(data, f)

New Code:

@classmethod
def save_commands_map(cls, commands: Iterable[str], default: Dict[str, str]):
    """Save command mappings with validation."""
    import os
    import stat
    
    # Load existing data
    data = cls.get_command_map()
    
    # Validate and update commands
    for command in commands:
        handler_value = default.get(command, "")
        
        # Validate command template
        if not cls.validate_command_template(handler_value):
            Settings.logger.warning(f"Skipping invalid command template for {command}")
            continue
        
        # Extract language key
        lang_key = command.removesuffix("_handler")
        
        # Additional validation for language key
        if len(lang_key) > 20 or not lang_key.replace('_', '').isalnum():
            Settings.logger.warning(f"Skipping invalid language key: {lang_key}")
            continue
        
        data[lang_key] = handler_value
    
    # Ensure directory exists with proper permissions
    config_dir = os.path.dirname(Settings.execute_command_extensions_map)
    if not os.path.exists(config_dir):
        os.makedirs(config_dir, mode=0o700)
    
    # Write with restricted permissions
    temp_file = f"{Settings.execute_command_extensions_map}.tmp"
    try:
        with open(temp_file, "w") as f:
            json.dump(data, f, indent=2)
        
        # Set restrictive permissions
        os.chmod(temp_file, stat.S_IRUSR | stat.S_IWUSR)
        
        # Atomic replace
        os.replace(temp_file, Settings.execute_command_extensions_map)
    except Exception as e:
        # Clean up temp file on error
        if os.path.exists(temp_file):
            os.unlink(temp_file)
        raise

2. src/pieces/commands/assets_command.py

Fix: check_assets_existence Decorator

Old Code:

def check_assets_existence(func):
    """Decorator to ensure user has assets."""
    def wrapper(*args, **kwargs):
        assets = Settings.pieces_client.assets()  # Check if there is an asset
        if not assets:
            return Settings.show_error(
                "No materials found", "Please create an material first."
            )
        return func(*args, **kwargs)
    return wrapper

New Code:

def check_assets_existence(func):
    """Decorator to ensure user has assets with error handling."""
    def wrapper(*args, **kwargs):
        try:
            assets = Settings.pieces_client.assets()
            if not assets:
                return Settings.show_error(
                    "No materials found", 
                    "Please create a material first using 'pieces create'"
                )
        except ConnectionError:
            return Settings.show_error(
                "Connection Error",
                "Unable to connect to Pieces OS. Please ensure it's running."
            )
        except Exception as e:
            Settings.logger.debug(f"Error checking assets: {e}")
            return Settings.show_error(
                "Error accessing materials",
                "Please try again or restart Pieces OS"
            )
        
        return func(*args, **kwargs)
    return wrapper

Fix: check_asset_selected Decorator

Old Code:

def check_asset_selected(func):
    """
    Decorator to check if there is a selected asset or not and if it is valid id.
    If valid id it returns the asset_data to the called function.
    """
    def wrapper(*args, **kwargs):
        from pieces.commands.list_command import ListCommand

        try:
            if AssetsCommands.current_asset is None:
                raise ValueError("No material selected")
            AssetsCommands.current_asset.asset  # Check if the current asset is vaild
        except (ValueError, NotFoundException):
            ListCommand.list_assets()
        return func(asset=AssetsCommands.current_asset, *args, **kwargs)
    return wrapper

New Code:

def check_asset_selected(func):
    """
    Decorator to check if there is a selected asset with comprehensive validation.
    """
    def wrapper(*args, **kwargs):
        from pieces.commands.list_command import ListCommand

        try:
            if AssetsCommands.current_asset is None:
                Settings.logger.print("No material currently selected.")
                if Settings.logger.confirm("Would you like to select one from the list?"):
                    ListCommand.list_assets()
                    return  # Don't proceed with the original function
                else:
                    return
            
            # Validate the asset is still valid
            try:
                # This will raise if asset was deleted
                _ = AssetsCommands.current_asset.asset
            except NotFoundException:
                Settings.logger.print("The selected material no longer exists.")
                AssetsCommands.current_asset = None
                if Settings.logger.confirm("Would you like to select another?"):
                    ListCommand.list_assets()
                return
            
        except ConnectionError:
            return Settings.show_error(
                "Connection Error",
                "Lost connection to Pieces OS"
            )
        except Exception as e:
            Settings.logger.debug(f"Error validating asset: {e}")
            return Settings.show_error(
                "Validation Error",
                "Unable to validate selected material"
            )
        
        return func(asset=AssetsCommands.current_asset, *args, **kwargs)
    return wrapper

Fix: print_code Method Error Handling

Old Code:

@staticmethod
def print_code(code_content, classification=None):
    try:
        if classification:
            lexer = get_lexer_by_name(classification, stripall=True)
        else:
            raise ClassNotFound
    except ClassNotFound:
        lexer = guess_lexer(code_content)
    formatted_code = highlight(code_content, lexer, TerminalFormatter())
    print(formatted_code)

New Code:

@staticmethod
def print_code(code_content, classification=None):
    """Print code with syntax highlighting and error handling."""
    MAX_DISPLAY_LENGTH = 10000  # Limit display length for performance
    
    if not code_content:
        Settings.logger.print("[No content to display]")
        return
    
    # Truncate very long content
    truncated = False
    if len(code_content) > MAX_DISPLAY_LENGTH:
        code_content = code_content[:MAX_DISPLAY_LENGTH]
        truncated = True
    
    try:
        # Try to get lexer by classification
        if classification:
            try:
                lexer = get_lexer_by_name(classification, stripall=True)
            except ClassNotFound:
                lexer = None
        else:
            lexer = None
        
        # Fallback to guessing if no lexer found
        if not lexer:
            try:
                lexer = guess_lexer(code_content)
            except ClassNotFound:
                # Final fallback - plain text
                Settings.logger.print(code_content)
                if truncated:
                    Settings.logger.print("\n[Content truncated for display]")
                return
        
        formatted_code = highlight(code_content, lexer, TerminalFormatter())
        print(formatted_code)
        
        if truncated:
            Settings.logger.print("\n[Content truncated for display]")
            
    except Exception as e:
        # Fallback to plain text on any error
        Settings.logger.debug(f"Syntax highlighting failed: {e}")
        Settings.logger.print(code_content)
        if truncated:
            Settings.logger.print("\n[Content truncated for display]")

Fix: save_asset Method

Old Code:

@classmethod
@check_asset_selected
def save_asset(cls, asset: BasicAsset, **kwargs):
    if not cls.check_editor()[0]:
        return

    file_path = os.path.join(
        Settings.open_snippet_dir,
        f"{(asset.id)}{get_file_extension(asset.classification)}",
    )
    data = None
    try:
        with open(file_path, "r") as f:
            data = f.read()
    except FileNotFoundError:
        cls.open_asset(asset.id, editor=True)
        Settings.logger.print(
            Markdown(
                "**Note:** Next time to open the material in your editor, use the `pieces list -e`"
            )
        )

    if data and asset.raw_content != data:
        Settings.logger.print(Markdown(f"Saving `{asset.name}` material"))
        asset.raw_content = data
    else:
        try:
            Settings.logger.input(
                f"Content not changed.\n"
                f"<Press enter when you finish editing {asset.name}>"
            )
            cls.save_asset(**kwargs)
        except KeyboardInterrupt:
            pass

New Code:

@classmethod
@check_asset_selected
def save_asset(cls, asset: BasicAsset, **kwargs):
    """Save asset with proper validation and error handling."""
    if not cls.check_editor()[0]:
        return

    # Sanitize asset ID for file path
    safe_id = re.sub(r'[^\w\-]', '_', asset.id)
    file_path = os.path.join(
        Settings.open_snippet_dir,
        f"{safe_id}{get_file_extension(asset.classification)}",
    )
    
    MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB limit
    data = None
    
    try:
        # Check file size before reading
        file_size = os.path.getsize(file_path)
        if file_size > MAX_FILE_SIZE:
            Settings.show_error(
                "File too large",
                f"File exceeds {MAX_FILE_SIZE // 1024 // 1024}MB limit"
            )
            return
        
        # Read with encoding detection
        with open(file_path, "rb") as f:
            raw_data = f.read()
        
        # Try to decode as UTF-8
        try:
            data = raw_data.decode('utf-8')
        except UnicodeDecodeError:
            # For binary files, work with raw bytes
            data = raw_data
            
    except FileNotFoundError:
        Settings.logger.print("File not found. Opening in editor...")
        cls.open_asset(asset.id, editor=True)
        Settings.logger.print(
            Markdown(
                "**Note:** The file is now open in your editor. "
                "Save the file and run this command again."
            )
        )
        return
    except PermissionError:
        Settings.show_error(
            "Permission denied",
            f"Cannot read file: {file_path}"
        )
        return
    except Exception as e:
        Settings.show_error(
            "Error reading file",
            str(e)
        )
        return

    # Check if content changed
    content_changed = False
    if isinstance(data, bytes) and isinstance(asset.raw_content, bytes):
        content_changed = data != asset.raw_content
    elif isinstance(data, str) and isinstance(asset.raw_content, str):
        # Normalize line endings for comparison
        content_changed = data.replace('\r\n', '\n') != asset.raw_content.replace('\r\n', '\n')
    else:
        # Type mismatch - content has changed
        content_changed = True

    if content_changed:
        try:
            Settings.logger.print(f"Saving changes to '{asset.name}'...")
            asset.raw_content = data
            Settings.logger.print("✓ Material saved successfully")
        except Exception as e:
            Settings.show_error(
                "Failed to save material",
                str(e)
            )
    else:
        Settings.logger.print("No changes detected in the file.")
        if Settings.logger.confirm("Would you like to reopen in editor?"):
            cls.open_asset(asset.id, editor=True)
            Settings.logger.print("After making changes, run 'pieces save' again.")

Fix: edit_asset Input Validation

Old Code:

@classmethod
@check_asset_selected
def edit_asset(cls, asset: BasicAsset, **kwargs):
    print_asset_details(asset)
    name = kwargs.get("name", "")
    classification = kwargs.get("classification", "")

    if (
        not name and not classification
    ):  # If no name or no classification is provided
        # Ask the user for a new name
        name = Settings.logger.input(
            "Enter the new name for the material[leave blank to keep the same]: "
        ).strip()
        classification = Settings.logger.input(
            "Enter the classification for the material[leave blank to keep the same]: "
        ).strip()

    # Check if the user actually entered a name
    if name:
        asset.name = name
    if classification:
        asset.classification = classification

New Code:

@classmethod
@check_asset_selected
def edit_asset(cls, asset: BasicAsset, **kwargs):
    """Edit asset with validation."""
    print_asset_details(asset)
    
    MAX_NAME_LENGTH = 255
    VALID_CLASSIFICATIONS = [c.value for c in ClassificationSpecificEnum]
    
    name = kwargs.get("name", "").strip()
    classification = kwargs.get("classification", "").strip()

    # Interactive mode if no arguments provided
    if not name and not classification:
        try:
            # Get new name with validation
            name = Settings.logger.input(
                "Enter the new name for the material [leave blank to keep the same]: "
            ).strip()
            
            # Get new classification with validation
            classification = Settings.logger.input(
                "Enter the classification for the material [leave blank to keep the same]: "
            ).strip()
        except KeyboardInterrupt:
            Settings.logger.print("\nEdit cancelled.")
            return
        except EOFError:
            return

    # Validate and apply name change
    if name:
        # Validate name length
        if len(name) > MAX_NAME_LENGTH:
            Settings.show_error(
                "Name too long",
                f"Maximum length is {MAX_NAME_LENGTH} characters"
            )
            return
        
        # Validate name characters (prevent control characters)
        if any(ord(c) < 32 for c in name):
            Settings.show_error(
                "Invalid name",
                "Name contains invalid characters"
            )
            return
        
        try:
            old_name = asset.name
            asset.name = name
            Settings.logger.print(f"✓ Name changed from '{old_name}' to '{name}'")
        except Exception as e:
            Settings.show_error(
                "Failed to update name",
                str(e)
            )
            return
    
    # Validate and apply classification change
    if classification:
        # Normalize classification
        classification_lower = classification.lower()
        
        # Find matching classification
        matching_classification = None
        for valid_class in VALID_CLASSIFICATIONS:
            if valid_class.lower() == classification_lower:
                matching_classification = valid_class
                break
        
        if not matching_classification:
            Settings.show_error(
                "Invalid classification",
                f"Must be one of: {', '.join(sorted(VALID_CLASSIFICATIONS))}"
            )
            # Show suggestions for close matches
            from difflib import get_close_matches
            suggestions = get_close_matches(
                classification_lower, 
                [c.lower() for c in VALID_CLASSIFICATIONS], 
                n=3, 
                cutoff=0.6
            )
            if suggestions:
                Settings.logger.print(f"Did you mean: {', '.join(suggestions)}?")
            return
        
        try:
            old_classification = asset.classification
            asset.classification = matching_classification
            Settings.logger.print(
                f"✓ Classification changed from '{old_classification}' to '{matching_classification}'"
            )
        except Exception as e:
            Settings.show_error(
                "Failed to update classification",
                str(e)
            )
            return
    
    if not name and not classification:
        Settings.logger.print("No changes made.")

Fix: share_asset Error Handling

Old Code:

@classmethod
@check_asset_selected
def share_asset(cls, asset: BasicAsset, **kwargs):
    Settings.logger.print("Generating shareable link")
    if asset.asset.shares:
        link = asset.asset.shares.iterable[0].link
    else:
        user = Settings.pieces_client.user_api.user_snapshot()
        # Update the local cache because the websockets are not running
        Settings.pieces_client.user.on_user_callback(user.user)
        try:
            share = asset.share()
        except PermissionError:
            Settings.logger.print(
                Markdown(
                    "Please login using `pieces login` command and make sure you are connected to the Pieces cloud"
                )
            )
            return
        link = share.iterable[0].link
    Settings.logger.print(f"Generated shareable link {link}")
    if Settings.logger.confirm("Do you want to open it in the browser?"):
        Settings.open_website(link)

New Code:

@classmethod
@check_asset_selected
def share_asset(cls, asset: BasicAsset, **kwargs):
    """Share asset with comprehensive error handling."""
    import time
    
    Settings.logger.print("Generating shareable link...")
    
    try:
        # Check if already shared
        if asset.asset.shares and asset.asset.shares.iterable:
            link = asset.asset.shares.iterable[0].link
            Settings.logger.print("✓ Using existing share link")
        else:
            # Check user authentication first
            try:
                user = Settings.pieces_client.user_api.user_snapshot()
                if not user or not user.user:
                    Settings.logger.print(
                        Markdown(
                            "**Authentication required**\n"
                            "Please login using `pieces login` command first"
                        )
                    )
                    return
                
                # Update the local cache
                Settings.pieces_client.user.on_user_callback(user.user)
                
            except Exception as e:
                Settings.logger.debug(f"Auth check failed: {e}")
                Settings.logger.print(
                    Markdown(
                        "**Unable to verify authentication**\n"
                        "Please ensure you're logged in with `pieces login`"
                    )
                )
                return
            
            # Generate share with retry logic
            max_retries = 3
            retry_delay = 1
            
            for attempt in range(max_retries):
                try:
                    share = asset.share()
                    if share and share.iterable:
                        link = share.iterable[0].link
                        break
                except PermissionError:
                    Settings.logger.print(
                        Markdown(
                            "**Permission denied**\n"
                            "Please ensure:\n"
                            "1. You're logged in with `pieces login`\n"
                            "2. You're connected to Pieces Cloud\n"
                            "3. Your account has sharing permissions"
                        )
                    )
                    return
                except ConnectionError:
                    if attempt < max_retries - 1:
                        Settings.logger.print(
                            f"Connection failed, retrying in {retry_delay}s..."
                        )
                        time.sleep(retry_delay)
                        retry_delay *= 2
                    else:
                        Settings.show_error(
                            "Connection failed",
                            "Unable to connect to Pieces Cloud"
                        )
                        return
                except Exception as e:
                    Settings.logger.debug(f"Share attempt {attempt + 1} failed: {e}")
                    if attempt < max_retries - 1:
                        time.sleep(retry_delay)
                        retry_delay *= 2
                    else:
                        Settings.show_error(
                            "Sharing failed",
                            "Unable to generate share link. Please try again."
                        )
                        return
            else:
                Settings.show_error(
                    "Sharing failed",
                    "Unable to generate share link after multiple attempts"
                )
                return
        
        # Validate link format
        if not link or not link.startswith(('http://', 'https://')):
            Settings.show_error(
                "Invalid share link",
                "Generated link appears to be invalid"
            )
            return
        
        # Display link with copy option
        Settings.logger.print(f"\n✓ Shareable link generated successfully:")
        Settings.logger.print(f"  {link}\n")
        
        # Copy to clipboard if available
        try:
            import pyperclip
            pyperclip.copy(link)
            Settings.logger.print("✓ Link copied to clipboard")
        except:
            pass  # Clipboard not available
        
        # Ask to open in browser
        if Settings.logger.confirm("Would you like to open it in the browser?"):
            try:
                Settings.open_website(link)
            except Exception as e:
                Settings.logger.debug(f"Failed to open browser: {e}")
                Settings.logger.print("Unable to open browser automatically")
                
    except Exception as e:
        Settings.logger.debug(f"Unexpected error in share_asset: {e}")
        Settings.show_error(
            "Sharing failed",
            "An unexpected error occurred"
        )

Fix: delete_asset Improvements

Old Code:

@classmethod
@check_asset_selected
def delete_asset(cls, asset: BasicAsset, **kwargs):
    print_asset_details(asset)

    confirm = Settings.logger.confirm(
        "Are you sure you really want to delete this material? This action cannot be undone."
    )
    if confirm:
        Settings.logger.print("Deleting material...")
        asset.delete()
        cls.current_asset = None
        space_below("Material Deleted")
    elif confirm == "n":
        Settings.logger.print("Deletion cancelled.")

New Code:

@classmethod
@check_asset_selected
def delete_asset(cls, asset: BasicAsset, **kwargs):
    """Delete asset with safety checks and cleanup."""
    print_asset_details(asset)
    
    # Show warning for shared assets
    try:
        if asset.asset.shares and asset.asset.shares.iterable:
            Settings.logger.print(
                Markdown("**Warning**: This material has been shared. "
                        "Deleting it will invalidate all share links.")
            )
    except:
        pass  # Ignore errors checking shares

    confirm = Settings.logger.confirm(
        "Are you sure you want to delete this material? This action cannot be undone."
    )
    
    if not confirm:
        Settings.logger.print("Deletion cancelled.")
        return
    
    try:
        # Clean up any temporary files
        safe_id = re.sub(r'[^\w\-]', '_', asset.id)
        temp_file = os.path.join(
            Settings.open_snippet_dir,
            f"{safe_id}{get_file_extension(asset.classification)}"
        )
        if os.path.exists(temp_file):
            try:
                os.unlink(temp_file)
            except:
                pass  # Ignore cleanup errors
        
        # Delete the asset
        Settings.logger.print("Deleting material...")
        asset.delete()
        
        # Clear current selection
        cls.current_asset = None
        
        space_below("✓ Material deleted successfully")
        
    except PermissionError:
        Settings.show_error(
            "Permission denied",
            "You don't have permission to delete this material"
        )
    except ConnectionError:
        Settings.show_error(
            "Connection error",
            "Unable to connect to Pieces OS"
        )
    except Exception as e:
        Settings.logger.debug(f"Delete failed: {e}")
        Settings.show_error(
            "Deletion failed",
            "Unable to delete the material. Please try again."
        )

3. src/pieces/app.py

Fix: Better Error Handling in run()

Old Code:

def run(self):
    config = ConfigCommands.load_config()
    Settings.logger = Logger(config.get("debug", False), Settings.pieces_data_dir)
    try:
        arg = sys.argv[1]
        if arg == "--ignore-onboarding":
            arg = sys.argv[2]
    except IndexError:  # No command provided
        self.parser.print_help()
        return
    # ... rest of method

New Code:

def run(self):
    """Run the CLI with comprehensive error handling."""
    try:
        # Load config with error handling
        try:
            config = ConfigCommands.load_config()
        except Exception as e:
            print(f"Error loading configuration: {e}")
            print("Using default configuration...")
            config = {}
        
        # Initialize logger
        Settings.logger = Logger(
            config.get("debug", False), 
            Settings.pieces_data_dir
        )
        
        # Validate command line arguments
        if len(sys.argv) > 100:  # Sanity check
            Settings.logger.error("Too many arguments provided")
            sys.exit(1)
        
        # Parse primary command
        try:
            primary_arg = None
            ignore_onboarding = "--ignore-onboarding" in sys.argv
            
            # Find the primary command (not a flag)
            for arg in sys.argv[1:]:
                if not arg.startswith("-"):
                    primary_arg = arg
                    break
            
            if not primary_arg and len(sys.argv) > 1:
                # Only flags provided, let argparse handle it
                primary_arg = sys.argv[1]
                
        except IndexError:
            self.parser.print_help()
            return

        # Check onboarding status
        onboarded = config.get("onboarded", False)
        
        if (
            not config.get("skip_onboarding", False)
            and not onboarded
            and not ignore_onboarding
            and primary_arg not in ["help", "--help", "-h", "onboarding"]
        ):
            res = Settings.logger.prompt(
                "It looks like this is your first time using the Pieces CLI."
                "\nWould you like to start onboarding?",
                choices=["y", "n", "skip"],
            )
            if res.lower() == "y":
                return onboarding_command()
            elif res.lower() == "skip":
                config["skip_onboarding"] = True
                try:
                    ConfigCommands.save_config(config)
                except Exception as e:
                    Settings.logger.warning(f"Failed to save config: {e}")

        # Parse arguments with error handling
        try:
            args = self.parser.parse_args()
        except SystemExit as e:
            if e.code != 0:  # Don't log successful exits
                Settings.logger.debug(f"Argument parsing failed with code {e.code}")
            raise
        except Exception as e:
            Settings.logger.error(f"Failed to parse arguments: {e}")
            self.parser.print_help()
            sys.exit(1)
        
        command = getattr(args, 'command', None)
        mcp_subcommand = getattr(args, "mcp", None)

        # Commands that don't need PiecesOS
        no_pieces_commands = [
            "help", "-v", "--version", "install", 
            "onboarding", "feedback", "contribute", "open"
        ]
        
        # Check if we need to start PiecesOS
        if command and command not in no_pieces_commands:
            if not (command == "mcp" and mcp_subcommand == "start"):
                try:
                    Settings.startup()
                except Exception as e:
                    Settings.logger.error(f"Failed to connect to Pieces OS: {e}")
                    if command != "install":
                        Settings.logger.print(
                            "\nPieces OS is not running. "
                            "Please start it or run 'pieces install'"
                        )
                    sys.exit(2)
        
        # Execute command
        Settings.logger.debug(f"Running command {command} with args: {args}")
        
        try:
            func = getattr(args, 'func', None)
            if func:
                func(**vars(args))
            else:
                self.parser.print_help()
        except KeyboardInterrupt:
            Settings.logger.print("\n\nOperation cancelled by user")
            sys.exit(130)  # Standard exit code for SIGINT
        except Exception as e:
            Settings.logger.error(f"Command execution failed: {e}")
            if config.get("debug", False):
                import traceback
                traceback.print_exc()
            sys.exit(1)
            
    except Exception as e:
        # Last resort error handler
        print(f"Fatal error: {e}")
        if "--debug" in sys.argv:
            import traceback
            traceback.print_exc()
        sys.exit(1)

Fix: Command Validation in add_subparsers

Add validation to subparser definitions:

New Code to add at the beginning of add_subparsers():

def add_subparsers(self):
    """Add subparsers with validation."""
    
    def validate_positive_int(value):
        """Validate positive integer input."""
        try:
            ivalue = int(value)
            if ivalue <= 0:
                raise argparse.ArgumentTypeError(
                    f"{value} is not a positive integer"
                )
            if ivalue > 10000:  # Reasonable upper limit
                raise argparse.ArgumentTypeError(
                    f"{value} is too large (max: 10000)"
                )
            return ivalue
        except ValueError:
            raise argparse.ArgumentTypeError(
                f"{value} is not a valid integer"
            )
    
    def validate_path(value):
        """Validate file/directory path."""
        if not value:
            return value
        
        # Prevent path traversal
        if ".." in value or value.startswith("/etc") or value.startswith("/sys"):
            raise argparse.ArgumentTypeError(
                f"Invalid path: {value}"
            )
        
        # Length check
        if len(value) > 4096:  # Linux PATH_MAX
            raise argparse.ArgumentTypeError(
                "Path too long"
            )
        
        return value
    
    def validate_query(value):
        """Validate search query."""
        if len(value) > 1000:
            raise argparse.ArgumentTypeError(
                "Query too long (max: 1000 characters)"
            )
        if '\x00' in value:
            raise argparse.ArgumentTypeError(
                "Query contains invalid characters"
            )
        return value
    
    # ... rest of the original method with modifications:
    
    # Update list_parser to use validation:
    list_parser.add_argument(
        "max_snippets",
        nargs="?",
        type=validate_positive_int,  # Use validator
        default=10,
        help="Max number of materials (1-10000)",
    )
    
    # Update ask_parser to use validation:
    ask_parser.add_argument(
        "query", 
        type=validate_query,  # Use validator
        help="Question to be asked to the Copilot"
    )
    ask_parser.add_argument(
        "--files",
        "-f",
        nargs="*",
        type=validate_path,  # Use validator
        dest="files",
        help="Folder or file as context (absolute or relative path)",
    )
    
    # Continue with rest of subparsers...

4. src/pieces/settings.py

Fix: startup() Method Improvements

Old Code:

@classmethod
def startup(cls):
    if cls.pieces_client.is_pieces_running():
        cls.version_check()  # Check the version first
    else:
        server_startup_failed()
        sys.exit(2)  # Exit the program

New Code:

@classmethod
def startup(cls):
    """Start up and verify Pieces OS connection with retries."""
    import time
    
    max_retries = 3
    retry_delay = 1
    
    for attempt in range(max_retries):
        try:
            if cls.pieces_client.is_pieces_running():
                # Verify we can actually communicate
                try:
                    cls.pieces_client.version  # Test API call
                    cls.version_check()
                    return  # Success
                except Exception as e:
                    cls.logger.debug(f"API test failed: {e}")
                    if attempt < max_retries - 1:
                        time.sleep(retry_delay)
                        retry_delay *= 2
                        continue
                    raise
            else:
                if attempt < max_retries - 1:
                    cls.logger.print(
                        f"Pieces OS not responding, retrying in {retry_delay}s..."
                    )
                    time.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    server_startup_failed()
                    sys.exit(2)
        except Exception as e:
            if attempt < max_retries - 1:
                cls.logger.debug(f"Startup attempt {attempt + 1} failed: {e}")
                time.sleep(retry_delay)
                retry_delay *= 2
            else:
                cls.logger.error(f"Failed to connect after {max_retries} attempts")
                server_startup_failed()
                sys.exit(2)

Fix: version_check() Improvements

Old Code:

@classmethod
def version_check(cls):
    """Check if the version of PiecesOS is compatible"""
    cls.pieces_os_version = cls.pieces_client.version
    result = VersionChecker(
        cls.PIECES_OS_MIN_VERSION, cls.PIECES_OS_MAX_VERSION, cls.pieces_os_version
    ).version_check()
    # ... rest of method

New Code:

@classmethod
def version_check(cls):
    """Check if the version of PiecesOS is compatible with error handling."""
    try:
        cls.pieces_os_version = cls.pieces_client.version
        
        # Validate version format
        if not cls.pieces_os_version or not isinstance(cls.pieces_os_version, str):
            cls.logger.warning("Unable to determine Pieces OS version")
            return  # Continue anyway
        
        # Sanitize version string
        import re
        if not re.match(r'^\d+\.\d+\.\d+', cls.pieces_os_version):
            cls.logger.warning(f"Invalid version format: {cls.pieces_os_version}")
            return
        
        result = VersionChecker(
            cls.PIECES_OS_MIN_VERSION, 
            cls.PIECES_OS_MAX_VERSION, 
            cls.pieces_os_version
        ).version_check()

        # Check compatibility
        if result.update == UpdateEnum.Plugin:
            print(
                "\n⚠️  CLI Update Required\n"
                "Your CLI version is not compatible with the current Pieces OS.\n"
                f"Please update: pip install --upgrade pieces-cli\n"
            )
            print_version_details(cls.pieces_os_version, __version__)
            
            if not cls.logger.confirm("Continue anyway? (not recommended)"):
                sys.exit(2)
                
        elif result.update == UpdateEnum.PiecesOS:
            print(
                "\n⚠️  Pieces OS Update Required\n"
                "Your Pieces OS version is not compatible with this CLI version.\n"
            )
            print_pieces_os_link()
            print_version_details(cls.pieces_os_version, __version__)
            
            if not cls.logger.confirm("Continue anyway? (not recommended)"):
                sys.exit(2)
                
    except Exception as e:
        cls.logger.warning(f"Version check failed: {e}")
        # Continue anyway - don't block the user

Fix: show_error() Thread Safety

Old Code:

@classmethod
def show_error(cls, error, error_message=None):
    print(f"\033[31m{error}\033[0m")
    print(f"\033[31m{error_message}\033[0m") if error_message else None
    if not cls.run_in_loop:
        sys.exit(2)

New Code:

@classmethod
def show_error(cls, error, error_message=None):
    """Show error message with thread safety and proper formatting."""
    import threading
    
    # Thread-safe print
    with threading.Lock():
        # Use logger if available, fallback to print
        if hasattr(cls, 'logger') and cls.logger:
            cls.logger.error(error)
            if error_message:
                cls.logger.error(error_message)
        else:
            # Fallback with color support detection
            if sys.stdout.isatty():
                print(f"\033[31m{error}\033[0m")
                if error_message:
                    print(f"\033[31m{error_message}\033[0m")
            else:
                print(f"ERROR: {error}")
                if error_message:
                    print(f"ERROR: {error_message}")
        
        # Ensure output is flushed
        sys.stdout.flush()
        sys.stderr.flush()
    
    # Exit if not in loop mode
    if not cls.run_in_loop:
        sys.exit(2)

Fix: get_os_id() Error Handling

Old Code:

@classmethod
def get_os_id(cls):
    from pieces_os_client.models.application_name_enum import ApplicationNameEnum

    if cls._os_id:
        return cls._os_id
    for app in cls.pieces_client.applications_api.applications_snapshot().iterable:
        if app.name == ApplicationNameEnum.OS_SERVER:
            cls._os_id = app.id
            return app.id

New Code:

@classmethod
def get_os_id(cls):
    """Get OS application ID with error handling."""
    from pieces_os_client.models.application_name_enum import ApplicationNameEnum

    if cls._os_id:
        return cls._os_id
    
    try:
        snapshot = cls.pieces_client.applications_api.applications_snapshot()
        if not snapshot or not snapshot.iterable:
            cls.logger.warning("No applications found")
            return None
        
        for app in snapshot.iterable:
            if app and hasattr(app, 'name') and app.name == ApplicationNameEnum.OS_SERVER:
                cls._os_id = app.id
                return app.id
        
        cls.logger.warning("OS Server application not found")
        return None
        
    except ConnectionError:
        cls.logger.debug("Unable to connect to Pieces OS")
        return None
    except Exception as e:
        cls.logger.debug(f"Error getting OS ID: {e}")
        return None

Fix: open_website() URL Validation

Old Code:

@classmethod
def open_website(cls, url: str):
    user_profile = cls.pieces_client.user_api.user_snapshot().user
    if (not cls.pieces_client.is_pieces_running) or ("pieces.app" not in url):
        return webbrowser.open(url)
    # ... rest of method

New Code:

@classmethod
def open_website(cls, url: str):
    """Open website with URL validation and error handling."""
    import re
    from urllib.parse import urlparse, parse_qsl, urlencode, urlunparse, quote
    
    # Validate URL format
    if not url or not isinstance(url, str):
        cls.logger.error("Invalid URL provided")
        return False
    
    # Basic URL validation
    url = url.strip()
    if not re.match(r'^https?://', url, re.IGNORECASE):
        cls.logger.error(f"Invalid URL format: {url}")
        return False
    
    # Length check
    if len(url) > 2048:  # Common browser limit
        cls.logger.error("URL too long")
        return False
    
    try:
        # Parse and validate URL components
        parsed = urlparse(url)
        if not parsed.netloc:
            cls.logger.error("Invalid URL: missing domain")
            return False
        
        # Check for Pieces app URL
        if "pieces.app" not in parsed.netloc.lower():
            # Non-Pieces URL, open directly
            return webbrowser.open(url)
        
        # For Pieces URLs, try to add user context
        try:
            if not cls.pieces_client.is_pieces_running():
                # Pieces not running, open URL as-is
                return webbrowser.open(url)
            
            # Get user profile safely
            user_profile = None
            try:
                user_snapshot = cls.pieces_client.user_api.user_snapshot()
                if user_snapshot and user_snapshot.user:
                    user_profile = user_snapshot.user
            except:
                pass  # Continue without user profile
            
            # Build parameters
            params = {}
            if user_profile and hasattr(user_profile, 'id'):
                params["user"] = user_profile.id
            
            os_id = cls.get_os_id()
            if os_id:
                params["os"] = os_id
            
            # Add parameters to URL
            if params:
                url_parts = list(parsed)
                query = dict(parse_qsl(url_parts[4]))
                query.update(params)
                url_parts[4] = urlencode(query, safe='', quote_via=quote)
                url = urlunparse(url_parts)
            
        except Exception as e:
            cls.logger.debug(f"Failed to enhance Pieces URL: {e}")
            # Continue with original URL
        
        # Open the URL
        return webbrowser.open(url)
        
    except Exception as e:
        cls.logger.error(f"Failed to open URL: {e}")
        cls.logger.debug(f"URL was: {url}")
        return False

5. tests/assets/execute_command_test.py

Fix: Add Security Tests

Add these test methods to the TestExecuteCommand class:

def test_command_injection_prevention(self):
    """Test that command injection is prevented."""
    # Test malicious content
    self.mock_asset.raw_content = "'; rm -rf /; echo '"
    self.mock_asset.classification = ClassificationSpecificEnum.BASH
    
    with patch('pieces.settings.Settings.show_error') as mock_error:
        # The new implementation should prevent this
        ExecuteCommand.execute_command()
        # Should not execute dangerous commands
        # Verify subprocess.run was called safely

def test_path_traversal_prevention(self):
    """Test that path traversal is prevented."""
    self.mock_asset.id = "../../../etc/passwd"
    
    with patch('pieces.commands.execute_command.AssetsCommands.create_asset_file') as mock_create:
        # The sanitization should prevent path traversal
        mock_create.return_value = "/tmp/safe_file.py"
        ExecuteCommand.execute_command()
        
        # Verify the ID was sanitized
        call_args = mock_create.call_args[0][0]
        self.assertNotIn("..", call_args.id)

def test_command_template_validation(self):
    """Test command template validation."""
    dangerous_templates = [
        "rm -rf {file}",
        "python {file} && rm -rf /",
        "echo {content} | bash",
        "eval {content}",
        "{content}; malicious_command",
    ]
    
    for template in dangerous_templates:
        self.assertFalse(
            ExecuteCommand.validate_command_template(template),
            f"Template should be rejected: {template}"
        )
    
    safe_templates = [
        "python {file}",
        "node {file}",
        "gcc {file} -o {file_no_extension}",
    ]
    
    for template in safe_templates:
        self.assertTrue(
            ExecuteCommand.validate_command_template(template),
            f"Template should be accepted: {template}"
        )

def test_timeout_handling(self):
    """Test that long-running commands are terminated."""
    import subprocess
    
    self.mock_asset.raw_content = "import time; time.sleep(60)"
    
    with patch('subprocess.run') as mock_run:
        mock_run.side_effect = subprocess.TimeoutExpired('cmd', 30)
        
        with patch('pieces.settings.Settings.logger') as mock_logger:
            ExecuteCommand.execute_command()
            
            # Verify timeout error was logged
            mock_logger.print.assert_any_call(
                "Error: Command execution timed out after 30 seconds"
            )

def test_resource_limits(self):
    """Test that resource limits are applied."""
    # This would test the resource limit implementation
    self.mock_asset.raw_content = "a = 'x' * (1024 * 1024 * 1024)"  # Try to allocate 1GB
    
    # The execution should fail due to memory limits
    # Implementation depends on how resource limits are enforced

def test_file_size_limits(self):
    """Test file size limits in asset creation."""
    large_content = "x" * (11 * 1024 * 1024)  # 11MB
    self.mock_asset.raw_content = large_content
    
    with patch('pieces.settings.Settings.show_error') as mock_error:
        # Should reject files that are too large
        ExecuteCommand.execute_command()

def test_concurrent_execution_safety(self):
    """Test thread safety of execution."""
    import threading
    import time
    
    results = []
    
    def execute_thread():
        try:
            ExecuteCommand.execute_command()
            results.append("success")
        except Exception as e:
            results.append(f"error: {e}")
    
    # Start multiple threads
    threads = []
    for _ in range(5):
        t = threading.Thread(target=execute_thread)
        threads.append(t)
        t.start()
    
    # Wait for completion
    for t in threads:
        t.join(timeout=5)
    
    # Should handle concurrent execution safely
    self.assertEqual(len(results), 5)

Summary

These fixes address:

1. Security vulnerabilities: Command injection, path traversal, pickle deserialization
2. Input validation: Size limits, character validation, path sanitization
3. Error handling: Connection failures, timeouts, permission errors
4. Resource management: File cleanup, memory limits, CPU limits
5. Thread safety: Concurrent access, race conditions
6. User experience: Better error messages, progress feedback, retry logic

The changes make the codebase significantly more secure and robust while maintaining backward compatibility where possible.