shotgun-sh 0.2.17__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -17,10 +17,9 @@ from tenacity import (
 )
 
 if TYPE_CHECKING:
-    from shotgun.agents.conversation_history import ConversationState
+    from shotgun.agents.conversation import ConversationState
 
 from pydantic_ai import (
-    Agent,
     RunContext,
     UsageLimits,
 )
@@ -38,6 +37,7 @@ from pydantic_ai.messages import (
     PartDeltaEvent,
     PartStartEvent,
     SystemPromptPart,
+    TextPartDelta,
     ToolCallPart,
     ToolCallPartDelta,
     UserPromptPart,
@@ -61,19 +61,24 @@ from shotgun.agents.context_analyzer import (
 from shotgun.agents.models import (
     AgentResponse,
     AgentType,
+    AnyAgent,
     FileOperation,
     FileOperationTracker,
+    RouterAgent,
+    ShotgunAgent,
 )
 from shotgun.posthog_telemetry import track_event
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage
 from shotgun.utils.source_detection import detect_source
 
+from .conversation.history.compaction import apply_persistent_compaction
 from .export import create_export_agent
-from .history.compaction import apply_persistent_compaction
 from .messages import AgentSystemPrompt
 from .models import AgentDeps, AgentRuntimeOptions
 from .plan import create_plan_agent
 from .research import create_research_agent
+from .router import create_router_agent
+from .router.models import RouterDeps
 from .specify import create_specify_agent
 from .tasks import create_tasks_agent
 
@@ -174,6 +179,67 @@ class CompactionCompletedMessage(Message):
     """Event posted when conversation compaction completes."""
 
 
+class ToolExecutionStartedMessage(Message):
+    """Event posted when a tool starts executing.
+
+    This allows the UI to update the spinner text to provide feedback
+    during long-running tool executions.
+    """
+
+    def __init__(self, spinner_text: str = "Processing...") -> None:
+        """Initialize the tool execution started message.
+
+        Args:
+            spinner_text: The spinner message to display
+        """
+        super().__init__()
+        self.spinner_text = spinner_text
+
+
+class ToolStreamingProgressMessage(Message):
+    """Event posted during tool call streaming to show progress.
+
+    This provides visual feedback while tool arguments are streaming,
+    especially useful for long-running writes like file content.
+    """
+
+    def __init__(self, streamed_tokens: int, spinner_text: str) -> None:
+        """Initialize the tool streaming progress message.
+
+        Args:
+            streamed_tokens: Approximate number of tokens streamed so far
+            spinner_text: The current spinner message to preserve
+        """
+        super().__init__()
+        self.streamed_tokens = streamed_tokens
+        self.spinner_text = spinner_text
+
+
+# Fun spinner messages to show during tool execution
+SPINNER_MESSAGES = [
+    "Pontificating...",
+    "Ruminating...",
+    "Cogitating...",
+    "Deliberating...",
+    "Contemplating...",
+    "Reticulating splines...",
+    "Consulting the oracle...",
+    "Gathering thoughts...",
+    "Processing neurons...",
+    "Summoning wisdom...",
+    "Brewing ideas...",
+    "Polishing pixels...",
+    "Herding electrons...",
+    "Warming up the flux capacitor...",
+    "Consulting ancient tomes...",
+    "Channeling the muses...",
+    "Percolating possibilities...",
+    "Untangling complexity...",
+    "Shuffling priorities...",
+    "Aligning the stars...",
+]
+
+
 class AgentStreamingStarted(Message):
     """Event posted when agent starts streaming responses."""
 
@@ -210,6 +276,11 @@ class _PartialStreamState:
 
     messages: list[ModelRequest | ModelResponse] = field(default_factory=list)
     current_response: ModelResponse | None = None
+    # Token counting for tool call streaming progress
+    streamed_tokens: int = 0
+    current_spinner_text: str = "Processing..."
+    # Track last reported tokens to throttle UI updates
+    last_reported_tokens: int = 0
 
 
 class AgentManager(Widget):
@@ -245,16 +316,18 @@ class AgentManager(Widget):
         )
 
         # Lazy initialization - agents created on first access
-        self._research_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._research_agent: ShotgunAgent | None = None
         self._research_deps: AgentDeps | None = None
-        self._plan_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._plan_agent: ShotgunAgent | None = None
         self._plan_deps: AgentDeps | None = None
-        self._tasks_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._tasks_agent: ShotgunAgent | None = None
         self._tasks_deps: AgentDeps | None = None
-        self._specify_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._specify_agent: ShotgunAgent | None = None
         self._specify_deps: AgentDeps | None = None
-        self._export_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._export_agent: ShotgunAgent | None = None
         self._export_deps: AgentDeps | None = None
+        self._router_agent: RouterAgent | None = None
+        self._router_deps: RouterDeps | None = None
         self._agents_initialized = False
 
         # Track current active agent
@@ -291,10 +364,13 @@ class AgentManager(Widget):
         self._export_agent, self._export_deps = await create_export_agent(
             agent_runtime_options=self._agent_runtime_options
         )
+        self._router_agent, self._router_deps = await create_router_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
         self._agents_initialized = True
 
     @property
-    def research_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def research_agent(self) -> ShotgunAgent:
         """Get research agent (must call _ensure_agents_initialized first)."""
         if self._research_agent is None:
             raise RuntimeError(
@@ -312,7 +388,7 @@ class AgentManager(Widget):
         return self._research_deps
 
     @property
-    def plan_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def plan_agent(self) -> ShotgunAgent:
         """Get plan agent (must call _ensure_agents_initialized first)."""
         if self._plan_agent is None:
             raise RuntimeError(
@@ -330,7 +406,7 @@ class AgentManager(Widget):
         return self._plan_deps
 
     @property
-    def tasks_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def tasks_agent(self) -> ShotgunAgent:
         """Get tasks agent (must call _ensure_agents_initialized first)."""
         if self._tasks_agent is None:
             raise RuntimeError(
@@ -348,7 +424,7 @@ class AgentManager(Widget):
         return self._tasks_deps
 
     @property
-    def specify_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def specify_agent(self) -> ShotgunAgent:
         """Get specify agent (must call _ensure_agents_initialized first)."""
         if self._specify_agent is None:
             raise RuntimeError(
@@ -366,7 +442,7 @@ class AgentManager(Widget):
         return self._specify_deps
 
     @property
-    def export_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def export_agent(self) -> ShotgunAgent:
         """Get export agent (must call _ensure_agents_initialized first)."""
         if self._export_agent is None:
             raise RuntimeError(
@@ -384,29 +460,48 @@ class AgentManager(Widget):
         return self._export_deps
 
     @property
-    def current_agent(self) -> Agent[AgentDeps, AgentResponse]:
+    def router_agent(self) -> RouterAgent:
+        """Get router agent (must call _ensure_agents_initialized first)."""
+        if self._router_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._router_agent
+
+    @property
+    def router_deps(self) -> RouterDeps:
+        """Get router deps (must call _ensure_agents_initialized first)."""
+        if self._router_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._router_deps
+
+    @property
+    def current_agent(self) -> AnyAgent:
         """Get the currently active agent.
 
         Returns:
-            The currently selected agent instance.
+            The currently selected agent instance (ShotgunAgent or RouterAgent).
         """
         return self._get_agent(self._current_agent_type)
 
-    def _get_agent(self, agent_type: AgentType) -> Agent[AgentDeps, AgentResponse]:
+    def _get_agent(self, agent_type: AgentType) -> AnyAgent:
         """Get agent by type.
 
         Args:
             agent_type: The type of agent to retrieve.
 
         Returns:
-            The requested agent instance.
+            The requested agent instance (ShotgunAgent or RouterAgent).
         """
-        agent_map = {
+        agent_map: dict[AgentType, AnyAgent] = {
             AgentType.RESEARCH: self.research_agent,
             AgentType.PLAN: self.plan_agent,
             AgentType.TASKS: self.tasks_agent,
             AgentType.SPECIFY: self.specify_agent,
             AgentType.EXPORT: self.export_agent,
+            AgentType.ROUTER: self.router_agent,
         }
         return agent_map[agent_type]
 
@@ -419,12 +514,13 @@ class AgentManager(Widget):
         Returns:
             The agent-specific dependencies.
         """
-        deps_map = {
+        deps_map: dict[AgentType, AgentDeps] = {
             AgentType.RESEARCH: self.research_deps,
             AgentType.PLAN: self.plan_deps,
             AgentType.TASKS: self.tasks_deps,
             AgentType.SPECIFY: self.specify_deps,
             AgentType.EXPORT: self.export_deps,
+            AgentType.ROUTER: self.router_deps,
         }
         return deps_map[agent_type]
 
@@ -433,6 +529,10 @@ class AgentManager(Widget):
 
         This preserves the agent's system_prompt_fn while using shared runtime state.
 
+        For Router agent, returns the shared deps directly (not a copy) because
+        Router state (pending_approval, current_plan, etc.) must be shared with
+        the TUI for features like plan approval widgets.
+
         Args:
             agent_type: The type of agent to create merged deps for.
 
@@ -445,8 +545,14 @@ class AgentManager(Widget):
         if self.deps is None:
             raise ValueError("Shared deps is None - this should not happen")
 
-        # Create new deps with shared runtime state but agent's system_prompt_fn
-        # Use a copy of the shared deps and update the system_prompt_fn
+        # For Router, use shared deps directly so state mutations are visible to TUI
+        # (e.g., pending_approval, current_plan need to be seen by ChatScreen)
+        if agent_type == AgentType.ROUTER:
+            # Update system_prompt_fn on shared deps in place
+            self.deps.system_prompt_fn = agent_deps.system_prompt_fn
+            return self.deps
+
+        # For other agents, create a copy with agent-specific system_prompt_fn
         merged_deps = self.deps.model_copy(
             update={"system_prompt_fn": agent_deps.system_prompt_fn}
         )
@@ -478,7 +584,7 @@ class AgentManager(Widget):
     )
     async def _run_agent_with_retry(
         self,
-        agent: Agent[AgentDeps, AgentResponse],
+        agent: AnyAgent,
         prompt: str | None,
         deps: AgentDeps,
         usage_limits: UsageLimits | None,
@@ -489,9 +595,9 @@ class AgentManager(Widget):
         """Run agent with automatic retry on transient errors.
 
         Args:
-            agent: The agent to run.
+            agent: The agent to run (ShotgunAgent or RouterAgent).
             prompt: Optional prompt to send to the agent.
-            deps: Agent dependencies.
+            deps: Agent dependencies (AgentDeps or RouterDeps).
             usage_limits: Optional usage limits.
             message_history: Message history to provide to agent.
             event_stream_handler: Event handler for streaming.
@@ -502,8 +608,16 @@ class AgentManager(Widget):
 
         Raises:
             Various exceptions if all retries fail.
+
+        Note:
+            Type safety for agent/deps pairing is maintained by AgentManager's
+            _get_agent_deps which ensures the correct deps type is used for each
+            agent type. The cast is needed because Agent is contravariant in deps.
         """
-        return await agent.run(
+        # Cast needed because Agent is contravariant in deps type parameter.
+        # The agent/deps pairing is ensured by _get_agent_deps returning the
+        # correct deps type for each agent type.
+        return await cast(ShotgunAgent, agent).run(
             prompt,
             deps=deps,
             usage_limits=usage_limits,
@@ -560,6 +674,11 @@ class AgentManager(Widget):
 
         deps.agent_mode = self._current_agent_type
 
+        # For router agent, set up the parent stream handler so sub-agents can stream
+        if self._current_agent_type == AgentType.ROUTER:
+            if isinstance(deps, RouterDeps):
+                deps.parent_stream_handler = self._handle_event_stream  # type: ignore[assignment]
+
         # Filter out system prompts from other agent types
         from pydantic_ai.messages import ModelRequestPart
 
@@ -615,19 +734,33 @@ class AgentManager(Widget):
         self._stream_state = _PartialStreamState()
 
         model_name = ""
+        supports_streaming = True  # Default to streaming enabled
+
         if hasattr(deps, "llm_model") and deps.llm_model is not None:
             model_name = deps.llm_model.name
+            supports_streaming = deps.llm_model.supports_streaming
 
-        # Check if it's a Shotgun account
-        is_shotgun_account = (
-            hasattr(deps, "llm_model")
-            and deps.llm_model is not None
-            and deps.llm_model.key_provider == KeyProvider.SHOTGUN
-        )
-
-        # Only disable streaming for GPT-5 if NOT a Shotgun account
-        # Shotgun accounts support streaming for GPT-5
-        is_gpt5_byok = "gpt-5" in model_name.lower() and not is_shotgun_account
+            # Add hint message if streaming is disabled for BYOK GPT-5 models
+            if (
+                not supports_streaming
+                and deps.llm_model.key_provider == KeyProvider.BYOK
+            ):
+                self.ui_message_history.append(
+                    HintMessage(
+                        message=(
+                            "⚠️ **Streaming not available for GPT-5**\n\n"
+                            "Your OpenAI organization doesn't have streaming enabled for this model.\n\n"
+                            "**Options:**\n"
+                            "- Get a [Shotgun Account](https://shotgun.sh) - streaming works out of the box\n"
+                            "- Complete [Biometric Verification](https://platform.openai.com/settings/organization/general) with OpenAI, then:\n"
+                            "  1. Press `Ctrl+P` → Open Provider Setup\n"
+                            "  2. Select OpenAI → Clear key\n"
+                            "  3. Re-add your OpenAI API key\n\n"
+                            "Continuing without streaming (responses will appear all at once)."
+                        )
+                    )
+                )
+                self._post_messages_updated()
 
         # Track message send event
         event_name = f"message_send_{self._current_agent_type.value}"
@@ -647,7 +780,7 @@ class AgentManager(Widget):
                 usage_limits=usage_limits,
                 message_history=message_history,
                 event_stream_handler=self._handle_event_stream
-                if not is_gpt5_byok
+                if supports_streaming
                 else None,
                 **kwargs,
             )
@@ -978,6 +1111,44 @@ class AgentManager(Widget):
                         )
                         continue
 
+                    # Count tokens from the delta for progress indication
+                    delta_len = 0
+                    is_tool_call_delta = False
+                    if isinstance(event.delta, ToolCallPartDelta):
+                        is_tool_call_delta = True
+                        # args_delta can be str or dict depending on provider
+                        args_delta = event.delta.args_delta
+                        if isinstance(args_delta, str):
+                            delta_len = len(args_delta)
+                        elif isinstance(args_delta, dict):
+                            # For dict deltas, estimate from JSON representation
+                            delta_len = len(json.dumps(args_delta))
+                        # Pick a spinner message when tool streaming starts
+                        if state.current_spinner_text == "Processing...":
+                            import random
+
+                            state.current_spinner_text = random.choice(  # noqa: S311
+                                SPINNER_MESSAGES
+                            )
+                    elif isinstance(event.delta, TextPartDelta):
+                        delta_len = len(event.delta.content_delta)
+
+                    if delta_len > 0:
+                        # Approximate tokens: len / 4 is a rough estimate
+                        state.streamed_tokens += delta_len // 4 + 1
+                        # Send progress update for tool call streaming
+                        # Throttle updates to every ~75 tokens to avoid flooding UI
+                        if is_tool_call_delta and (
+                            state.streamed_tokens - state.last_reported_tokens >= 75
+                        ):
+                            state.last_reported_tokens = state.streamed_tokens
+                            self.post_message(
+                                ToolStreamingProgressMessage(
+                                    state.streamed_tokens,
+                                    state.current_spinner_text,
+                                )
+                            )
+
                     try:
                         updated_part = event.delta.apply(
                             cast(ModelResponsePart, partial_parts[index])
@@ -1073,6 +1244,17 @@ class AgentManager(Widget):
                     if partial_message is not None:
                         state.current_response = partial_message
                         self._post_partial_message(False)
+
+                    # Notify UI that a tool is about to execute
+                    # This updates the spinner with a fun message during tool execution
+                    # Pick a random spinner message and store it for progress updates
+                    import random
+
+                    spinner_text = random.choice(SPINNER_MESSAGES)  # noqa: S311
+                    state.current_spinner_text = spinner_text
+                    state.streamed_tokens = 0  # Reset token count for new tool
+                    self.post_message(ToolExecutionStartedMessage(spinner_text))
+
                 elif isinstance(event, FunctionToolResultEvent):
                     # Track tool completion event
 
@@ -1300,7 +1482,7 @@ class AgentManager(Widget):
         Returns:
             ConversationState object containing UI and agent messages and current type
         """
-        from shotgun.agents.conversation_history import ConversationState
+        from shotgun.agents.conversation import ConversationState
 
         return ConversationState(
             agent_messages=self.message_history.copy(),