shotgun-sh 0.1.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

shotgun/agents/plan.py

@@ -4,7 +4,6 @@ from functools import partial
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
 )
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
@@ -19,14 +18,14 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions, AgentType
+from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
 
-def create_plan_agent(
+async def create_plan_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a plan agent with artifact management capabilities.
 
     Args:
@@ -40,7 +39,7 @@ def create_plan_agent(
     # Use partial to create system prompt function for plan agent
     system_prompt_fn = partial(build_agent_system_prompt, "plan")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,
@@ -52,11 +51,11 @@ def create_plan_agent(
 
 
 async def run_plan_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     goal: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     """Create or update a plan based on the given goal using artifacts.
 
     Args:

shotgun/agents/research.py

@@ -4,7 +4,6 @@ from functools import partial
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
 )
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import (
@@ -21,15 +20,15 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions, AgentType
+from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 from .tools import get_available_web_search_tools
 
 logger = get_logger(__name__)
 
 
-def create_research_agent(
+async def create_research_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a research agent with web search and artifact management capabilities.
 
     Args:
@@ -42,7 +41,7 @@ def create_research_agent(
     logger.debug("Initializing research agent")
 
     # Get available web search tools based on configured API keys
-    web_search_tools = get_available_web_search_tools()
+    web_search_tools = await get_available_web_search_tools()
     if web_search_tools:
         logger.info(
             "Research agent configured with %d web search tool(s)",
@@ -54,7 +53,7 @@ def create_research_agent(
     # Use partial to create system prompt function for research agent
     system_prompt_fn = partial(build_agent_system_prompt, "research")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,
@@ -66,11 +65,11 @@ def create_research_agent(
 
 
 async def run_research_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     query: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     """Perform research on the given query and update research artifacts.
 
     Args:

shotgun/agents/specify.py

@@ -4,7 +4,6 @@ from functools import partial
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
 )
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
@@ -19,14 +18,14 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions, AgentType
+from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
 
-def create_specify_agent(
+async def create_specify_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a specify agent with artifact management capabilities.
 
     Args:
@@ -40,7 +39,7 @@ def create_specify_agent(
     # Use partial to create system prompt function for specify agent
     system_prompt_fn = partial(build_agent_system_prompt, "specify")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         load_codebase_understanding_tools=True,
@@ -52,11 +51,11 @@ def create_specify_agent(
 
 
 async def run_specify_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     requirement: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     """Create or update specifications based on the given requirement.
 
     Args:

shotgun/agents/tasks.py

@@ -4,7 +4,6 @@ from functools import partial
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
 )
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
@@ -19,14 +18,14 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions, AgentType
+from .models import AgentDeps, AgentResponse, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
 
-def create_tasks_agent(
+async def create_tasks_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a tasks agent with file management capabilities.
 
     Args:
@@ -40,7 +39,7 @@ def create_tasks_agent(
     # Use partial to create system prompt function for tasks agent
     system_prompt_fn = partial(build_agent_system_prompt, "tasks")
 
-    agent, deps = create_base_agent(
+    agent, deps = await create_base_agent(
         system_prompt_fn,
         agent_runtime_options,
         provider=provider,
@@ -50,11 +49,11 @@ def create_tasks_agent(
 
 
 async def run_tasks_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     instruction: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     """Create or update tasks based on the given instruction.
 
     Args:

shotgun/agents/tools/__init__.py

@@ -8,7 +8,6 @@ from .codebase import (
     retrieve_code,
 )
 from .file_management import append_file, read_file, write_file
-from .user_interaction import ask_user
 from .web_search import (
     anthropic_web_search_tool,
     gemini_web_search_tool,
@@ -21,7 +20,6 @@ __all__ = [
     "anthropic_web_search_tool",
     "gemini_web_search_tool",
     "get_available_web_search_tools",
-    "ask_user",
     "read_file",
     "write_file",
     "append_file",

shotgun/agents/tools/codebase/codebase_shell.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 
 from .models import ShellCommandResult
@@ -48,6 +49,11 @@ DANGEROUS_PATTERNS = [
 ]
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Running shell",
+    key_arg="command",
+)
 async def codebase_shell(
     ctx: RunContext[AgentDeps],
     command: str,

shotgun/agents/tools/codebase/directory_lister.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 
 from .models import DirectoryListResult
@@ -12,6 +13,11 @@ from .models import DirectoryListResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Listing directory",
+    key_arg="directory",
+)
 async def directory_lister(
     ctx: RunContext[AgentDeps], graph_id: str, directory: str = "."
 ) -> DirectoryListResult:

shotgun/agents/tools/codebase/file_read.py

@@ -2,9 +2,11 @@
 
 from pathlib import Path
 
+import aiofiles
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.core.language_config import get_language_config
 from shotgun.logging_config import get_logger
 
@@ -13,6 +15,11 @@ from .models import FileReadResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Reading file",
+    key_arg="file_path",
+)
 async def file_read(
     ctx: RunContext[AgentDeps], graph_id: str, file_path: str
 ) -> FileReadResult:
@@ -87,7 +94,8 @@ async def file_read(
         # Read file contents
         encoding_used = "utf-8"
         try:
-            content = full_file_path.read_text(encoding="utf-8")
+            async with aiofiles.open(full_file_path, encoding="utf-8") as f:
+                content = await f.read()
             size_bytes = full_file_path.stat().st_size
 
             logger.debug(
@@ -113,7 +121,8 @@ async def file_read(
             try:
                 # Try with different encoding
                 encoding_used = "latin-1"
-                content = full_file_path.read_text(encoding="latin-1")
+                async with aiofiles.open(full_file_path, encoding="latin-1") as f:
+                    content = await f.read()
                 size_bytes = full_file_path.stat().st_size
 
                 # Detect language from file extension

shotgun/agents/tools/codebase/query_graph.py

@@ -3,6 +3,7 @@
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.codebase.models import QueryType
 from shotgun.logging_config import get_logger
 
@@ -11,6 +12,11 @@ from .models import QueryGraphResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Querying code",
+    key_arg="query",
+)
 async def query_graph(
     ctx: RunContext[AgentDeps], graph_id: str, query: str
 ) -> QueryGraphResult:

shotgun/agents/tools/codebase/retrieve_code.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from pydantic_ai import RunContext
 
 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.logging_config import get_logger
@@ -14,6 +15,11 @@ from .models import CodeSnippetResult
 logger = get_logger(__name__)
 
 
+@register_tool(
+    category=ToolCategory.CODEBASE_UNDERSTANDING,
+    display_text="Retrieving code",
+    key_arg="qualified_name",
+)
 async def retrieve_code(
     ctx: RunContext[AgentDeps], graph_id: str, qualified_name: str
 ) -> CodeSnippetResult:

shotgun/agents/tools/file_management.py

@@ -6,20 +6,30 @@ These tools are restricted to the .shotgun directory for security.
 from pathlib import Path
 from typing import Literal
 
+import aiofiles
+import aiofiles.os
 from pydantic_ai import RunContext
 
 from shotgun.agents.models import AgentDeps, AgentType, FileOperationType
+from shotgun.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
 from shotgun.utils.file_system_utils import get_shotgun_base_path
 
 logger = get_logger(__name__)
 
 # Map agent modes to their allowed directories/files (in workflow order)
-AGENT_DIRECTORIES = {
-    AgentType.RESEARCH: "research.md",
-    AgentType.SPECIFY: "specification.md",
-    AgentType.PLAN: "plan.md",
-    AgentType.TASKS: "tasks.md",
+# Values can be:
+# - A Path: exact file (e.g., Path("research.md"))
+# - A list of Paths: multiple allowed files/directories (e.g., [Path("specification.md"), Path("contracts")])
+# - "*": any file except protected files (for export agent)
+AGENT_DIRECTORIES: dict[AgentType, str | Path | list[Path]] = {
+    AgentType.RESEARCH: Path("research.md"),
+    AgentType.SPECIFY: [
+        Path("specification.md"),
+        Path("contracts"),
+    ],  # Specify can write specs and contract files
+    AgentType.PLAN: Path("plan.md"),
+    AgentType.TASKS: Path("tasks.md"),
     AgentType.EXPORT: "*",  # Export agent can write anywhere except protected files
 }
 
@@ -60,13 +70,52 @@ def _validate_agent_scoped_path(filename: str, agent_mode: AgentType | None) ->
             # Allow writing anywhere else in .shotgun directory
             full_path = (base_path / filename).resolve()
         else:
-            # For other agents, only allow writing to their specific file
-            allowed_file = AGENT_DIRECTORIES[agent_mode]
-            if filename != allowed_file:
+            # For other agents, check if they have access to the requested file
+            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
+                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 if filename matches any allowed path
+            is_allowed = False
+            for allowed_path in allowed_paths:
+                allowed_str = str(allowed_path)
+
+                # Check if it's a directory (no .md extension or suffix)
+                # Directories: Path("contracts") has no suffix, files: Path("spec.md") has .md suffix
+                if not allowed_path.suffix or (
+                    allowed_path.suffix and not allowed_str.endswith(".md")
+                ):
+                    # Directory - allow any file within this directory
+                    # Check both "contracts/file.py" and "contracts" prefix
+                    if (
+                        filename.startswith(allowed_str + "/")
+                        or filename == allowed_str
+                    ):
+                        is_allowed = True
+                        break
+                else:
+                    # Exact file match
+                    if filename == allowed_str:
+                        is_allowed = True
+                        break
+
+            if not is_allowed:
+                allowed_display = ", ".join(f"'{p}'" for p in allowed_paths)
                 raise ValueError(
-                    f"{agent_mode.value.capitalize()} agent can only write to '{allowed_file}'. "
+                    f"{agent_mode.value.capitalize()} agent can only write to {allowed_display}. "
                     f"Attempted to write to '{filename}'"
                 )
+
             full_path = (base_path / filename).resolve()
     else:
         # No agent mode specified, fall back to old validation
@@ -111,6 +160,11 @@ def _validate_shotgun_path(filename: str) -> Path:
     return full_path
 
 
+@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.
 
@@ -129,10 +183,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
     try:
         file_path = _validate_shotgun_path(filename)
 
-        if not file_path.exists():
+        if not await aiofiles.os.path.exists(file_path):
             raise FileNotFoundError(f"File not found: {filename}")
 
-        content = file_path.read_text(encoding="utf-8")
+        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)
         return content
 
@@ -142,6 +197,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
         return error_msg
 
 
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Writing file",
+    key_arg="filename",
+)
 async def write_file(
     ctx: RunContext[AgentDeps],
     filename: str,
@@ -176,21 +236,22 @@ async def write_file(
         else:
             operation = (
                 FileOperationType.CREATED
-                if not file_path.exists()
+                if not await aiofiles.os.path.exists(file_path)
                 else FileOperationType.UPDATED
             )
 
         # Ensure parent directory exists
-        file_path.parent.mkdir(parents=True, exist_ok=True)
+        await aiofiles.os.makedirs(file_path.parent, exist_ok=True)
 
         # Write content
         if mode == "a":
-            with open(file_path, "a", encoding="utf-8") as f:
-                f.write(content)
+            async with aiofiles.open(file_path, "a", encoding="utf-8") as f:
+                await f.write(content)
             logger.debug("📄 Appended %d characters to %s", len(content), filename)
             result = f"Successfully appended {len(content)} characters to {filename}"
         else:
-            file_path.write_text(content, encoding="utf-8")
+            async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
+                await f.write(content)
             logger.debug("📄 Wrote %d characters to %s", len(content), filename)
             result = f"Successfully wrote {len(content)} characters to {filename}"
 
@@ -205,6 +266,11 @@ async def write_file(
         return error_msg
 
 
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Appending to file",
+    key_arg="filename",
+)
 async def append_file(ctx: RunContext[AgentDeps], filename: str, content: str) -> str:
     """Append content to a file in the .shotgun directory.