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

shotgun/llm_proxy/client.py

@@ -0,0 +1,215 @@
+"""HTTP client for LiteLLM Proxy API."""
+
+import logging
+from typing import Any
+
+import httpx
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+
+from shotgun.api_endpoints import LITELLM_PROXY_BASE_URL
+from shotgun.logging_config import get_logger
+
+from .models import BudgetInfo, KeyInfoResponse, TeamInfoResponse
+
+logger = get_logger(__name__)
+
+
+def _is_retryable_http_error(exception: BaseException) -> bool:
+    """Check if HTTP exception should trigger a retry.
+
+    Args:
+        exception: The exception to check
+
+    Returns:
+        True if the exception is a transient error that should be retried
+    """
+    # Retry on network errors and timeouts
+    if isinstance(exception, (httpx.RequestError, httpx.TimeoutException)):
+        return True
+
+    # Retry on server errors (5xx) and rate limits (429)
+    if isinstance(exception, httpx.HTTPStatusError):
+        status_code = exception.response.status_code
+        return status_code >= 500 or status_code == 429
+
+    # Don't retry on other errors (e.g., 4xx client errors)
+    return False
+
+
+class LiteLLMProxyClient:
+    """HTTP client for LiteLLM Proxy API.
+
+    Provides methods to query budget information and key/team metadata
+    from a LiteLLM proxy server.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str | None = None,
+        timeout: float = 10.0,
+    ):
+        """Initialize LiteLLM Proxy client.
+
+        Args:
+            api_key: LiteLLM API key for authentication
+            base_url: Base URL for LiteLLM proxy. If None, uses LITELLM_PROXY_BASE_URL
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.base_url = base_url or LITELLM_PROXY_BASE_URL
+        self.timeout = timeout
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential_jitter(initial=1, max=8),
+        retry=retry_if_exception(_is_retryable_http_error),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )
+    async def _request_with_retry(
+        self,
+        method: str,
+        url: str,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        """Make async HTTP request with exponential backoff retry and jitter.
+
+        Uses tenacity to retry on transient errors (5xx, 429, network errors)
+        with exponential backoff and jitter. Client errors (4xx except 429)
+        are not retried.
+
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            url: Request URL
+            **kwargs: Additional arguments to pass to httpx request
+
+        Returns:
+            HTTP response
+
+        Raises:
+            httpx.HTTPError: If request fails after all retries
+        """
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.request(method, url, **kwargs)
+            response.raise_for_status()
+            return response
+
+    async def get_key_info(self) -> KeyInfoResponse:
+        """Get key information from LiteLLM proxy.
+
+        Returns:
+            Key information including spend, budget, and team_id
+
+        Raises:
+            httpx.HTTPError: If request fails
+        """
+        url = f"{self.base_url}/key/info"
+        params = {"key": self.api_key}
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        logger.debug("Fetching key info from %s", url)
+
+        response = await self._request_with_retry(
+            "GET", url, params=params, headers=headers
+        )
+
+        data = response.json()
+        result = KeyInfoResponse.model_validate(data)
+
+        logger.info(
+            "Successfully fetched key info: key_alias=%s, team_id=%s",
+            result.info.key_alias,
+            result.info.team_id,
+        )
+        return result
+
+    async def get_team_info(self, team_id: str) -> TeamInfoResponse:
+        """Get team information from LiteLLM proxy.
+
+        Args:
+            team_id: Team identifier
+
+        Returns:
+            Team information including spend and budget
+
+        Raises:
+            httpx.HTTPError: If request fails
+        """
+        url = f"{self.base_url}/team/info"
+        params = {"team_id": team_id}
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        logger.debug("Fetching team info from %s for team_id=%s", url, team_id)
+
+        response = await self._request_with_retry(
+            "GET", url, params=params, headers=headers
+        )
+
+        data = response.json()
+        result = TeamInfoResponse.model_validate(data)
+
+        logger.info(
+            "Successfully fetched team info: team_alias=%s",
+            result.team_info.team_alias,
+        )
+        return result
+
+    async def get_budget_info(self) -> BudgetInfo:
+        """Get team-level budget information for this key.
+
+        Budget is always configured at the team level, never at the key level.
+        This method fetches the team_id from the key info, then retrieves
+        the team's budget information.
+
+        Returns:
+            Team-level budget information
+
+        Raises:
+            httpx.HTTPError: If request fails
+            ValueError: If team has no budget configured
+        """
+        logger.debug("Fetching budget info")
+
+        # Get key info to retrieve team_id
+        key_response = await self.get_key_info()
+        key_info = key_response.info
+
+        # Fetch team budget (budget is always at team level)
+        logger.debug(
+            "Fetching team budget for team_id=%s",
+            key_info.team_id,
+        )
+        team_response = await self.get_team_info(key_info.team_id)
+        team_info = team_response.team_info
+
+        if team_info.max_budget is None:
+            raise ValueError(
+                f"Team (team_id={key_info.team_id}) has no max_budget configured"
+            )
+
+        logger.debug("Using team-level budget: $%.6f", team_info.max_budget)
+        return BudgetInfo.from_team_info(team_info)
+
+
+# Convenience function for standalone use
+async def get_budget_info(api_key: str, base_url: str | None = None) -> BudgetInfo:
+    """Get budget information for an API key.
+
+    Convenience function that creates a client and calls get_budget_info.
+
+    Args:
+        api_key: LiteLLM API key
+        base_url: Optional base URL for LiteLLM proxy
+
+    Returns:
+        Budget information
+    """
+    client = LiteLLMProxyClient(api_key, base_url=base_url)
+    return await client.get_budget_info()

shotgun/llm_proxy/models.py

@@ -0,0 +1,137 @@
+"""Pydantic models for LiteLLM Proxy API."""
+
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+
+class BudgetSource(StrEnum):
+    """Source of budget information."""
+
+    KEY = "key"
+    TEAM = "team"
+
+
+class KeyInfoData(BaseModel):
+    """Key information data from /key/info endpoint."""
+
+    key_name: str = Field(description="Key name/identifier")
+    key_alias: str | None = Field(default=None, description="Human-readable key alias")
+    spend: float = Field(description="Current spend for this key in USD")
+    max_budget: float | None = Field(
+        default=None, description="Maximum budget for this key in USD"
+    )
+    team_id: str = Field(description="Team ID associated with this key")
+    user_id: str = Field(description="User ID associated with this key")
+    models: list[str] = Field(
+        default_factory=list, description="List of models available to this key"
+    )
+
+
+class KeyInfoResponse(BaseModel):
+    """Response from /key/info endpoint."""
+
+    key: str = Field(description="The API key")
+    info: KeyInfoData = Field(description="Key information data")
+
+
+class TeamInfoData(BaseModel):
+    """Team information data from /team/info endpoint."""
+
+    team_id: str = Field(description="Team identifier")
+    team_alias: str | None = Field(
+        default=None, description="Human-readable team alias"
+    )
+    max_budget: float | None = Field(
+        default=None, description="Maximum budget for this team in USD"
+    )
+    spend: float = Field(description="Current spend for this team in USD")
+    models: list[str] = Field(
+        default_factory=list, description="List of models available to this team"
+    )
+
+
+class TeamInfoResponse(BaseModel):
+    """Response from /team/info endpoint."""
+
+    team_id: str = Field(description="Team identifier")
+    team_info: TeamInfoData = Field(description="Team information data")
+
+
+class BudgetInfo(BaseModel):
+    """Unified budget information.
+
+    Combines key and team budget information to provide a single view
+    of budget status. Budget can come from either key-level or team-level,
+    with key-level taking priority if set.
+    """
+
+    max_budget: float = Field(description="Maximum budget in USD")
+    spend: float = Field(description="Current spend in USD")
+    remaining: float = Field(description="Remaining budget in USD")
+    source: BudgetSource = Field(
+        description="Source of budget information (key or team)"
+    )
+    percentage_used: float = Field(description="Percentage of budget used (0-100)")
+
+    @classmethod
+    def from_key_info(cls, key_info: KeyInfoData) -> "BudgetInfo":
+        """Create BudgetInfo from key-level budget.
+
+        Args:
+            key_info: Key information containing budget data
+
+        Returns:
+            BudgetInfo instance with key-level budget
+
+        Raises:
+            ValueError: If key does not have max_budget set
+        """
+        if key_info.max_budget is None:
+            raise ValueError("Key does not have max_budget set")
+
+        remaining = key_info.max_budget - key_info.spend
+        percentage_used = (
+            (key_info.spend / key_info.max_budget * 100)
+            if key_info.max_budget > 0
+            else 0.0
+        )
+
+        return cls(
+            max_budget=key_info.max_budget,
+            spend=key_info.spend,
+            remaining=remaining,
+            source=BudgetSource.KEY,
+            percentage_used=percentage_used,
+        )
+
+    @classmethod
+    def from_team_info(cls, team_info: TeamInfoData) -> "BudgetInfo":
+        """Create BudgetInfo from team-level budget.
+
+        Args:
+            team_info: Team information containing budget data
+
+        Returns:
+            BudgetInfo instance with team-level budget
+
+        Raises:
+            ValueError: If team does not have max_budget set
+        """
+        if team_info.max_budget is None:
+            raise ValueError("Team does not have max_budget set")
+
+        remaining = team_info.max_budget - team_info.spend
+        percentage_used = (
+            (team_info.spend / team_info.max_budget * 100)
+            if team_info.max_budget > 0
+            else 0.0
+        )
+
+        return cls(
+            max_budget=team_info.max_budget,
+            spend=team_info.spend,
+            remaining=remaining,
+            source=BudgetSource.TEAM,
+            percentage_used=percentage_used,
+        )

shotgun/sdk/codebase.py

@@ -93,6 +93,19 @@ class CodebaseSDK:
         if indexed_from_cwd is None:
             indexed_from_cwd = str(Path.cwd().resolve())
 
+        # Track codebase indexing started event
+        source = detect_source()
+        logger.debug(
+            "Tracking codebase_index_started event: source=%s",
+            source,
+        )
+        track_event(
+            "codebase_index_started",
+            {
+                "source": source,
+            },
+        )
+
         graph = await self.service.create_graph(
             resolved_path,
             name,
@@ -101,9 +114,7 @@ class CodebaseSDK:
         )
         file_count = sum(graph.language_stats.values()) if graph.language_stats else 0
 
-        # Track codebase indexing event
-        # Detect if called from TUI by checking the call stack
-        source = detect_source()
+        # Track codebase indexing completion event (reuse source from start event)
 
         logger.debug(
             "Tracking codebase_indexed event: file_count=%d, node_count=%d, relationship_count=%d, source=%s",

shotgun/tui/app.py

@@ -14,7 +14,10 @@ from shotgun.agents.models import AgentType
 from shotgun.logging_config import get_logger
 from shotgun.tui.containers import TUIContainer
 from shotgun.tui.screens.splash import SplashScreen
-from shotgun.utils.file_system_utils import get_shotgun_base_path
+from shotgun.utils.file_system_utils import (
+    ensure_shotgun_directory_exists,
+    get_shotgun_base_path,
+)
 from shotgun.utils.update_checker import (
     detect_installation_method,
     perform_auto_update_async,
@@ -34,10 +37,10 @@ logger = get_logger(__name__)
 class ShotgunApp(App[None]):
     # ChatScreen removed from SCREENS dict since it requires dependency injection
     # and is instantiated manually in refresh_startup_screen()
+    # DirectorySetupScreen also removed since it requires error_message parameter
     SCREENS = {
         "provider_config": ProviderConfigScreen,
         "model_picker": ModelPickerScreen,
-        "directory_setup": DirectorySetupScreen,
         "github_issue": GitHubIssueScreen,
     }
     BINDINGS = [
@@ -117,16 +120,32 @@ class ShotgunApp(App[None]):
                 )
                 return
 
+            # Try to create .shotgun directory if it doesn't exist
             if not self.check_local_shotgun_directory_exists():
-                if isinstance(self.screen, DirectorySetupScreen):
+                try:
+                    path = ensure_shotgun_directory_exists()
+                    # Verify directory was created successfully
+                    if not path.is_dir():
+                        # Show error screen if creation failed
+                        if isinstance(self.screen, DirectorySetupScreen):
+                            return
+                        self.push_screen(
+                            DirectorySetupScreen(
+                                error_message="Unable to create .shotgun directory due to filesystem conflict."
+                            ),
+                            callback=lambda _arg: self.refresh_startup_screen(),
+                        )
+                        return
+                except Exception as exc:
+                    # Show error screen if creation failed with exception
+                    if isinstance(self.screen, DirectorySetupScreen):
+                        return
+                    self.push_screen(
+                        DirectorySetupScreen(error_message=str(exc)),
+                        callback=lambda _arg: self.refresh_startup_screen(),
+                    )
                     return
 
-                self.push_screen(
-                    DirectorySetupScreen(),
-                    callback=lambda _arg: self.refresh_startup_screen(),
-                )
-                return
-
             if isinstance(self.screen, ChatScreen):
                 return
 

shotgun/tui/screens/chat/chat_screen.py

@@ -44,12 +44,17 @@ from shotgun.agents.models import (
     AgentType,
     FileOperationTracker,
 )
+from shotgun.agents.runner import AgentRunner
 from shotgun.codebase.core.manager import (
     CodebaseAlreadyIndexedError,
     CodebaseGraphManager,
 )
 from shotgun.codebase.models import IndexProgress, ProgressPhase
-from shotgun.exceptions import ContextSizeLimitExceeded
+from shotgun.exceptions import (
+    SHOTGUN_CONTACT_EMAIL,
+    ErrorNotPickedUpBySentry,
+    ShotgunAccountException,
+)
 from shotgun.posthog_telemetry import track_event
 from shotgun.sdk.codebase import CodebaseSDK
 from shotgun.sdk.exceptions import CodebaseNotFoundError, InvalidPathError
@@ -59,6 +64,8 @@ from shotgun.tui.components.mode_indicator import ModeIndicator
 from shotgun.tui.components.prompt_input import PromptInput
 from shotgun.tui.components.spinner import Spinner
 from shotgun.tui.components.status_bar import StatusBar
+
+# TUIErrorHandler removed - exceptions now caught directly
 from shotgun.tui.screens.chat.codebase_index_prompt_screen import (
     CodebaseIndexPromptScreen,
 )
@@ -301,9 +308,75 @@ class ChatScreen(Screen[None]):
         # Re-focus input after mode change
         self.call_later(lambda: self.widget_coordinator.update_prompt_input(focus=True))
 
-    def action_show_usage(self) -> None:
+    async def action_show_usage(self) -> None:
         usage_hint = self.agent_manager.get_usage_hint()
         logger.info(f"Usage hint: {usage_hint}")
+
+        # Add budget info for Shotgun Account users
+        if self.deps.llm_model.is_shotgun_account:
+            try:
+                from shotgun.llm_proxy import LiteLLMProxyClient
+
+                logger.debug("Fetching budget info for Shotgun Account")
+                client = LiteLLMProxyClient(self.deps.llm_model.api_key)
+                budget_info = await client.get_budget_info()
+
+                # Format budget section
+                source_label = "Key" if budget_info.source == "key" else "Team"
+                budget_section = 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
+
+**Questions or need help?**"""
+
+                # Build markdown_before (usage + budget info before email)
+                if usage_hint:
+                    markdown_before = f"{usage_hint}\n\n{budget_section}"
+                else:
+                    markdown_before = budget_section
+
+                markdown_after = (
+                    "\n\n_Reach out anytime for billing questions "
+                    "or to increase your budget._"
+                )
+
+                # Mount with email copy button
+                self.mount_hint_with_email(
+                    markdown_before=markdown_before,
+                    email="contact@shotgun.sh",
+                    markdown_after=markdown_after,
+                )
+                logger.debug("Successfully added budget info to usage hint")
+                return  # Exit early since we've already mounted
+
+            except Exception as e:
+                logger.warning(f"Failed to fetch budget info: {e}")
+                # For Shotgun Account, show budget fetch error
+                # If we have usage data, still show it
+                if usage_hint:
+                    # Show usage even though budget fetch failed
+                    self.mount_hint(usage_hint)
+                else:
+                    # No usage and budget fetch failed - show specific error with email
+                    markdown_before = (
+                        "⚠️ **Unable to fetch budget information**\n\n"
+                        "There was an error retrieving your budget data."
+                    )
+                    markdown_after = (
+                        "\n\n_Try the command again in a moment. "
+                        "If the issue persists, reach out for help._"
+                    )
+                    self.mount_hint_with_email(
+                        markdown_before=markdown_before,
+                        email="contact@shotgun.sh",
+                        markdown_after=markdown_after,
+                    )
+                return  # Exit early
+
+        # Fallback for non-Shotgun Account users
         if usage_hint:
             self.mount_hint(usage_hint)
         else:
@@ -582,6 +655,21 @@ class ChatScreen(Screen[None]):
         hint = HintMessage(message=markdown)
         self.agent_manager.add_hint_message(hint)
 
+    def mount_hint_with_email(
+        self, markdown_before: str, email: str, markdown_after: str = ""
+    ) -> None:
+        """Mount a hint with inline email copy button.
+
+        Args:
+            markdown_before: Markdown content to display before the email line
+            email: Email address to display with copy button
+            markdown_after: Optional markdown content to display after the email line
+        """
+        hint = HintMessage(
+            message=markdown_before, email=email, markdown_after=markdown_after
+        )
+        self.agent_manager.add_hint_message(hint)
+
     @on(PartialResponseMessage)
     def handle_partial_response(self, event: PartialResponseMessage) -> None:
         # Filter event.messages to exclude ModelRequest with only ToolReturnPart
@@ -768,6 +856,19 @@ class ChatScreen(Screen[None]):
             # Update the agent manager's model configuration
             self.agent_manager.deps.llm_model = result.model_config
 
+            # Reset agents so they get recreated with new model
+            self.agent_manager._agents_initialized = False
+            self.agent_manager._research_agent = None
+            self.agent_manager._plan_agent = None
+            self.agent_manager._tasks_agent = None
+            self.agent_manager._specify_agent = None
+            self.agent_manager._export_agent = None
+            self.agent_manager._research_deps = None
+            self.agent_manager._plan_deps = None
+            self.agent_manager._tasks_deps = None
+            self.agent_manager._specify_deps = None
+            self.agent_manager._export_deps = None
+
             # Get current analysis and update context indicator via coordinator
             analysis = await self.agent_manager.get_context_analysis()
             self.widget_coordinator.update_context_indicator(analysis, result.new_model)
@@ -1139,8 +1240,6 @@ class ChatScreen(Screen[None]):
 
     @work
     async def run_agent(self, message: str) -> None:
-        prompt = None
-
         # Start processing with spinner
         from textual.worker import get_current_worker
 
@@ -1150,60 +1249,31 @@ class ChatScreen(Screen[None]):
         # Start context indicator animation immediately
         self.widget_coordinator.set_context_streaming(True)
 
-        prompt = message
-
         try:
-            await self.agent_manager.run(
-                prompt=prompt,
-            )
-        except asyncio.CancelledError:
-            # Handle cancellation gracefully - DO NOT re-raise
-            self.mount_hint("⚠️ Operation cancelled by user")
-        except ContextSizeLimitExceeded as e:
-            # User-friendly error with actionable options
-            hint = (
-                f"⚠️ **Context too large for {e.model_name}**\n\n"
-                f"Your conversation history exceeds this model's limit ({e.max_tokens:,} tokens).\n\n"
-                f"**Choose an action:**\n\n"
-                f"1. Switch to a larger model (`Ctrl+P` → Change Model)\n"
-                f"2. Switch to a larger model, compact (`/compact`), then switch back to {e.model_name}\n"
-                f"3. Clear conversation (`/clear`)\n"
-            )
-
-            self.mount_hint(hint)
-
-            # Log for debugging (won't send to Sentry due to ErrorNotPickedUpBySentry)
-            logger.info(
-                "Context size limit exceeded",
-                extra={
-                    "max_tokens": e.max_tokens,
-                    "model_name": e.model_name,
-                },
-            )
-        except Exception as e:
-            # Log with full stack trace to shotgun.log
-            logger.exception(
-                "Agent run failed",
-                extra={
-                    "agent_mode": self.mode.value,
-                    "error_type": type(e).__name__,
-                },
-            )
-
-            # Determine user-friendly message based on error type
-            error_name = type(e).__name__
-            error_message = str(e)
-
-            if "APIStatusError" in error_name and "overload" in error_message.lower():
-                hint = "⚠️ The AI service is temporarily overloaded. Please wait a moment and try again."
-            elif "APIStatusError" in error_name and "rate" in error_message.lower():
-                hint = "⚠️ Rate limit reached. Please wait before trying again."
-            elif "APIStatusError" in error_name:
-                hint = f"⚠️ AI service error: {error_message}"
+            # Use unified agent runner - exceptions propagate for handling
+            runner = AgentRunner(self.agent_manager)
+            await runner.run(message)
+        except ShotgunAccountException as e:
+            # Shotgun Account errors show contact email UI
+            message_parts = e.to_markdown().split("**Need help?**")
+            if len(message_parts) == 2:
+                markdown_before = message_parts[0] + "**Need help?**"
+                markdown_after = message_parts[1].strip()
+                self.mount_hint_with_email(
+                    markdown_before=markdown_before,
+                    email=SHOTGUN_CONTACT_EMAIL,
+                    markdown_after=markdown_after,
+                )
             else:
-                hint = f"⚠️ An error occurred: {error_message}\n\nCheck logs at ~/.shotgun-sh/logs/shotgun.log"
-
-            self.mount_hint(hint)
+                # Fallback if message format is unexpected
+                self.mount_hint(e.to_markdown())
+        except ErrorNotPickedUpBySentry as e:
+            # All other user-actionable errors - display with markdown
+            self.mount_hint(e.to_markdown())
+        except Exception as e:
+            # Unexpected errors that weren't wrapped (shouldn't happen)
+            logger.exception("Unexpected error in run_agent")
+            self.mount_hint(f"⚠️ An unexpected error occurred: {str(e)}")
         finally:
             self.processing_state.stop_processing()
             # Stop context indicator animation