shotgun-sh 0.2.3.dev2__py3-none-any.whl → 0.2.11.dev1__py3-none-any.whl

shotgun/agents/history/history_processors.py

@@ -127,6 +127,7 @@ calculate_max_summarization_tokens = _calculate_max_summarization_tokens
 async def token_limit_compactor(
     ctx: ContextProtocol,
     messages: list[ModelMessage],
+    force: bool = False,
 ) -> list[ModelMessage]:
     """Compact message history based on token limits with incremental processing.
 
@@ -139,6 +140,7 @@ async def token_limit_compactor(
     Args:
         ctx: Run context with usage information and dependencies
         messages: Current conversation history
+        force: If True, force compaction even if below token threshold
 
     Returns:
         Compacted list of messages within token limits
@@ -169,7 +171,7 @@ async def token_limit_compactor(
         )
 
         # Only do incremental compaction if post-summary conversation exceeds threshold
-        if post_summary_tokens < max_tokens:
+        if post_summary_tokens < max_tokens and not force:
             logger.debug(
                 f"Post-summary conversation under threshold ({post_summary_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -340,6 +342,7 @@ async def token_limit_compactor(
             else 0
         )
 
+        # Track incremental compaction with simple metrics (fast, no token counting)
         track_event(
             "context_compaction_triggered",
             {
@@ -352,6 +355,10 @@ async def token_limit_compactor(
                 "agent_mode": deps.agent_mode.value
                 if hasattr(deps, "agent_mode") and deps.agent_mode
                 else "unknown",
+                # Model and provider info (no computation needed)
+                "model_name": deps.llm_model.name.value,
+                "provider": deps.llm_model.provider.value,
+                "key_provider": deps.llm_model.key_provider.value,
             },
         )
 
@@ -368,7 +375,7 @@ async def token_limit_compactor(
         )
 
         # Only do full compaction if total conversation exceeds threshold
-        if total_tokens < max_tokens:
+        if total_tokens < max_tokens and not force:
             logger.debug(
                 f"Total conversation under threshold ({total_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -468,6 +475,7 @@ async def _full_compaction(
     tokens_before = current_tokens  # Already calculated above
     tokens_after = summary_usage.output_tokens if summary_usage else 0
 
+    # Track full compaction with simple metrics (fast, no token counting)
     track_event(
         "context_compaction_triggered",
         {
@@ -480,6 +488,10 @@ async def _full_compaction(
             "agent_mode": deps.agent_mode.value
             if hasattr(deps, "agent_mode") and deps.agent_mode
             else "unknown",
+            # Model and provider info (no computation needed)
+            "model_name": deps.llm_model.name.value,
+            "provider": deps.llm_model.provider.value,
+            "key_provider": deps.llm_model.key_provider.value,
         },
     )
 

shotgun/agents/history/token_counting/anthropic.py

@@ -1,9 +1,10 @@
 """Anthropic token counting using official client."""
 
+import logfire
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config.models import KeyProvider
-from shotgun.llm_proxy import create_anthropic_proxy_client
+from shotgun.llm_proxy import create_anthropic_proxy_provider
 from shotgun.logging_config import get_logger
 
 from .base import TokenCounter, extract_text_from_messages
@@ -36,19 +37,28 @@ class AnthropicTokenCounter(TokenCounter):
         try:
             if key_provider == KeyProvider.SHOTGUN:
                 # Use LiteLLM proxy for Shotgun Account
-                # Proxies to Anthropic's token counting API
-                self.client = create_anthropic_proxy_client(api_key)
+                # Get async client from AnthropicProvider
+                provider = create_anthropic_proxy_provider(api_key)
+                self.client = provider.client
                 logger.debug(
-                    f"Initialized Anthropic token counter for {model_name} via LiteLLM proxy"
+                    f"Initialized async Anthropic token counter for {model_name} via LiteLLM proxy"
                 )
             else:
-                # Direct Anthropic API for BYOK
-                self.client = anthropic.Anthropic(api_key=api_key)
+                # Direct Anthropic API for BYOK - use async client
+                self.client = anthropic.AsyncAnthropic(api_key=api_key)
                 logger.debug(
-                    f"Initialized Anthropic token counter for {model_name} via direct API"
+                    f"Initialized async Anthropic token counter for {model_name} via direct API"
                 )
         except Exception as e:
-            raise RuntimeError("Failed to initialize Anthropic client") from e
+            logfire.exception(
+                f"Failed to initialize Anthropic token counter for {model_name}",
+                model_name=model_name,
+                key_provider=key_provider.value,
+                exception_type=type(e).__name__,
+            )
+            raise RuntimeError(
+                f"Failed to initialize Anthropic async client for {model_name}: {type(e).__name__}: {str(e)}"
+            ) from e
 
     async def count_tokens(self, text: str) -> int:
         """Count tokens using Anthropic's official API (async).
@@ -64,13 +74,25 @@ class AnthropicTokenCounter(TokenCounter):
         """
         try:
             # Anthropic API expects messages format and model parameter
-            result = self.client.messages.count_tokens(
+            # Use await with async client
+            result = await self.client.messages.count_tokens(
                 messages=[{"role": "user", "content": text}], model=self.model_name
             )
             return result.input_tokens
         except Exception as e:
+            # Create a preview of the text for logging (truncated to avoid huge logs)
+            text_preview = text[:100] + "..." if len(text) > 100 else text
+
+            logfire.exception(
+                f"Anthropic token counting failed for {self.model_name}",
+                model_name=self.model_name,
+                text_length=len(text),
+                text_preview=text_preview,
+                exception_type=type(e).__name__,
+                exception_message=str(e),
+            )
             raise RuntimeError(
-                f"Anthropic token counting API failed for {self.model_name}"
+                f"Anthropic token counting API failed for {self.model_name}: {type(e).__name__}: {str(e)}"
             ) from e
 
     async def count_message_tokens(self, messages: list[ModelMessage]) -> int:

shotgun/agents/models.py

@@ -19,6 +19,30 @@ if TYPE_CHECKING:
     from shotgun.codebase.service import CodebaseService
 
 
+class AgentResponse(BaseModel):
+    """Structured response from an agent with optional clarifying questions.
+
+    This model provides a consistent response format for all agents:
+    - response: The main response text (can be empty if only asking questions)
+    - clarifying_questions: Optional list of questions to ask the user
+
+    When clarifying_questions is provided, the agent expects to receive
+    answers before continuing its work. This replaces the ask_questions tool.
+    """
+
+    response: str = Field(
+        description="The agent's response text. Always respond with some text summarizing what happened, whats next, etc.",
+    )
+    clarifying_questions: list[str] | None = Field(
+        default=None,
+        description="""
+Optional list of clarifying questions to ask the user.
+- Single question: Shown as a non-blocking suggestion (user can answer or continue with other prompts)
+- Multiple questions (2+): Asked sequentially in Q&A mode (blocks input until all answered or cancelled)
+""",
+    )
+
+
 class AgentType(StrEnum):
     """Enumeration for available agent types."""
 
@@ -73,6 +97,30 @@ class UserQuestion(BaseModel):
     )
 
 
+class MultipleUserQuestions(BaseModel):
+    """Multiple questions to ask the user sequentially."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    questions: list[str] = Field(
+        description="List of questions to ask the user",
+    )
+    current_index: int = Field(
+        default=0,
+        description="Current question index being asked",
+    )
+    answers: list[str] = Field(
+        default_factory=list,
+        description="Accumulated answers from the user",
+    )
+    tool_call_id: str = Field(
+        description="Tool call id",
+    )
+    result: Future[UserAnswer] = Field(
+        description="Future that will contain all answers formatted as Q&A pairs"
+    )
+
+
 class AgentRuntimeOptions(BaseModel):
     """User interface options for agents."""
 
@@ -100,9 +148,9 @@ class AgentRuntimeOptions(BaseModel):
         description="Maximum number of iterations for agent loops",
     )
 
-    queue: Queue[UserQuestion] = Field(
+    queue: Queue[UserQuestion | MultipleUserQuestions] = Field(
         default_factory=Queue,
-        description="Queue for storing user responses",
+        description="Queue for storing user questions (single or multiple)",
     )
 
     tasks: list[Future[UserAnswer]] = Field(

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(
     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:
@@ -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,7 +20,7 @@ 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__)
@@ -29,7 +28,7 @@ logger = get_logger(__name__)
 
 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:
@@ -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(
     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:
@@ -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(
     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:
@@ -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

@@ -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.language_config import get_language_config
 from shotgun.logging_config import get_logger
 
@@ -13,6 +14,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:

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

@@ -9,17 +9,25 @@ from typing import Literal
 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 +68,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 +158,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.
 
@@ -142,6 +194,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,
@@ -205,6 +262,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.