shotgun-sh 0.2.19__py3-none-any.whl → 0.2.23.dev1__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -633,7 +633,10 @@ class AgentManager(Widget):
                             "Your OpenAI organization doesn't have streaming enabled for this model.\n\n"
                             "**Options:**\n"
                             "- Get a [Shotgun Account](https://shotgun.sh) - streaming works out of the box\n"
-                            "- Complete [Biometric Verification](https://platform.openai.com/settings/organization/billing/payment-methods) with OpenAI\n\n"
+                            "- Complete [Biometric Verification](https://platform.openai.com/settings/organization/general) with OpenAI, then:\n"
+                            "  1. Press `Ctrl+P` → Open Provider Setup\n"
+                            "  2. Select OpenAI → Clear key\n"
+                            "  3. Re-add your OpenAI API key\n\n"
                             "Continuing without streaming (responses will appear all at once)."
                         )
                     )

shotgun/agents/config/models.py

@@ -86,6 +86,15 @@ class ModelConfig(BaseModel):
         }
         return f"{provider_prefix[self.provider]}:{self.name}"
 
+    @property
+    def is_shotgun_account(self) -> bool:
+        """Check if this model is using Shotgun Account authentication.
+
+        Returns:
+            True if using Shotgun Account, False if BYOK
+        """
+        return self.key_provider == KeyProvider.SHOTGUN
+
 
 # Model specifications registry (static metadata)
 MODEL_SPECS: dict[ModelName, ModelSpec] = {

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/history/token_counting/anthropic.py

@@ -1,6 +1,7 @@
 """Anthropic token counting using official client."""
 
 import logfire
+from anthropic import APIStatusError
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config.models import KeyProvider
@@ -103,6 +104,13 @@ class AnthropicTokenCounter(TokenCounter):
                 exception_type=type(e).__name__,
                 exception_message=str(e),
             )
+
+            # Re-raise API errors directly so they can be classified by the runner
+            # This allows proper error classification for BYOK users (authentication, rate limits, etc.)
+            if isinstance(e, APIStatusError):
+                raise
+
+            # Only wrap library-level errors in RuntimeError
             raise RuntimeError(
                 f"Anthropic token counting API failed for {self.model_name}: {type(e).__name__}: {str(e)}"
             ) from e

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

@@ -12,8 +12,8 @@ POSTHOG_API_KEY = 'phc_KKnChzZUKeNqZDOTJ6soCBWNQSx3vjiULdwTR9H5Mcr'
 POSTHOG_PROJECT_ID = '191396'
 
 # Logfire configuration embedded at build time (only for dev builds)
-LOGFIRE_ENABLED = ''
-LOGFIRE_TOKEN = ''
+LOGFIRE_ENABLED = 'true'
+LOGFIRE_TOKEN = 'pylf_v1_us_RwZMlJm1tX6j0PL5RWWbmZpzK2hLBNtFWStNKlySfjh8'
 
 # Build metadata
 BUILD_TIME_ENV = "production" if SENTRY_DSN else "development"

shotgun/cli/context.py

@@ -5,6 +5,7 @@ import json
 from pathlib import Path
 from typing import Annotated
 
+import httpx
 import typer
 from rich.console import Console
 
@@ -16,6 +17,7 @@ from shotgun.agents.context_analyzer import (
 )
 from shotgun.agents.conversation_manager import ConversationManager
 from shotgun.cli.models import OutputFormat
+from shotgun.llm_proxy import BudgetInfo, LiteLLMProxyClient
 from shotgun.logging_config import get_logger
 
 app = typer.Typer(
@@ -108,4 +110,45 @@ async def analyze_context() -> ContextAnalysisOutput:
     markdown = ContextFormatter.format_markdown(analysis)
     json_data = ContextFormatter.format_json(analysis)
 
+    # Add budget info for Shotgun Account users
+    if model_config.is_shotgun_account:
+        try:
+            logger.debug("Fetching budget info for Shotgun Account")
+            client = LiteLLMProxyClient(model_config.api_key)
+            budget_info = await client.get_budget_info()
+
+            # Format budget section for markdown
+            budget_markdown = _format_budget_markdown(budget_info)
+            markdown = f"{markdown}\n\n{budget_markdown}"
+
+            # Add budget info to JSON using Pydantic model
+            json_data["budget"] = budget_info.model_dump()
+            logger.debug("Successfully added budget info to context output")
+
+        except httpx.HTTPError as e:
+            logger.warning(f"Failed to fetch budget info: {e}")
+            # Don't fail the entire command if budget fetch fails
+        except Exception as e:
+            logger.warning(f"Unexpected error fetching budget info: {e}")
+            # Don't fail the entire command if budget fetch fails
+
     return ContextAnalysisOutput(markdown=markdown, json_data=json_data)
+
+
+def _format_budget_markdown(budget_info: BudgetInfo) -> str:
+    """Format budget information as markdown.
+
+    Args:
+        budget_info: BudgetInfo instance
+
+    Returns:
+        Formatted markdown string
+    """
+    source_label = "Key" if budget_info.source == "key" else "Team"
+
+    return f"""## Shotgun Account Budget
+
+* Max Budget:     ${budget_info.max_budget:.2f}
+* Current Spend:  ${budget_info.spend:.2f}
+* Remaining:      ${budget_info.remaining:.2f} ({100 - budget_info.percentage_used:.1f}%)
+* Budget Source:  {source_label}-level"""

shotgun/cli/error_handler.py

@@ -0,0 +1,24 @@
+"""CLI-specific error handling utilities.
+
+This module provides utilities for displaying agent errors in the CLI
+by printing formatted messages to the console.
+"""
+
+from rich.console import Console
+
+from shotgun.exceptions import ErrorNotPickedUpBySentry
+
+console = Console(stderr=True)
+
+
+def print_agent_error(exception: ErrorNotPickedUpBySentry) -> None:
+    """Print an agent error to the console in yellow.
+
+    Args:
+        exception: The error exception with formatting methods
+    """
+    # Get plain text version for CLI
+    message = exception.to_plain_text()
+
+    # Print with yellow styling
+    console.print(message, style="yellow")

shotgun/cli/export.py

@@ -11,7 +11,10 @@ from shotgun.agents.export import (
     run_export_agent,
 )
 from shotgun.agents.models import AgentRuntimeOptions
+from shotgun.cli.error_handler import print_agent_error
+from shotgun.exceptions import ErrorNotPickedUpBySentry
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 
 app = typer.Typer(
     name="export", help="Export artifacts to various formats with agentic approach"
@@ -45,37 +48,34 @@ def export(
 
     logger.info("📤 Export Instruction: %s", instruction)
 
-    try:
-        # Track export command usage
-        from shotgun.posthog_telemetry import track_event
-
-        track_event(
-            "export_command",
-            {
-                "non_interactive": non_interactive,
-                "provider": provider.value if provider else "default",
-            },
-        )
-
-        # Create agent dependencies
-        agent_runtime_options = AgentRuntimeOptions(
-            interactive_mode=not non_interactive
-        )
-
-        # Create the export agent with deps and provider
-        agent, deps = asyncio.run(create_export_agent(agent_runtime_options, provider))
-
-        # Start export process
-        logger.info("🎯 Starting export...")
-        result = asyncio.run(run_export_agent(agent, instruction, deps))
-
-        # Display results
-        logger.info("✅ Export Complete!")
-        logger.info("📤 Results:")
-        logger.info("%s", result.output)
-
-    except Exception as e:
-        logger.error("❌ Error during export: %s", str(e))
-        import traceback
-
-        logger.debug("Full traceback:\n%s", traceback.format_exc())
+    # Track export command usage
+    track_event(
+        "export_command",
+        {
+            "non_interactive": non_interactive,
+            "provider": provider.value if provider else "default",
+        },
+    )
+
+    # Create agent dependencies
+    agent_runtime_options = AgentRuntimeOptions(interactive_mode=not non_interactive)
+
+    # Create the export agent with deps and provider
+    agent, deps = asyncio.run(create_export_agent(agent_runtime_options, provider))
+
+    # Start export process with error handling
+    logger.info("🎯 Starting export...")
+
+    async def async_export() -> None:
+        try:
+            result = await run_export_agent(agent, instruction, deps)
+            logger.info("✅ Export Complete!")
+            logger.info("📤 Results:")
+            logger.info("%s", result.output)
+        except ErrorNotPickedUpBySentry as e:
+            print_agent_error(e)
+        except Exception as e:
+            logger.exception("Unexpected error in export command")
+            print(f"⚠️  An unexpected error occurred: {str(e)}")
+
+    asyncio.run(async_export())

shotgun/cli/plan.py

@@ -8,7 +8,10 @@ import typer
 from shotgun.agents.config import ProviderType
 from shotgun.agents.models import AgentRuntimeOptions
 from shotgun.agents.plan import create_plan_agent, run_plan_agent
+from shotgun.cli.error_handler import print_agent_error
+from shotgun.exceptions import ErrorNotPickedUpBySentry
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 
 app = typer.Typer(name="plan", help="Generate structured plans", no_args_is_help=True)
 logger = get_logger(__name__)
@@ -37,37 +40,34 @@ def plan(
 
     logger.info("📋 Planning Goal: %s", goal)
 
-    try:
-        # Track plan command usage
-        from shotgun.posthog_telemetry import track_event
-
-        track_event(
-            "plan_command",
-            {
-                "non_interactive": non_interactive,
-                "provider": provider.value if provider else "default",
-            },
-        )
-
-        # Create agent dependencies
-        agent_runtime_options = AgentRuntimeOptions(
-            interactive_mode=not non_interactive
-        )
-
-        # Create the plan agent with deps and provider
-        agent, deps = asyncio.run(create_plan_agent(agent_runtime_options, provider))
-
-        # Start planning process
-        logger.info("🎯 Starting planning...")
-        result = asyncio.run(run_plan_agent(agent, goal, deps))
-
-        # Display results
-        logger.info("✅ Planning Complete!")
-        logger.info("📋 Results:")
-        logger.info("%s", result.output)
-
-    except Exception as e:
-        logger.error("❌ Error during planning: %s", str(e))
-        import traceback
-
-        logger.debug("Full traceback:\n%s", traceback.format_exc())
+    # Track plan command usage
+    track_event(
+        "plan_command",
+        {
+            "non_interactive": non_interactive,
+            "provider": provider.value if provider else "default",
+        },
+    )
+
+    # Create agent dependencies
+    agent_runtime_options = AgentRuntimeOptions(interactive_mode=not non_interactive)
+
+    # Create the plan agent with deps and provider
+    agent, deps = asyncio.run(create_plan_agent(agent_runtime_options, provider))
+
+    # Start planning process with error handling
+    logger.info("🎯 Starting planning...")
+
+    async def async_plan() -> None:
+        try:
+            result = await run_plan_agent(agent, goal, deps)
+            logger.info("✅ Planning Complete!")
+            logger.info("📋 Results:")
+            logger.info("%s", result.output)
+        except ErrorNotPickedUpBySentry as e:
+            print_agent_error(e)
+        except Exception as e:
+            logger.exception("Unexpected error in plan command")
+            print(f"⚠️  An unexpected error occurred: {str(e)}")
+
+    asyncio.run(async_plan())

shotgun/cli/research.py

@@ -11,7 +11,10 @@ from shotgun.agents.research import (
     create_research_agent,
     run_research_agent,
 )
+from shotgun.cli.error_handler import print_agent_error
+from shotgun.exceptions import ErrorNotPickedUpBySentry
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 
 app = typer.Typer(
     name="research", help="Perform research with agentic loops", no_args_is_help=True
@@ -59,8 +62,6 @@ async def async_research(
 ) -> None:
     """Async wrapper for research process."""
     # Track research command usage
-    from shotgun.posthog_telemetry import track_event
-
     track_event(
         "research_command",
         {
@@ -75,11 +76,18 @@ async def async_research(
     # Create the research agent with deps and provider
     agent, deps = await create_research_agent(agent_runtime_options, provider)
 
-    # Start research process
+    # Start research process with error handling
     logger.info("🔬 Starting research...")
-    result = await run_research_agent(agent, query, deps)
-
-    # Display results
-    print("✅ Research Complete!")
-    print("📋 Findings:")
-    print(result.output)
+    try:
+        result = await run_research_agent(agent, query, deps)
+        # Display results
+        print("✅ Research Complete!")
+        print("📋 Findings:")
+        print(result.output)
+    except ErrorNotPickedUpBySentry as e:
+        # All user-actionable errors - display with plain text
+        print_agent_error(e)
+    except Exception as e:
+        # Unexpected errors that weren't wrapped (shouldn't happen)
+        logger.exception("Unexpected error in research command")
+        print(f"⚠️  An unexpected error occurred: {str(e)}")