shotgun-sh 0.2.1.dev1__tar.gz → 0.2.1.dev2__tar.gz

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.2.1.dev1
+Version: 0.2.1.dev2
 Summary: AI-powered research, planning, and task management CLI tool
 Project-URL: Homepage, https://shotgun.sh/
 Project-URL: Repository, https://github.com/shotgun-sh/shotgun
@@ -22,6 +22,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Utilities
 Requires-Python: >=3.11
 Requires-Dist: anthropic>=0.39.0
+Requires-Dist: genai-prices>=0.0.27
 Requires-Dist: httpx>=0.27.0
 Requires-Dist: jinja2>=3.1.0
 Requires-Dist: kuzu>=0.7.0

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shotgun-sh"
-version = "0.2.1.dev1"
+version = "0.2.1.dev2"
 description = "AI-powered research, planning, and task management CLI tool"
 readme = "README.md"
 license = { text = "MIT" }
@@ -45,6 +45,7 @@ dependencies = [
     "tiktoken>=0.7.0",
     "sentencepiece>=0.2.0",
     "packaging>=23.0",
+    "genai-prices>=0.0.27",
 ]
 
 [project.urls]

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/agent_manager.py

@@ -115,12 +115,12 @@ class AgentManager(Widget):
         super().__init__()
         self.display = False
 
+        if deps is None:
+            raise ValueError("AgentDeps must be provided to AgentManager")
+
         # Use provided deps or create default with interactive mode
         self.deps = deps
 
-        if self.deps is None:
-            raise ValueError("AgentDeps must be provided to AgentManager")
-
         # Create AgentRuntimeOptions from deps for agent creation
         agent_runtime_options = AgentRuntimeOptions(
             interactive_mode=self.deps.interactive_mode,
@@ -269,6 +269,7 @@ class AgentManager(Widget):
         Returns:
             The agent run result.
         """
+        logger.info(f"Running agent {self._current_agent_type.value}")
         # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
             deps = self._create_merged_deps(self._current_agent_type)
@@ -395,6 +396,10 @@ class AgentManager(Widget):
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
         self.message_history = await apply_persistent_compaction(all_messages, deps)
+        usage = result.usage()
+        deps.usage_manager.add_usage(
+            usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
+        )
 
         # Log file operations summary if any files were modified
         file_operations = deps.file_tracker.operations.copy()
@@ -641,6 +646,9 @@ class AgentManager(Widget):
                 filtered_messages.append(msg)
         return filtered_messages
 
+    def get_usage_hint(self) -> str | None:
+        return self.deps.usage_manager.build_usage_hint()
+
     def get_conversation_state(self) -> "ConversationState":
         """Get the current conversation state.
 

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/config/manager.py

@@ -1,7 +1,6 @@
 """Configuration manager for Shotgun CLI."""
 
 import json
-import os
 import uuid
 from pathlib import Path
 from typing import Any
@@ -12,10 +11,7 @@ from shotgun.logging_config import get_logger
 from shotgun.utils import get_shotgun_home
 
 from .constants import (
-    ANTHROPIC_API_KEY_ENV,
     API_KEY_FIELD,
-    GEMINI_API_KEY_ENV,
-    OPENAI_API_KEY_ENV,
     ConfigSection,
 )
 from .models import (
@@ -245,25 +241,13 @@ class ConfigManager:
     def has_provider_key(self, provider: ProviderType | str) -> bool:
         """Check if the given provider has a non-empty API key configured.
 
-        This checks both the configuration file and environment variables.
+        This checks only the configuration file.
         """
         config = self.load()
         provider_enum = self._ensure_provider_enum(provider)
         provider_config = self._get_provider_config(config, provider_enum)
 
-        # Check config first
-        if self._provider_has_api_key(provider_config):
-            return True
-
-        # Check environment variable
-        if provider_enum == ProviderType.OPENAI:
-            return bool(os.getenv(OPENAI_API_KEY_ENV))
-        elif provider_enum == ProviderType.ANTHROPIC:
-            return bool(os.getenv(ANTHROPIC_API_KEY_ENV))
-        elif provider_enum == ProviderType.GOOGLE:
-            return bool(os.getenv(GEMINI_API_KEY_ENV))
-
-        return False
+        return self._provider_has_api_key(provider_config)
 
     def has_any_provider_key(self) -> bool:
         """Determine whether any provider has a configured API key."""

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/config/models.py

@@ -27,6 +27,7 @@ class ModelName(StrEnum):
     GPT_5 = "gpt-5"
     GPT_5_MINI = "gpt-5-mini"
     CLAUDE_OPUS_4_1 = "claude-opus-4-1"
+    CLAUDE_SONNET_4_5 = "claude-sonnet-4-5"
     GEMINI_2_5_PRO = "gemini-2.5-pro"
     GEMINI_2_5_FLASH = "gemini-2.5-flash"
 
@@ -102,6 +103,13 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         max_output_tokens=32_000,
         litellm_proxy_model_name="anthropic/claude-opus-4-1",
     ),
+    ModelName.CLAUDE_SONNET_4_5: ModelSpec(
+        name=ModelName.CLAUDE_SONNET_4_5,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=16_000,
+        litellm_proxy_model_name="anthropic/claude-sonnet-4-5",
+    ),
     ModelName.GEMINI_2_5_PRO: ModelSpec(
         name=ModelName.GEMINI_2_5_PRO,
         provider=ProviderType.GOOGLE,

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/config/provider.py

@@ -1,7 +1,5 @@
 """Provider management for LLM configuration."""
 
-import os
-
 from pydantic import SecretStr
 from pydantic_ai.models import Model
 from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
@@ -15,12 +13,6 @@ from pydantic_ai.settings import ModelSettings
 from shotgun.llm_proxy import create_litellm_provider
 from shotgun.logging_config import get_logger
 
-from .constants import (
-    ANTHROPIC_API_KEY_ENV,
-    GEMINI_API_KEY_ENV,
-    OPENAI_API_KEY_ENV,
-    SHOTGUN_API_KEY_ENV,
-)
 from .manager import get_config_manager
 from .models import (
     MODEL_SPECS,
@@ -150,10 +142,10 @@ def get_provider_model(
     config = config_manager.load()
 
     # Priority 1: Check if Shotgun key exists - if so, use it for ANY model
-    shotgun_api_key = _get_api_key(config.shotgun.api_key, SHOTGUN_API_KEY_ENV)
+    shotgun_api_key = _get_api_key(config.shotgun.api_key)
     if shotgun_api_key:
-        # Use selected model or default to claude-opus-4-1
-        model_name = config.selected_model or ModelName.CLAUDE_OPUS_4_1
+        # Use selected model or default to claude-sonnet-4-5
+        model_name = config.selected_model or ModelName.CLAUDE_SONNET_4_5
         if model_name not in MODEL_SPECS:
             raise ValueError(f"Model '{model_name.value}' not found")
         spec = MODEL_SPECS[model_name]
@@ -202,11 +194,9 @@ def get_provider_model(
         requested_model = None  # Will use provider's default model
 
     if provider_enum == ProviderType.OPENAI:
-        api_key = _get_api_key(config.openai.api_key, OPENAI_API_KEY_ENV)
+        api_key = _get_api_key(config.openai.api_key)
         if not api_key:
-            raise ValueError(
-                f"OpenAI API key not configured. Set via environment variable {OPENAI_API_KEY_ENV} or config."
-            )
+            raise ValueError("OpenAI API key not configured. Set via config.")
 
         # Use requested model or default to gpt-5
         model_name = requested_model if requested_model else ModelName.GPT_5
@@ -225,11 +215,9 @@ def get_provider_model(
         )
 
     elif provider_enum == ProviderType.ANTHROPIC:
-        api_key = _get_api_key(config.anthropic.api_key, ANTHROPIC_API_KEY_ENV)
+        api_key = _get_api_key(config.anthropic.api_key)
         if not api_key:
-            raise ValueError(
-                f"Anthropic API key not configured. Set via environment variable {ANTHROPIC_API_KEY_ENV} or config."
-            )
+            raise ValueError("Anthropic API key not configured. Set via config.")
 
         # Use requested model or default to claude-opus-4-1
         model_name = requested_model if requested_model else ModelName.CLAUDE_OPUS_4_1
@@ -248,11 +236,9 @@ def get_provider_model(
         )
 
     elif provider_enum == ProviderType.GOOGLE:
-        api_key = _get_api_key(config.google.api_key, GEMINI_API_KEY_ENV)
+        api_key = _get_api_key(config.google.api_key)
         if not api_key:
-            raise ValueError(
-                f"Gemini API key not configured. Set via environment variable {GEMINI_API_KEY_ENV} or config."
-            )
+            raise ValueError("Gemini API key not configured. Set via config.")
 
         # Use requested model or default to gemini-2.5-pro
         model_name = requested_model if requested_model else ModelName.GEMINI_2_5_PRO
@@ -285,20 +271,19 @@ def _has_provider_key(config: "ShotgunConfig", provider: ProviderType) -> bool:
         True if provider has a configured API key
     """
     if provider == ProviderType.OPENAI:
-        return bool(_get_api_key(config.openai.api_key, OPENAI_API_KEY_ENV))
+        return bool(_get_api_key(config.openai.api_key))
     elif provider == ProviderType.ANTHROPIC:
-        return bool(_get_api_key(config.anthropic.api_key, ANTHROPIC_API_KEY_ENV))
+        return bool(_get_api_key(config.anthropic.api_key))
     elif provider == ProviderType.GOOGLE:
-        return bool(_get_api_key(config.google.api_key, GEMINI_API_KEY_ENV))
+        return bool(_get_api_key(config.google.api_key))
     return False
 
 
-def _get_api_key(config_key: SecretStr | None, env_var: str) -> str | None:
-    """Get API key from config or environment variable.
+def _get_api_key(config_key: SecretStr | None) -> str | None:
+    """Get API key from config.
 
     Args:
         config_key: API key from configuration
-        env_var: Environment variable name to check
 
     Returns:
         API key string or None
@@ -306,4 +291,4 @@ def _get_api_key(config_key: SecretStr | None, env_var: str) -> str | None:
     if config_key is not None:
         return config_key.get_secret_value()
 
-    return os.getenv(env_var)
+    return None

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/models.py

@@ -11,6 +11,8 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai import RunContext
 
+from shotgun.agents.usage_manager import SessionUsageManager, get_session_usage_manager
+
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
@@ -108,6 +110,11 @@ class AgentRuntimeOptions(BaseModel):
         description="Tasks for storing deferred tool results",
     )
 
+    usage_manager: SessionUsageManager = Field(
+        default_factory=get_session_usage_manager,
+        description="Usage manager for tracking usage",
+    )
+
 
 class FileOperationType(StrEnum):
     """Types of file operations that can be tracked."""

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/agents/tools/web_search/anthropic.py

@@ -94,7 +94,6 @@ async def anthropic_web_search_tool(query: str) -> str:
 
 async def main() -> None:
     """Main function for testing the Anthropic web search tool."""
-    import os
     import sys
 
     from shotgun.logging_config import setup_logger
@@ -119,15 +118,14 @@ async def main() -> None:
     print("=" * 60)
 
     # Check if API key is available
-    if not (
-        os.getenv("ANTHROPIC_API_KEY")
-        or (
-            callable(get_provider_model)
-            and get_provider_model(ProviderType.ANTHROPIC).api_key
-        )
-    ):
-        print("❌ Error: ANTHROPIC_API_KEY environment variable not set")
-        print("   Please set it with: export ANTHROPIC_API_KEY=your_key_here")
+    try:
+        if callable(get_provider_model):
+            model_config = get_provider_model(ProviderType.ANTHROPIC)
+            if not model_config.api_key:
+                raise ValueError("No API key configured")
+    except (ValueError, Exception):
+        print("❌ Error: Anthropic API key not configured")
+        print("   Please set it in your config file")
         sys.exit(1)
 
     try:

shotgun_sh-0.2.1.dev2/src/shotgun/agents/usage_manager.py

@@ -0,0 +1,159 @@
+import json
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime
+from logging import getLogger
+from pathlib import Path
+from typing import TypeAlias
+
+from genai_prices import calc_price
+from pydantic import BaseModel, Field
+from pydantic_ai import RunUsage
+
+from shotgun.agents.config.models import ProviderType
+from shotgun.utils import get_shotgun_home
+
+logger = getLogger(__name__)
+ModelName: TypeAlias = str
+
+
+@dataclass(frozen=True)
+class UsageSummaryEntry:
+    model_name: ModelName
+    provider: ProviderType
+    usage: RunUsage
+
+
+class UsageLogEntry(BaseModel):
+    timestamp: datetime = Field(default_factory=datetime.now)
+    model_name: ModelName
+    usage: RunUsage
+    provider: ProviderType
+
+
+class SessionUsage(BaseModel):
+    usage: RunUsage
+    log: list[UsageLogEntry]
+
+
+class UsageState(BaseModel):
+    usage: dict[ModelName, RunUsage] = Field(default_factory=dict)
+    model_providers: dict[ModelName, ProviderType] = Field(default_factory=dict)
+    usage_log: list[UsageLogEntry] = Field(default_factory=list)
+
+
+class SessionUsageManager:
+    def __init__(self) -> None:
+        self.usage: defaultdict[ModelName, RunUsage] = defaultdict(RunUsage)
+        self._model_providers: dict[ModelName, ProviderType] = {}
+        self._usage_log: list[UsageLogEntry] = []
+        self._usage_path: Path = get_shotgun_home() / "usage.json"
+        self.restore_usage_state()
+
+    def add_usage(
+        self, usage: RunUsage, *, model_name: ModelName, provider: ProviderType
+    ) -> None:
+        self.usage[model_name] += usage
+        self._model_providers[model_name] = provider
+        self._usage_log.append(
+            UsageLogEntry(model_name=model_name, usage=usage, provider=provider)
+        )
+        self.persist_usage_state()
+
+    def get_usage_report(self) -> dict[ModelName, RunUsage]:
+        return self.usage.copy()
+
+    def get_usage_breakdown(self) -> list[UsageSummaryEntry]:
+        breakdown: list[UsageSummaryEntry] = []
+        for model_name, usage in self.usage.items():
+            provider = self._model_providers.get(model_name)
+            if provider is None:
+                continue
+            breakdown.append(
+                UsageSummaryEntry(model_name=model_name, provider=provider, usage=usage)
+            )
+        breakdown.sort(key=lambda entry: entry.model_name.lower())
+        return breakdown
+
+    def build_usage_hint(self) -> str | None:
+        return format_usage_hint(self.get_usage_breakdown())
+
+    def persist_usage_state(self) -> None:
+        state = UsageState(
+            usage=dict(self.usage.items()),
+            model_providers=self._model_providers.copy(),
+            usage_log=self._usage_log.copy(),
+        )
+
+        try:
+            self._usage_path.parent.mkdir(parents=True, exist_ok=True)
+            with self._usage_path.open("w", encoding="utf-8") as f:
+                json.dump(state.model_dump(mode="json"), f, indent=2)
+            logger.debug("Usage state persisted to %s", self._usage_path)
+        except Exception as exc:
+            logger.error(
+                "Failed to persist usage state to %s: %s", self._usage_path, exc
+            )
+
+    def restore_usage_state(self) -> None:
+        if not self._usage_path.exists():
+            logger.debug("No usage state file found at %s", self._usage_path)
+            return
+
+        try:
+            with self._usage_path.open(encoding="utf-8") as f:
+                data = json.load(f)
+
+            state = UsageState.model_validate(data)
+        except Exception as exc:
+            logger.error(
+                "Failed to restore usage state from %s: %s", self._usage_path, exc
+            )
+            return
+
+        self.usage = defaultdict(RunUsage)
+        for model_name, usage in state.usage.items():
+            self.usage[model_name] = usage
+
+        self._model_providers = state.model_providers.copy()
+        self._usage_log = state.usage_log.copy()
+
+
+def format_usage_hint(breakdown: list[UsageSummaryEntry]) -> str | None:
+    if not breakdown:
+        return None
+
+    lines = ["# Token usage by model"]
+
+    for entry in breakdown:
+        usage = entry.usage
+        input_tokens = usage.input_tokens
+        output_tokens = usage.output_tokens
+        cached_tokens = usage.cache_read_tokens
+
+        cost = calc_price(usage=usage, model_ref=entry.model_name)
+        input_line = f"* Input: {input_tokens:,}"
+        if cached_tokens > 0:
+            input_line += f" (+ {cached_tokens:,} cached)"
+        input_line += " tokens"
+        section = f"""
+### {entry.model_name}
+
+{input_line}
+* Output: {output_tokens:,} tokens
+* Total: {input_tokens + output_tokens:,} tokens
+* Cost: ${cost.total_price:,.2f}
+""".strip()
+        lines.append(section)
+
+    return "\n\n".join(lines)
+
+
+_usage_manager = None
+
+
+def get_session_usage_manager() -> SessionUsageManager:
+    global _usage_manager
+    if _usage_manager is None:
+        _usage_manager = SessionUsageManager()
+    return _usage_manager

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/cli/feedback.py

@@ -43,4 +43,4 @@ def send_feedback(
 
     submit_feedback_survey(feedback)
 
-    console.print("✅ Feedback sent. Thank you!")
+    console.print("Feedback sent. Thank you!")

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/tui/app.py

@@ -108,7 +108,7 @@ class ShotgunApp(App[None]):
         def handle_feedback(feedback: Feedback | None) -> None:
             if feedback is not None:
                 submit_feedback_survey(feedback)
-                self.notify("✅ Feedback sent. Thank you!")
+                self.notify("Feedback sent. Thank you!")
 
         self.push_screen("feedback", callback=handle_feedback)
 

{shotgun_sh-0.2.1.dev1 → shotgun_sh-0.2.1.dev2}/src/shotgun/tui/screens/chat.py

@@ -58,6 +58,7 @@ from .chat_screen.command_providers import (
     CodebaseCommandProvider,
     DeleteCodebasePaletteProvider,
     ProviderSetupProvider,
+    UsageProvider,
 )
 
 logger = logging.getLogger(__name__)
@@ -228,9 +229,15 @@ class ChatScreen(Screen[None]):
     BINDINGS = [
         ("ctrl+p", "command_palette", "Command Palette"),
         ("shift+tab", "toggle_mode", "Toggle mode"),
+        ("ctrl+u", "show_usage", "Show usage"),
     ]
 
-    COMMANDS = {AgentModeProvider, ProviderSetupProvider, CodebaseCommandProvider}
+    COMMANDS = {
+        AgentModeProvider,
+        ProviderSetupProvider,
+        CodebaseCommandProvider,
+        UsageProvider,
+    }
 
     value = reactive("")
     mode = reactive(AgentType.RESEARCH)
@@ -401,6 +408,14 @@ class ChatScreen(Screen[None]):
         # whoops it actually changes focus. Let's be brutal for now
         self.call_later(lambda: self.query_one(PromptInput).focus())
 
+    def action_show_usage(self) -> None:
+        usage_hint = self.agent_manager.get_usage_hint()
+        logger.info(f"Usage hint: {usage_hint}")
+        if usage_hint:
+            self.mount_hint(usage_hint)
+        else:
+            self.notify("No usage hint available", severity="error")
+
     @work
     async def add_question_listener(self) -> None:
         while True:
@@ -589,7 +604,7 @@ class ChatScreen(Screen[None]):
     async def index_codebase(self, selection: CodebaseIndexSelection) -> None:
         label = self.query_one("#indexing-job-display", Static)
         label.update(
-            f"[$foreground-muted]Indexing [bold $text-accent]{selection.name}[/]...[/]"
+            f"[$foreground-muted]Indexing codebase: [bold $text-accent]{selection.name}[/][/]"
         )
         label.refresh()
 
@@ -599,33 +614,33 @@ class ChatScreen(Screen[None]):
             empty = width - filled
             return "▓" * filled + "░" * empty
 
-        # Spinner animation state (shared between timer and progress callback)
+        # Spinner animation frames
         spinner_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
-        spinner_state: dict[str, int | str | float] = {
+
+        # Progress state (shared between timer and progress callback)
+        progress_state: dict[str, int | float] = {
             "frame_index": 0,
-            "phase_name": "Starting...",
             "percentage": 0.0,
         }
 
-        def update_spinner_display() -> None:
-            """Update spinner frame on timer - runs every 100ms."""
+        def update_progress_display() -> None:
+            """Update progress bar on timer - runs every 100ms."""
             # Advance spinner frame
-            frame_idx = int(spinner_state["frame_index"])
-            spinner_state["frame_index"] = (frame_idx + 1) % len(spinner_frames)
+            frame_idx = int(progress_state["frame_index"])
+            progress_state["frame_index"] = (frame_idx + 1) % len(spinner_frames)
             spinner = spinner_frames[frame_idx]
 
             # Get current state
-            phase = str(spinner_state["phase_name"])
-            pct = float(spinner_state["percentage"])
+            pct = float(progress_state["percentage"])
             bar = create_progress_bar(pct)
 
             # Update label
             label.update(
-                f"[$foreground-muted]Indexing codebase: {spinner} {phase}... {bar} {pct:.0f}%[/]"
+                f"[$foreground-muted]Indexing codebase: {spinner} {bar} {pct:.0f}%[/]"
             )
 
         def progress_callback(progress_info: IndexProgress) -> None:
-            """Update progress state (spinner animates independently on timer)."""
+            """Update progress state (timer renders it independently)."""
             # Calculate overall percentage (0-95%, reserve 95-100% for finalization)
             if progress_info.phase == ProgressPhase.STRUCTURE:
                 # Phase 1: 0-10%, always show 5% while running, 10% when complete
@@ -648,11 +663,10 @@ class ChatScreen(Screen[None]):
                 overall_pct = 0.0
 
             # Update shared state (timer will render it)
-            spinner_state["phase_name"] = progress_info.phase_name
-            spinner_state["percentage"] = overall_pct
+            progress_state["percentage"] = overall_pct
 
-        # Start spinner animation timer (10 fps = 100ms interval)
-        spinner_timer = self.set_interval(0.1, update_spinner_display)
+        # Start progress animation timer (10 fps = 100ms interval)
+        progress_timer = self.set_interval(0.1, update_progress_display)
 
         try:
             # Pass the current working directory as the indexed_from_cwd
@@ -667,14 +681,12 @@ class ChatScreen(Screen[None]):
                 progress_callback=progress_callback,
             )
 
-            # Stop spinner animation
-            spinner_timer.stop()
+            # Stop progress animation
+            progress_timer.stop()
 
             # Show 100% completion after indexing finishes
             final_bar = create_progress_bar(100.0)
-            label.update(
-                f"[$foreground-muted]Indexing codebase: ✓ Complete {final_bar} 100%[/]"
-            )
+            label.update(f"[$foreground-muted]Indexing codebase: {final_bar} 100%[/]")
             label.refresh()
 
             logger.info(
@@ -687,12 +699,12 @@ class ChatScreen(Screen[None]):
             )
 
         except CodebaseAlreadyIndexedError as exc:
-            spinner_timer.stop()
+            progress_timer.stop()
             logger.warning(f"Codebase already indexed: {exc}")
             self.notify(str(exc), severity="warning")
             return
         except InvalidPathError as exc:
-            spinner_timer.stop()
+            progress_timer.stop()
             logger.error(f"Invalid path error: {exc}")
             self.notify(str(exc), severity="error")
 
@@ -704,8 +716,8 @@ class ChatScreen(Screen[None]):
             )
             self.notify(f"Failed to index codebase: {exc}", severity="error")
         finally:
-            # Always stop the spinner timer
-            spinner_timer.stop()
+            # Always stop the progress timer
+            progress_timer.stop()
             label.update("")
             label.refresh()
 
@@ -774,6 +786,7 @@ class ChatScreen(Screen[None]):
 
         # Update the current mode
         self.mode = AgentType(conversation.last_agent_model)
+        self.deps.usage_manager.restore_usage_state()
 
 
 def help_text_with_codebase(already_indexed: bool = False) -> str: