ShellToolsMixin

src/gaia/agents/tools/shell_tools.py

from typing import Any, Dict, Optional
from collections import deque

class ShellToolsMixin:
    """
    Mixin providing shell command execution with rate limiting.

    Attributes:
        shell_command_times: Deque tracking command timestamps
        max_commands_per_minute: Sustained rate limit (default: 10)
        max_commands_per_10_seconds: Burst rate limit (default: 3)

    Tools provided:
    - run_shell_command: Execute safe shell commands with security checks
    - wait_for_condition: Poll a predicate until it succeeds, or hit a deadline
    - get_shell_state: Report the session's current working directory
    - reset_shell_session: Return the session to its starting directory
    """

    def __init__(self, *args, **kwargs):
        """Initialize shell tools with rate limiting."""
        super().__init__(*args, **kwargs)
        self.shell_command_times = deque(maxlen=100)
        self.max_commands_per_minute = 10
        self.max_commands_per_10_seconds = 3

    def register_shell_tools(self) -> None:
        """Register shell command execution tools."""
        pass

    def _check_rate_limit(self) -> tuple:
        """
        Check if rate limit allows another command.

        Returns:
            (allowed: bool, reason: str, wait_time: float)
        """
        pass

    def _record_command_execution(self):
        """Record command execution timestamp for rate limiting."""
        pass

    @tool
    def run_shell_command(
        command: str,
        working_directory: Optional[str] = None,
        timeout: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Execute a shell command and return output.

        Security Layers:
        1. Rate limit check (burst + sustained) — one line costs one slot
        2. Working directory validation
        3. Command parsing (connectors, pipes, shlex for safety)
        4. Command whitelist check, for every segment of the line
        5. Git subcommand filtering
        6. Argument path validation, against each segment's own directory

        Args:
            command: Shell command to execute (e.g., 'ls -la', 'cd src && ls')
            working_directory: Optional directory to run in
            timeout: Max execution time in seconds for the whole line. None
                (the default) picks the default for the command's class — see
                Adaptive Timeouts.

        Returns:
            {
                "status": "success" | "error",   # "error" when a command failed
                "error": str,    # On failure: "'<step>' exited with code N: <last stderr line>"
                "expected_exit": str,  # A documented answer exit, e.g. grep's 1 (no match)
                "command": str,
                "stdout": str,   # every command's output, in order
                "stderr": str,
                "return_code": int,   # the last command that ran
                "has_errors": bool,
                "steps": list,   # [{"command": str, "return_code": int, "expected_exit"?: str}, ...]
                "duration_seconds": float,
                "timeout": int,           # The applied timeout
                "timeout_class": str,     # Which class it came from
                "cwd": str,
                "output_truncated": bool,
                "check_result": dict | None,  # CheckResult when a test/lint/build runner ran
                "rate_limited": bool,  # If refused rather than paced
                "wait_time_seconds": float,  # Time to wait
                "waited_seconds": float,  # Present only when it paced first
                "timed_out": bool  # If command timed out
            }
        """
        pass

    @tool
    def wait_for_condition(
        command: str,
        working_directory: Optional[str] = None,
        timeout: int = 120,
        poll_interval: int = 5
    ) -> Dict[str, Any]:
        """
        Block until `command` exits 0, or until the deadline passes.

        Args:
            command: The predicate — exits 0 once the condition holds. Runs
                through run_shell_command, so the same guardrails apply.
            working_directory: Optional directory to run the predicate in
            timeout: Give up after this many seconds (max 600)
            poll_interval: Seconds between checks (5-60)

        Returns:
            {
                "status": "success" | "error",
                "condition_met": bool,
                "command": str,
                "elapsed_seconds": float,
                "polls": int,
                "poll_interval_seconds": int,
                "timeout": int,
                "stdout": str,
                "stderr": str,
                "timed_out": bool,          # Deadline expired
                "last_return_code": int,    # On expiry
                "cancelled": bool           # Agent cancel signal fired
            }
        """
        pass

def _check_rate_limit(self) -> tuple:
    current_time = time.time()
    minute_ago = current_time - 60
    ten_sec_ago = current_time - 10

    recent_minute = sum(1 for t in self.shell_command_times if t > minute_ago)
    recent_10_sec = sum(1 for t in self.shell_command_times if t > ten_sec_ago)

    # Check burst limit first
    if recent_10_sec >= self.max_commands_per_10_seconds:
        wait_time = 10 - (current_time - min(recent_times))
        return False, f"Rate limit: max 3 commands per 10 seconds", wait_time

    # Check sustained limit
    if recent_minute >= self.max_commands_per_minute:
        wait_time = 60 - (current_time - min(recent_times))
        return False, f"Rate limit: max 10 commands per minute", wait_time

    return True, "", 0.0

ALLOWED_COMMANDS = {
    # File listing (READ-ONLY)
    "ls", "dir", "pwd", "cd",
    # File content (READ-ONLY)
    "cat", "head", "tail", "more", "less",
    # Text processing (READ-ONLY)
    "grep", "find", "wc", "sort", "uniq", "diff",
    # File info (READ-ONLY)
    "file", "stat", "du", "df",
    # System info (READ-ONLY)
    "whoami", "hostname", "uname", "date", "uptime",
    # Git (filtered separately)
    "git"
}

cmd_base = cmd_parts[0].lower()
if cmd_base not in ALLOWED_COMMANDS:
    return {
        "status": "error",
        "tier": TIER_CONFIRM,  # asks the user; not a refusal
        "error": f"Command '{cmd_base}' is not in the allowed list for security reasons",
        "hint": "Commands outside the no-prompt list run once the user approves them; they are not refused.",
        "examples": "ls, cat, grep, find, etc."
    }

DEVELOPER_COMMANDS = frozenset({
    "node", "npm", "make", "cmake", "go", "cargo",
    "sed", "awk", "curl", "sleep", "timeout", "export", "cp", "mv",
    # consolidation, not new reach — these already have their own paths
    "python", "python3", "pytest", "gh", "git",
})   # "rm" deliberately excluded

if cmd_base == "git":
    cmd_parts, error_message = _git_policy_argv(cmd_parts)
    if error_message is not None:
        return {"status": "error", "error": error_message, "has_errors": True}

policy = BINARY_POLICIES.get(normalize_binary(cmd_base))
if policy is not None:
    granted = binary in granted_binaries
    classify = classify_invocation if granted else classify_ungranted_invocation
    decision = classify(policy, cmd_parts)
    if decision.outcome == REFUSE:
        return {"status": "error", "error": decision.message, ...}

# Validate all file-like arguments
for arg in cmd_parts[1:]:
    candidate_path = arg
    if arg.startswith("-") and "=" in arg:
        _, candidate_path = arg.split("=", 1)

    # Check for path indicators
    if os.sep in candidate_path or "/" in candidate_path or ".." in candidate_path:
        # Ignore URLs
        if candidate_path.startswith(("http://", "https://", "git://")):
            continue

        # Resolve and validate path
        resolved_path = str(Path(cwd).joinpath(candidate_path).resolve())
        if not self.path_validator.is_path_allowed(resolved_path):
            return {
                "status": "error",
                "error": f"Access denied: '{arg}' resolves to forbidden path"
            }

def test_run_shell_command_success(agent):
    """Test successful command execution."""
    result = agent.run_shell_command("ls -la")
    assert result["status"] == "success"
    assert result["return_code"] == 0

def test_run_shell_command_rate_limit_burst(agent):
    """Test burst rate limiting (3 per 10 seconds)."""
    # Execute 3 commands quickly
    for i in range(3):
        result = agent.run_shell_command("pwd")
        assert result["status"] == "success"

    # The 4th waits out the burst window, then runs
    result = agent.run_shell_command("pwd")
    assert result["status"] == "success"
    assert result["waited_seconds"] > 0

def test_run_shell_command_rate_limit_cap(agent):
    """A wait past max_rate_limit_wait_seconds is refused, not paced."""
    agent.max_rate_limit_wait_seconds = 1
    for i in range(4):
        result = agent.run_shell_command("pwd")
    assert result["rate_limited"] is True

def test_run_shell_command_rate_limit_sustained(agent):
    """Test sustained rate limiting (10 per minute)."""
    # Would require time manipulation or mocking

def test_a_write_is_confirmable_not_refused():
    """Writes reach the confirmation prompt instead of dying in front of it.

    Checked on the validator, never by running the command: a direct tool call
    is what an approved call looks like, so it would execute.
    """
    for cmd in ["rm notes.txt", "chmod 644 a.txt", "git commit -m test"]:
        parts = cmd.split()
        assert ShellToolsMixin._validate_command(parts[0], parts, cmd)["tier"] == TIER_CONFIRM

def test_an_undescribable_invocation_is_refused():
    cmd = "git -c core.pager=x status"
    parts = cmd.split()
    assert ShellToolsMixin._validate_command(parts[0], parts, cmd)["tier"] == TIER_REFUSE

def test_run_shell_command_git_read_allowed(agent):
    """Test git read operations are allowed."""
    result = agent.run_shell_command("git status")
    assert result["status"] == "success" or "not a git repository" in result["error"]

def test_run_shell_command_path_traversal_blocked(agent, monkeypatch):
    """Test path traversal prevention."""
    # Mock path_validator to deny parent directory
    def mock_denied(path):
        return False

    monkeypatch.setattr(agent.path_validator, "is_path_allowed", mock_denied)

    result = agent.run_shell_command("cat ../secret.txt")
    assert result["status"] == "error"
    assert "denied" in result["error"].lower()

def test_run_shell_command_timeout(agent):
    """Test command timeout handling."""
    result = agent.run_shell_command("sleep 100", timeout=1)
    assert result["status"] == "error"
    assert result["timed_out"] is True

def test_run_shell_command_output_truncation(agent):
    """Test large output truncation."""
    # Command that produces >10,000 chars
    result = agent.run_shell_command("yes | head -n 1000")
    assert result["output_truncated"] is True
    assert len(result["stdout"]) <= 10100  # 10,000 + truncation message

import logging
import os
import shlex
import subprocess
import time
from collections import deque
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

# List directory
result = agent.run_shell_command("ls -la /home/user/Documents")
print(result["stdout"])

# Search for files
result = agent.run_shell_command("find . -name '*.pdf'")
files = result["stdout"].splitlines()

# Read file
result = agent.run_shell_command("cat README.md")
content = result["stdout"]

# Allowed: git status
result = agent.run_shell_command("git status")
print(result["stdout"])

# Allowed: git log
result = agent.run_shell_command("git log --oneline -10")

# Asks first: git commit is TIER_CONFIRM — shown to the user, runs if approved
# (and refused under GAIA_AUTO_APPROVE_TOOLS alone)

# Quick burst of commands. The mixin paces itself: a command over the limit
# waits for the window and then runs, so the caller writes no retry loop. The
# wait is sliced and stops early if the agent cancels the call.
for i in range(5):
    result = agent.run_shell_command("pwd")
    if result.get("waited_seconds"):
        print(f"Paced {result['waited_seconds']:.1f}s before running")

# Only a wait longer than max_rate_limit_wait_seconds is refused outright.
if result.get("rate_limited"):
    print(f"Refused: the window is {result['wait_time_seconds']:.1f}s away")