shotgun-sh 0.2.11.dev7__py3-none-any.whl → 0.2.23.dev1__py3-none-any.whl

shotgun/agents/config/streaming_test.py

@@ -0,0 +1,119 @@
+"""Utility for testing streaming capability of OpenAI models."""
+
+import logging
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Maximum number of attempts to test streaming capability
+MAX_STREAMING_TEST_ATTEMPTS = 3
+
+# Timeout for each streaming test attempt (in seconds)
+STREAMING_TEST_TIMEOUT = 10.0
+
+
+async def check_streaming_capability(
+    api_key: str, model_name: str, max_attempts: int = MAX_STREAMING_TEST_ATTEMPTS
+) -> bool:
+    """Check if the given OpenAI model supports streaming with this API key.
+
+    Retries multiple times to handle transient network issues. Only returns False
+    if streaming definitively fails after all retry attempts.
+
+    Args:
+        api_key: The OpenAI API key to test
+        model_name: The model name (e.g., "gpt-5", "gpt-5-mini")
+        max_attempts: Maximum number of attempts (default: 3)
+
+    Returns:
+        True if streaming is supported, False if it definitively fails
+    """
+    url = "https://api.openai.com/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+    # GPT-5 family uses max_completion_tokens instead of max_tokens
+    payload = {
+        "model": model_name,
+        "messages": [{"role": "user", "content": "test"}],
+        "stream": True,
+        "max_completion_tokens": 10,
+    }
+
+    last_error = None
+
+    for attempt in range(1, max_attempts + 1):
+        logger.debug(
+            f"Streaming test attempt {attempt}/{max_attempts} for {model_name}"
+        )
+
+        try:
+            async with httpx.AsyncClient(timeout=STREAMING_TEST_TIMEOUT) as client:
+                async with client.stream(
+                    "POST", url, json=payload, headers=headers
+                ) as response:
+                    # Check if we get a successful response
+                    if response.status_code != 200:
+                        last_error = f"HTTP {response.status_code}"
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name}: {last_error}"
+                        )
+
+                        # For definitive errors (403 Forbidden, 404 Not Found), don't retry
+                        if response.status_code in (403, 404):
+                            logger.info(
+                                f"Streaming definitively unsupported for {model_name} (HTTP {response.status_code})"
+                            )
+                            return False
+
+                        # For other errors, retry
+                        continue
+
+                    # Try to read at least one chunk from the stream
+                    try:
+                        async for _ in response.aiter_bytes():
+                            # Successfully received streaming data
+                            logger.info(
+                                f"Streaming test passed for {model_name} (attempt {attempt})"
+                            )
+                            return True
+                    except Exception as e:
+                        last_error = str(e)
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name} while reading stream: {e}"
+                        )
+                        continue
+
+        except httpx.TimeoutException:
+            last_error = "timeout"
+            logger.warning(
+                f"Streaming test attempt {attempt} timed out for {model_name}"
+            )
+            continue
+        except httpx.HTTPStatusError as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name}: {e}"
+            )
+            continue
+        except Exception as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name} with unexpected error: {e}"
+            )
+            continue
+
+        # If we got here without reading any chunks, streaming didn't work
+        last_error = "no data received"
+        logger.warning(
+            f"Streaming test attempt {attempt} failed for {model_name}: no data received"
+        )
+
+    # All attempts exhausted
+    logger.error(
+        f"Streaming test failed for {model_name} after {max_attempts} attempts. "
+        f"Last error: {last_error}. Assuming streaming is NOT supported."
+    )
+    return False

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

@@ -8,7 +8,7 @@ DO NOT EDIT MANUALLY.
 SENTRY_DSN = 'https://2818a6d165c64eccc94cfd51ce05d6aa@o4506813296738304.ingest.us.sentry.io/4510045952409600'
 
 # PostHog configuration embedded at build time (empty strings if not provided)
-POSTHOG_API_KEY = ''
+POSTHOG_API_KEY = 'phc_KKnChzZUKeNqZDOTJ6soCBWNQSx3vjiULdwTR9H5Mcr'
 POSTHOG_PROJECT_ID = '191396'
 
 # Logfire configuration embedded at build time (only for dev builds)

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())