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

shotgun/agents/router/tools/delegation_tools.py

@@ -6,6 +6,7 @@ These tools allow the Router to delegate work to specialized sub-agents
 Sub-agents run with isolated message histories to prevent context window bloat.
 """
 
+import time
 from collections.abc import Awaitable, Callable
 from typing import Any
 
@@ -33,6 +34,7 @@ from shotgun.agents.specify import create_specify_agent, run_specify_agent
 from shotgun.agents.tasks import create_tasks_agent, run_tasks_agent
 from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 
 logger = get_logger(__name__)
 
@@ -253,6 +255,9 @@ async def _run_sub_agent(
     # Set up SubAgentContext so sub-agent knows it's being orchestrated
     sub_agent_deps.sub_agent_context = _build_sub_agent_context(deps)
 
+    # Propagate cancellation event for responsive ESC handling in sub-agents
+    sub_agent_deps.cancellation_event = deps.cancellation_event
+
     # Clear sub-agent's file tracker for fresh tracking
     sub_agent_deps.file_tracker.clear()
 
@@ -260,11 +265,23 @@ async def _run_sub_agent(
     deps.active_sub_agent = agent_type
     logger.info("Delegating to %s agent: %s", agent_type.value, task[:100])
 
+    # Track delegation start time and event
+    start_time = time.time()
+    track_event(
+        "delegation_started",
+        {
+            "target_agent": agent_type.value,
+            "task_length": len(task),
+            "has_context_hint": context_hint is not None,
+        },
+    )
+
     # Get the run function for this agent type
     _, run_fn = AGENT_FACTORIES[agent_type]
 
     # Retry loop for transient errors
     last_error: BaseException | None = None
+    retries_attempted = 0
     for attempt in range(MAX_RETRIES + 1):
         try:
             # Run sub-agent with isolated message history and streaming support
@@ -286,6 +303,11 @@ async def _run_sub_agent(
                 op.file_path for op in sub_agent_deps.file_tracker.operations
             ]
 
+            # Extract files found (used by FileReadAgent)
+            files_found: list[str] = []
+            if result and result.output and result.output.files_found:
+                files_found = result.output.files_found
+
             # Check for clarifying questions
             has_questions = False
             questions: list[str] = []
@@ -294,9 +316,21 @@ async def _run_sub_agent(
                 questions = result.output.clarifying_questions
 
             logger.info(
-                "Sub-agent %s completed. Files modified: %s",
+                "Sub-agent %s completed. Files modified: %s, files found: %s",
                 agent_type.value,
                 files_modified,
+                files_found,
+            )
+
+            # Track delegation completion metric
+            track_event(
+                "delegation_completed",
+                {
+                    "target_agent": agent_type.value,
+                    "files_modified_count": len(files_modified),
+                    "has_questions": has_questions,
+                    "duration_seconds": round(time.time() - start_time, 2),
+                },
             )
 
             # Clear active_sub_agent
@@ -306,12 +340,14 @@ async def _run_sub_agent(
                 success=True,
                 response=response_text,
                 files_modified=files_modified,
+                files_found=files_found,
                 has_questions=has_questions,
                 questions=questions,
             )
 
         except Exception as e:
             last_error = e
+            retries_attempted = attempt
             if _is_retryable_error(e) and attempt < MAX_RETRIES:
                 logger.warning(
                     "Sub-agent %s failed (attempt %d/%d), retrying: %s",
@@ -329,8 +365,26 @@ async def _run_sub_agent(
                 attempt + 1,
                 str(e),
             )
+            # Clear the agent from cache on failure so next request gets a fresh agent
+            # This is especially important for request_limit errors
+            if agent_type in deps.sub_agent_cache:
+                del deps.sub_agent_cache[agent_type]
+                logger.debug(
+                    "Cleared %s from sub_agent_cache after failure", agent_type.value
+                )
             break
 
+    # Track delegation failure metric
+    track_event(
+        "delegation_failed",
+        {
+            "target_agent": agent_type.value,
+            "error_type": type(last_error).__name__ if last_error else "Unknown",
+            "retries_attempted": retries_attempted,
+            "duration_seconds": round(time.time() - start_time, 2),
+        },
+    )
+
     # Clear active_sub_agent on failure
     deps.active_sub_agent = None
 

shotgun/agents/router/tools/plan_tools.py

@@ -25,6 +25,7 @@ from shotgun.agents.router.models import (
 )
 from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 
 logger = get_logger(__name__)
 
@@ -44,7 +45,7 @@ def _notify_plan_changed(deps: RouterDeps) -> None:
 
 @register_tool(
     category=ToolCategory.PLANNING,
-    display_text="Creating execution plan",
+    display_text="Creating a Drafting plan",
     key_arg="input",
 )
 async def create_plan(
@@ -62,7 +63,7 @@ async def create_plan(
     Returns:
         ToolResult indicating success or failure
     """
-    logger.debug("Creating execution plan with goal: %s", input.goal)
+    logger.debug("Creating a Drafting plan with goal: %s", input.goal)
 
     # Convert step inputs to ExecutionStep objects
     steps = [
@@ -107,6 +108,17 @@ async def create_plan(
         input.goal,
     )
 
+    # Track plan creation metric
+    track_event(
+        "plan_created",
+        {
+            "step_count": len(steps),
+            "goal_preview": input.goal[:100],
+            "requires_approval": plan.needs_approval(),
+            "router_mode": ctx.deps.router_mode.value,
+        },
+    )
+
     _notify_plan_changed(ctx.deps)
 
     # Return different message based on whether approval is needed
@@ -151,11 +163,22 @@ async def mark_step_done(
         )
 
     # Find the step by ID
-    for _i, step in enumerate(plan.steps):
+    for step_index, step in enumerate(plan.steps):
         if step.id == input.step_id:
             step.done = True
             logger.info("Marked step '%s' as done", input.step_id)
 
+            # Track step completion metric
+            completed_count = sum(1 for s in plan.steps if s.done)
+            track_event(
+                "plan_step_completed",
+                {
+                    "step_position": step_index + 1,  # 1-indexed for human readability
+                    "total_steps": len(plan.steps),
+                    "steps_remaining": len(plan.steps) - completed_count,
+                },
+            )
+
             # Advance current_step_index to next incomplete step
             while (
                 plan.current_step_index < len(plan.steps)
@@ -167,6 +190,20 @@ async def mark_step_done(
             if plan.is_complete():
                 ctx.deps.is_executing = False
                 logger.debug("Plan complete, is_executing=False")
+
+                # Track plan completion metric
+                track_event(
+                    "plan_completed",
+                    {
+                        "step_count": len(plan.steps),
+                    },
+                )
+
+                # Set pending completion for Drafting mode
+                # The TUI will detect this and show the completion message
+                if ctx.deps.router_mode == RouterMode.DRAFTING:
+                    ctx.deps.pending_completion = True
+                    logger.debug("Set pending_completion=True for drafting mode")
             # Set pending checkpoint for Planning mode
             # The TUI will detect this and show the StepCheckpointWidget
             elif ctx.deps.router_mode == RouterMode.PLANNING:
@@ -184,10 +221,28 @@ async def mark_step_done(
 
             _notify_plan_changed(ctx.deps)
 
-            return ToolResult(
-                success=True,
-                message=f"Marked step '{step.title}' as complete.",
-            )
+            # Return different messages based on mode and plan state
+            if plan.is_complete():
+                return ToolResult(
+                    success=True,
+                    message=f"Marked step '{step.title}' as complete.\n\n"
+                    "All plan steps are now complete. You may return your final response.",
+                )
+            elif ctx.deps.router_mode == RouterMode.DRAFTING:
+                next_step = plan.current_step()
+                return ToolResult(
+                    success=True,
+                    message=f"Marked step '{step.title}' as complete.\n\n"
+                    f"NEXT STEP: {next_step.title if next_step else 'None'}\n"
+                    "IMPORTANT: Do NOT call final_result. "
+                    "Immediately delegate the next step.",
+                )
+            else:
+                # Planning mode - checkpoint will be shown
+                return ToolResult(
+                    success=True,
+                    message=f"Marked step '{step.title}' as complete.",
+                )
 
     return ToolResult(
         success=False,
@@ -242,6 +297,15 @@ async def add_step(ctx: RunContext[RouterDeps], input: AddStepInput) -> ToolResu
         plan.steps.append(new_step)
         logger.info("Appended step '%s' to end of plan", input.step.id)
 
+        # Track step added metric
+        track_event(
+            "plan_step_added",
+            {
+                "new_step_count": len(plan.steps),
+                "position": len(plan.steps),  # Appended at end
+            },
+        )
+
         _notify_plan_changed(ctx.deps)
 
         return ToolResult(
@@ -259,6 +323,15 @@ async def add_step(ctx: RunContext[RouterDeps], input: AddStepInput) -> ToolResu
                 input.after_step_id,
             )
 
+            # Track step added metric
+            track_event(
+                "plan_step_added",
+                {
+                    "new_step_count": len(plan.steps),
+                    "position": i + 2,  # 1-indexed, inserted after position i+1
+                },
+            )
+
             _notify_plan_changed(ctx.deps)
 
             return ToolResult(
@@ -303,6 +376,14 @@ async def remove_step(
             removed_step = plan.steps.pop(i)
             logger.info("Removed step '%s' from plan", input.step_id)
 
+            # Track step removed metric
+            track_event(
+                "plan_step_removed",
+                {
+                    "new_step_count": len(plan.steps),
+                },
+            )
+
             # Adjust current_step_index if needed
             if plan.current_step_index > i:
                 plan.current_step_index -= 1

shotgun/agents/runner.py

@@ -13,6 +13,7 @@ from openai import APIStatusError as OpenAIAPIStatusError
 from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior
 
 from shotgun.agents.error.models import AgentErrorContext
+from shotgun.attachments import FileAttachment
 from shotgun.exceptions import (
     AgentCancelledException,
     BudgetExceededException,
@@ -29,6 +30,8 @@ from shotgun.exceptions import (
 )
 
 if TYPE_CHECKING:
+    from pydantic_ai import BinaryContent
+
     from shotgun.agents.agent_manager import AgentManager
 
 logger = logging.getLogger(__name__)
@@ -66,11 +69,19 @@ class AgentRunner:
         """
         self.agent_manager = agent_manager
 
-    async def run(self, prompt: str) -> None:
+    async def run(
+        self,
+        prompt: str,
+        attachment: FileAttachment | None = None,
+        file_contents: list[tuple[str, "BinaryContent"]] | None = None,
+    ) -> None:
         """Run the agent with the given prompt.
 
         Args:
             prompt: The user's prompt/query
+            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.
 
         Raises:
             Custom exceptions for different error types:
@@ -88,7 +99,11 @@ class AgentRunner:
             - UnknownAgentException: Unknown/unclassified error
         """
         try:
-            await self.agent_manager.run(prompt=prompt)
+            await self.agent_manager.run(
+                prompt=prompt,
+                attachment=attachment,
+                file_contents=file_contents,
+            )
 
         except asyncio.CancelledError as e:
             # User cancelled - wrap and re-raise as our custom exception

shotgun/agents/tools/__init__.py

@@ -8,6 +8,11 @@ from .codebase import (
     retrieve_code,
 )
 from .file_management import append_file, read_file, write_file
+from .markdown_tools import (
+    insert_markdown_section,
+    remove_markdown_section,
+    replace_markdown_section,
+)
 from .web_search import (
     anthropic_web_search_tool,
     gemini_web_search_tool,
@@ -23,6 +28,9 @@ __all__ = [
     "read_file",
     "write_file",
     "append_file",
+    "replace_markdown_section",
+    "insert_markdown_section",
+    "remove_markdown_section",
     # Codebase understanding tools
     "query_graph",
     "retrieve_code",

shotgun/agents/tools/codebase/directory_lister.py

@@ -19,60 +19,48 @@ logger = get_logger(__name__)
     key_arg="directory",
 )
 async def directory_lister(
-    ctx: RunContext[AgentDeps], graph_id: str, directory: str = "."
+    ctx: RunContext[AgentDeps], directory: str = ".", graph_id: str = ""
 ) -> DirectoryListResult:
-    """List directory contents in codebase.
+    """List directory contents in codebase or current working directory.
 
     Args:
         ctx: RunContext containing AgentDeps with codebase service
-        graph_id: Graph ID to identify the repository
-        directory: Path to directory relative to repository root (default: ".")
+        directory: Path to directory relative to repository root or CWD (default: ".")
+        graph_id: Graph ID to identify the repository (optional - uses CWD if not provided)
 
     Returns:
         DirectoryListResult with formatted output via __str__
     """
-    logger.debug("🔧 Listing directory: %s in graph %s", directory, graph_id)
+    logger.debug("🔧 Listing directory: %s in graph %s", directory, graph_id or "(CWD)")
 
     try:
-        if not ctx.deps.codebase_service:
-            return DirectoryListResult(
-                success=False,
-                directory=directory,
-                full_path="",
-                error="No codebase indexed",
-            )
-
-        # Get the graph to find the repository path
-        try:
-            graphs = await ctx.deps.codebase_service.list_graphs()
-            graph = next((g for g in graphs if g.graph_id == graph_id), None)
-        except Exception as e:
-            logger.error("Error getting graph: %s", e)
-            return DirectoryListResult(
-                success=False,
-                directory=directory,
-                full_path="",
-                error=f"Could not find graph with ID '{graph_id}'",
-            )
-
-        if not graph:
-            return DirectoryListResult(
-                success=False,
-                directory=directory,
-                full_path="",
-                error=f"Graph '{graph_id}' not found",
-            )
-
-        # Validate the directory path is within the repository
-        repo_path = Path(graph.repo_path).resolve()
+        # Determine the root path - either from indexed codebase or CWD
+        repo_path: Path | None = None
+
+        if graph_id and ctx.deps.codebase_service:
+            # Try to get the graph to find the repository path
+            try:
+                graphs = await ctx.deps.codebase_service.list_graphs()
+                graph = next((g for g in graphs if g.graph_id == graph_id), None)
+                if graph:
+                    repo_path = Path(graph.repo_path).resolve()
+            except Exception as e:
+                logger.debug("Could not find graph '%s': %s", graph_id, e)
+
+        # Fall back to CWD if no graph found or no graph_id provided
+        if repo_path is None:
+            repo_path = Path.cwd().resolve()
+            logger.debug("📂 Using CWD as root: %s", repo_path)
+
+        # Validate the directory path is within the root
         full_dir_path = (repo_path / directory).resolve()
 
-        # Security check: ensure the resolved path is within the repository
+        # Security check: ensure the resolved path is within the root directory
         try:
             full_dir_path.relative_to(repo_path)
         except ValueError:
             error_msg = (
-                f"Access denied: Path '{directory}' is outside repository bounds"
+                f"Access denied: Path '{directory}' is outside allowed directory bounds"
             )
             logger.warning("🚨 Security violation attempt: %s", error_msg)
             return DirectoryListResult(
@@ -88,7 +76,7 @@ async def directory_lister(
                 success=False,
                 directory=directory,
                 full_path=str(full_dir_path),
-                error=f"Directory '{directory}' not found in repository",
+                error=f"Directory not found: {directory}",
             )
 
         if not full_dir_path.is_dir():

shotgun/agents/tools/codebase/file_read.py

@@ -21,57 +21,48 @@ logger = get_logger(__name__)
     key_arg="file_path",
 )
 async def file_read(
-    ctx: RunContext[AgentDeps], graph_id: str, file_path: str
+    ctx: RunContext[AgentDeps], file_path: str, graph_id: str = ""
 ) -> FileReadResult:
-    """Read file contents from codebase.
+    """Read file contents from codebase or current working directory.
 
     Args:
         ctx: RunContext containing AgentDeps with codebase service
-        graph_id: Graph ID to identify the repository
-        file_path: Path to file relative to repository root
+        file_path: Path to file relative to repository root or CWD
+        graph_id: Graph ID to identify the repository (optional - uses CWD if not provided)
 
     Returns:
         FileReadResult with formatted output via __str__
     """
-    logger.debug("🔧 Reading file: %s in graph %s", file_path, graph_id)
+    logger.debug("🔧 Reading file: %s in graph %s", file_path, graph_id or "(CWD)")
 
     try:
-        if not ctx.deps.codebase_service:
-            return FileReadResult(
-                success=False,
-                file_path=file_path,
-                error="No codebase indexed",
-            )
-
-        # Get the graph to find the repository path
-        try:
-            graphs = await ctx.deps.codebase_service.list_graphs()
-            graph = next((g for g in graphs if g.graph_id == graph_id), None)
-        except Exception as e:
-            logger.error("Error getting graph: %s", e)
-            return FileReadResult(
-                success=False,
-                file_path=file_path,
-                error=f"Could not find graph with ID '{graph_id}'",
-            )
-
-        if not graph:
-            return FileReadResult(
-                success=False,
-                file_path=file_path,
-                error=f"Graph '{graph_id}' not found",
-            )
+        # Determine the root path - either from indexed codebase or CWD
+        repo_path: Path | None = None
 
-        # Validate the file path is within the repository
-        repo_path = Path(graph.repo_path).resolve()
+        if graph_id and ctx.deps.codebase_service:
+            # Try to get the graph to find the repository path
+            try:
+                graphs = await ctx.deps.codebase_service.list_graphs()
+                graph = next((g for g in graphs if g.graph_id == graph_id), None)
+                if graph:
+                    repo_path = Path(graph.repo_path).resolve()
+            except Exception as e:
+                logger.debug("Could not find graph '%s': %s", graph_id, e)
+
+        # Fall back to CWD if no graph found or no graph_id provided
+        if repo_path is None:
+            repo_path = Path.cwd().resolve()
+            logger.debug("📂 Using CWD as root: %s", repo_path)
+
+        # Validate the file path is within the root
         full_file_path = (repo_path / file_path).resolve()
 
-        # Security check: ensure the resolved path is within the repository
+        # Security check: ensure the resolved path is within the root directory
         try:
             full_file_path.relative_to(repo_path)
         except ValueError:
             error_msg = (
-                f"Access denied: Path '{file_path}' is outside repository bounds"
+                f"Access denied: Path '{file_path}' is outside allowed directory bounds"
             )
             logger.warning("🚨 Security violation attempt: %s", error_msg)
             return FileReadResult(success=False, file_path=file_path, error=error_msg)
@@ -81,7 +72,7 @@ async def file_read(
             return FileReadResult(
                 success=False,
                 file_path=file_path,
-                error=f"File '{file_path}' not found in repository",
+                error=f"File not found: {file_path}",
             )
 
         if full_file_path.is_dir():

shotgun/agents/tools/codebase/query_graph.py

@@ -4,6 +4,7 @@ from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
 from shotgun.agents.tools.registry import ToolCategory, register_tool
+from shotgun.codebase.indexing_state import IndexingState
 from shotgun.codebase.models import QueryType
 from shotgun.logging_config import get_logger
 
@@ -40,6 +41,14 @@ async def query_graph(
                 error="No codebase indexed",
             )
 
+        # Check if graph is currently being indexed
+        if ctx.deps.codebase_service.indexing.is_active(graph_id):
+            return QueryGraphResult(
+                success=False,
+                query=query,
+                error=IndexingState.INDEXING_IN_PROGRESS_ERROR,
+            )
+
         # Execute natural language query
         result = await ctx.deps.codebase_service.execute_query(
             graph_id=graph_id,

shotgun/agents/tools/codebase/retrieve_code.py

@@ -8,6 +8,7 @@ from shotgun.agents.models import AgentDeps
 from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.core.code_retrieval import retrieve_code_by_qualified_name
 from shotgun.codebase.core.language_config import get_language_config
+from shotgun.codebase.indexing_state import IndexingState
 from shotgun.logging_config import get_logger
 
 from .models import CodeSnippetResult
@@ -43,6 +44,14 @@ async def retrieve_code(
                 error="No codebase indexed",
             )
 
+        # Check if graph is currently being indexed
+        if ctx.deps.codebase_service.indexing.is_active(graph_id):
+            return CodeSnippetResult(
+                found=False,
+                qualified_name=qualified_name,
+                error=IndexingState.INDEXING_IN_PROGRESS_ERROR,
+            )
+
         # Use the existing code retrieval functionality
         code_snippet = await retrieve_code_by_qualified_name(
             manager=ctx.deps.codebase_service.manager,

shotgun/agents/tools/file_management.py

@@ -163,19 +163,33 @@ def _validate_shotgun_path(filename: str) -> Path:
     return full_path
 
 
+# Binary file extensions that should be loaded via file_requests
+BINARY_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".gif", ".webp"}
+
+
 @register_tool(
     category=ToolCategory.ARTIFACT_MANAGEMENT,
     display_text="Reading file",
     key_arg="filename",
 )
 async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
-    """Read a file from the .shotgun directory.
+    """Read a TEXT file from the .shotgun directory.
+
+    IMPORTANT: This tool is for TEXT files only (.md, .txt, .json, etc.).
+
+    For BINARY files (PDFs, images), DO NOT use this tool. Instead:
+    - Use file_requests in your response to load binary files
+    - Example: {"response": "Let me check that.", "file_requests": ["/path/to/file.pdf"]}
+
+    Binary file extensions that require file_requests instead:
+    - .pdf, .png, .jpg, .jpeg, .gif, .webp
 
     Args:
         filename: Relative path to file within .shotgun directory
 
     Returns:
-        File contents as string
+        File contents as string. For binary files, returns instructions
+        with the absolute path to use in file_requests.
 
     Raises:
         ValueError: If path is outside .shotgun directory
@@ -189,6 +203,22 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
         if not await aiofiles.os.path.exists(file_path):
             raise FileNotFoundError(f"File not found: {filename}")
 
+        # Check if it's a binary file type (PDF, image)
+        suffix = file_path.suffix.lower()
+        if suffix in BINARY_EXTENSIONS:
+            # Return info for the agent to use file_requests
+            logger.debug(
+                "📎 Binary file detected (%s), returning path for file_requests: %s",
+                suffix,
+                file_path,
+            )
+            return (
+                f"This is a binary file ({suffix}) that cannot be read as text. "
+                f"To view its contents, include the absolute path in your "
+                f"`file_requests` response field:\n\n"
+                f"Absolute path: {file_path}"
+            )
+
         async with aiofiles.open(file_path, encoding="utf-8") as f:
             content = await f.read()
         logger.debug("📄 Read %d characters from %s", len(content), filename)

shotgun/agents/tools/file_read_tools/__init__.py

@@ -0,0 +1,7 @@
+"""File reading tools for the FileRead agent."""
+
+from shotgun.agents.tools.file_read_tools.multimodal_file_read import (
+    multimodal_file_read,
+)
+
+__all__ = ["multimodal_file_read"]