shotgun-sh 0.1.0__py3-none-any.whl

shotgun/agents/tools/codebase/__init__.py

@@ -0,0 +1,28 @@
+"""Codebase understanding tools for Pydantic AI agents."""
+
+from .codebase_shell import codebase_shell
+from .directory_lister import directory_lister
+from .file_read import file_read
+from .models import (
+    CodeSnippetResult,
+    DirectoryListResult,
+    FileReadResult,
+    QueryGraphResult,
+    ShellCommandResult,
+)
+from .query_graph import query_graph
+from .retrieve_code import retrieve_code
+
+__all__ = [
+    "query_graph",
+    "retrieve_code",
+    "file_read",
+    "directory_lister",
+    "codebase_shell",
+    # Result models
+    "QueryGraphResult",
+    "CodeSnippetResult",
+    "FileReadResult",
+    "DirectoryListResult",
+    "ShellCommandResult",
+]

shotgun/agents/tools/codebase/codebase_shell.py

@@ -0,0 +1,256 @@
+"""Execute safe shell commands in codebase context."""
+
+import asyncio
+import re
+import time
+from pathlib import Path
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps
+from shotgun.logging_config import get_logger
+
+from .models import ShellCommandResult
+
+# Output size limits
+MAX_OUTPUT_SIZE = 50000  # Maximum characters allowed in combined stdout/stderr
+
+logger = get_logger(__name__)
+
+# Whitelist of safe read-only commands
+ALLOWED_COMMANDS = {
+    "ls",
+    "grep",
+    "find",
+    "git",
+    "cat",
+    "head",
+    "tail",
+    "wc",
+    "tree",
+    "rg",
+    "fd",
+    "ag",
+    "awk",
+    "sed",
+    "sort",
+    "uniq",
+    "cut",
+    "pwd",
+}
+
+# Patterns that indicate command injection attempts
+DANGEROUS_PATTERNS = [
+    r"[|&;`$]",  # Pipes, background, command termination, backticks, variable expansion
+    r"[<>]",  # Redirections
+    r"\$\(",  # Command substitution
+    r"^\s*\w+\s*=",  # Variable assignment
+]
+
+
+async def codebase_shell(
+    ctx: RunContext[AgentDeps],
+    command: str,
+    args: list[str],
+    graph_id: str | None = None,
+) -> ShellCommandResult:
+    """Execute safe shell commands in codebase context.
+
+    Example: Use grep patterns like this so you limit the
+    number of results while also getting the total count
+    in one command. So as not to exceed output limits.
+      `command`:
+    ```
+    # first 10 hits + grand total
+    grep -m 10 -nH "foo" src/main.cpp
+    echo "-----"
+    echo "total: $(grep -c 'foo' src/main.cpp)"
+
+    # case-insensitive, whole word, with totals
+    grep -iw -nH "foo" src/*.cpp | tee /dev/tty | wc -l
+    ```
+
+    Args:
+        ctx: RunContext containing AgentDeps with codebase service
+        command: Command to execute (must be in whitelist)
+        args: List of command arguments
+        graph_id: Optional graph ID to use (defaults to first available graph)
+
+    Returns:
+        ShellCommandResult with formatted output via __str__
+    """
+    logger.debug("🔧 Executing shell command: %s with args: %s", command, args)
+
+    try:
+        if not ctx.deps.codebase_service:
+            return ShellCommandResult(
+                success=False,
+                command=command,
+                args=args,
+                error="No codebase indexed",
+            )
+
+        # Security validation
+        if command not in ALLOWED_COMMANDS:
+            return ShellCommandResult(
+                success=False,
+                command=command,
+                args=args,
+                error=f"Command '{command}' is not allowed. Allowed commands: {', '.join(sorted(ALLOWED_COMMANDS))}",
+            )
+
+        # Validate arguments for dangerous patterns
+        full_command_str = f"{command} {' '.join(args)}"
+        for pattern in DANGEROUS_PATTERNS:
+            if re.search(pattern, full_command_str):
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error="Command contains dangerous patterns. No piping, redirection, or command substitution allowed.",
+                )
+
+        # Validate each argument individually
+        for arg in args:
+            if any(re.search(pattern, arg) for pattern in DANGEROUS_PATTERNS):
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error=f"Argument '{arg}' contains dangerous patterns.",
+                )
+
+        # Get repository path from specified graph or first available graph
+        try:
+            graphs = await ctx.deps.codebase_service.list_graphs()
+
+            if not graphs:
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error="No codebase indexed. Index a codebase first.",
+                )
+
+            # Select the appropriate graph
+            if graph_id:
+                # Find specific graph by ID
+                graph = next((g for g in graphs if g.graph_id == graph_id), None)
+                if not graph:
+                    return ShellCommandResult(
+                        success=False,
+                        command=command,
+                        args=args,
+                        error=f"Graph '{graph_id}' not found",
+                    )
+            else:
+                # Use the first available graph
+                graph = graphs[0]
+
+            repo_path = Path(graph.repo_path)
+            if not repo_path.exists():
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error=f"Repository path '{repo_path}' does not exist",
+                )
+
+        except Exception as e:
+            logger.error("Error getting graphs: %s", e)
+            return ShellCommandResult(
+                success=False,
+                command=command,
+                args=args,
+                error="Could not access codebase information",
+            )
+
+        # Execute command asynchronously
+        start_time = time.time()
+        try:
+            # Use asyncio subprocess for proper async execution
+            process = await asyncio.create_subprocess_exec(
+                command,
+                *args,
+                cwd=repo_path,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+
+            try:
+                stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                    process.communicate(), timeout=30.0
+                )
+                stdout = stdout_bytes.decode("utf-8", errors="replace")
+                stderr = stderr_bytes.decode("utf-8", errors="replace")
+                return_code = process.returncode or 0
+            except asyncio.TimeoutError:
+                # Kill the process and return timeout error
+                process.kill()
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error="Command timed out after 30 seconds",
+                )
+
+            execution_time_ms = (time.time() - start_time) * 1000
+            success = return_code == 0
+
+            logger.debug(
+                "📄 Command completed: %s with exit code %d in %.1fms",
+                "success" if success else "failed",
+                return_code,
+                execution_time_ms,
+            )
+
+            # Check if output is too large
+            combined_output_size = len(stdout) + len(stderr)
+            if combined_output_size > MAX_OUTPUT_SIZE:
+                # Format size info
+                if combined_output_size < 1024 * 1024:
+                    size_str = f"{combined_output_size / 1024:.1f}KB"
+                else:
+                    size_str = f"{combined_output_size / (1024 * 1024):.1f}MB"
+
+                guidance_msg = (
+                    f"Command output is very large ({size_str}). "
+                    "Consider using more targeted commands:\n"
+                    "• Use 'head' or 'tail' to limit lines: `head -50 file.txt`\n"
+                    "• Add filters to grep: `grep -n 'pattern' file.txt`\n"
+                    "• Use find with specific criteria: `find . -name '*.py' -type f`\n"
+                    "• Limit directory depth: `find . -maxdepth 2 -type f`\n"
+                    "• Use wc to get counts: `wc -l *.py`"
+                )
+
+                return ShellCommandResult(
+                    success=False,
+                    command=command,
+                    args=args,
+                    error=guidance_msg,
+                )
+
+            return ShellCommandResult(
+                success=success,
+                command=command,
+                args=args,
+                stdout=stdout,
+                stderr=stderr,
+                return_code=return_code,
+                execution_time_ms=execution_time_ms,
+            )
+
+        except FileNotFoundError:
+            return ShellCommandResult(
+                success=False,
+                command=command,
+                args=args,
+                error=f"Command '{command}' not found on system",
+            )
+
+    except Exception as e:
+        error_msg = f"Error executing command: {str(e)}"
+        logger.error("❌ Shell command failed: %s", str(e))
+        return ShellCommandResult(
+            success=False, command=command, args=args, error=error_msg
+        )

shotgun/agents/tools/codebase/directory_lister.py

@@ -0,0 +1,141 @@
+"""List directory contents in codebase."""
+
+from pathlib import Path
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps
+from shotgun.logging_config import get_logger
+
+from .models import DirectoryListResult
+
+logger = get_logger(__name__)
+
+
+async def directory_lister(
+    ctx: RunContext[AgentDeps], graph_id: str, directory: str = "."
+) -> DirectoryListResult:
+    """List directory contents in codebase.
+
+    Args:
+        ctx: RunContext containing AgentDeps with codebase service
+        graph_id: Graph ID to identify the repository
+        directory: Path to directory relative to repository root (default: ".")
+
+    Returns:
+        DirectoryListResult with formatted output via __str__
+    """
+    logger.debug("🔧 Listing directory: %s in graph %s", directory, graph_id)
+
+    try:
+        if not ctx.deps.codebase_service:
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path="",
+                error="No codebase indexed",
+            )
+
+        # Get the graph to find the repository path
+        try:
+            graphs = await ctx.deps.codebase_service.list_graphs()
+            graph = next((g for g in graphs if g.graph_id == graph_id), None)
+        except Exception as e:
+            logger.error("Error getting graph: %s", e)
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path="",
+                error=f"Could not find graph with ID '{graph_id}'",
+            )
+
+        if not graph:
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path="",
+                error=f"Graph '{graph_id}' not found",
+            )
+
+        # Validate the directory path is within the repository
+        repo_path = Path(graph.repo_path).resolve()
+        full_dir_path = (repo_path / directory).resolve()
+
+        # Security check: ensure the resolved path is within the repository
+        try:
+            full_dir_path.relative_to(repo_path)
+        except ValueError:
+            error_msg = (
+                f"Access denied: Path '{directory}' is outside repository bounds"
+            )
+            logger.warning("🚨 Security violation attempt: %s", error_msg)
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path=str(full_dir_path),
+                error=error_msg,
+            )
+
+        # Check if directory exists
+        if not full_dir_path.exists():
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path=str(full_dir_path),
+                error=f"Directory '{directory}' not found in repository",
+            )
+
+        if not full_dir_path.is_dir():
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path=str(full_dir_path),
+                error=f"'{directory}' is not a directory",
+            )
+
+        # List directory contents
+        try:
+            entries = list(full_dir_path.iterdir())
+            entries.sort(key=lambda p: (not p.is_dir(), p.name.lower()))
+
+            directories = []
+            files = []
+
+            for entry in entries:
+                if entry.is_dir():
+                    directories.append(entry.name)
+                elif entry.is_file():
+                    try:
+                        size = entry.stat().st_size
+                        files.append((entry.name, size))
+                    except OSError:
+                        files.append((entry.name, 0))
+
+            logger.debug(
+                "📄 Listed directory: %d directories, %d files",
+                len(directories),
+                len(files),
+            )
+
+            return DirectoryListResult(
+                success=True,
+                directory=directory,
+                full_path=str(full_dir_path),
+                directories=directories,
+                files=files,
+            )
+
+        except PermissionError:
+            return DirectoryListResult(
+                success=False,
+                directory=directory,
+                full_path=str(full_dir_path),
+                error=f"Permission denied accessing directory '{directory}'",
+            )
+
+    except Exception as e:
+        error_msg = f"Error listing directory: {str(e)}"
+        logger.error("❌ Directory listing failed: %s", str(e))
+        return DirectoryListResult(
+            success=False, directory=directory, full_path="", error=error_msg
+        )

shotgun/agents/tools/codebase/file_read.py

@@ -0,0 +1,144 @@
+"""Read file contents from codebase."""
+
+from pathlib import Path
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps
+from shotgun.codebase.core.language_config import get_language_config
+from shotgun.logging_config import get_logger
+
+from .models import FileReadResult
+
+logger = get_logger(__name__)
+
+
+async def file_read(
+    ctx: RunContext[AgentDeps], graph_id: str, file_path: str
+) -> FileReadResult:
+    """Read file contents from codebase.
+
+    Args:
+        ctx: RunContext containing AgentDeps with codebase service
+        graph_id: Graph ID to identify the repository
+        file_path: Path to file relative to repository root
+
+    Returns:
+        FileReadResult with formatted output via __str__
+    """
+    logger.debug("🔧 Reading file: %s in graph %s", file_path, graph_id)
+
+    try:
+        if not ctx.deps.codebase_service:
+            return FileReadResult(
+                success=False,
+                file_path=file_path,
+                error="No codebase indexed",
+            )
+
+        # Get the graph to find the repository path
+        try:
+            graphs = await ctx.deps.codebase_service.list_graphs()
+            graph = next((g for g in graphs if g.graph_id == graph_id), None)
+        except Exception as e:
+            logger.error("Error getting graph: %s", e)
+            return FileReadResult(
+                success=False,
+                file_path=file_path,
+                error=f"Could not find graph with ID '{graph_id}'",
+            )
+
+        if not graph:
+            return FileReadResult(
+                success=False,
+                file_path=file_path,
+                error=f"Graph '{graph_id}' not found",
+            )
+
+        # Validate the file path is within the repository
+        repo_path = Path(graph.repo_path).resolve()
+        full_file_path = (repo_path / file_path).resolve()
+
+        # Security check: ensure the resolved path is within the repository
+        try:
+            full_file_path.relative_to(repo_path)
+        except ValueError:
+            error_msg = (
+                f"Access denied: Path '{file_path}' is outside repository bounds"
+            )
+            logger.warning("🚨 Security violation attempt: %s", error_msg)
+            return FileReadResult(success=False, file_path=file_path, error=error_msg)
+
+        # Check if file exists
+        if not full_file_path.exists():
+            return FileReadResult(
+                success=False,
+                file_path=file_path,
+                error=f"File '{file_path}' not found in repository",
+            )
+
+        if full_file_path.is_dir():
+            return FileReadResult(
+                success=False,
+                file_path=file_path,
+                error=f"'{file_path}' is a directory, not a file. Use directory_lister instead.",
+            )
+
+        # Read file contents
+        encoding_used = "utf-8"
+        try:
+            content = full_file_path.read_text(encoding="utf-8")
+            size_bytes = full_file_path.stat().st_size
+
+            logger.debug(
+                "📄 Read file: %d characters, %d bytes", len(content), size_bytes
+            )
+
+            # Detect language from file extension
+            language = ""
+            file_extension = Path(file_path).suffix
+            language_config = get_language_config(file_extension)
+            if language_config:
+                language = language_config.name
+
+            return FileReadResult(
+                success=True,
+                file_path=file_path,
+                content=content,
+                encoding=encoding_used,
+                size_bytes=size_bytes,
+                language=language,
+            )
+        except UnicodeDecodeError:
+            try:
+                # Try with different encoding
+                encoding_used = "latin-1"
+                content = full_file_path.read_text(encoding="latin-1")
+                size_bytes = full_file_path.stat().st_size
+
+                # Detect language from file extension
+                language = ""
+                file_extension = Path(file_path).suffix
+                language_config = get_language_config(file_extension)
+                if language_config:
+                    language = language_config.name
+
+                return FileReadResult(
+                    success=True,
+                    file_path=file_path,
+                    content=content,
+                    encoding=encoding_used,
+                    size_bytes=size_bytes,
+                    language=language,
+                )
+            except Exception:
+                return FileReadResult(
+                    success=False,
+                    file_path=file_path,
+                    error=f"Unable to read file '{file_path}' - binary or unsupported encoding",
+                )
+
+    except Exception as e:
+        error_msg = f"Error reading file: {str(e)}"
+        logger.error("❌ File read failed: %s", str(e))
+        return FileReadResult(success=False, file_path=file_path, error=error_msg)