shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.3.3.dev1__py3-none-any.whl

shotgun/agents/{conversation_manager.py → conversation/manager.py}

@@ -1,13 +1,17 @@
 """Manager for handling conversation persistence operations."""
 
+import asyncio
 import json
-import shutil
 from pathlib import Path
 
+import aiofiles
+import aiofiles.os
+
 from shotgun.logging_config import get_logger
 from shotgun.utils import get_shotgun_home
+from shotgun.utils.file_system_utils import async_copy_file
 
-from .conversation_history import ConversationHistory
+from .models import ConversationHistory
 
 logger = get_logger(__name__)
 
@@ -27,14 +31,14 @@ class ConversationManager:
         else:
             self.conversation_path = conversation_path
 
-    def save(self, conversation: ConversationHistory) -> None:
+    async def save(self, conversation: ConversationHistory) -> None:
         """Save conversation history to file.
 
         Args:
             conversation: ConversationHistory to save
         """
         # Ensure directory exists
-        self.conversation_path.parent.mkdir(parents=True, exist_ok=True)
+        await aiofiles.os.makedirs(self.conversation_path.parent, exist_ok=True)
 
         try:
             # Update timestamp
@@ -42,11 +46,17 @@ class ConversationManager:
 
             conversation.updated_at = datetime.now()
 
-            # Serialize to JSON using Pydantic's model_dump
-            data = conversation.model_dump(mode="json")
+            # Serialize to JSON in background thread to avoid blocking event loop
+            # This is crucial for large conversations (5k+ tokens)
+            data = await asyncio.to_thread(conversation.model_dump, mode="json")
+            json_content = await asyncio.to_thread(
+                json.dumps, data, indent=2, ensure_ascii=False
+            )
 
-            with open(self.conversation_path, "w", encoding="utf-8") as f:
-                json.dump(data, f, indent=2, ensure_ascii=False)
+            async with aiofiles.open(
+                self.conversation_path, "w", encoding="utf-8"
+            ) as f:
+                await f.write(json_content)
 
             logger.debug("Conversation saved to %s", self.conversation_path)
 
@@ -56,21 +66,26 @@ class ConversationManager:
             )
             # Don't raise - we don't want to interrupt the user's session
 
-    def load(self) -> ConversationHistory | None:
+    async def load(self) -> ConversationHistory | None:
         """Load conversation history from file.
 
         Returns:
             ConversationHistory if file exists and is valid, None otherwise
         """
-        if not self.conversation_path.exists():
+        if not await aiofiles.os.path.exists(self.conversation_path):
             logger.debug("No conversation history found at %s", self.conversation_path)
             return None
 
         try:
-            with open(self.conversation_path, encoding="utf-8") as f:
-                data = json.load(f)
-
-            conversation = ConversationHistory.model_validate(data)
+            async with aiofiles.open(self.conversation_path, encoding="utf-8") as f:
+                content = await f.read()
+                # Deserialize JSON in background thread to avoid blocking
+                data = await asyncio.to_thread(json.loads, content)
+
+            # Validate model in background thread for large conversations
+            conversation = await asyncio.to_thread(
+                ConversationHistory.model_validate, data
+            )
             logger.debug(
                 "Conversation loaded from %s with %d agent messages",
                 self.conversation_path,
@@ -89,7 +104,7 @@ class ConversationManager:
             # Create a backup of the corrupted file for debugging
             backup_path = self.conversation_path.with_suffix(".json.backup")
             try:
-                shutil.copy2(self.conversation_path, backup_path)
+                await async_copy_file(self.conversation_path, backup_path)
                 logger.info("Backed up corrupted conversation to %s", backup_path)
             except Exception as backup_error:  # pragma: no cover
                 logger.warning("Failed to backup corrupted file: %s", backup_error)
@@ -105,11 +120,12 @@ class ConversationManager:
             )
             return None
 
-    def clear(self) -> None:
+    async def clear(self) -> None:
         """Delete the conversation history file."""
-        if self.conversation_path.exists():
+        if await aiofiles.os.path.exists(self.conversation_path):
             try:
-                self.conversation_path.unlink()
+                # Use asyncio.to_thread for unlink operation
+                await asyncio.to_thread(self.conversation_path.unlink)
                 logger.debug(
                     "Conversation history cleared at %s", self.conversation_path
                 )
@@ -118,10 +134,10 @@ class ConversationManager:
                     "Failed to clear conversation at %s: %s", self.conversation_path, e
                 )
 
-    def exists(self) -> bool:
+    async def exists(self) -> bool:
         """Check if a conversation history file exists.
 
         Returns:
             True if conversation file exists, False otherwise
         """
-        return self.conversation_path.exists()
+        return await aiofiles.os.path.exists(str(self.conversation_path))

shotgun/agents/{conversation_history.py → conversation/models.py}

@@ -1,7 +1,5 @@
-"""Models and utilities for persisting TUI conversation history."""
+"""Models for persisting TUI conversation history."""
 
-import json
-import logging
 from datetime import datetime
 from typing import Any, cast
 
@@ -16,99 +14,15 @@ from pydantic_core import to_jsonable_python
 
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage
 
-logger = logging.getLogger(__name__)
+from .filters import (
+    filter_incomplete_messages,
+    filter_orphaned_tool_responses,
+    is_tool_call_complete,
+)
 
 SerializedMessage = dict[str, Any]
 
 
-def is_tool_call_complete(tool_call: ToolCallPart) -> bool:
-    """Check if a tool call has valid, complete JSON arguments.
-
-    Args:
-        tool_call: The tool call part to validate
-
-    Returns:
-        True if the tool call args are valid JSON, False otherwise
-    """
-    if tool_call.args is None:
-        return True  # No args is valid
-
-    if isinstance(tool_call.args, dict):
-        return True  # Already parsed dict is valid
-
-    if not isinstance(tool_call.args, str):
-        return False
-
-    # Try to parse the JSON string
-    try:
-        json.loads(tool_call.args)
-        return True
-    except (json.JSONDecodeError, ValueError) as e:
-        # Log incomplete tool call detection
-        args_preview = (
-            tool_call.args[:100] + "..."
-            if len(tool_call.args) > 100
-            else tool_call.args
-        )
-        logger.info(
-            "Detected incomplete tool call in validation",
-            extra={
-                "tool_name": tool_call.tool_name,
-                "tool_call_id": tool_call.tool_call_id,
-                "args_preview": args_preview,
-                "error": str(e),
-            },
-        )
-        return False
-
-
-def filter_incomplete_messages(messages: list[ModelMessage]) -> list[ModelMessage]:
-    """Filter out messages with incomplete tool calls.
-
-    Args:
-        messages: List of messages to filter
-
-    Returns:
-        List of messages with only complete tool calls
-    """
-    filtered: list[ModelMessage] = []
-    filtered_count = 0
-    filtered_tool_names: list[str] = []
-
-    for message in messages:
-        # Only check ModelResponse messages for tool calls
-        if not isinstance(message, ModelResponse):
-            filtered.append(message)
-            continue
-
-        # Check if any tool calls are incomplete
-        has_incomplete_tool_call = False
-        for part in message.parts:
-            if isinstance(part, ToolCallPart) and not is_tool_call_complete(part):
-                has_incomplete_tool_call = True
-                filtered_tool_names.append(part.tool_name)
-                break
-
-        # Only include messages without incomplete tool calls
-        if not has_incomplete_tool_call:
-            filtered.append(message)
-        else:
-            filtered_count += 1
-
-    # Log if any messages were filtered
-    if filtered_count > 0:
-        logger.info(
-            "Filtered incomplete messages before saving",
-            extra={
-                "filtered_count": filtered_count,
-                "total_messages": len(messages),
-                "filtered_tool_names": filtered_tool_names,
-            },
-        )
-
-    return filtered
-
-
 class ConversationState(BaseModel):
     """Represents the complete state of a conversation in memory."""
 
@@ -142,6 +56,8 @@ class ConversationHistory(BaseModel):
         """
         # Filter out messages with incomplete tool calls to prevent corruption
         filtered_messages = filter_incomplete_messages(messages)
+        # Filter out orphaned tool responses (tool responses without tool calls)
+        filtered_messages = filter_orphaned_tool_responses(filtered_messages)
 
         # Serialize ModelMessage list to JSON-serializable format
         self.agent_history = to_jsonable_python(

shotgun/agents/error/__init__.py

@@ -0,0 +1,11 @@
+"""Agent error handling module.
+
+This module provides the AgentErrorContext model used by AgentRunner
+for error classification.
+"""
+
+from shotgun.agents.error.models import AgentErrorContext
+
+__all__ = [
+    "AgentErrorContext",
+]

shotgun/agents/error/models.py

@@ -0,0 +1,19 @@
+"""Pydantic models for agent error handling."""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AgentErrorContext(BaseModel):
+    """Context information needed to classify and handle agent errors.
+
+    Attributes:
+        exception: The exception that was raised
+        is_shotgun_account: Whether the user is using a Shotgun Account
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    exception: Any = Field(...)
+    is_shotgun_account: bool

shotgun/agents/export.py

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

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/runner.py

@@ -0,0 +1,230 @@
+"""Unified agent execution with consistent error handling.
+
+This module provides a reusable agent runner that wraps agent execution exceptions
+in user-friendly custom exceptions that can be caught and displayed by TUI or CLI.
+"""
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, NoReturn
+
+from anthropic import APIStatusError as AnthropicAPIStatusError
+from openai import APIStatusError as OpenAIAPIStatusError
+from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior
+
+from shotgun.agents.error.models import AgentErrorContext
+from shotgun.exceptions import (
+    AgentCancelledException,
+    BudgetExceededException,
+    BYOKAuthenticationException,
+    BYOKGenericAPIException,
+    BYOKQuotaBillingException,
+    BYOKRateLimitException,
+    BYOKServiceOverloadException,
+    ContextSizeLimitExceeded,
+    GenericAPIStatusException,
+    ShotgunRateLimitException,
+    ShotgunServiceOverloadException,
+    UnknownAgentException,
+)
+
+if TYPE_CHECKING:
+    from shotgun.agents.agent_manager import AgentManager
+
+logger = logging.getLogger(__name__)
+
+
+class AgentRunner:
+    """Unified agent execution wrapper with consistent error handling.
+
+    This class wraps agent execution and converts any exceptions into
+    user-friendly custom exceptions that can be caught and displayed by the
+    calling interface (TUI or CLI).
+
+    The runner:
+    - Executes the agent
+    - Logs errors for debugging
+    - Wraps exceptions in custom exception types (AgentCancelledException,
+      BYOKRateLimitException, etc.)
+    - Lets exceptions propagate to caller for display
+
+    Example:
+        >>> runner = AgentRunner(agent_manager)
+        >>> try:
+        >>>     await runner.run("Write a hello world function")
+        >>> except ContextSizeLimitExceeded as e:
+        >>>     print(e.to_markdown())
+        >>> except BYOKRateLimitException as e:
+        >>>     print(e.to_plain_text())
+    """
+
+    def __init__(self, agent_manager: "AgentManager"):
+        """Initialize the agent runner.
+
+        Args:
+            agent_manager: The agent manager to execute
+        """
+        self.agent_manager = agent_manager
+
+    async def run(self, prompt: str) -> None:
+        """Run the agent with the given prompt.
+
+        Args:
+            prompt: The user's prompt/query
+
+        Raises:
+            Custom exceptions for different error types:
+            - AgentCancelledException: User cancelled the operation
+            - ContextSizeLimitExceeded: Context too large for model
+            - BudgetExceededException: Shotgun Account budget exceeded
+            - BYOKRateLimitException: BYOK rate limit hit
+            - BYOKQuotaBillingException: BYOK quota/billing issue
+            - BYOKAuthenticationException: BYOK authentication failed
+            - BYOKServiceOverloadException: BYOK service overloaded
+            - BYOKGenericAPIException: Generic BYOK API error
+            - ShotgunServiceOverloadException: Shotgun service overloaded
+            - ShotgunRateLimitException: Shotgun rate limit hit
+            - GenericAPIStatusException: Generic API error
+            - UnknownAgentException: Unknown/unclassified error
+        """
+        try:
+            await self.agent_manager.run(prompt=prompt)
+
+        except asyncio.CancelledError as e:
+            # User cancelled - wrap and re-raise as our custom exception
+            context = self._create_error_context(e)
+            self._classify_and_raise(context)
+
+        except ContextSizeLimitExceeded as e:
+            # Already a custom exception - log and re-raise
+            logger.info(
+                "Context size limit exceeded",
+                extra={
+                    "max_tokens": e.max_tokens,
+                    "model_name": e.model_name,
+                },
+            )
+            raise
+
+        except Exception as e:
+            # Log with full stack trace to shotgun.log
+            logger.exception(
+                "Agent run failed",
+                extra={
+                    "agent_mode": self.agent_manager._current_agent_type.value,
+                    "error_type": type(e).__name__,
+                },
+            )
+
+            # Create error context and wrap/raise custom exception
+            context = self._create_error_context(e)
+            self._classify_and_raise(context)
+
+    def _create_error_context(self, exception: BaseException) -> AgentErrorContext:
+        """Create error context from exception and agent state.
+
+        Args:
+            exception: The exception that was raised
+
+        Returns:
+            AgentErrorContext with all necessary information for classification
+        """
+        return AgentErrorContext(
+            exception=exception,
+            is_shotgun_account=self.agent_manager.deps.llm_model.is_shotgun_account,
+        )
+
+    def _classify_and_raise(self, context: AgentErrorContext) -> NoReturn:
+        """Classify an exception and raise the appropriate custom exception.
+
+        Args:
+            context: Context information about the error
+
+        Raises:
+            Custom exception based on the error type
+        """
+        exception = context.exception
+        error_name = type(exception).__name__
+        error_message = str(exception)
+
+        # Check for cancellation
+        if isinstance(exception, asyncio.CancelledError):
+            raise AgentCancelledException() from exception
+
+        # Check for context size limit exceeded
+        if isinstance(exception, ContextSizeLimitExceeded):
+            # Already the right exception type, re-raise it
+            raise exception
+
+        # Check for budget exceeded (Shotgun Account only)
+        if (
+            context.is_shotgun_account
+            and "apistatuserror" in error_name.lower()
+            and "budget" in error_message.lower()
+            and "exceeded" in error_message.lower()
+        ):
+            raise BudgetExceededException(message=error_message) from exception
+
+        # Check for empty model response (e.g., model unavailable or misconfigured)
+        if isinstance(exception, UnexpectedModelBehavior):
+            raise GenericAPIStatusException(
+                "The model returned an empty response. This may indicate:\n"
+                "- The model is unavailable or misconfigured\n"
+                "- A temporary service issue\n\n"
+                "Try switching to a different model or try again later."
+            ) from exception
+
+        # Detect API errors
+        is_api_error = False
+        if isinstance(exception, OpenAIAPIStatusError):
+            is_api_error = True
+        elif isinstance(exception, AnthropicAPIStatusError):
+            is_api_error = True
+        elif isinstance(exception, ModelHTTPError):
+            # pydantic_ai wraps API errors in ModelHTTPError
+            # Check for HTTP error status codes (4xx client errors)
+            if 400 <= exception.status_code < 500:
+                is_api_error = True
+
+        # BYOK user API errors
+        if not context.is_shotgun_account and is_api_error:
+            self._raise_byok_api_error(error_message, exception)
+
+        # Shotgun Account specific errors
+        if "APIStatusError" in error_name:
+            if "overload" in error_message.lower():
+                raise ShotgunServiceOverloadException(error_message) from exception
+            elif "rate" in error_message.lower():
+                raise ShotgunRateLimitException(error_message) from exception
+            else:
+                raise GenericAPIStatusException(error_message) from exception
+
+        # Unknown error - wrap in our custom exception
+        raise UnknownAgentException(exception) from exception
+
+    def _raise_byok_api_error(
+        self, error_message: str, original_exception: Exception
+    ) -> NoReturn:
+        """Classify and raise API errors for BYOK users into specific types.
+
+        Args:
+            error_message: The error message from the API
+            original_exception: The original exception
+
+        Raises:
+            Specific BYOK exception type
+        """
+        error_lower = error_message.lower()
+
+        if "rate" in error_lower:
+            raise BYOKRateLimitException(error_message) from original_exception
+        elif "quota" in error_lower or "billing" in error_lower:
+            raise BYOKQuotaBillingException(error_message) from original_exception
+        elif "authentication" in error_lower or (
+            "invalid" in error_lower and "key" in error_lower
+        ):
+            raise BYOKAuthenticationException(error_message) from original_exception
+        elif "overload" in error_lower:
+            raise BYOKServiceOverloadException(error_message) from original_exception
+        else:
+            raise BYOKGenericAPIException(error_message) from original_exception

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: