shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.2.17__py3-none-any.whl

shotgun/agents/history/token_counting/openai.py

@@ -57,9 +57,15 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary encoding
+        if not text or not text.strip():
+            return 0
+
         try:
             return len(self.encoding.encode(text))
-        except Exception as e:
+        except BaseException as e:
+            # Must catch BaseException to handle PanicException from tiktoken's Rust layer
+            # which can occur with extremely long texts. Regular Exception won't catch it.
             raise RuntimeError(
                 f"Failed to count tokens for OpenAI model {self.model_name}"
             ) from e
@@ -76,5 +82,9 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/sentencepiece_counter.py

@@ -88,6 +88,10 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary tokenization
+        if not text or not text.strip():
+            return 0
+
         await self._ensure_tokenizer()
 
         if self.sp is None:
@@ -115,5 +119,9 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/tokenizer_cache.py

@@ -3,6 +3,7 @@
 import hashlib
 from pathlib import Path
 
+import aiofiles
 import httpx
 
 from shotgun.logging_config import get_logger
@@ -78,7 +79,8 @@ async def download_gemini_tokenizer() -> Path:
 
         # Atomic write: write to temp file first, then rename
         temp_path = cache_path.with_suffix(".tmp")
-        temp_path.write_bytes(content)
+        async with aiofiles.open(temp_path, "wb") as f:
+            await f.write(content)
         temp_path.rename(cache_path)
 
         logger.info(f"Gemini tokenizer downloaded and cached at {cache_path}")

shotgun/agents/history/token_counting/utils.py

@@ -44,9 +44,6 @@ def get_token_counter(model_config: ModelConfig) -> TokenCounter:
 
     # Return cached instance if available
     if cache_key in _token_counter_cache:
-        logger.debug(
-            f"Reusing cached token counter for {model_config.provider.value}:{model_config.name}"
-        )
         return _token_counter_cache[cache_key]
 
     # Create new instance and cache it

shotgun/agents/plan.py

@@ -23,7 +23,7 @@ from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 logger = get_logger(__name__)
 
 
-def create_plan_agent(
+async def create_plan_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
 ) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a plan agent with artifact management capabilities.
@@ -39,7 +39,7 @@ def create_plan_agent(
     # Use partial to create system prompt function for plan agent
     system_prompt_fn = partial(build_agent_system_prompt, "plan")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,

shotgun/agents/research.py

@@ -26,7 +26,7 @@ from .tools import get_available_web_search_tools
 logger = get_logger(__name__)
 
 
-def create_research_agent(
+async def create_research_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
 ) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a research agent with web search and artifact management capabilities.
@@ -41,7 +41,7 @@ def create_research_agent(
     logger.debug("Initializing research agent")
 
     # Get available web search tools based on configured API keys
-    web_search_tools = get_available_web_search_tools()
+    web_search_tools = await get_available_web_search_tools()
     if web_search_tools:
         logger.info(
             "Research agent configured with %d web search tool(s)",
@@ -53,7 +53,7 @@ def create_research_agent(
     # Use partial to create system prompt function for research agent
     system_prompt_fn = partial(build_agent_system_prompt, "research")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,

shotgun/agents/specify.py

@@ -23,7 +23,7 @@ from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 logger = get_logger(__name__)
 
 
-def create_specify_agent(
+async def create_specify_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
 ) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a specify agent with artifact management capabilities.
@@ -39,7 +39,7 @@ def create_specify_agent(
     # Use partial to create system prompt function for specify agent
     system_prompt_fn = partial(build_agent_system_prompt, "specify")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,

shotgun/agents/tasks.py

@@ -23,7 +23,7 @@ from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 logger = get_logger(__name__)
 
 
-def create_tasks_agent(
+async def create_tasks_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
 ) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a tasks agent with file management capabilities.
@@ -39,7 +39,7 @@ def create_tasks_agent(
     # Use partial to create system prompt function for tasks agent
     system_prompt_fn = partial(build_agent_system_prompt, "tasks")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         provider=provider,

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,9 +6,12 @@ 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
 
@@ -157,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.
 
@@ -175,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
 
@@ -188,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,
@@ -222,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}"
 
@@ -251,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

@@ -26,7 +26,7 @@ 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.
 
     Works with both Shotgun Account (via LiteLLM proxy) and BYOK (individual provider keys).
@@ -43,25 +43,25 @@ 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:
         logger.debug("🔑 Shotgun Account - only Gemini web search available")
 
         # Gemini: Only search tool available for Shotgun Account
-        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)
 
         # Anthropic: Not available for Shotgun Account (Gemini-only for Shotgun)
-        if is_provider_available(ProviderType.ANTHROPIC):
+        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 is_provider_available(ProviderType.OPENAI):
+        if await is_provider_available(ProviderType.OPENAI):
             logger.debug(
                 "⚠️  OpenAI web search requires BYOK (Responses API not supported via proxy)"
             )
@@ -69,15 +69,15 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
         # 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)