shotgun-sh 0.3.3.dev1__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -20,7 +20,6 @@ if TYPE_CHECKING:
     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,8 +61,11 @@ 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
@@ -74,6 +77,8 @@ 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
 
@@ -992,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])
@@ -1087,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
 

shotgun/agents/common.py

@@ -1,6 +1,6 @@
 """Common utilities for agent creation and management."""
 
-from collections.abc import Callable
+from collections.abc import AsyncIterable, Awaitable, Callable
 from pathlib import Path
 from typing import Any
 
@@ -17,7 +17,12 @@ from pydantic_ai.messages import (
 )
 
 from shotgun.agents.config import ProviderType, get_provider_model
-from shotgun.agents.models import AgentResponse, AgentType
+from shotgun.agents.models import (
+    AgentResponse,
+    AgentSystemPromptContext,
+    AgentType,
+    ShotgunAgent,
+)
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
 from shotgun.sdk.services import get_codebase_service
@@ -38,7 +43,6 @@ from .tools import (
     retrieve_code,
     write_file,
 )
-from .tools.file_management import AGENT_DIRECTORIES
 
 logger = get_logger(__name__)
 
@@ -74,6 +78,19 @@ async def add_system_status_message(
     # Get current datetime with timezone information
     dt_context = get_datetime_context()
 
+    # Get execution plan and pending approval state if this is the Router agent
+    execution_plan = None
+    pending_approval = False
+    if deps.agent_mode == AgentType.ROUTER:
+        # Import here to avoid circular imports
+        from shotgun.agents.router.models import RouterDeps
+
+        if isinstance(deps, RouterDeps):
+            if deps.current_plan is not None:
+                execution_plan = deps.current_plan.format_for_display()
+            # Check if plan is pending approval (multi-step plan in Planning mode)
+            pending_approval = deps.pending_approval is not None
+
     system_state = prompt_loader.render(
         "agents/state/system_state.j2",
         codebase_understanding_graphs=codebase_understanding_graphs,
@@ -83,6 +100,8 @@ async def add_system_status_message(
         current_datetime=dt_context.datetime_formatted,
         timezone_name=dt_context.timezone_name,
         utc_offset=dt_context.utc_offset,
+        execution_plan=execution_plan,
+        pending_approval=pending_approval,
     )
 
     message_history.append(
@@ -102,7 +121,7 @@ async def create_base_agent(
     additional_tools: list[Any] | None = None,
     provider: ProviderType | None = None,
     agent_mode: AgentType | None = None,
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create a base agent with common configuration.
 
     Args:
@@ -352,86 +371,33 @@ async def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
 
 
 def get_agent_existing_files(agent_mode: AgentType | None = None) -> list[str]:
-    """Get list of existing files for the given agent mode.
+    """Get list of all existing files in .shotgun directory.
+
+    All agents can read any file in .shotgun/, so we list all files regardless
+    of agent mode. This includes user-added files that agents should be aware of.
 
     Args:
-        agent_mode: The agent mode to check files for. If None, lists all files.
+        agent_mode: Unused, kept for backwards compatibility.
 
     Returns:
         List of existing file paths relative to .shotgun directory
     """
     base_path = get_shotgun_base_path()
-    existing_files = []
-
-    # If no agent mode, list all files in base path and first level subdirectories
-    if agent_mode is None:
-        # List files in the root .shotgun directory
-        for item in base_path.iterdir():
-            if item.is_file():
-                existing_files.append(item.name)
-            elif item.is_dir():
-                # List files in first-level subdirectories
-                for subitem in item.iterdir():
-                    if subitem.is_file():
-                        relative_path = subitem.relative_to(base_path)
-                        existing_files.append(str(relative_path))
+    existing_files: list[str] = []
+
+    if not base_path.exists():
         return existing_files
 
-    # Handle specific agent modes
-    if agent_mode not in AGENT_DIRECTORIES:
-        return []
-
-    if agent_mode == AgentType.EXPORT:
-        # For export agent, list all files in exports directory
-        exports_dir = base_path / "exports"
-        if exports_dir.exists():
-            for file_path in exports_dir.rglob("*"):
-                if file_path.is_file():
-                    relative_path = file_path.relative_to(base_path)
+    # List all files in .shotgun directory and subdirectories
+    for item in base_path.iterdir():
+        if item.is_file():
+            existing_files.append(item.name)
+        elif item.is_dir():
+            # List files in subdirectories (one level deep to avoid too much noise)
+            for subitem in item.iterdir():
+                if subitem.is_file():
+                    relative_path = subitem.relative_to(base_path)
                     existing_files.append(str(relative_path))
-    else:
-        # For other agents, check files/directories they have access to
-        allowed_paths_raw = AGENT_DIRECTORIES[agent_mode]
-
-        # Convert single Path/string to list of Paths for uniform handling
-        if isinstance(allowed_paths_raw, str):
-            # Special case: "*" means export agent (shouldn't reach here but handle it)
-            allowed_paths = (
-                [Path(allowed_paths_raw)] if allowed_paths_raw != "*" else []
-            )
-        elif isinstance(allowed_paths_raw, Path):
-            allowed_paths = [allowed_paths_raw]
-        else:
-            # Already a list
-            allowed_paths = allowed_paths_raw
-
-        # Check each allowed path
-        for allowed_path in allowed_paths:
-            allowed_str = str(allowed_path)
-
-            # Check if it's a directory (no .md suffix)
-            if not allowed_path.suffix or not allowed_str.endswith(".md"):
-                # It's a directory - list all files within it
-                dir_path = base_path / allowed_str
-                if dir_path.exists() and dir_path.is_dir():
-                    for file_path in dir_path.rglob("*"):
-                        if file_path.is_file():
-                            relative_path = file_path.relative_to(base_path)
-                            existing_files.append(str(relative_path))
-            else:
-                # It's a file - check if it exists
-                file_path = base_path / allowed_str
-                if file_path.exists():
-                    existing_files.append(allowed_str)
-
-                # Also check for associated directory (e.g., research/ for research.md)
-                base_name = allowed_str.replace(".md", "")
-                dir_path = base_path / base_name
-                if dir_path.exists() and dir_path.is_dir():
-                    for file_path in dir_path.rglob("*"):
-                        if file_path.is_file():
-                            relative_path = file_path.relative_to(base_path)
-                            existing_files.append(str(relative_path))
 
     return existing_files
 
@@ -458,10 +424,24 @@ def build_agent_system_prompt(
         logger.debug("🔧 Building research agent system prompt...")
         logger.debug("Interactive mode: %s", ctx.deps.interactive_mode)
 
-    result = prompt_loader.render(
-        f"agents/{agent_type}.j2",
+    # Build template context using Pydantic model for type safety and testability
+    # Import here to avoid circular imports (same pattern as add_system_status_message)
+    from shotgun.agents.router.models import RouterDeps
+
+    router_mode = None
+    if isinstance(ctx.deps, RouterDeps):
+        router_mode = ctx.deps.router_mode.value
+
+    template_context = AgentSystemPromptContext(
         interactive_mode=ctx.deps.interactive_mode,
         mode=agent_type,
+        sub_agent_context=ctx.deps.sub_agent_context,
+        router_mode=router_mode,
+    )
+
+    result = prompt_loader.render(
+        f"agents/{agent_type}.j2",
+        **template_context.model_dump(),
     )
 
     if agent_type == "research":
@@ -525,13 +505,33 @@ async def add_system_prompt_message(
     return message_history
 
 
+EventStreamHandler = Callable[
+    [RunContext[AgentDeps], AsyncIterable[Any]], Awaitable[None]
+]
+
+
 async def run_agent(
-    agent: Agent[AgentDeps, AgentResponse],
+    agent: ShotgunAgent,
     prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
     usage_limits: UsageLimits | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
+    """Run an agent with optional streaming support.
+
+    Args:
+        agent: The agent to run.
+        prompt: The prompt to send to the agent.
+        deps: Agent dependencies.
+        message_history: Optional message history to continue from.
+        usage_limits: Optional usage limits for the run.
+        event_stream_handler: Optional callback for streaming events.
+            When provided, enables real-time streaming of agent responses.
+
+    Returns:
+        The agent run result.
+    """
     # Clear file tracker for new run
     deps.file_tracker.clear()
     logger.debug("🔧 Cleared file tracker for new agent run")
@@ -544,6 +544,7 @@ async def run_agent(
         deps=deps,
         usage_limits=usage_limits,
         message_history=message_history,
+        event_stream_handler=event_stream_handler,
     )
 
     # Log file operations summary if any files were modified