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

shotgun/agents/{history → conversation/history}/history_processors.py

@@ -13,6 +13,7 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
+from shotgun.agents.conversation.filters import filter_orphaned_tool_responses
 from shotgun.agents.llm import shotgun_model_request
 from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.agents.models import AgentDeps
@@ -21,7 +22,7 @@ from shotgun.logging_config import get_logger
 from shotgun.posthog_telemetry import track_event
 from shotgun.prompts import PromptLoader
 
-from .constants import SUMMARY_MARKER, TOKEN_LIMIT_RATIO
+from .constants import CHUNK_SAFE_RATIO, SUMMARY_MARKER, TOKEN_LIMIT_RATIO
 from .context_extraction import extract_context_from_messages
 from .history_building import ensure_ends_with_model_request
 from .message_utils import (
@@ -38,7 +39,7 @@ from .token_estimation import (
 )
 
 if TYPE_CHECKING:
-    pass
+    from . import chunking
 
 
 class ContextProtocol(Protocol):
@@ -416,6 +417,9 @@ async def token_limit_compactor(
             compacted_messages, messages
         )
 
+        # Filter out orphaned tool responses (tool responses without tool calls)
+        compacted_messages = filter_orphaned_tool_responses(compacted_messages)
+
         logger.debug(
             f"Incremental compaction complete: {len(messages)} -> {len(compacted_messages)} messages"
         )
@@ -489,10 +493,32 @@ async def _full_compaction(
     deps: AgentDeps,
     messages: list[ModelMessage],
 ) -> list[ModelMessage]:
-    """Perform full compaction for first-time summarization."""
+    """Perform full compaction for first-time summarization.
+
+    If the conversation is too large for single-pass compaction, delegates
+    to chunked compaction which breaks the conversation into logical chunks.
+    """
     # Extract context from all messages
     context = extract_context_from_messages(messages)
 
+    # Check if context would exceed model limit for compaction request
+    # We use CHUNK_SAFE_RATIO (70%) to leave room for prompt overhead
+    max_safe_input = int(deps.llm_model.max_input_tokens * CHUNK_SAFE_RATIO)
+
+    # Estimate context tokens
+    context_request: list[ModelMessage] = [ModelRequest.user_text_prompt(context)]
+    context_tokens = await estimate_tokens_from_messages(
+        context_request, deps.llm_model
+    )
+
+    if context_tokens > max_safe_input:
+        # Context too large for single-pass compaction - use chunked approach
+        logger.info(
+            f"Context ({context_tokens:,} tokens) exceeds safe limit "
+            f"({max_safe_input:,} tokens), using chunked compaction"
+        )
+        return await _chunked_compaction(deps, messages)
+
     # Use regular summarization prompt
     summarization_prompt = prompt_loader.render("history/summarization.j2")
     request_messages: list[ModelMessage] = [
@@ -565,6 +591,9 @@ async def _full_compaction(
     # Ensure history ends with ModelRequest for PydanticAI compatibility
     compacted_messages = ensure_ends_with_model_request(compacted_messages, messages)
 
+    # Filter out orphaned tool responses (tool responses without tool calls)
+    compacted_messages = filter_orphaned_tool_responses(compacted_messages)
+
     # Track full compaction event
     messages_before = len(messages)
     messages_after = len(compacted_messages)
@@ -592,3 +621,238 @@ async def _full_compaction(
     )
 
     return compacted_messages
+
+
+async def _chunked_compaction(
+    deps: AgentDeps,
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Perform chunked compaction for oversized conversations.
+
+    Breaks the conversation into logical chunks, summarizes each sequentially,
+    then combines the summaries into a master summary.
+    """
+    from .chunking import chunk_messages_for_compaction
+
+    # Split into chunks and retention window
+    chunks, retained_messages = await chunk_messages_for_compaction(
+        messages, deps.llm_model
+    )
+
+    if not chunks:
+        # No chunks to summarize (conversation too small), return retained messages
+        logger.debug("No chunks to summarize, returning retained messages")
+        return retained_messages
+
+    # Track chunked compaction
+    total_chunks = len(chunks)
+    logger.info(f"Starting chunked compaction: {total_chunks} chunks to process")
+
+    # Summarize each chunk sequentially
+    chunk_summaries: list[str] = []
+    for chunk in chunks:
+        try:
+            summary = await _summarize_chunk(chunk, total_chunks, deps)
+            chunk_summaries.append(summary)
+            logger.debug(
+                f"Chunk {chunk.chunk_index + 1}/{total_chunks} summarized successfully"
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to summarize chunk {chunk.chunk_index + 1}/{total_chunks}: {e}"
+            )
+            # Continue with other chunks - we'll note the gap in fusion
+            chunk_summaries.append(
+                f"[Chunk {chunk.chunk_index + 1} summary unavailable]"
+            )
+
+    # Combine summaries into master summary
+    if len(chunk_summaries) == 1:
+        final_summary = chunk_summaries[0]
+    else:
+        final_summary = await _combine_chunk_summaries(chunk_summaries, deps)
+
+    # Build final compacted history
+    compacted = _build_chunked_compaction_result(
+        final_summary, messages, retained_messages, deps
+    )
+
+    # Track chunked compaction event
+    track_event(
+        "chunked_compaction_triggered",
+        {
+            "num_chunks": total_chunks,
+            "chunks_succeeded": sum(
+                1 for s in chunk_summaries if not s.startswith("[Chunk")
+            ),
+            "retention_window_size": len(retained_messages),
+            "model_name": deps.llm_model.name.value,
+            "provider": deps.llm_model.provider.value,
+        },
+    )
+
+    return compacted
+
+
+async def _summarize_chunk(
+    chunk: "chunking.Chunk",
+    total_chunks: int,
+    deps: AgentDeps,
+) -> str:
+    """Summarize a single chunk of messages."""
+    chunk_messages = chunk.get_all_messages()
+    context = extract_context_from_messages(chunk_messages)
+
+    # Use chunk summarization template
+    chunk_prompt = prompt_loader.render(
+        "history/chunk_summarization.j2",
+        chunk_index=chunk.chunk_index + 1,
+        total_chunks=total_chunks,
+        chunk_content=context,
+    )
+
+    request_messages: list[ModelMessage] = [
+        ModelRequest.user_text_prompt(context, instructions=chunk_prompt)
+    ]
+
+    max_tokens = await calculate_max_summarization_tokens(
+        deps.llm_model, request_messages
+    )
+
+    log_summarization_request(
+        deps.llm_model,
+        max_tokens,
+        chunk_prompt,
+        context[:500] + "..." if len(context) > 500 else context,
+        f"CHUNK_{chunk.chunk_index + 1}",
+    )
+
+    response = await shotgun_model_request(
+        model_config=deps.llm_model,
+        messages=request_messages,
+        model_settings=ModelSettings(max_tokens=max_tokens),
+    )
+
+    log_summarization_response(response, f"CHUNK_{chunk.chunk_index + 1}")
+
+    if response.parts and isinstance(response.parts[0], TextPart):
+        return response.parts[0].content
+    return ""
+
+
+async def _combine_chunk_summaries(
+    summaries: list[str],
+    deps: AgentDeps,
+) -> str:
+    """Combine multiple chunk summaries into a unified summary."""
+    # Check if combined summaries exceed limit (may need recursive combination)
+    combined_text = "\n\n".join(summaries)
+    combined_request: list[ModelMessage] = [
+        ModelRequest.user_text_prompt(combined_text)
+    ]
+    combined_tokens = await estimate_tokens_from_messages(
+        combined_request, deps.llm_model
+    )
+
+    max_safe_input = int(deps.llm_model.max_input_tokens * CHUNK_SAFE_RATIO)
+
+    if combined_tokens > max_safe_input:
+        # Recursive: split summaries in half and combine each half first
+        logger.warning(
+            f"Combined summaries too large ({combined_tokens:,} tokens), "
+            f"applying recursive combination"
+        )
+        mid = len(summaries) // 2
+        first_half = await _combine_chunk_summaries(summaries[:mid], deps)
+        second_half = await _combine_chunk_summaries(summaries[mid:], deps)
+        summaries = [first_half, second_half]
+
+    # Use combination template
+    combine_prompt = prompt_loader.render(
+        "history/combine_summaries.j2",
+        num_summaries=len(summaries),
+        chunk_summaries=summaries,
+    )
+
+    request_messages: list[ModelMessage] = [
+        ModelRequest.user_text_prompt(
+            "\n\n---\n\n".join(summaries), instructions=combine_prompt
+        )
+    ]
+
+    max_tokens = await calculate_max_summarization_tokens(
+        deps.llm_model, request_messages
+    )
+
+    log_summarization_request(
+        deps.llm_model,
+        max_tokens,
+        combine_prompt,
+        f"[{len(summaries)} summaries to combine]",
+        "COMBINE",
+    )
+
+    response = await shotgun_model_request(
+        model_config=deps.llm_model,
+        messages=request_messages,
+        model_settings=ModelSettings(max_tokens=max_tokens),
+    )
+
+    log_summarization_response(response, "COMBINE")
+
+    if response.parts and isinstance(response.parts[0], TextPart):
+        return response.parts[0].content
+    return ""
+
+
+def _build_chunked_compaction_result(
+    final_summary: str,
+    original_messages: list[ModelMessage],
+    retained_messages: list[ModelMessage],
+    deps: AgentDeps,
+) -> list[ModelMessage]:
+    """Build the final compacted history from chunked compaction."""
+    from pydantic_ai.messages import ModelRequestPart
+
+    # Extract system context from original messages
+    agent_prompt = get_agent_system_prompt(original_messages) or ""
+    system_status = get_latest_system_status(original_messages) or ""
+    first_user = get_first_user_request(original_messages) or ""
+
+    # Create marked summary
+    summary_part = TextPart(content=f"{SUMMARY_MARKER} {final_summary}")
+    summary_message = ModelResponse(parts=[summary_part])
+
+    # Build compacted structure
+    compacted: list[ModelMessage] = []
+
+    # Initial request with system context
+    parts: list[ModelRequestPart] = []
+    if agent_prompt:
+        parts.append(AgentSystemPrompt(content=agent_prompt))
+    if system_status:
+        parts.append(SystemStatusPrompt(content=system_status))
+    if first_user:
+        parts.append(UserPromptPart(content=first_user))
+
+    if parts:
+        compacted.append(ModelRequest(parts=parts))
+
+    # Add summary
+    compacted.append(summary_message)
+
+    # Add retained messages (recent context)
+    compacted.extend(retained_messages)
+
+    # Ensure ends with ModelRequest for PydanticAI compatibility
+    compacted = ensure_ends_with_model_request(compacted, original_messages)
+
+    # Filter orphaned tool responses
+    compacted = filter_orphaned_tool_responses(compacted)
+
+    logger.info(
+        f"Chunked compaction complete: {len(original_messages)} messages -> "
+        f"{len(compacted)} messages (retained {len(retained_messages)} recent)"
+    )
+
+    return compacted

shotgun/agents/{history → conversation/history}/token_counting/anthropic.py

@@ -1,6 +1,7 @@
 """Anthropic token counting using official client."""
 
 import logfire
+from anthropic import APIStatusError
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config.models import KeyProvider
@@ -103,6 +104,13 @@ class AnthropicTokenCounter(TokenCounter):
                 exception_type=type(e).__name__,
                 exception_message=str(e),
             )
+
+            # Re-raise API errors directly so they can be classified by the runner
+            # This allows proper error classification for BYOK users (authentication, rate limits, etc.)
+            if isinstance(e, APIStatusError):
+                raise
+
+            # Only wrap library-level errors in RuntimeError
             raise RuntimeError(
                 f"Anthropic token counting API failed for {self.model_name}: {type(e).__name__}: {str(e)}"
             ) from e

shotgun/agents/{conversation_manager.py → conversation/manager.py}

@@ -11,7 +11,7 @@ from shotgun.logging_config import get_logger
 from shotgun.utils import get_shotgun_home
 from shotgun.utils.file_system_utils import async_copy_file
 
-from .conversation_history import ConversationHistory
+from .models import ConversationHistory
 
 logger = get_logger(__name__)
 

shotgun/agents/{conversation_history.py → conversation/models.py}

@@ -1,7 +1,5 @@
-"""Models and utilities for persisting TUI conversation history."""
+"""Models for persisting TUI conversation history."""
 
-import json
-import logging
 from datetime import datetime
 from typing import Any, cast
 
@@ -16,101 +14,15 @@ from pydantic_core import to_jsonable_python
 
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage
 
-__all__ = ["HintMessage", "ConversationHistory"]
-
-logger = logging.getLogger(__name__)
+from .filters import (
+    filter_incomplete_messages,
+    filter_orphaned_tool_responses,
+    is_tool_call_complete,
+)
 
 SerializedMessage = dict[str, Any]
 
 
-def is_tool_call_complete(tool_call: ToolCallPart) -> bool:
-    """Check if a tool call has valid, complete JSON arguments.
-
-    Args:
-        tool_call: The tool call part to validate
-
-    Returns:
-        True if the tool call args are valid JSON, False otherwise
-    """
-    if tool_call.args is None:
-        return True  # No args is valid
-
-    if isinstance(tool_call.args, dict):
-        return True  # Already parsed dict is valid
-
-    if not isinstance(tool_call.args, str):
-        return False
-
-    # Try to parse the JSON string
-    try:
-        json.loads(tool_call.args)
-        return True
-    except (json.JSONDecodeError, ValueError) as e:
-        # Log incomplete tool call detection
-        args_preview = (
-            tool_call.args[:100] + "..."
-            if len(tool_call.args) > 100
-            else tool_call.args
-        )
-        logger.info(
-            "Detected incomplete tool call in validation",
-            extra={
-                "tool_name": tool_call.tool_name,
-                "tool_call_id": tool_call.tool_call_id,
-                "args_preview": args_preview,
-                "error": str(e),
-            },
-        )
-        return False
-
-
-def filter_incomplete_messages(messages: list[ModelMessage]) -> list[ModelMessage]:
-    """Filter out messages with incomplete tool calls.
-
-    Args:
-        messages: List of messages to filter
-
-    Returns:
-        List of messages with only complete tool calls
-    """
-    filtered: list[ModelMessage] = []
-    filtered_count = 0
-    filtered_tool_names: list[str] = []
-
-    for message in messages:
-        # Only check ModelResponse messages for tool calls
-        if not isinstance(message, ModelResponse):
-            filtered.append(message)
-            continue
-
-        # Check if any tool calls are incomplete
-        has_incomplete_tool_call = False
-        for part in message.parts:
-            if isinstance(part, ToolCallPart) and not is_tool_call_complete(part):
-                has_incomplete_tool_call = True
-                filtered_tool_names.append(part.tool_name)
-                break
-
-        # Only include messages without incomplete tool calls
-        if not has_incomplete_tool_call:
-            filtered.append(message)
-        else:
-            filtered_count += 1
-
-    # Log if any messages were filtered
-    if filtered_count > 0:
-        logger.info(
-            "Filtered incomplete messages before saving",
-            extra={
-                "filtered_count": filtered_count,
-                "total_messages": len(messages),
-                "filtered_tool_names": filtered_tool_names,
-            },
-        )
-
-    return filtered
-
-
 class ConversationState(BaseModel):
     """Represents the complete state of a conversation in memory."""
 
@@ -144,6 +56,8 @@ class ConversationHistory(BaseModel):
         """
         # Filter out messages with incomplete tool calls to prevent corruption
         filtered_messages = filter_incomplete_messages(messages)
+        # Filter out orphaned tool responses (tool responses without tool calls)
+        filtered_messages = filter_orphaned_tool_responses(filtered_messages)
 
         # Serialize ModelMessage list to JSON-serializable format
         self.agent_history = to_jsonable_python(

shotgun/agents/error/__init__.py

@@ -0,0 +1,11 @@
+"""Agent error handling module.
+
+This module provides the AgentErrorContext model used by AgentRunner
+for error classification.
+"""
+
+from shotgun.agents.error.models import AgentErrorContext
+
+__all__ = [
+    "AgentErrorContext",
+]

shotgun/agents/error/models.py

@@ -0,0 +1,19 @@
+"""Pydantic models for agent error handling."""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AgentErrorContext(BaseModel):
+    """Context information needed to classify and handle agent errors.
+
+    Attributes:
+        exception: The exception that was raised
+        is_shotgun_account: Whether the user is using a Shotgun Account
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    exception: Any = Field(...)
+    is_shotgun_account: bool

shotgun/agents/export.py

@@ -2,16 +2,15 @@
 
 from functools import partial
 
-from pydantic_ai import (
-    Agent,
-)
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config import ProviderType
+from shotgun.agents.models import ShotgunAgent
 from shotgun.logging_config import get_logger
 
 from .common import (
+    EventStreamHandler,
     add_system_status_message,
     build_agent_system_prompt,
     create_base_agent,
@@ -25,7 +24,7 @@ logger = get_logger(__name__)
 
 async def create_export_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create an export agent with file management capabilities.
 
     Args:
@@ -49,39 +48,39 @@ async def create_export_agent(
 
 
 async def run_export_agent(
-    agent: Agent[AgentDeps, AgentResponse],
-    instruction: str,
+    agent: ShotgunAgent,
+    prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
-    """Export artifacts based on the given instruction.
+    """Export artifacts based on the given prompt.
 
     Args:
         agent: The configured export agent
-        instruction: The export instruction
+        prompt: The export prompt
         deps: Agent dependencies
         message_history: Optional message history for conversation continuity
+        event_stream_handler: Optional callback for streaming events
 
     Returns:
         AgentRunResult containing the export process output
     """
-    logger.debug("📤 Starting export for instruction: %s", instruction)
+    logger.debug("📤 Starting export for prompt: %s", prompt)
 
     message_history = await add_system_status_message(deps, message_history)
 
-    # Let the agent use its tools to read existing artifacts and export them
-    full_prompt = f"Export artifacts or findings based on: {instruction}"
-
     try:
         # Create usage limits for responsible API usage
         usage_limits = create_usage_limits()
 
         result = await run_agent(
             agent=agent,
-            prompt=full_prompt,
+            prompt=prompt,
             deps=deps,
             message_history=message_history,
             usage_limits=usage_limits,
+            event_stream_handler=event_stream_handler,
         )
 
         logger.debug("✅ Export completed successfully")

shotgun/agents/models.py

@@ -6,7 +6,7 @@ from collections.abc import Callable
 from datetime import datetime
 from enum import StrEnum
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai import RunContext
@@ -16,9 +16,57 @@ from shotgun.agents.usage_manager import SessionUsageManager, get_session_usage_
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
+    from pydantic_ai import Agent
+
+    from shotgun.agents.router.models import RouterDeps
     from shotgun.codebase.service import CodebaseService
 
 
+class SubAgentContext(BaseModel):
+    """
+    Context passed to sub-agents so they know they're being orchestrated.
+
+    When sub-agents receive this context, they should:
+    - Be more concise (router handles user communication)
+    - Focus on their specific task
+    - Return structured results
+    """
+
+    is_router_delegated: bool = Field(
+        default=True, description="Always True when passed to sub-agent"
+    )
+    plan_goal: str = Field(
+        default="", description="High-level goal from execution plan"
+    )
+    current_step_id: str | None = Field(
+        default=None, description="ID of the current execution step"
+    )
+    current_step_title: str | None = Field(
+        default=None, description="Title of the current execution step"
+    )
+
+
+class AgentSystemPromptContext(BaseModel):
+    """Context passed to agent system prompt templates.
+
+    This model standardizes the context variables passed to Jinja2 templates
+    when rendering agent system prompts. Using a model makes it easier to
+    test template context construction and ensures type safety.
+    """
+
+    interactive_mode: bool = Field(
+        description="Whether the agent is running in interactive mode"
+    )
+    mode: str = Field(description="The agent type (research, specify, plan, etc.)")
+    sub_agent_context: SubAgentContext | None = Field(
+        default=None, description="Context when running as a sub-agent of the router"
+    )
+    router_mode: str | None = Field(
+        default=None,
+        description="Router mode value (planning/drafting) if router agent",
+    )
+
+
 class AgentResponse(BaseModel):
     """Structured response from an agent with optional clarifying questions.
 
@@ -51,6 +99,7 @@ class AgentType(StrEnum):
     PLAN = "plan"
     TASKS = "tasks"
     EXPORT = "export"
+    ROUTER = "router"
 
 
 class PipelineConfigEntry(BaseModel):
@@ -319,6 +368,11 @@ class AgentDeps(AgentRuntimeOptions):
         description="Current agent mode for file scoping",
     )
 
+    sub_agent_context: SubAgentContext | None = Field(
+        default=None,
+        description="Context when agent is delegated to by router",
+    )
+
 
 # Rebuild model to resolve forward references after imports are available
 try:
@@ -328,3 +382,14 @@ try:
 except ImportError:
     # Services may not be available in all contexts
     pass
+
+
+# Type alias for the standard agent type used throughout the codebase
+ShotgunAgent: TypeAlias = "Agent[AgentDeps, AgentResponse]"
+
+# Type alias for router agent (uses RouterDeps which extends AgentDeps)
+# Note: Agent is contravariant in deps, so RouterAgent is NOT a subtype of ShotgunAgent
+RouterAgent: TypeAlias = "Agent[RouterDeps, AgentResponse]"
+
+# Union type for any agent type (used in AgentManager)
+AnyAgent: TypeAlias = "ShotgunAgent | RouterAgent"