shotgun-sh 0.2.3.dev2__py3-none-any.whl → 0.2.11.dev5__py3-none-any.whl

shotgun/tui/screens/provider_config.py

@@ -97,7 +97,7 @@ class ProviderConfigScreen(Screen[None]):
                 "Don't have an API Key? Use these links to get one: [OpenAI](https://platform.openai.com/api-keys) | [Anthropic](https://console.anthropic.com) | [Google Gemini](https://aistudio.google.com)",
                 id="provider-links",
             )
-        yield ListView(*self._build_provider_items(), id="provider-list")
+        yield ListView(*self._build_provider_items_sync(), id="provider-list")
         yield Input(
             placeholder=self._input_placeholder(self.selected_provider),
             password=True,
@@ -110,8 +110,6 @@ class ProviderConfigScreen(Screen[None]):
             yield Button("Done \\[ESC]", id="done")
 
     def on_mount(self) -> None:
-        self.refresh_provider_status()
-        self._update_done_button_visibility()
         list_view = self.query_one(ListView)
         if list_view.children:
             list_view.index = 0
@@ -121,13 +119,20 @@ class ProviderConfigScreen(Screen[None]):
         self.query_one("#authenticate", Button).display = False
         self.set_focus(self.query_one("#api-key", Input))
 
+        # Refresh UI asynchronously
+        self.run_worker(self._refresh_ui(), exclusive=False)
+
     def on_screenresume(self) -> None:
         """Refresh provider status when screen is resumed.
 
         This ensures the UI reflects any provider changes made elsewhere.
         """
-        self.refresh_provider_status()
-        self._update_done_button_visibility()
+        self.run_worker(self._refresh_ui(), exclusive=False)
+
+    async def _refresh_ui(self) -> None:
+        """Refresh provider status and button visibility."""
+        await self.refresh_provider_status()
+        await self._update_done_button_visibility()
 
     def action_done(self) -> None:
         self.dismiss()
@@ -170,7 +175,11 @@ class ProviderConfigScreen(Screen[None]):
         if not self.is_mounted:
             return
 
-        # Show/hide UI elements based on provider type
+        # Show/hide UI elements based on provider type asynchronously
+        self.run_worker(self._update_provider_ui(provider), exclusive=False)
+
+    async def _update_provider_ui(self, provider: ProviderType) -> None:
+        """Update UI elements based on selected provider."""
         is_shotgun = provider == "shotgun"
 
         input_widget = self.query_one("#api-key", Input)
@@ -183,7 +192,7 @@ class ProviderConfigScreen(Screen[None]):
             save_button.display = False
 
             # Only show Authenticate button if shotgun is NOT already configured
-            if self._has_provider_key("shotgun"):
+            if await self._has_provider_key("shotgun"):
                 auth_button.display = False
             else:
                 auth_button.display = True
@@ -200,22 +209,29 @@ class ProviderConfigScreen(Screen[None]):
         app = cast("ShotgunApp", self.app)
         return app.config_manager
 
-    def refresh_provider_status(self) -> None:
+    async def refresh_provider_status(self) -> None:
         """Update the list view entries to reflect configured providers."""
         for provider_id in get_configurable_providers():
             label = self.query_one(f"#label-{provider_id}", Label)
-            label.update(self._provider_label(provider_id))
+            label.update(await self._provider_label(provider_id))
 
-    def _update_done_button_visibility(self) -> None:
+    async def _update_done_button_visibility(self) -> None:
         """Show/hide Done button based on whether any provider keys are configured."""
         done_button = self.query_one("#done", Button)
-        has_keys = self.config_manager.has_any_provider_key()
+        has_keys = await self.config_manager.has_any_provider_key()
         done_button.display = has_keys
 
-    def _build_provider_items(self) -> list[ListItem]:
+    def _build_provider_items_sync(self) -> list[ListItem]:
+        """Build provider items synchronously for compose().
+
+        Labels will be populated with status asynchronously in on_mount().
+        """
         items: list[ListItem] = []
         for provider_id in get_configurable_providers():
-            label = Label(self._provider_label(provider_id), id=f"label-{provider_id}")
+            # Create labels with placeholder text - will be updated in on_mount()
+            label = Label(
+                self._provider_display_name(provider_id), id=f"label-{provider_id}"
+            )
             items.append(ListItem(label, id=f"provider-{provider_id}"))
         return items
 
@@ -225,11 +241,10 @@ class ProviderConfigScreen(Screen[None]):
         provider_id = item.id.removeprefix("provider-")
         return provider_id if provider_id in get_configurable_providers() else None
 
-    def _provider_label(self, provider_id: str) -> str:
+    async def _provider_label(self, provider_id: str) -> str:
         display = self._provider_display_name(provider_id)
-        status = (
-            "Configured" if self._has_provider_key(provider_id) else "Not configured"
-        )
+        has_key = await self._has_provider_key(provider_id)
+        status = "Configured" if has_key else "Not configured"
         return f"{display} · {status}"
 
     def _provider_display_name(self, provider_id: str) -> str:
@@ -244,21 +259,25 @@ class ProviderConfigScreen(Screen[None]):
     def _input_placeholder(self, provider_id: str) -> str:
         return f"{self._provider_display_name(provider_id)} API key"
 
-    def _has_provider_key(self, provider_id: str) -> bool:
+    async def _has_provider_key(self, provider_id: str) -> bool:
         """Check if provider has a configured API key."""
         if provider_id == "shotgun":
             # Check shotgun key directly
-            config = self.config_manager.load()
+            config = await self.config_manager.load()
             return self.config_manager._provider_has_api_key(config.shotgun)
         else:
             # Check LLM provider key
             try:
                 provider = ProviderType(provider_id)
-                return self.config_manager.has_provider_key(provider)
+                return await self.config_manager.has_provider_key(provider)
             except ValueError:
                 return False
 
     def _save_api_key(self) -> None:
+        self.run_worker(self._do_save_api_key(), exclusive=True)
+
+    async def _do_save_api_key(self) -> None:
+        """Async implementation of API key saving."""
         input_widget = self.query_one("#api-key", Input)
         api_key = input_widget.value.strip()
 
@@ -267,7 +286,7 @@ class ProviderConfigScreen(Screen[None]):
             return
 
         try:
-            self.config_manager.update_provider(
+            await self.config_manager.update_provider(
                 self.selected_provider,
                 api_key=api_key,
             )
@@ -276,21 +295,25 @@ class ProviderConfigScreen(Screen[None]):
             return
 
         input_widget.value = ""
-        self.refresh_provider_status()
-        self._update_done_button_visibility()
+        await self.refresh_provider_status()
+        await self._update_done_button_visibility()
         self.notify(
             f"Saved API key for {self._provider_display_name(self.selected_provider)}."
         )
 
     def _clear_api_key(self) -> None:
+        self.run_worker(self._do_clear_api_key(), exclusive=True)
+
+    async def _do_clear_api_key(self) -> None:
+        """Async implementation of API key clearing."""
         try:
-            self.config_manager.clear_provider_key(self.selected_provider)
+            await self.config_manager.clear_provider_key(self.selected_provider)
         except Exception as exc:  # pragma: no cover - defensive; textual path
             self.notify(f"Failed to clear key: {exc}", severity="error")
             return
 
-        self.refresh_provider_status()
-        self._update_done_button_visibility()
+        await self.refresh_provider_status()
+        await self._update_done_button_visibility()
         self.query_one("#api-key", Input).value = ""
 
         # If we just cleared shotgun, show the Authenticate button
@@ -311,5 +334,5 @@ class ProviderConfigScreen(Screen[None]):
 
         # Refresh provider status after auth completes
         if result:
-            self.refresh_provider_status()
+            await self.refresh_provider_status()
             # Notify handled by auth screen

shotgun/tui/screens/shotgun_auth.py

@@ -135,7 +135,7 @@ class ShotgunAuthScreen(Screen[bool]):
         """Start the authentication flow."""
         try:
             # Get shotgun instance ID from config
-            shotgun_instance_id = self.config_manager.get_shotgun_instance_id()
+            shotgun_instance_id = await self.config_manager.get_shotgun_instance_id()
             logger.info("Starting auth flow with instance ID: %s", shotgun_instance_id)
 
             # Update status
@@ -215,7 +215,7 @@ class ShotgunAuthScreen(Screen[bool]):
                     logger.info("Authentication completed successfully")
 
                     if status_response.litellm_key and status_response.supabase_key:
-                        self.config_manager.update_shotgun_account(
+                        await self.config_manager.update_shotgun_account(
                             api_key=status_response.litellm_key,
                             supabase_jwt=status_response.supabase_key,
                         )

shotgun/tui/screens/welcome.py

@@ -137,6 +137,15 @@ class WelcomeScreen(Screen[None]):
     def on_mount(self) -> None:
         """Focus the first button on mount."""
         self.query_one("#shotgun-button", Button).focus()
+        # Update BYOK button text asynchronously
+        self.run_worker(self._update_byok_button_text(), exclusive=False)
+
+    async def _update_byok_button_text(self) -> None:
+        """Update BYOK button text based on whether user has existing providers."""
+        byok_button = self.query_one("#byok-button", Button)
+        app = cast("ShotgunApp", self.app)
+        if await app.config_manager.has_any_provider_key():
+            byok_button.label = "I'll stick with my BYOK setup"
 
     @on(Button.Pressed, "#shotgun-button")
     def _on_shotgun_pressed(self) -> None:
@@ -146,21 +155,34 @@ class WelcomeScreen(Screen[None]):
     @on(Button.Pressed, "#byok-button")
     def _on_byok_pressed(self) -> None:
         """Handle BYOK button press."""
-        self._mark_welcome_shown()
-        # Push provider config screen before dismissing
+        self.run_worker(self._start_byok_config(), exclusive=True)
+
+    async def _start_byok_config(self) -> None:
+        """Launch BYOK provider configuration flow."""
+        await self._mark_welcome_shown()
+
+        app = cast("ShotgunApp", self.app)
+
+        # If user already has providers, just dismiss and continue to chat
+        if await app.config_manager.has_any_provider_key():
+            self.dismiss()
+            return
+
+        # Otherwise, push provider config screen and wait for result
         from .provider_config import ProviderConfigScreen
 
-        self.app.push_screen(
-            ProviderConfigScreen(),
-            callback=lambda _arg: self.dismiss(),
-        )
+        await self.app.push_screen_wait(ProviderConfigScreen())
+
+        # Dismiss welcome screen after config if providers are now configured
+        if await app.config_manager.has_any_provider_key():
+            self.dismiss()
 
     async def _start_shotgun_auth(self) -> None:
         """Launch Shotgun Account authentication flow."""
         from .shotgun_auth import ShotgunAuthScreen
 
         # Mark welcome screen as shown before auth
-        self._mark_welcome_shown()
+        await self._mark_welcome_shown()
 
         # Push the auth screen and wait for result
         await self.app.push_screen_wait(ShotgunAuthScreen())
@@ -168,9 +190,9 @@ class WelcomeScreen(Screen[None]):
         # Dismiss welcome screen after auth
         self.dismiss()
 
-    def _mark_welcome_shown(self) -> None:
+    async def _mark_welcome_shown(self) -> None:
         """Mark the welcome screen as shown in config."""
         app = cast("ShotgunApp", self.app)
-        config = app.config_manager.load()
+        config = await app.config_manager.load()
         config.shown_welcome_screen = True
-        app.config_manager.save(config)
+        await app.config_manager.save(config)

shotgun/tui/services/__init__.py

@@ -0,0 +1,5 @@
+"""Services for TUI business logic."""
+
+from shotgun.tui.services.conversation_service import ConversationService
+
+__all__ = ["ConversationService"]

shotgun/tui/services/conversation_service.py

@@ -0,0 +1,184 @@
+"""Service for managing conversation persistence and restoration.
+
+This service extracts conversation save/load/restore logic from ChatScreen,
+making it testable and reusable.
+"""
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import aiofiles.os
+
+from shotgun.agents.conversation_history import ConversationHistory, ConversationState
+from shotgun.agents.conversation_manager import ConversationManager
+from shotgun.agents.models import AgentType
+
+if TYPE_CHECKING:
+    from shotgun.agents.agent_manager import AgentManager
+    from shotgun.agents.usage_manager import SessionUsageManager
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationService:
+    """Handles conversation persistence and restoration.
+
+    This service provides:
+    - Save current conversation to disk
+    - Load conversation from disk
+    - Restore conversation state to agent manager
+    - Handle corrupted conversations gracefully
+    """
+
+    def __init__(
+        self,
+        conversation_manager: ConversationManager | None = None,
+        conversation_path: Path | None = None,
+    ):
+        """Initialize the conversation service.
+
+        Args:
+            conversation_manager: Optional conversation manager. If not provided,
+                                creates a default one.
+            conversation_path: Optional custom path for conversation storage.
+        """
+        if conversation_manager:
+            self.conversation_manager = conversation_manager
+        elif conversation_path:
+            self.conversation_manager = ConversationManager(conversation_path)
+        else:
+            self.conversation_manager = ConversationManager()
+
+    async def save_conversation(self, agent_manager: "AgentManager") -> bool:
+        """Save the current conversation to persistent storage.
+
+        Args:
+            agent_manager: The agent manager containing conversation state.
+
+        Returns:
+            True if save was successful, False otherwise.
+        """
+        try:
+            # Get conversation state from agent manager
+            state = agent_manager.get_conversation_state()
+
+            # Create conversation history object
+            conversation = ConversationHistory(
+                last_agent_model=state.agent_type,
+            )
+            conversation.set_agent_messages(state.agent_messages)
+            conversation.set_ui_messages(state.ui_messages)
+
+            # Save to file (now async)
+            await self.conversation_manager.save(conversation)
+            logger.debug("Conversation saved successfully")
+            return True
+        except Exception as e:
+            logger.exception(f"Failed to save conversation: {e}")
+            return False
+
+    async def load_conversation(self) -> ConversationHistory | None:
+        """Load conversation from persistent storage.
+
+        Returns:
+            The loaded conversation history, or None if no conversation exists
+            or if loading failed.
+        """
+        try:
+            conversation = await self.conversation_manager.load()
+            if conversation is None:
+                logger.debug("No conversation file found")
+                return None
+
+            logger.debug("Conversation loaded successfully")
+            return conversation
+        except Exception as e:
+            logger.exception(f"Failed to load conversation: {e}")
+            return None
+
+    async def check_for_corrupted_conversation(self) -> bool:
+        """Check if a conversation backup exists (indicating corruption).
+
+        Returns:
+            True if a backup exists (conversation was corrupted), False otherwise.
+        """
+        backup_path = self.conversation_manager.conversation_path.with_suffix(
+            ".json.backup"
+        )
+        return await aiofiles.os.path.exists(str(backup_path))
+
+    async def restore_conversation(
+        self,
+        agent_manager: "AgentManager",
+        usage_manager: "SessionUsageManager | None" = None,
+    ) -> tuple[bool, str | None, AgentType | None]:
+        """Restore conversation state from disk.
+
+        Args:
+            agent_manager: The agent manager to restore state to.
+            usage_manager: Optional usage manager to restore usage state.
+
+        Returns:
+            Tuple of (success, error_message, restored_agent_type)
+            - success: True if restoration succeeded
+            - error_message: Error message if restoration failed, None otherwise
+            - restored_agent_type: The agent type from restored conversation
+        """
+        conversation = await self.load_conversation()
+
+        if conversation is None:
+            # Check for corruption
+            if await self.check_for_corrupted_conversation():
+                return (
+                    False,
+                    "⚠️ Previous session was corrupted and has been backed up. Starting fresh conversation.",
+                    None,
+                )
+            return True, None, None  # No conversation to restore is not an error
+
+        try:
+            # Restore agent state
+            agent_messages = conversation.get_agent_messages()
+            ui_messages = conversation.get_ui_messages()
+
+            # Create ConversationState for restoration
+            state = ConversationState(
+                agent_messages=agent_messages,
+                ui_messages=ui_messages,
+                agent_type=conversation.last_agent_model,
+            )
+
+            agent_manager.restore_conversation_state(state)
+
+            # Restore usage state if manager provided
+            if usage_manager:
+                await usage_manager.restore_usage_state()
+
+            restored_type = AgentType(conversation.last_agent_model)
+            logger.info(f"Conversation restored successfully (mode: {restored_type})")
+            return True, None, restored_type
+
+        except Exception as e:
+            logger.exception(f"Failed to restore conversation state: {e}")
+            return (
+                False,
+                "⚠️ Could not restore previous session. Starting fresh conversation.",
+                None,
+            )
+
+    async def clear_conversation(self) -> bool:
+        """Clear the saved conversation file.
+
+        Returns:
+            True if clearing succeeded, False otherwise.
+        """
+        try:
+            conversation_path = self.conversation_manager.conversation_path
+            if await aiofiles.os.path.exists(str(conversation_path)):
+                await aiofiles.os.unlink(str(conversation_path))
+                logger.info("Conversation file cleared")
+            return True
+        except Exception as e:
+            logger.exception(f"Failed to clear conversation: {e}")
+            return False

shotgun/tui/state/__init__.py

@@ -0,0 +1,7 @@
+"""State management utilities for TUI."""
+
+from .processing_state import ProcessingStateManager
+
+__all__ = [
+    "ProcessingStateManager",
+]

shotgun/tui/state/processing_state.py

@@ -0,0 +1,185 @@
+"""Processing state management for TUI operations.
+
+This module provides centralized management of processing state including:
+- Tracking whether operations are in progress
+- Managing worker references for cancellation
+- Coordinating spinner widget updates
+- Providing clean cancellation API
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
+
+if TYPE_CHECKING:
+    from textual.screen import Screen
+    from textual.worker import Worker
+
+    from shotgun.tui.components.spinner import Spinner
+
+logger = get_logger(__name__)
+
+
+class ProcessingStateManager:
+    """Manages processing state and spinner coordination for async operations.
+
+    This class centralizes the logic for tracking whether the TUI is processing
+    an operation, managing the current worker for cancellation, and updating
+    spinner text.
+
+    Example:
+        ```python
+        # In ChatScreen
+        self.processing_state = ProcessingStateManager(self)
+
+        # Start processing
+        @work
+        async def some_operation(self) -> None:
+            self.processing_state.start_processing("Doing work...")
+            self.processing_state.bind_worker(get_current_worker())
+            try:
+                # ... do work ...
+            finally:
+                self.processing_state.stop_processing()
+        ```
+    """
+
+    def __init__(
+        self, screen: "Screen[Any]", telemetry_context: dict[str, Any] | None = None
+    ) -> None:
+        """Initialize the processing state manager.
+
+        Args:
+            screen: The Textual screen this manager is attached to
+            telemetry_context: Optional context to include in telemetry events
+                (e.g., {"agent_mode": "research"})
+        """
+        self.screen = screen
+        self._working = False
+        self._current_worker: Worker[Any] | None = None
+        self._spinner_widget: Spinner | None = None
+        self._default_spinner_text = "Processing..."
+        self._telemetry_context = telemetry_context or {}
+
+    @property
+    def is_working(self) -> bool:
+        """Check if an operation is currently in progress.
+
+        Returns:
+            True if processing, False if idle
+        """
+        return self._working
+
+    def bind_spinner(self, spinner: "Spinner") -> None:
+        """Bind a spinner widget for state coordination.
+
+        Should be called during screen mount after the spinner widget is available.
+
+        Args:
+            spinner: The Spinner widget to coordinate with
+        """
+        self._spinner_widget = spinner
+        logger.debug(f"Spinner widget bound: {spinner}")
+
+    def start_processing(self, spinner_text: str | None = None) -> None:
+        """Start processing state with optional custom spinner text.
+
+        Args:
+            spinner_text: Custom text to display in spinner. If None, uses default.
+        """
+        if self._working:
+            logger.warning("Attempted to start processing while already processing")
+            return
+
+        self._working = True
+        text = spinner_text or self._default_spinner_text
+
+        # Update screen's reactive working state
+        if hasattr(self.screen, "working"):
+            self.screen.working = True
+
+        if self._spinner_widget:
+            self._spinner_widget.text = text
+            logger.debug(f"Processing started with spinner text: {text}")
+        else:
+            logger.warning("Processing started but no spinner widget bound")
+
+    def stop_processing(self) -> None:
+        """Stop processing state and reset to default."""
+        if not self._working:
+            logger.debug("stop_processing called when not working (no-op)")
+            return
+
+        self._working = False
+        self._current_worker = None
+
+        # Update screen's reactive working state
+        if hasattr(self.screen, "working"):
+            self.screen.working = False
+
+        # Reset spinner to default text
+        if self._spinner_widget:
+            self._spinner_widget.text = self._default_spinner_text
+            logger.debug("Processing stopped, spinner reset to default")
+
+    def bind_worker(self, worker: "Worker[Any]") -> None:
+        """Bind a worker for cancellation tracking.
+
+        Should be called immediately after starting a @work decorated method
+        using get_current_worker().
+
+        Args:
+            worker: The Worker instance to track for cancellation
+        """
+        self._current_worker = worker
+        logger.debug(f"Worker bound: {worker}")
+
+    def cancel_current_operation(self, cancel_key: str | None = None) -> bool:
+        """Attempt to cancel the current operation if one is running.
+
+        Automatically tracks cancellation telemetry with context from initialization.
+
+        Args:
+            cancel_key: Optional key that triggered cancellation (e.g., "Escape")
+
+        Returns:
+            True if an operation was cancelled, False if no operation was running
+        """
+        if not self._working or not self._current_worker:
+            logger.debug("No operation to cancel")
+            return False
+
+        try:
+            self._current_worker.cancel()
+            logger.info("Operation cancelled successfully")
+
+            # Track cancellation event with context
+            event_data = {**self._telemetry_context}
+            if cancel_key:
+                event_data["cancel_key"] = cancel_key
+
+            track_event("agent_cancelled", event_data)
+
+            return True
+        except Exception as e:
+            logger.error(f"Failed to cancel operation: {e}", exc_info=True)
+            return False
+
+    def update_spinner_text(self, text: str) -> None:
+        """Update spinner text during processing.
+
+        Args:
+            text: New text to display in spinner
+        """
+        if not self._working:
+            logger.warning(
+                f"Attempted to update spinner text while not working: {text}"
+            )
+            return
+
+        if self._spinner_widget:
+            self._spinner_widget.text = text
+            logger.debug(f"Spinner text updated to: {text}")
+        else:
+            logger.warning(f"Cannot update spinner text, widget not bound: {text}")