shotgun-sh 0.3.3.dev1__py3-none-any.whl → 0.6.2__py3-none-any.whl

shotgun/agents/cancellation.py

@@ -0,0 +1,103 @@
+"""Cancellation utilities for agent execution.
+
+This module provides utilities for responsive cancellation of agent operations,
+particularly for handling ESC key presses during LLM streaming.
+"""
+
+import asyncio
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import TypeVar
+
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+T = TypeVar("T")
+
+CANCELLATION_CHECK_INTERVAL = 0.5  # Check every 500ms
+CANCELLATION_MESSAGE = "Operation cancelled by user"
+
+
+class CancellableStreamIterator(AsyncIterator[T]):
+    """Wraps an async iterable to check for cancellation periodically.
+
+    This allows ESC cancellation to be responsive even when the underlying
+    stream (LLM chunks) is slow to produce events. Instead of blocking
+    indefinitely on the next chunk, we timeout periodically and check
+    if cancellation was requested.
+
+    Example:
+        ```python
+        cancellation_event = asyncio.Event()
+
+        async def process_stream(stream):
+            wrapped = CancellableStreamIterator(stream, cancellation_event)
+            async for event in wrapped:
+                process(event)
+
+        # In another task, set the event to cancel:
+        cancellation_event.set()
+        ```
+    """
+
+    def __init__(
+        self,
+        stream: AsyncIterable[T],
+        cancellation_event: asyncio.Event | None = None,
+        check_interval: float = CANCELLATION_CHECK_INTERVAL,
+    ) -> None:
+        """Initialize the cancellable stream iterator.
+
+        Args:
+            stream: The underlying async iterable to wrap
+            cancellation_event: Event that signals cancellation when set
+            check_interval: How often to check for cancellation (seconds)
+        """
+        self._stream = stream
+        self._iterator: AsyncIterator[T] | None = None
+        self._cancellation_event = cancellation_event
+        self._check_interval = check_interval
+        self._pending_task: asyncio.Task[T] | None = None
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        return self
+
+    async def __anext__(self) -> T:
+        if self._iterator is None:
+            self._iterator = self._stream.__aiter__()
+
+        # Create a task for the next item if we don't have one pending
+        if self._pending_task is None:
+            # Capture iterator reference for the coroutine
+            iterator = self._iterator
+
+            async def get_next() -> T:
+                return await iterator.__anext__()
+
+            self._pending_task = asyncio.create_task(get_next())
+
+        while True:
+            # Check if cancellation was requested
+            if self._cancellation_event and self._cancellation_event.is_set():
+                logger.debug("Cancellation detected in stream iterator")
+                # Cancel the pending task and raise
+                self._pending_task.cancel()
+                self._pending_task = None
+                raise asyncio.CancelledError(CANCELLATION_MESSAGE)
+
+            # Wait for the task with a short timeout
+            # Using asyncio.wait instead of wait_for to avoid cancelling the task on timeout
+            done, _ = await asyncio.wait(
+                [self._pending_task],
+                timeout=self._check_interval,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if done:
+                # Task completed - get result and clear pending task
+                task = done.pop()
+                self._pending_task = None
+                # Re-raise StopAsyncIteration or return the result
+                return task.result()
+
+            # Task not done yet, loop and check cancellation again

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
@@ -33,12 +38,14 @@ from .tools import (
     codebase_shell,
     directory_lister,
     file_read,
+    insert_markdown_section,
     query_graph,
     read_file,
+    remove_markdown_section,
+    replace_markdown_section,
     retrieve_code,
     write_file,
 )
-from .tools.file_management import AGENT_DIRECTORIES
 
 logger = get_logger(__name__)
 
@@ -65,6 +72,11 @@ async def add_system_status_message(
         await deps.codebase_service.list_graphs_for_directory()
     )
 
+    # Get graphs currently being indexed
+    indexing_graph_ids: set[str] = set()
+    if deps.codebase_service:
+        indexing_graph_ids = deps.codebase_service.indexing.get_active_ids()
+
     # Get existing files for the agent
     existing_files = get_agent_existing_files(deps.agent_mode)
 
@@ -74,15 +86,31 @@ 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,
+        indexing_graph_ids=indexing_graph_ids,
         is_tui_context=deps.is_tui_context,
         existing_files=existing_files,
         markdown_toc=markdown_toc,
         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 +130,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:
@@ -179,6 +207,9 @@ async def create_base_agent(
     agent.tool(write_file)
     agent.tool(append_file)
     agent.tool(read_file)
+    agent.tool(replace_markdown_section)
+    agent.tool(insert_markdown_section)
+    agent.tool(remove_markdown_section)
 
     # Register codebase understanding tools (conditional)
     if load_codebase_understanding_tools:
@@ -352,86 +383,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 +436,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 +517,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 +556,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

shotgun/agents/config/README.md

@@ -41,7 +41,6 @@ This directory contains the configuration management system for Shotgun, includi
 - **Title**: "feat: add config migration for streaming capability field (v4->v5)"
 - **Key Changes**:
   - Added `supports_streaming` field to OpenAI config
-  - Added `shown_onboarding_popup` timestamp field
   - Added `supabase_jwt` to Shotgun Account config
 
 ## Migration System

shotgun/agents/config/manager.py

@@ -51,7 +51,7 @@ class ConfigMigrationError(Exception):
 ProviderConfig = OpenAIConfig | AnthropicConfig | GoogleConfig | ShotgunAccountConfig
 
 # Current config version
-CURRENT_CONFIG_VERSION = 5
+CURRENT_CONFIG_VERSION = 6
 
 # Backup directory name
 BACKUP_DIR_NAME = "backup"
@@ -183,6 +183,26 @@ def _migrate_v4_to_v5(data: dict[str, Any]) -> dict[str, Any]:
     return data
 
 
+def _migrate_v5_to_v6(data: dict[str, Any]) -> dict[str, Any]:
+    """Migrate config from version 5 to version 6.
+
+    Changes:
+    - Add 'router_mode' field with default 'planning'
+
+    Args:
+        data: Config data dict at version 5
+
+    Returns:
+        Modified config data dict at version 6
+    """
+    if "router_mode" not in data:
+        data["router_mode"] = "planning"
+        logger.info("Migrated config v5->v6: added router_mode field")
+
+    data["config_version"] = 6
+    return data
+
+
 def _apply_migrations(data: dict[str, Any]) -> dict[str, Any]:
     """Apply all necessary migrations to bring config to current version.
 
@@ -203,6 +223,7 @@ def _apply_migrations(data: dict[str, Any]) -> dict[str, Any]:
         2: _migrate_v2_to_v3,
         3: _migrate_v3_to_v4,
         4: _migrate_v4_to_v5,
+        5: _migrate_v5_to_v6,
     }
 
     # Apply migrations sequentially
@@ -356,9 +377,9 @@ class ConfigManager:
 
                             # Find default model for this provider
                             provider_models = {
-                                ProviderType.OPENAI: ModelName.GPT_5_1,
-                                ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
-                                ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
+                                ProviderType.OPENAI: ModelName.GPT_5_2,
+                                ProviderType.ANTHROPIC: ModelName.CLAUDE_SONNET_4_5,
+                                ProviderType.GOOGLE: ModelName.GEMINI_3_PRO_PREVIEW,
                             }
 
                             if provider in provider_models:
@@ -500,15 +521,18 @@ class ConfigManager:
             if provider_enum is None:
                 raise RuntimeError("Provider enum should not be None for LLM providers")
             other_providers = [p for p in ProviderType if p != provider_enum]
-            has_other_keys = any(self.has_provider_key(p) for p in other_providers)
+            has_other_keys = any(
+                self._provider_has_api_key(self._get_provider_config(config, p))
+                for p in other_providers
+            )
             if not has_other_keys:
                 # Set selected_model to this provider's default model
                 from .models import ModelName
 
                 provider_models = {
-                    ProviderType.OPENAI: ModelName.GPT_5_1,
-                    ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
-                    ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
+                    ProviderType.OPENAI: ModelName.GPT_5_2,
+                    ProviderType.ANTHROPIC: ModelName.CLAUDE_SONNET_4_5,
+                    ProviderType.GOOGLE: ModelName.GEMINI_3_PRO_PREVIEW,
                 }
                 if provider_enum in provider_models:
                     config.selected_model = provider_models[provider_enum]
@@ -772,6 +796,26 @@ class ConfigManager:
         await self.save(config)
         logger.info("Updated Shotgun Account configuration")
 
+    async def get_router_mode(self) -> str:
+        """Get the saved router mode.
+
+        Returns:
+            The router mode string ('planning' or 'drafting')
+        """
+        config = await self.load()
+        return config.router_mode
+
+    async def set_router_mode(self, mode: str) -> None:
+        """Save the router mode.
+
+        Args:
+            mode: Router mode to save ('planning' or 'drafting')
+        """
+        config = await self.load()
+        config.router_mode = mode
+        await self.save(config)
+        logger.debug("Router mode saved: %s", mode)
+
 
 # Global singleton instance
 _config_manager_instance: ConfigManager | None = None

shotgun/agents/config/models.py

@@ -26,10 +26,8 @@ class ModelName(StrEnum):
     """Available AI model names."""
 
     GPT_5_1 = "gpt-5.1"
-    GPT_5_1_CODEX = "gpt-5.1-codex"
-    GPT_5_1_CODEX_MINI = "gpt-5.1-codex-mini"
+    GPT_5_2 = "gpt-5.2"
     CLAUDE_OPUS_4_5 = "claude-opus-4-5"
-    CLAUDE_SONNET_4 = "claude-sonnet-4"
     CLAUDE_SONNET_4_5 = "claude-sonnet-4-5"
     CLAUDE_HAIKU_4_5 = "claude-haiku-4-5"
     GEMINI_2_5_PRO = "gemini-2.5-pro"
@@ -110,21 +108,13 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         litellm_proxy_model_name="openai/gpt-5.1",
         short_name="GPT-5.1",
     ),
-    ModelName.GPT_5_1_CODEX: ModelSpec(
-        name=ModelName.GPT_5_1_CODEX,
+    ModelName.GPT_5_2: ModelSpec(
+        name=ModelName.GPT_5_2,
         provider=ProviderType.OPENAI,
         max_input_tokens=272_000,
         max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5.1-codex",
-        short_name="GPT-5.1 Codex",
-    ),
-    ModelName.GPT_5_1_CODEX_MINI: ModelSpec(
-        name=ModelName.GPT_5_1_CODEX_MINI,
-        provider=ProviderType.OPENAI,
-        max_input_tokens=272_000,
-        max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5.1-codex-mini",
-        short_name="GPT-5.1 Codex Mini",
+        litellm_proxy_model_name="openai/gpt-5.2",
+        short_name="GPT-5.2",
     ),
     ModelName.CLAUDE_SONNET_4_5: ModelSpec(
         name=ModelName.CLAUDE_SONNET_4_5,
@@ -166,14 +156,6 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         litellm_proxy_model_name="anthropic/claude-opus-4-5",
         short_name="Opus 4.5",
     ),
-    ModelName.CLAUDE_SONNET_4: ModelSpec(
-        name=ModelName.CLAUDE_SONNET_4,
-        provider=ProviderType.ANTHROPIC,
-        max_input_tokens=200_000,
-        max_output_tokens=64_000,
-        litellm_proxy_model_name="anthropic/claude-sonnet-4",
-        short_name="Sonnet 4",
-    ),
     ModelName.GEMINI_2_5_FLASH_LITE: ModelSpec(
         name=ModelName.GEMINI_2_5_FLASH_LITE,
         provider=ProviderType.GOOGLE,
@@ -226,6 +208,18 @@ class ShotgunAccountConfig(BaseModel):
         default=None, description="Default workspace ID for shared specs"
     )
 
+    @property
+    def has_valid_account(self) -> bool:
+        """Check if the user has a valid Shotgun Account configured.
+
+        Returns:
+            True if api_key is set and non-empty, False otherwise
+        """
+        if self.api_key is None:
+            return False
+        value = self.api_key.get_secret_value()
+        return bool(value and value.strip())
+
 
 class MarketingMessageRecord(BaseModel):
     """Record of when a marketing message was shown to the user."""
@@ -261,10 +255,6 @@ class ShotgunConfig(BaseModel):
         default=False,
         description="Whether the welcome screen has been shown to the user",
     )
-    shown_onboarding_popup: datetime | None = Field(
-        default=None,
-        description="Timestamp when the onboarding popup was shown to the user (ISO8601 format)",
-    )
     marketing: MarketingConfig = Field(
         default_factory=MarketingConfig,
         description="Marketing messages configuration and tracking",
@@ -277,3 +267,7 @@ class ShotgunConfig(BaseModel):
         default=None,
         description="Path to the backup file created when migration failed",
     )
+    router_mode: str = Field(
+        default="planning",
+        description="Router execution mode: 'planning' or 'drafting'",
+    )