shotgun-sh 0.2.1.dev3__py3-none-any.whl → 0.2.1.dev5__py3-none-any.whl

shotgun/tui/app.py

@@ -8,6 +8,7 @@ from textual.screen import Screen
 from shotgun.agents.config import ConfigManager, get_config_manager
 from shotgun.logging_config import get_logger
 from shotgun.tui.screens.splash import SplashScreen
+from shotgun.utils.env_utils import is_shotgun_account_enabled
 from shotgun.utils.file_system_utils import get_shotgun_base_path
 from shotgun.utils.update_checker import perform_auto_update_async
 
@@ -16,6 +17,7 @@ from .screens.directory_setup import DirectorySetupScreen
 from .screens.feedback import FeedbackScreen
 from .screens.model_picker import ModelPickerScreen
 from .screens.provider_config import ProviderConfigScreen
+from .screens.welcome import WelcomeScreen
 
 logger = get_logger(__name__)
 
@@ -60,20 +62,34 @@ class ShotgunApp(App[None]):
     def refresh_startup_screen(self) -> None:
         """Push the appropriate screen based on configured providers."""
         if not self.config_manager.has_any_provider_key():
-            if isinstance(self.screen, ProviderConfigScreen):
+            # If Shotgun Account is enabled, show welcome screen with choice
+            # Otherwise, go directly to provider config (BYOK only)
+            if is_shotgun_account_enabled():
+                if isinstance(self.screen, WelcomeScreen):
+                    return
+
+                self.push_screen(
+                    WelcomeScreen(),
+                    callback=lambda _arg: self.refresh_startup_screen(),
+                )
+                return
+            else:
+                if isinstance(self.screen, ProviderConfigScreen):
+                    return
+
+                self.push_screen(
+                    ProviderConfigScreen(),
+                    callback=lambda _arg: self.refresh_startup_screen(),
+                )
                 return
-
-            self.push_screen(
-                "provider_config", callback=lambda _arg: self.refresh_startup_screen()
-            )
-            return
 
         if not self.check_local_shotgun_directory_exists():
             if isinstance(self.screen, DirectorySetupScreen):
                 return
 
             self.push_screen(
-                "directory_setup", callback=lambda _arg: self.refresh_startup_screen()
+                DirectorySetupScreen(),
+                callback=lambda _arg: self.refresh_startup_screen(),
             )
             return
 
@@ -110,7 +126,7 @@ class ShotgunApp(App[None]):
                 submit_feedback_survey(feedback)
                 self.notify("Feedback sent. Thank you!")
 
-        self.push_screen("feedback", callback=handle_feedback)
+        self.push_screen(FeedbackScreen(), callback=handle_feedback)
 
 
 def run(no_update_check: bool = False, continue_session: bool = False) -> None:

shotgun/tui/screens/chat_screen/command_providers.py

@@ -5,6 +5,8 @@ from textual.command import DiscoveryHit, Hit, Provider
 
 from shotgun.agents.models import AgentType
 from shotgun.codebase.models import CodebaseGraph
+from shotgun.tui.screens.model_picker import ModelPickerScreen
+from shotgun.tui.screens.provider_config import ProviderConfigScreen
 
 if TYPE_CHECKING:
     from shotgun.tui.screens.chat import ChatScreen
@@ -139,11 +141,11 @@ class ProviderSetupProvider(Provider):
 
     def open_provider_config(self) -> None:
         """Show the provider configuration screen."""
-        self.chat_screen.app.push_screen("provider_config")
+        self.chat_screen.app.push_screen(ProviderConfigScreen())
 
     def open_model_picker(self) -> None:
         """Show the model picker screen."""
-        self.chat_screen.app.push_screen("model_picker")
+        self.chat_screen.app.push_screen(ModelPickerScreen())
 
     async def discover(self) -> AsyncGenerator[DiscoveryHit, None]:
         yield DiscoveryHit(
@@ -282,11 +284,11 @@ class UnifiedCommandProvider(Provider):
 
     def open_provider_config(self) -> None:
         """Show the provider configuration screen."""
-        self.chat_screen.app.push_screen("provider_config")
+        self.chat_screen.app.push_screen(ProviderConfigScreen())
 
     def open_model_picker(self) -> None:
         """Show the model picker screen."""
-        self.chat_screen.app.push_screen("model_picker")
+        self.chat_screen.app.push_screen(ModelPickerScreen())
 
     async def discover(self) -> AsyncGenerator[DiscoveryHit, None]:
         """Provide commands in alphabetical order when palette opens."""

shotgun/tui/screens/feedback.py

@@ -182,12 +182,12 @@ class FeedbackScreen(Screen[Feedback | None]):
             return
 
         app = cast("ShotgunApp", self.app)
-        user_id = app.config_manager.get_user_id()
+        shotgun_instance_id = app.config_manager.get_shotgun_instance_id()
 
         feedback = Feedback(
             kind=self.selected_kind,
             description=description,
-            user_id=user_id,
+            shotgun_instance_id=shotgun_instance_id,
         )
 
         self.dismiss(feedback)

shotgun/tui/screens/model_picker.py

@@ -12,11 +12,14 @@ from textual.screen import Screen
 from textual.widgets import Button, Label, ListItem, ListView, Static
 
 from shotgun.agents.config import ConfigManager
-from shotgun.agents.config.models import MODEL_SPECS, ModelName
+from shotgun.agents.config.models import MODEL_SPECS, ModelName, ShotgunConfig
+from shotgun.logging_config import get_logger
 
 if TYPE_CHECKING:
     from ..app import ShotgunApp
 
+logger = get_logger(__name__)
+
 
 # Available models for selection
 AVAILABLE_MODELS = list(ModelName)
@@ -82,20 +85,52 @@ class ModelPickerScreen(Screen[None]):
                 "Select the AI model you want to use for your tasks.",
                 id="model-picker-summary",
             )
-        yield ListView(*self._build_model_items(), id="model-list")
+        yield ListView(id="model-list")
         with Horizontal(id="model-actions"):
             yield Button("Select \\[ENTER]", variant="primary", id="select")
             yield Button("Done \\[ESC]", id="done")
 
-    def on_mount(self) -> None:
-        # Load current selection
+    def _rebuild_model_list(self) -> None:
+        """Rebuild the model list from current config.
+
+        This method is called both on first show and when screen is resumed
+        to ensure the list always reflects the current configuration.
+        """
+        logger.debug("Rebuilding model list from current config")
+
+        # Load current config with force_reload to get latest API keys
         config_manager = self.config_manager
-        config = config_manager.load()
+        config = config_manager.load(force_reload=True)
+
+        # Log provider key status
+        logger.debug(
+            "Provider keys: openai=%s, anthropic=%s, google=%s, shotgun=%s",
+            config_manager._provider_has_api_key(config.openai),
+            config_manager._provider_has_api_key(config.anthropic),
+            config_manager._provider_has_api_key(config.google),
+            config_manager._provider_has_api_key(config.shotgun),
+        )
+
         current_model = config.selected_model or ModelName.CLAUDE_SONNET_4_5
         self.selected_model = current_model
+        logger.debug("Current selected model: %s", current_model)
 
-        # Find and highlight current selection (if it's in the filtered list)
+        # Rebuild the model list with current available models
         list_view = self.query_one(ListView)
+
+        # Remove all existing items
+        old_count = len(list(list_view.children))
+        for child in list(list_view.children):
+            child.remove()
+        logger.debug("Removed %d existing model items from list", old_count)
+
+        # Add new items (labels already have correct text including current indicator)
+        new_items = self._build_model_items(config)
+        for item in new_items:
+            list_view.append(item)
+        logger.debug("Added %d available model items to list", len(new_items))
+
+        # Find and highlight current selection (if it's in the filtered list)
         if list_view.children:
             for i, child in enumerate(list_view.children):
                 if isinstance(child, ListItem) and child.id:
@@ -106,7 +141,20 @@ class ModelPickerScreen(Screen[None]):
                             if model_name == current_model:
                                 list_view.index = i
                                 break
-        self.refresh_model_labels()
+
+    def on_show(self) -> None:
+        """Rebuild model list when screen is first shown."""
+        logger.debug("ModelPickerScreen.on_show() called")
+        self._rebuild_model_list()
+
+    def on_screenresume(self) -> None:
+        """Rebuild model list when screen is resumed (subsequent visits).
+
+        This is called when returning to the screen after it was suspended,
+        ensuring the model list reflects any config changes made while away.
+        """
+        logger.debug("ModelPickerScreen.on_screenresume() called")
+        self._rebuild_model_list()
 
     def action_done(self) -> None:
         self.dismiss()
@@ -138,13 +186,19 @@ class ModelPickerScreen(Screen[None]):
         return app.config_manager
 
     def refresh_model_labels(self) -> None:
-        """Update the list view entries to reflect current selection."""
-        current_model = (
-            self.config_manager.load().selected_model or ModelName.CLAUDE_SONNET_4_5
-        )
+        """Update the list view entries to reflect current selection.
+
+        Note: This method only updates labels for currently displayed models.
+        To rebuild the entire list after provider changes, on_show() should be used.
+        """
+        # Load config once with force_reload
+        config = self.config_manager.load(force_reload=True)
+        current_model = config.selected_model or ModelName.CLAUDE_SONNET_4_5
+
         # Update labels for available models only
         for model_name in AVAILABLE_MODELS:
-            if not self._is_model_available(model_name):
+            # Pass config to avoid multiple force reloads
+            if not self._is_model_available(model_name, config):
                 continue
             label = self.query_one(
                 f"#label-{_sanitize_model_name_for_id(model_name)}", Label
@@ -153,12 +207,15 @@ class ModelPickerScreen(Screen[None]):
                 self._model_label(model_name, is_current=model_name == current_model)
             )
 
-    def _build_model_items(self) -> list[ListItem]:
+    def _build_model_items(self, config: ShotgunConfig | None = None) -> list[ListItem]:
+        if config is None:
+            config = self.config_manager.load(force_reload=True)
+
         items: list[ListItem] = []
         current_model = self.selected_model
         for model_name in AVAILABLE_MODELS:
             # Only add models that are available
-            if not self._is_model_available(model_name):
+            if not self._is_model_available(model_name, config):
                 continue
 
             label = Label(
@@ -181,7 +238,9 @@ class ModelPickerScreen(Screen[None]):
                 return model_name
         return None
 
-    def _is_model_available(self, model_name: ModelName) -> bool:
+    def _is_model_available(
+        self, model_name: ModelName, config: ShotgunConfig | None = None
+    ) -> bool:
         """Check if a model is available based on provider key configuration.
 
         A model is available if:
@@ -190,22 +249,38 @@ class ModelPickerScreen(Screen[None]):
 
         Args:
             model_name: The model to check availability for
+            config: Optional pre-loaded config to avoid multiple reloads
 
         Returns:
             True if the model can be used, False otherwise
         """
-        config = self.config_manager.load()
+        if config is None:
+            config = self.config_manager.load(force_reload=True)
 
         # If Shotgun Account is configured, all models are available
         if self.config_manager._provider_has_api_key(config.shotgun):
+            logger.debug("Model %s available (Shotgun Account configured)", model_name)
             return True
 
         # In BYOK mode, check if the model's provider has a key
         if model_name not in MODEL_SPECS:
+            logger.debug("Model %s not available (not in MODEL_SPECS)", model_name)
             return False
 
         spec = MODEL_SPECS[model_name]
-        return self.config_manager.has_provider_key(spec.provider)
+        # Check provider key directly using the loaded config to avoid stale cache
+        provider_config = self.config_manager._get_provider_config(
+            config, spec.provider
+        )
+        has_key = self.config_manager._provider_has_api_key(provider_config)
+        logger.debug(
+            "Model %s available=%s (provider=%s, has_key=%s)",
+            model_name,
+            has_key,
+            spec.provider,
+            has_key,
+        )
+        return has_key
 
     def _model_label(self, model_name: ModelName, is_current: bool) -> str:
         """Generate label for model with specs and current indicator."""

shotgun/tui/screens/provider_config.py

@@ -85,6 +85,7 @@ class ProviderConfigScreen(Screen[None]):
 
     BINDINGS = [
         ("escape", "done", "Back"),
+        ("ctrl+c", "app.quit", "Quit"),
     ]
 
     selected_provider: reactive[str] = reactive("openai")
@@ -108,17 +109,30 @@ class ProviderConfigScreen(Screen[None]):
         )
         with Horizontal(id="provider-actions"):
             yield Button("Save key \\[ENTER]", variant="primary", id="save")
+            yield Button("Authenticate", variant="success", id="authenticate")
             yield Button("Clear key", id="clear", variant="warning")
             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
         self.selected_provider = "openai"
+
+        # Hide authenticate button by default (shown only for shotgun)
+        self.query_one("#authenticate", Button).display = False
         self.set_focus(self.query_one("#api-key", Input))
 
+    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()
+
     def action_done(self) -> None:
         self.dismiss()
 
@@ -139,6 +153,10 @@ class ProviderConfigScreen(Screen[None]):
     def _on_save_pressed(self) -> None:
         self._save_api_key()
 
+    @on(Button.Pressed, "#authenticate")
+    def _on_authenticate_pressed(self) -> None:
+        self.run_worker(self._start_shotgun_auth(), exclusive=True)
+
     @on(Button.Pressed, "#clear")
     def _on_clear_pressed(self) -> None:
         self._clear_api_key()
@@ -155,9 +173,31 @@ class ProviderConfigScreen(Screen[None]):
     def watch_selected_provider(self, provider: ProviderType) -> None:
         if not self.is_mounted:
             return
+
+        # Show/hide UI elements based on provider type
+        is_shotgun = provider == "shotgun"
+
         input_widget = self.query_one("#api-key", Input)
-        input_widget.placeholder = self._input_placeholder(provider)
-        input_widget.value = ""
+        save_button = self.query_one("#save", Button)
+        auth_button = self.query_one("#authenticate", Button)
+
+        if is_shotgun:
+            # Hide API key input and save button
+            input_widget.display = False
+            save_button.display = False
+
+            # Only show Authenticate button if shotgun is NOT already configured
+            if self._has_provider_key("shotgun"):
+                auth_button.display = False
+            else:
+                auth_button.display = True
+        else:
+            # Show API key input and save button, hide authenticate button
+            input_widget.display = True
+            save_button.display = True
+            auth_button.display = False
+            input_widget.placeholder = self._input_placeholder(provider)
+            input_widget.value = ""
 
     @property
     def config_manager(self) -> ConfigManager:
@@ -170,6 +210,12 @@ class ProviderConfigScreen(Screen[None]):
             label = self.query_one(f"#label-{provider_id}", Label)
             label.update(self._provider_label(provider_id))
 
+    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()
+        done_button.display = has_keys
+
     def _build_provider_items(self) -> list[ListItem]:
         items: list[ListItem] = []
         for provider_id in get_configurable_providers():
@@ -235,6 +281,7 @@ class ProviderConfigScreen(Screen[None]):
 
         input_widget.value = ""
         self.refresh_provider_status()
+        self._update_done_button_visibility()
         self.notify(
             f"Saved API key for {self._provider_display_name(self.selected_provider)}."
         )
@@ -247,7 +294,26 @@ class ProviderConfigScreen(Screen[None]):
             return
 
         self.refresh_provider_status()
+        self._update_done_button_visibility()
         self.query_one("#api-key", Input).value = ""
+
+        # If we just cleared shotgun, show the Authenticate button
+        if self.selected_provider == "shotgun":
+            auth_button = self.query_one("#authenticate", Button)
+            auth_button.display = True
+
         self.notify(
             f"Cleared API key for {self._provider_display_name(self.selected_provider)}."
         )
+
+    async def _start_shotgun_auth(self) -> None:
+        """Launch Shotgun Account authentication flow."""
+        from .shotgun_auth import ShotgunAuthScreen
+
+        # Push the auth screen and wait for result
+        result = await self.app.push_screen_wait(ShotgunAuthScreen())
+
+        # Refresh provider status after auth completes
+        if result:
+            self.refresh_provider_status()
+            # Notify handled by auth screen