shotgun-sh 0.2.3.dev2__py3-none-any.whl → 0.2.11.dev5__py3-none-any.whl

shotgun/agents/tools/__init__.py

@@ -8,7 +8,6 @@ from .codebase import (
     retrieve_code,
 )
 from .file_management import append_file, read_file, write_file
-from .user_interaction import ask_user
 from .web_search import (
     anthropic_web_search_tool,
     gemini_web_search_tool,
@@ -21,7 +20,6 @@ __all__ = [
     "anthropic_web_search_tool",
     "gemini_web_search_tool",
     "get_available_web_search_tools",
-    "ask_user",
     "read_file",
     "write_file",
     "append_file",

shotgun/agents/tools/codebase/codebase_shell.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 
 from .models import ShellCommandResult
@@ -48,6 +49,11 @@ DANGEROUS_PATTERNS = [
 ]
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Running shell",
+    key_arg="command",
+)
 async def codebase_shell(
     ctx: RunContext[AgentDeps],
     command: str,

shotgun/agents/tools/codebase/directory_lister.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 
 from .models import DirectoryListResult
@@ -12,6 +13,11 @@ from .models import DirectoryListResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Listing directory",
+    key_arg="directory",
+)
 async def directory_lister(
     ctx: RunContext[AgentDeps], graph_id: str, directory: str = "."
 ) -> DirectoryListResult:

shotgun/agents/tools/codebase/file_read.py

@@ -2,9 +2,11 @@
 
 from pathlib import Path
 
+import aiofiles
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.core.language_config import get_language_config
 from shotgun.logging_config import get_logger
 
@@ -13,6 +15,11 @@ from .models import FileReadResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Reading file",
+    key_arg="file_path",
+)
 async def file_read(
     ctx: RunContext[AgentDeps], graph_id: str, file_path: str
 ) -> FileReadResult:
@@ -87,7 +94,8 @@ async def file_read(
         # Read file contents
         encoding_used = "utf-8"
         try:
-            content = full_file_path.read_text(encoding="utf-8")
+            async with aiofiles.open(full_file_path, encoding="utf-8") as f:
+                content = await f.read()
             size_bytes = full_file_path.stat().st_size
 
             logger.debug(
@@ -113,7 +121,8 @@ async def file_read(
             try:
                 # Try with different encoding
                 encoding_used = "latin-1"
-                content = full_file_path.read_text(encoding="latin-1")
+                async with aiofiles.open(full_file_path, encoding="latin-1") as f:
+                    content = await f.read()
                 size_bytes = full_file_path.stat().st_size
 
                 # Detect language from file extension

shotgun/agents/tools/codebase/query_graph.py

@@ -3,6 +3,7 @@
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.models import QueryType
 from shotgun.logging_config import get_logger
 
@@ -11,6 +12,11 @@ from .models import QueryGraphResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Querying code",
+    key_arg="query",
+)
 async def query_graph(
     ctx: RunContext[AgentDeps], graph_id: str, query: str
 ) -> QueryGraphResult:

shotgun/agents/tools/codebase/retrieve_code.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.core.code_retrieval import retrieve_code_by_qualified_name
 from shotgun.codebase.core.language_config import get_language_config
 from shotgun.logging_config import get_logger
@@ -14,6 +15,11 @@ from .models import CodeSnippetResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Retrieving code",
+    key_arg="qualified_name",
+)
 async def retrieve_code(
     ctx: RunContext[AgentDeps], graph_id: str, qualified_name: str
 ) -> CodeSnippetResult:

shotgun/agents/tools/file_management.py

@@ -6,20 +6,30 @@ These tools are restricted to the .shotgun directory for security.
 from pathlib import Path
 from typing import Literal
 
+import aiofiles
+import aiofiles.os
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps, AgentType, FileOperationType
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 from shotgun.utils.file_system_utils import get_shotgun_base_path
 
 logger = get_logger(__name__)
 
 # Map agent modes to their allowed directories/files (in workflow order)
-AGENT_DIRECTORIES = {
-    AgentType.RESEARCH: "research.md",
-    AgentType.SPECIFY: "specification.md",
-    AgentType.PLAN: "plan.md",
-    AgentType.TASKS: "tasks.md",
+# Values can be:
+# - A Path: exact file (e.g., Path("research.md"))
+# - A list of Paths: multiple allowed files/directories (e.g., [Path("specification.md"), Path("contracts")])
+# - "*": any file except protected files (for export agent)
+AGENT_DIRECTORIES: dict[AgentType, str | Path | list[Path]] = {
+    AgentType.RESEARCH: Path("research.md"),
+    AgentType.SPECIFY: [
+        Path("specification.md"),
+        Path("contracts"),
+    ],  # Specify can write specs and contract files
+    AgentType.PLAN: Path("plan.md"),
+    AgentType.TASKS: Path("tasks.md"),
     AgentType.EXPORT: "*",  # Export agent can write anywhere except protected files
 }
 
@@ -60,13 +70,52 @@ def _validate_agent_scoped_path(filename: str, agent_mode: AgentType | None) ->
             # Allow writing anywhere else in .shotgun directory
             full_path = (base_path / filename).resolve()
         else:
-            # For other agents, only allow writing to their specific file
-            allowed_file = AGENT_DIRECTORIES[agent_mode]
-            if filename != allowed_file:
+            # For other agents, check if they have access to the requested file
+            allowed_paths_raw = AGENT_DIRECTORIES[agent_mode]
+
+            # Convert single Path/string to list of Paths for uniform handling
+            if isinstance(allowed_paths_raw, str):
+                # Special case: "*" means export agent
+                allowed_paths = (
+                    [Path(allowed_paths_raw)] if allowed_paths_raw != "*" else []
+                )
+            elif isinstance(allowed_paths_raw, Path):
+                allowed_paths = [allowed_paths_raw]
+            else:
+                # Already a list
+                allowed_paths = allowed_paths_raw
+
+            # Check if filename matches any allowed path
+            is_allowed = False
+            for allowed_path in allowed_paths:
+                allowed_str = str(allowed_path)
+
+                # Check if it's a directory (no .md extension or suffix)
+                # Directories: Path("contracts") has no suffix, files: Path("spec.md") has .md suffix
+                if not allowed_path.suffix or (
+                    allowed_path.suffix and not allowed_str.endswith(".md")
+                ):
+                    # Directory - allow any file within this directory
+                    # Check both "contracts/file.py" and "contracts" prefix
+                    if (
+                        filename.startswith(allowed_str + "/")
+                        or filename == allowed_str
+                    ):
+                        is_allowed = True
+                        break
+                else:
+                    # Exact file match
+                    if filename == allowed_str:
+                        is_allowed = True
+                        break
+
+            if not is_allowed:
+                allowed_display = ", ".join(f"'{p}'" for p in allowed_paths)
                 raise ValueError(
-                    f"{agent_mode.value.capitalize()} agent can only write to '{allowed_file}'. "
+                    f"{agent_mode.value.capitalize()} agent can only write to {allowed_display}. "
                     f"Attempted to write to '{filename}'"
                 )
+
             full_path = (base_path / filename).resolve()
     else:
         # No agent mode specified, fall back to old validation
@@ -111,6 +160,11 @@ def _validate_shotgun_path(filename: str) -> Path:
     return full_path
 
 
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Reading file",
+    key_arg="filename",
+)
 async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
     """Read a file from the .shotgun directory.
 
@@ -129,10 +183,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
     try:
         file_path = _validate_shotgun_path(filename)
 
-        if not file_path.exists():
+        if not await aiofiles.os.path.exists(file_path):
             raise FileNotFoundError(f"File not found: {filename}")
 
-        content = file_path.read_text(encoding="utf-8")
+        async with aiofiles.open(file_path, encoding="utf-8") as f:
+            content = await f.read()
         logger.debug("📄 Read %d characters from %s", len(content), filename)
         return content
 
@@ -142,6 +197,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
         return error_msg
 
 
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Writing file",
+    key_arg="filename",
+)
 async def write_file(
     ctx: RunContext[AgentDeps],
     filename: str,
@@ -176,21 +236,22 @@ async def write_file(
         else:
             operation = (
                 FileOperationType.CREATED
-                if not file_path.exists()
+                if not await aiofiles.os.path.exists(file_path)
                 else FileOperationType.UPDATED
             )
 
         # Ensure parent directory exists
-        file_path.parent.mkdir(parents=True, exist_ok=True)
+        await aiofiles.os.makedirs(file_path.parent, exist_ok=True)
 
         # Write content
         if mode == "a":
-            with open(file_path, "a", encoding="utf-8") as f:
-                f.write(content)
+            async with aiofiles.open(file_path, "a", encoding="utf-8") as f:
+                await f.write(content)
             logger.debug("📄 Appended %d characters to %s", len(content), filename)
             result = f"Successfully appended {len(content)} characters to {filename}"
         else:
-            file_path.write_text(content, encoding="utf-8")
+            async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
+                await f.write(content)
             logger.debug("📄 Wrote %d characters to %s", len(content), filename)
             result = f"Successfully wrote {len(content)} characters to {filename}"
 
@@ -205,6 +266,11 @@ async def write_file(
         return error_msg
 
 
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Appending to file",
+    key_arg="filename",
+)
 async def append_file(ctx: RunContext[AgentDeps], filename: str, content: str) -> str:
     """Append content to a file in the .shotgun directory.
 

shotgun/agents/tools/registry.py

@@ -0,0 +1,217 @@
+"""Tool category registry using decorators for automatic registration.
+
+This module provides a decorator-based system for categorizing tools used by agents.
+Tools can be decorated with @register_tool to automatically register their category,
+which is then used by the context analyzer to break down token usage by tool type.
+
+It also provides a display registry system for tool formatting in the TUI, allowing
+tools to declare how they should be displayed when streaming.
+"""
+
+from collections.abc import Callable
+from enum import StrEnum
+from typing import TypeVar, overload
+
+import sentry_sdk
+from pydantic import BaseModel
+
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+# Type variable for decorated functions
+F = TypeVar("F", bound=Callable[..., object])
+
+
+class ToolCategory(StrEnum):
+    """Categories for agent tools used in context analysis."""
+
+    CODEBASE_UNDERSTANDING = "codebase_understanding"
+    ARTIFACT_MANAGEMENT = "artifact_management"
+    WEB_RESEARCH = "web_research"
+    AGENT_RESPONSE = "agent_response"
+    UNKNOWN = "unknown"
+
+
+class ToolDisplayConfig(BaseModel):
+    """Configuration for how a tool should be displayed in the TUI.
+
+    Attributes:
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument to extract from tool args for display
+        hide: Whether to completely hide this tool call from the UI
+    """
+
+    display_text: str
+    key_arg: str
+    hide: bool = False
+
+
+# Global registry mapping tool names to categories
+_TOOL_REGISTRY: dict[str, ToolCategory] = {}
+
+# Global registry mapping tool names to display configs
+_TOOL_DISPLAY_REGISTRY: dict[str, ToolDisplayConfig] = {}
+
+
+@overload
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+) -> Callable[[F], F]: ...
+
+
+@overload
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool,
+) -> Callable[[F], F]: ...
+
+
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool = False,
+) -> Callable[[F], F]:
+    """Decorator to register a tool's category and display configuration.
+
+    Args:
+        category: The ToolCategory enum value for this tool
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument name to extract for display (e.g., "query", "filename")
+        hide: Whether to hide this tool call completely from the UI (default: False)
+
+    Returns:
+        Decorator function that registers the tool and returns it unchanged
+
+    Display Format:
+        - When key_arg value is missing: Shows just display_text (e.g., "Reading file")
+        - When key_arg value is present: Shows "display_text: key_arg_value" (e.g., "Reading file: foo.py")
+
+    Example:
+        @register_tool(
+            category=ToolCategory.CODEBASE_UNDERSTANDING,
+            display_text="Querying code",
+            key_arg="query",
+        )
+        async def query_graph(ctx: RunContext[AgentDeps], query: str) -> str:
+            ...
+    """
+
+    def decorator(func: F) -> F:
+        tool_name = func.__name__
+        _TOOL_REGISTRY[tool_name] = category
+        logger.debug(f"Registered tool '{tool_name}' as category '{category.value}'")
+
+        # Register display config
+        config = ToolDisplayConfig(
+            display_text=display_text,
+            key_arg=key_arg,
+            hide=hide,
+        )
+        _TOOL_DISPLAY_REGISTRY[tool_name] = config
+        logger.debug(f"Registered display config for tool '{tool_name}'")
+
+        return func
+
+    return decorator
+
+
+# Backwards compatibility alias
+tool_category = register_tool
+
+
+def get_tool_category(tool_name: str) -> ToolCategory:
+    """Get category for a tool, logging unknown tools to Sentry.
+
+    Args:
+        tool_name: Name of the tool to look up
+
+    Returns:
+        ToolCategory enum value for the tool, or UNKNOWN if not registered
+    """
+    category = _TOOL_REGISTRY.get(tool_name)
+
+    if category is None:
+        logger.warning(f"Unknown tool encountered in context analysis: {tool_name}")
+        sentry_sdk.capture_message(
+            f"Unknown tool in context analysis: {tool_name}",
+            level="warning",
+            extras={"tool_name": tool_name},
+        )
+        return ToolCategory.UNKNOWN
+
+    return category
+
+
+def register_special_tool(tool_name: str, category: ToolCategory) -> None:
+    """Register a special tool that doesn't have a decorator.
+
+    Used for tools like 'final_result' that aren't actual Python functions
+    but need to be categorized.
+
+    Args:
+        tool_name: Name of the special tool
+        category: Category to assign to this tool
+    """
+    _TOOL_REGISTRY[tool_name] = category
+    logger.debug(
+        f"Registered special tool '{tool_name}' as category '{category.value}'"
+    )
+
+
+def get_tool_display_config(tool_name: str) -> ToolDisplayConfig | None:
+    """Get display configuration for a tool.
+
+    Args:
+        tool_name: Name of the tool to look up
+
+    Returns:
+        ToolDisplayConfig for the tool, or None if not registered
+    """
+    return _TOOL_DISPLAY_REGISTRY.get(tool_name)
+
+
+def register_tool_display(
+    tool_name: str,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool = False,
+) -> None:
+    """Register a display config for a special tool that doesn't have a decorator.
+
+    Used for tools like 'final_result' or builtin tools that aren't actual Python functions.
+
+    Args:
+        tool_name: Name of the special tool
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument name to extract for display
+        hide: Whether to hide this tool call completely
+    """
+    config = ToolDisplayConfig(
+        display_text=display_text,
+        key_arg=key_arg,
+        hide=hide,
+    )
+    _TOOL_DISPLAY_REGISTRY[tool_name] = config
+    logger.debug(f"Registered display config for special tool '{tool_name}'")
+
+
+# Register special tools that don't have decorators
+register_special_tool("final_result", ToolCategory.AGENT_RESPONSE)
+register_tool_display("final_result", display_text="", key_arg="", hide=True)
+
+# Register builtin tools (tools that come from Pydantic AI or model providers)
+# These don't have Python function definitions but need display formatting
+register_tool_display(
+    "web_search",
+    display_text="Searching",
+    key_arg="query",
+)

shotgun/agents/tools/web_search/__init__.py

@@ -3,7 +3,10 @@
 Provides web search capabilities for multiple LLM providers:
 - OpenAI: Uses Responses API with web_search tool (BYOK only)
 - Anthropic: Uses Messages API with web_search_20250305 tool (BYOK only)
-- Gemini: Uses grounding with Google Search via Pydantic AI (works with Shotgun Account)
+- Gemini: Uses grounding with Google Search via Pydantic AI (Shotgun Account and BYOK)
+
+Shotgun Account: Only Gemini web search is available
+BYOK: All tools work with direct provider API keys
 """
 
 from collections.abc import Awaitable, Callable
@@ -23,14 +26,15 @@ logger = get_logger(__name__)
 WebSearchTool = Callable[[str], Awaitable[str]]
 
 
-def get_available_web_search_tools() -> list[WebSearchTool]:
+async def get_available_web_search_tools() -> list[WebSearchTool]:
     """Get list of available web search tools based on configured API keys.
 
-    When using Shotgun Account (via LiteLLM proxy):
-        Only Gemini web search is available (others use provider-specific APIs)
+    Works with both Shotgun Account (via LiteLLM proxy) and BYOK (individual provider keys).
 
-    When using BYOK (individual provider keys):
-        All provider tools are available based on their respective keys
+    Available tools:
+    - Gemini: Available for both Shotgun Account and BYOK
+    - Anthropic: BYOK only (uses Messages API with web search)
+    - OpenAI: BYOK only (uses Responses API not compatible with LiteLLM proxy)
 
     Returns:
         List of web search tool functions that have API keys configured
@@ -39,33 +43,41 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
 
     # Check if using Shotgun Account
     config_manager = get_config_manager()
-    config = config_manager.load()
+    config = await config_manager.load()
     has_shotgun_key = config.shotgun.api_key is not None
 
     if has_shotgun_key:
-        # Shotgun Account mode: Only Gemini supports web search via LiteLLM
-        if is_provider_available(ProviderType.GOOGLE):
-            logger.info("🔑 Shotgun Account detected - using Gemini web search only")
-            logger.debug("   OpenAI and Anthropic web search require direct API keys")
+        logger.debug("🔑 Shotgun Account - only Gemini web search available")
+
+        # Gemini: Only search tool available for Shotgun Account
+        if await is_provider_available(ProviderType.GOOGLE):
+            logger.debug("✅ Gemini web search tool available")
             tools.append(gemini_web_search_tool)
-        else:
-            logger.warning(
-                "⚠️ Shotgun Account configured but no Gemini key - "
-                "web search unavailable"
+
+        # Anthropic: Not available for Shotgun Account (Gemini-only for Shotgun)
+        if await is_provider_available(ProviderType.ANTHROPIC):
+            logger.debug(
+                "⚠️  Anthropic web search requires BYOK (Shotgun Account uses Gemini only)"
+            )
+
+        # OpenAI: Not available for Shotgun Account (Responses API incompatible with proxy)
+        if await is_provider_available(ProviderType.OPENAI):
+            logger.debug(
+                "⚠️  OpenAI web search requires BYOK (Responses API not supported via proxy)"
             )
     else:
         # BYOK mode: Load all available tools based on individual provider keys
         logger.debug("🔑 BYOK mode - checking all provider web search tools")
 
-        if is_provider_available(ProviderType.OPENAI):
+        if await is_provider_available(ProviderType.OPENAI):
             logger.debug("✅ OpenAI web search tool available")
             tools.append(openai_web_search_tool)
 
-        if is_provider_available(ProviderType.ANTHROPIC):
+        if await is_provider_available(ProviderType.ANTHROPIC):
             logger.debug("✅ Anthropic web search tool available")
             tools.append(anthropic_web_search_tool)
 
-        if is_provider_available(ProviderType.GOOGLE):
+        if await is_provider_available(ProviderType.GOOGLE):
             logger.debug("✅ Gemini web search tool available")
             tools.append(gemini_web_search_tool)
 

shotgun/agents/tools/web_search/anthropic.py

@@ -8,11 +8,22 @@ from shotgun.agents.config import get_provider_model
 from shotgun.agents.config.constants import MEDIUM_TEXT_8K_TOKENS
 from shotgun.agents.config.models import ProviderType
 from shotgun.agents.llm import shotgun_model_request
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
+from shotgun.prompts import PromptLoader
+from shotgun.utils.datetime_utils import get_datetime_context
 
 logger = get_logger(__name__)
 
+# Global prompt loader instance
+prompt_loader = PromptLoader()
 
+
+@register_tool(
+    category=ToolCategory.WEB_RESEARCH,
+    display_text="Searching web",
+    key_arg="query",
+)
 async def anthropic_web_search_tool(query: str) -> str:
     """Perform a web search using Anthropic's Claude API.
 
@@ -35,17 +46,27 @@ async def anthropic_web_search_tool(query: str) -> str:
 
     # Get model configuration (supports both Shotgun and BYOK)
     try:
-        model_config = get_provider_model(ProviderType.ANTHROPIC)
+        model_config = await get_provider_model(ProviderType.ANTHROPIC)
     except ValueError as e:
         error_msg = f"Anthropic API key not configured: {str(e)}"
         logger.error("❌ %s", error_msg)
         span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
         return error_msg
 
+    # Get datetime context for the search prompt
+    dt_context = get_datetime_context()
+
+    # Render search prompt from template
+    search_prompt = prompt_loader.render(
+        "tools/web_search.j2",
+        query=query,
+        current_datetime=dt_context.datetime_formatted,
+        timezone_name=dt_context.timezone_name,
+        utc_offset=dt_context.utc_offset,
+    )
+
     # Build the request messages
-    messages: list[ModelMessage] = [
-        ModelRequest.user_text_prompt(f"Search for: {query}")
-    ]
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
 
     # Use the Messages API with web search tool
     try:
@@ -120,7 +141,7 @@ async def main() -> None:
     # Check if API key is available
     try:
         if callable(get_provider_model):
-            model_config = get_provider_model(ProviderType.ANTHROPIC)
+            model_config = await get_provider_model(ProviderType.ANTHROPIC)
             if not model_config.api_key:
                 raise ValueError("No API key configured")
     except (ValueError, Exception):