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

shotgun/agents/agent_manager.py

@@ -19,7 +19,10 @@ from tenacity import (
 if TYPE_CHECKING:
     from shotgun.agents.conversation import ConversationState
 
+import base64
+
 from pydantic_ai import (
+    BinaryContent,
     RunContext,
     UsageLimits,
 )
@@ -40,11 +43,13 @@ from pydantic_ai.messages import (
     TextPartDelta,
     ToolCallPart,
     ToolCallPartDelta,
+    UserContent,
     UserPromptPart,
 )
 from textual.message import Message
 from textual.widget import Widget
 
+from shotgun.agents.cancellation import CancellableStreamIterator
 from shotgun.agents.common import add_system_prompt_message, add_system_status_message
 from shotgun.agents.config.models import (
     KeyProvider,
@@ -67,18 +72,19 @@ from shotgun.agents.models import (
     RouterAgent,
     ShotgunAgent,
 )
+from shotgun.attachments import FileAttachment
 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 .messages import AgentSystemPrompt
+from .messages import AgentSystemPrompt, InternalPromptPart
 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 .router.models import RouterDeps, RouterMode
 from .specify import create_specify_agent
 from .tasks import create_tasks_agent
 
@@ -171,6 +177,29 @@ class ClarifyingQuestionsMessage(Message):
         self.response_text = response_text
 
 
+class FileRequestPendingMessage(Message):
+    """Event posted when agent requests files to be loaded.
+
+    This triggers the TUI to load the requested files and resume
+    the agent with the file contents in the next prompt.
+    """
+
+    def __init__(
+        self,
+        file_paths: list[str],
+        response_text: str,
+    ) -> None:
+        """Initialize the file request pending message.
+
+        Args:
+            file_paths: List of file paths the agent wants to read
+            response_text: The agent's response text before requesting files
+        """
+        super().__init__()
+        self.file_paths = file_paths
+        self.response_text = response_text
+
+
 class CompactionStartedMessage(Message):
     """Event posted when conversation compaction starts."""
 
@@ -343,6 +372,10 @@ class AgentManager(Widget):
         self._qa_questions: list[str] | None = None
         self._qa_mode_active: bool = False
 
+        # File request state for structured output file loading
+        self._file_request_pending: bool = False
+        self._pending_file_requests: list[str] = []
+
     async def _ensure_agents_initialized(self) -> None:
         """Ensure all agents are initialized (lazy initialization)."""
         if self._agents_initialized:
@@ -486,6 +519,72 @@ class AgentManager(Widget):
         """
         return self._get_agent(self._current_agent_type)
 
+    @property
+    def file_request_pending(self) -> bool:
+        """Check if there's a pending file request."""
+        return self._file_request_pending
+
+    @property
+    def pending_file_requests(self) -> list[str]:
+        """Get the list of pending file requests."""
+        return self._pending_file_requests
+
+    def process_file_requests(self) -> list[tuple[str, BinaryContent]]:
+        """Process pending file requests and return loaded content.
+
+        This method is called by the TUI after FileRequestPendingMessage is received.
+        It loads the requested files as BinaryContent and clears the pending state.
+
+        Returns:
+            List of (file_path, BinaryContent) tuples for files that were successfully loaded.
+        """
+        if not self._file_request_pending:
+            return []
+
+        # MIME type mapping for supported file types
+        mime_types: dict[str, str] = {
+            ".pdf": "application/pdf",
+            ".png": "image/png",
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".gif": "image/gif",
+            ".webp": "image/webp",
+        }
+
+        loaded_files: list[tuple[str, BinaryContent]] = []
+        for file_path_str in self._pending_file_requests:
+            try:
+                path = Path(file_path_str).expanduser().resolve()
+                if not path.exists():
+                    logger.warning(f"Requested file not found: {path}")
+                    continue
+
+                # Get MIME type
+                suffix = path.suffix.lower()
+                mime_type = mime_types.get(suffix)
+                if mime_type is None:
+                    logger.warning(f"Unsupported file type: {suffix} for {path}")
+                    continue
+
+                # Read file and create BinaryContent
+                data = path.read_bytes()
+                loaded_files.append(
+                    (str(path), BinaryContent(data=data, media_type=mime_type))
+                )
+                logger.debug(f"Loaded file: {path} ({len(data)} bytes)")
+
+            except Exception as e:
+                logger.error(f"Error loading file {file_path_str}: {e}")
+
+        # Clear pending state
+        self._file_request_pending = False
+        self._pending_file_requests = []
+
+        logger.info(
+            f"Loaded {len(loaded_files)} of {len(self._pending_file_requests)} requested files"
+        )
+        return loaded_files
+
     def _get_agent(self, agent_type: AgentType) -> AnyAgent:
         """Get agent by type.
 
@@ -585,7 +684,7 @@ class AgentManager(Widget):
     async def _run_agent_with_retry(
         self,
         agent: AnyAgent,
-        prompt: str | None,
+        prompt: str | Sequence[UserContent] | None,
         deps: AgentDeps,
         usage_limits: UsageLimits | None,
         message_history: list[ModelMessage],
@@ -596,7 +695,8 @@ class AgentManager(Widget):
 
         Args:
             agent: The agent to run (ShotgunAgent or RouterAgent).
-            prompt: Optional prompt to send to the agent.
+            prompt: Optional prompt to send to the agent. Can be a string,
+                a sequence of UserContent (for multimodal), or None.
             deps: Agent dependencies (AgentDeps or RouterDeps).
             usage_limits: Optional usage limits.
             message_history: Message history to provide to agent.
@@ -630,6 +730,8 @@ class AgentManager(Widget):
         self,
         prompt: str | None = None,
         *,
+        attachment: FileAttachment | None = None,
+        file_contents: list[tuple[str, BinaryContent]] | None = None,
         deps: AgentDeps | None = None,
         usage_limits: UsageLimits | None = None,
         **kwargs: Any,
@@ -641,6 +743,9 @@ class AgentManager(Widget):
 
         Args:
             prompt: Optional prompt to send to the agent.
+            attachment: Optional file attachment to include as multimodal content.
+            file_contents: Optional list of (file_path, BinaryContent) tuples to include
+                          as multimodal content. Used when resuming after file_requests.
             deps: Optional dependencies override (defaults to manager's deps).
             usage_limits: Optional usage limits for the agent run.
             **kwargs: Additional keyword arguments to pass to the agent.
@@ -753,7 +858,7 @@ class AgentManager(Widget):
                             "**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"
+                            "  1. Press `/` → 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)."
@@ -769,13 +874,46 @@ class AgentManager(Widget):
             {
                 "has_prompt": prompt is not None,
                 "model_name": model_name,
+                "has_attachment": attachment is not None,
             },
         )
 
+        # Construct multimodal prompt if attachment or file_contents is provided
+        user_prompt: str | Sequence[UserContent] | None = prompt
+
+        if file_contents:
+            # File contents from file_requests - construct multimodal prompt with files
+            content_parts: list[UserContent] = [
+                prompt or "Here are the files you requested:"
+            ]
+            for file_path, binary in file_contents:
+                content_parts.append(f"\n\n--- File: {file_path} ---")
+                content_parts.append(binary)
+            user_prompt = content_parts
+            logger.debug(
+                "Constructed multimodal prompt with requested files",
+                extra={"num_files": len(file_contents)},
+            )
+        elif attachment and attachment.content_base64:
+            # Use BinaryContent which is supported by all providers (OpenAI, Anthropic, Google)
+            binary_data = base64.b64decode(attachment.content_base64)
+            binary_content = BinaryContent(
+                data=binary_data,
+                media_type=attachment.mime_type,
+            )
+            user_prompt = [prompt or "", binary_content]
+            logger.debug(
+                "Constructed multimodal prompt with attachment",
+                extra={
+                    "attachment_type": attachment.file_type.value,
+                    "attachment_size": attachment.file_size_bytes,
+                },
+            )
+
         try:
             result: AgentRunResult[AgentResponse] = await self._run_agent_with_retry(
                 agent=self.current_agent,
-                prompt=prompt,
+                prompt=user_prompt,
                 deps=deps,
                 usage_limits=usage_limits,
                 message_history=message_history,
@@ -861,17 +999,38 @@ class AgentManager(Widget):
         )
 
         # Deduplicate: skip user prompts that are already in original_messages
+        # Note: We compare content only, not timestamps, since UserPromptPart
+        # has a timestamp field that differs between instances
+        def get_user_prompt_text(
+            request: ModelRequest,
+        ) -> str | None:
+            """Extract just the text content from a ModelRequest for deduplication.
+
+            When content is multimodal (list with text + binary), extract just the text.
+            This ensures text-only and multimodal versions of the same prompt match.
+            """
+            for part in request.parts:
+                if isinstance(part, UserPromptPart):
+                    content = part.content
+                    if isinstance(content, str):
+                        return content
+                    elif isinstance(content, list):
+                        # Multimodal content - extract text strings only
+                        text_parts = [item for item in content if isinstance(item, str)]
+                        return text_parts[0] if text_parts else None
+            return None
+
         deduplicated_new_messages = []
         for msg in new_messages:
             # Check if this is a user prompt that's already in original_messages
             if isinstance(msg, ModelRequest) and any(
                 isinstance(part, UserPromptPart) for part in msg.parts
             ):
+                msg_text = get_user_prompt_text(msg)
                 # Check if an identical user prompt is already in original_messages
                 already_exists = any(
                     isinstance(existing, ModelRequest)
-                    and any(isinstance(p, UserPromptPart) for p in existing.parts)
-                    and existing.parts == msg.parts
+                    and get_user_prompt_text(existing) == msg_text
                     for existing in original_messages[
                         -5:
                     ]  # Check last 5 messages for efficiency
@@ -881,6 +1040,13 @@ class AgentManager(Widget):
 
             deduplicated_new_messages.append(msg)
 
+        # Mark file resume prompts as internal (hidden from UI)
+        # When file_contents is provided, the prompt is system-generated, not user input
+        if file_contents:
+            deduplicated_new_messages = self._mark_as_internal_prompts(
+                deduplicated_new_messages
+            )
+
         self.ui_message_history = original_messages + deduplicated_new_messages
 
         # Get file operations early so we can use them for contextual messages
@@ -895,6 +1061,48 @@ class AgentManager(Widget):
             },
         )
 
+        # Check if there are file requests (takes priority over clarifying questions)
+        # But ignore file_requests if we just provided file_contents (prevents infinite loops)
+        if agent_response.file_requests and not file_contents:
+            logger.info(
+                f"Agent requested {len(agent_response.file_requests)} files to be loaded"
+            )
+
+            # Set pending state
+            self._file_request_pending = True
+            self._pending_file_requests = agent_response.file_requests
+
+            # Add agent's response as hint if present
+            if agent_response.response:
+                self.ui_message_history.append(
+                    HintMessage(message=agent_response.response)
+                )
+
+            # Add file loading indicator
+            files_list = "\n".join(f"- `{p}`" for p in agent_response.file_requests)
+            self.ui_message_history.append(
+                HintMessage(message=f"📁 Loading requested files:\n{files_list}")
+            )
+
+            # Post UI update with hint messages
+            self._post_messages_updated([])
+
+            # Post event to TUI to load files and resume
+            self.post_message(
+                FileRequestPendingMessage(
+                    file_paths=agent_response.file_requests,
+                    response_text=agent_response.response,
+                )
+            )
+
+            return result
+        elif agent_response.file_requests and file_contents:
+            # We just provided files, ignore any new file_requests to prevent loops
+            logger.debug(
+                "Ignoring file_requests (files were just provided): %s",
+                agent_response.file_requests,
+            )
+
         # Check if there are clarifying questions
         if agent_response.clarifying_questions:
             logger.info(
@@ -918,11 +1126,16 @@ class AgentManager(Widget):
                 self.ui_message_history.append(
                     HintMessage(message=f"💡 {agent_response.clarifying_questions[0]}")
                 )
+                # Add plan hint for Drafting mode (Planning mode uses PlanPanelWidget)
+                self._maybe_add_plan_hint_drafting_mode(deps)
             else:
                 # Multiple questions (2+) - enter Q&A mode
                 self._qa_questions = agent_response.clarifying_questions
                 self._qa_mode_active = True
 
+                # In Drafting mode, show plan BEFORE Q&A questions (without "Shall I continue?")
+                self._maybe_add_plan_hint_drafting_mode(deps, in_qa_mode=True)
+
                 # Show intro with list, then first question
                 questions_list_with_intro = (
                     f"I have {len(agent_response.clarifying_questions)} questions:\n\n"
@@ -984,6 +1197,9 @@ class AgentManager(Widget):
                         HintMessage(message="✅ Task completed")
                     )
 
+            # Add plan hint for Drafting mode (Planning mode uses PlanPanelWidget)
+            self._maybe_add_plan_hint_drafting_mode(deps)
+
             # Post UI update immediately so user sees the response without delay
             # (file operations will be posted after compaction to avoid duplicates)
             logger.debug("Posting immediate UI update with hint message")
@@ -1082,6 +1298,11 @@ class AgentManager(Widget):
         else:
             partial_parts = []
 
+        # Wrap stream with cancellable iterator for responsive ESC handling
+        deps = _ctx.deps
+        if deps.cancellation_event:
+            stream = CancellableStreamIterator(stream, deps.cancellation_event)
+
         async for event in stream:
             try:
                 if isinstance(event, PartStartEvent):
@@ -1359,6 +1580,47 @@ class AgentManager(Widget):
                 # Common path is a file, show parent directory
                 return f"📁 Modified {num_files} files in: `{path_obj.parent}`"
 
+    def _maybe_add_plan_hint_drafting_mode(
+        self, deps: AgentDeps, in_qa_mode: bool = False
+    ) -> None:
+        """Add execution plan hint for router agent in Drafting mode only.
+
+        In Drafting mode, there's no PlanPanelWidget, so we show the plan
+        in the chat history with a "Shall I continue?" prompt (unless in Q&A mode).
+
+        In Planning mode, the PlanPanelWidget handles plan display.
+
+        Args:
+            deps: Agent dependencies (may be RouterDeps for router agent)
+            in_qa_mode: If True, skip the "Shall I continue?" prompt since user
+                       needs to answer Q&A questions first.
+        """
+        if self._current_agent_type != AgentType.ROUTER:
+            return
+
+        if not isinstance(deps, RouterDeps):
+            return
+
+        # Only show plan hints in Drafting mode
+        # Planning mode uses PlanPanelWidget instead
+        if deps.router_mode != RouterMode.DRAFTING:
+            return
+
+        if deps.current_plan is None:
+            return
+
+        plan_display = deps.current_plan.format_for_display()
+
+        # In drafting mode, if plan is not complete and NOT in Q&A mode,
+        # prompt user to continue
+        if not deps.current_plan.is_complete() and not in_qa_mode:
+            plan_display += "\n\n**Shall I continue?**"
+
+        logger.debug("Adding plan hint to UI history (Drafting mode)")
+        self.ui_message_history.append(
+            HintMessage(message=f"**Current Plan**\n\n{plan_display}")
+        )
+
     def _post_messages_updated(
         self, file_operations: list[FileOperation] | None = None
     ) -> None:
@@ -1371,6 +1633,43 @@ class AgentManager(Widget):
             )
         )
 
+    def _mark_as_internal_prompts(
+        self,
+        messages: list[ModelRequest | ModelResponse | HintMessage],
+    ) -> list[ModelRequest | ModelResponse | HintMessage]:
+        """Mark UserPromptPart as InternalPromptPart for system-generated prompts.
+
+        Used when file_contents is provided - the resume prompt is system-generated,
+        not actual user input, and should be hidden from the UI.
+
+        Args:
+            messages: List of messages that may contain user prompts to mark as internal
+
+        Returns:
+            List of messages with UserPromptPart converted to InternalPromptPart
+        """
+        result: list[ModelRequest | ModelResponse | HintMessage] = []
+        for msg in messages:
+            if isinstance(msg, ModelRequest):
+                new_parts: list[ModelRequestPart] = []
+                for part in msg.parts:
+                    if isinstance(part, UserPromptPart) and not isinstance(
+                        part, InternalPromptPart
+                    ):
+                        # Convert to InternalPromptPart
+                        new_parts.append(
+                            InternalPromptPart(
+                                content=part.content,
+                                timestamp=part.timestamp,
+                            )
+                        )
+                    else:
+                        new_parts.append(part)
+                result.append(ModelRequest(parts=new_parts))
+            else:
+                result.append(msg)
+        return result
+
     def _filter_system_prompts(
         self, messages: list[ModelMessage | HintMessage]
     ) -> list[ModelMessage | HintMessage]:

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

@@ -38,8 +38,11 @@ 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,
 )
@@ -69,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)
 
@@ -94,6 +102,7 @@ async def add_system_status_message(
     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,
@@ -198,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:

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

@@ -377,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:
@@ -521,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]