shotgun-sh 0.1.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

shotgun/tui/screens/provider_config.py

@@ -9,7 +9,7 @@ from textual.app import ComposeResult
 from textual.containers import Horizontal, Vertical
 from textual.reactive import reactive
 from textual.screen import Screen
-from textual.widgets import Button, Input, Label, ListItem, ListView, Static
+from textual.widgets import Button, Input, Label, ListItem, ListView, Markdown, Static
 
 from shotgun.agents.config import ConfigManager, ProviderType
 
@@ -17,6 +17,16 @@ if TYPE_CHECKING:
     from ..app import ShotgunApp
 
 
+def get_configurable_providers() -> list[str]:
+    """Get list of configurable providers.
+
+    Returns:
+        List of provider identifiers that can be configured.
+        Includes all providers: openai, anthropic, google, and shotgun.
+    """
+    return ["openai", "anthropic", "google", "shotgun"]
+
+
 class ProviderConfigScreen(Screen[None]):
     """Collect API keys for available providers."""
 
@@ -47,6 +57,10 @@ class ProviderConfigScreen(Screen[None]):
             color: $text-accent;
         }
 
+        #provider-links {
+            padding: 1 0;
+        }
+
         #provider-list {
             margin: 2 0;
             height: auto;
@@ -67,9 +81,10 @@ class ProviderConfigScreen(Screen[None]):
 
     BINDINGS = [
         ("escape", "done", "Back"),
+        ("ctrl+c", "app.quit", "Quit"),
     ]
 
-    selected_provider: reactive[ProviderType] = reactive(ProviderType.OPENAI)
+    selected_provider: reactive[str] = reactive("openai")
 
     def compose(self) -> ComposeResult:
         with Vertical(id="titlebox"):
@@ -78,7 +93,11 @@ class ProviderConfigScreen(Screen[None]):
                 "Select a provider and enter the API key needed to activate it.",
                 id="provider-config-summary",
             )
-        yield ListView(*self._build_provider_items(), id="provider-list")
+            yield Markdown(
+                "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_sync(), id="provider-list")
         yield Input(
             placeholder=self._input_placeholder(self.selected_provider),
             password=True,
@@ -86,17 +105,35 @@ 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()
         list_view = self.query_one(ListView)
         if list_view.children:
             list_view.index = 0
-        self.selected_provider = ProviderType.OPENAI
+        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))
 
+        # 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.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()
 
@@ -117,6 +154,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()
@@ -133,58 +174,110 @@ 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 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)
-        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 await 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:
         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 in ProviderType:
-            label = self.query_one(f"#label-{provider.value}", Label)
-            label.update(self._provider_label(provider))
+        for provider_id in get_configurable_providers():
+            label = self.query_one(f"#label-{provider_id}", Label)
+            label.update(await self._provider_label(provider_id))
+
+    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 = await self.config_manager.has_any_provider_key()
+        done_button.display = has_keys
+
+    def _build_provider_items_sync(self) -> list[ListItem]:
+        """Build provider items synchronously for compose().
 
-    def _build_provider_items(self) -> list[ListItem]:
+        Labels will be populated with status asynchronously in on_mount().
+        """
         items: list[ListItem] = []
-        for provider in ProviderType:
-            label = Label(self._provider_label(provider), id=f"label-{provider.value}")
-            items.append(ListItem(label, id=f"provider-{provider.value}"))
+        for provider_id in get_configurable_providers():
+            # 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
 
-    def _provider_from_item(self, item: ListItem | None) -> ProviderType | None:
+    def _provider_from_item(self, item: ListItem | None) -> str | None:
         if item is None or item.id is None:
             return None
         provider_id = item.id.removeprefix("provider-")
-        try:
-            return ProviderType(provider_id)
-        except ValueError:
-            return None
+        return provider_id if provider_id in get_configurable_providers() else None
 
-    def _provider_label(self, provider: ProviderType) -> str:
-        display = self._provider_display_name(provider)
-        status = (
-            "Configured"
-            if self.config_manager.has_provider_key(provider)
-            else "Not configured"
-        )
+    async def _provider_label(self, provider_id: str) -> str:
+        display = self._provider_display_name(provider_id)
+        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: ProviderType) -> str:
+    def _provider_display_name(self, provider_id: str) -> str:
         names = {
-            ProviderType.OPENAI: "OpenAI",
-            ProviderType.ANTHROPIC: "Anthropic",
-            ProviderType.GOOGLE: "Google Gemini",
+            "openai": "OpenAI",
+            "anthropic": "Anthropic",
+            "google": "Google Gemini",
+            "shotgun": "Shotgun Account",
         }
-        return names.get(provider, provider.value.title())
-
-    def _input_placeholder(self, provider: ProviderType) -> str:
-        return f"{self._provider_display_name(provider)} API key"
+        return names.get(provider_id, provider_id.title())
+
+    def _input_placeholder(self, provider_id: str) -> str:
+        return f"{self._provider_display_name(provider_id)} API key"
+
+    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 = 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 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()
 
@@ -193,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,
             )
@@ -202,20 +295,44 @@ class ProviderConfigScreen(Screen[None]):
             return
 
         input_widget.value = ""
-        self.refresh_provider_status()
+        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()
+        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
+        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:
+            await self.refresh_provider_status()
+            # Notify handled by auth screen

shotgun/tui/screens/shotgun_auth.py

@@ -0,0 +1,295 @@
+"""Shotgun Account authentication screen."""
+
+import asyncio
+import webbrowser
+from typing import TYPE_CHECKING, cast
+
+import httpx
+from textual import on
+from textual.app import ComposeResult
+from textual.containers import Vertical
+from textual.screen import Screen
+from textual.widgets import Button, Label, Markdown, Static
+from textual.worker import Worker, WorkerState
+
+from shotgun.agents.config import ConfigManager
+from shotgun.logging_config import get_logger
+from shotgun.shotgun_web import (
+    ShotgunWebClient,
+    TokenStatus,
+)
+from shotgun.shotgun_web.constants import DEFAULT_POLL_INTERVAL_SECONDS
+
+if TYPE_CHECKING:
+    from ..app import ShotgunApp
+
+logger = get_logger(__name__)
+
+
+class ShotgunAuthScreen(Screen[bool]):
+    """Screen for Shotgun Account authentication flow.
+
+    Returns True if authentication was successful, False otherwise.
+    """
+
+    CSS = """
+        ShotgunAuth {
+            layout: vertical;
+        }
+
+        #titlebox {
+            height: auto;
+            margin: 2 0;
+            padding: 1;
+            border: hkey $border;
+            content-align: center middle;
+
+            & > * {
+                text-align: center;
+            }
+        }
+
+        #auth-title {
+            padding: 1 0;
+            text-style: bold;
+            color: $text-accent;
+        }
+
+        #content {
+            padding: 2;
+            height: auto;
+        }
+
+        #status {
+            padding: 1 0;
+            text-align: center;
+        }
+
+        #auth-url {
+            padding: 1;
+            border: solid $primary;
+            background: $surface;
+            text-align: center;
+        }
+
+        #actions {
+            padding: 1;
+            align: center middle;
+        }
+    """
+
+    BINDINGS = [
+        ("escape", "cancel", "Cancel"),
+    ]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.token: str | None = None
+        self.auth_url: str | None = None
+        self.poll_worker: Worker[None] | None = None
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="titlebox"):
+            yield Static("Shotgun Account Setup", id="auth-title")
+            yield Static(
+                "Authenticate with your Shotgun Account to get started",
+                id="auth-subtitle",
+            )
+
+        with Vertical(id="content"):
+            yield Label("Initializing...", id="status")
+            yield Markdown("", id="auth-url")
+            yield Markdown(
+                "**Instructions:**\n"
+                "1. A browser window will open automatically\n"
+                "2. Sign in or create a Shotgun Account\n"
+                "3. Complete payment if required\n"
+                "4. This window will automatically detect completion",
+                id="instructions",
+            )
+
+        with Vertical(id="actions"):
+            yield Button("Cancel", variant="default", id="cancel")
+
+    def on_mount(self) -> None:
+        """Start authentication flow when screen is mounted."""
+        self.run_worker(self._start_auth_flow(), exclusive=True)
+
+    def action_cancel(self) -> None:
+        """Cancel authentication and close screen."""
+        if self.poll_worker and self.poll_worker.state == WorkerState.RUNNING:
+            self.poll_worker.cancel()
+        self.dismiss(False)
+
+    @on(Button.Pressed, "#cancel")
+    def _on_cancel_pressed(self) -> None:
+        """Handle cancel button press."""
+        self.action_cancel()
+
+    @property
+    def config_manager(self) -> ConfigManager:
+        app = cast("ShotgunApp", self.app)
+        return app.config_manager
+
+    async def _start_auth_flow(self) -> None:
+        """Start the authentication flow."""
+        try:
+            # Get shotgun instance ID from config
+            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
+            self.query_one("#status", Label).update(
+                "🔄 Creating authentication token..."
+            )
+
+            # Create unification token
+            client = ShotgunWebClient()
+            response = client.create_unification_token(shotgun_instance_id)
+
+            self.token = response.token
+            self.auth_url = response.auth_url
+
+            logger.info("Auth URL: %s", self.auth_url)
+
+            # Update UI with auth URL
+            self.query_one("#status", Label).update("✅ Authentication URL ready")
+            self.query_one("#auth-url", Markdown).update(
+                f"**Authentication URL:**\n\n[{self.auth_url}]({self.auth_url})"
+            )
+
+            # Try to open browser
+            try:
+                self.query_one("#status", Label).update("🌐 Opening browser...")
+                webbrowser.open(self.auth_url)
+                await asyncio.sleep(1)
+                self.query_one("#status", Label).update(
+                    "⏳ Waiting for authentication... (opened in browser)"
+                )
+            except Exception as e:
+                logger.warning("Failed to open browser: %s", e)
+                self.query_one("#status", Label).update(
+                    "⚠️ Please click the link above to authenticate"
+                )
+
+            # Start polling for status
+            self.poll_worker = self.run_worker(
+                self._poll_token_status(), exclusive=False
+            )
+
+        except httpx.HTTPError as e:
+            logger.error("Failed to create auth token: %s", e)
+            self.query_one("#status", Label).update(
+                f"❌ Error: Failed to create authentication token\n{e}"
+            )
+            self.notify("Failed to start authentication", severity="error")
+
+        except Exception as e:
+            logger.error("Unexpected error during auth flow: %s", e)
+            self.query_one("#status", Label).update(f"❌ Unexpected error: {e}")
+            self.notify("Authentication failed", severity="error")
+
+    async def _poll_token_status(self) -> None:
+        """Poll token status until completed or expired."""
+        if not self.token:
+            logger.error("No token available for polling")
+            return
+
+        client = ShotgunWebClient()
+        poll_count = 0
+        max_polls = 600  # 30 minutes with 3 second intervals
+
+        while poll_count < max_polls:
+            try:
+                await asyncio.sleep(DEFAULT_POLL_INTERVAL_SECONDS)
+                poll_count += 1
+
+                logger.debug(
+                    "Polling token status (attempt %d/%d)", poll_count, max_polls
+                )
+
+                status_response = client.check_token_status(self.token)
+
+                if status_response.status == TokenStatus.COMPLETED:
+                    # Success! Save keys and dismiss
+                    logger.info("Authentication completed successfully")
+
+                    if status_response.litellm_key and status_response.supabase_key:
+                        await self.config_manager.update_shotgun_account(
+                            api_key=status_response.litellm_key,
+                            supabase_jwt=status_response.supabase_key,
+                        )
+
+                        self.query_one("#status", Label).update(
+                            "✅ Authentication successful! Saving credentials..."
+                        )
+                        await asyncio.sleep(1)
+                        self.notify(
+                            "Shotgun Account configured successfully!",
+                            severity="information",
+                        )
+                        self.dismiss(True)
+                    else:
+                        logger.error("Completed but missing keys")
+                        self.query_one("#status", Label).update(
+                            "❌ Error: Authentication completed but keys are missing"
+                        )
+                        self.notify("Authentication failed", severity="error")
+                        await asyncio.sleep(3)
+                        self.dismiss(False)
+                    return
+
+                elif status_response.status == TokenStatus.AWAITING_PAYMENT:
+                    self.query_one("#status", Label).update(
+                        "💳 Waiting for payment completion..."
+                    )
+
+                elif status_response.status == TokenStatus.EXPIRED:
+                    logger.error("Token expired")
+                    self.query_one("#status", Label).update(
+                        "❌ Authentication token expired (30 minutes)\n"
+                        "Please try again."
+                    )
+                    self.notify("Authentication token expired", severity="error")
+                    await asyncio.sleep(3)
+                    self.dismiss(False)
+                    return
+
+                elif status_response.status == TokenStatus.PENDING:
+                    # Still waiting, update status message
+                    elapsed_minutes = (poll_count * DEFAULT_POLL_INTERVAL_SECONDS) // 60
+                    self.query_one("#status", Label).update(
+                        f"⏳ Waiting for authentication... ({elapsed_minutes}m elapsed)"
+                    )
+
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 410:
+                    # Token expired
+                    logger.error("Token expired (410)")
+                    self.query_one("#status", Label).update(
+                        "❌ Authentication token expired"
+                    )
+                    self.notify("Authentication token expired", severity="error")
+                    await asyncio.sleep(3)
+                    self.dismiss(False)
+                    return
+                else:
+                    logger.error("HTTP error polling status: %s", e)
+                    self.query_one("#status", Label).update(
+                        f"❌ Error checking status: {e}"
+                    )
+                    await asyncio.sleep(5)  # Wait a bit longer on error
+
+            except Exception as e:
+                logger.error("Error polling token status: %s", e)
+                self.query_one("#status", Label).update(f"⚠️ Error checking status: {e}")
+                await asyncio.sleep(5)  # Wait a bit longer on error
+
+        # Timeout reached
+        logger.error("Polling timeout reached")
+        self.query_one("#status", Label).update(
+            "❌ Authentication timeout (30 minutes)\nPlease try again."
+        )
+        self.notify("Authentication timeout", severity="error")
+        await asyncio.sleep(3)
+        self.dismiss(False)