shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.3.3.dev1__py3-none-any.whl

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

@@ -1,7 +1,9 @@
 """History processors for managing conversation history in Shotgun agents."""
 
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Protocol
 
+from anthropic import APIStatusError
 from pydantic_ai import ModelSettings
 from pydantic_ai.messages import (
     ModelMessage,
@@ -11,14 +13,16 @@ 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
+from shotgun.exceptions import ContextSizeLimitExceeded
 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 (
@@ -35,7 +39,7 @@ from .token_estimation import (
 )
 
 if TYPE_CHECKING:
-    pass
+    from . import chunking
 
 
 class ContextProtocol(Protocol):
@@ -51,6 +55,86 @@ logger = get_logger(__name__)
 prompt_loader = PromptLoader()
 
 
+async def _safe_token_estimation(
+    estimation_func: Callable[..., Awaitable[int]],
+    model_name: str,
+    max_tokens: int,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Safely estimate tokens with proper error handling.
+
+    Wraps token estimation functions to handle failures gracefully.
+    Only RuntimeError (from token counters) is wrapped in ContextSizeLimitExceeded.
+    Other errors (network, auth) are allowed to bubble up.
+
+    Args:
+        estimation_func: Async function that estimates tokens
+        model_name: Name of the model for error messages
+        max_tokens: Maximum tokens for the model
+        *args: Arguments to pass to estimation_func
+        **kwargs: Keyword arguments to pass to estimation_func
+
+    Returns:
+        Token count from estimation_func
+
+    Raises:
+        ContextSizeLimitExceeded: If token counting fails with RuntimeError
+        Exception: Any other exceptions from estimation_func
+    """
+    try:
+        return await estimation_func(*args, **kwargs)
+    except Exception as e:
+        # Log the error with full context
+        logger.warning(
+            f"Token counting failed for {model_name}",
+            extra={
+                "error_type": type(e).__name__,
+                "error_message": str(e),
+                "model": model_name,
+            },
+        )
+
+        # Token counting behavior with oversized context (verified via testing):
+        #
+        # 1. OpenAI/tiktoken:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: ValueError, KeyError, AttributeError, SSLError (file/cache issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 2. Gemini/SentencePiece:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: RuntimeError, IOError, TypeError (file/model loading issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 3. Anthropic API:
+        #    - Successfully counts large token counts (tested with 752K tokens, no error)
+        #    - Only enforces 32 MB request size limit (not token count)
+        #    - Raises: APIStatusError(413) with error type 'request_too_large' for 32MB+ requests
+        #    - Other API errors: APIConnectionError, RateLimitError, APIStatusError (4xx/5xx)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # IMPORTANT: No provider raises errors for "too many tokens" during counting.
+        # Token count validation happens separately by comparing count to max_input_tokens.
+        #
+        # We wrap RuntimeError (library-level failures from tiktoken/sentencepiece).
+        # We also wrap Anthropic's 413 error (request exceeds 32 MB) as it indicates
+        # context is effectively too large and needs user action to reduce it.
+        if isinstance(e, RuntimeError):
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Check for Anthropic's 32 MB request size limit (APIStatusError with status 413)
+        if isinstance(e, APIStatusError) and e.status_code == 413:
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Re-raise other exceptions (network errors, auth failures, etc.)
+        raise
+
+
 def is_summary_part(part: Any) -> bool:
     """Check if a message part is a compacted summary."""
     return isinstance(part, TextPart) and part.content.startswith(SUMMARY_MARKER)
@@ -127,6 +211,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 +224,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
@@ -155,9 +241,15 @@ async def token_limit_compactor(
 
     if last_summary_index is not None:
         # Check if post-summary conversation exceeds threshold for incremental compaction
-        post_summary_tokens = await estimate_post_summary_tokens(
-            messages, last_summary_index, deps.llm_model
+        post_summary_tokens = await _safe_token_estimation(
+            estimate_post_summary_tokens,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            last_summary_index,
+            deps.llm_model,
         )
+
         post_summary_percentage = (
             (post_summary_tokens / max_tokens) * 100 if max_tokens > 0 else 0
         )
@@ -169,7 +261,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"
@@ -325,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"
         )
@@ -340,6 +435,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 +448,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,
             },
         )
 
@@ -359,7 +459,14 @@ async def token_limit_compactor(
 
     else:
         # Check if total conversation exceeds threshold for full compaction
-        total_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
+        total_tokens = await _safe_token_estimation(
+            estimate_tokens_from_messages,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            deps.llm_model,
+        )
+
         total_percentage = (total_tokens / max_tokens) * 100 if max_tokens > 0 else 0
 
         logger.debug(
@@ -368,7 +475,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"
@@ -386,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] = [
@@ -462,12 +591,16 @@ 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)
     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,7 +613,246 @@ 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,
         },
     )
 
     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
@@ -72,11 +73,23 @@ class AnthropicTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If API call fails
         """
+        # Handle empty text to avoid unnecessary API calls
+        # Anthropic API requires non-empty content, so we need a strict check
+        if not text or not text.strip():
+            return 0
+
+        # Additional validation: ensure the text has actual content
+        # Some edge cases might have only whitespace or control characters
+        cleaned_text = text.strip()
+        if not cleaned_text:
+            return 0
+
         try:
             # Anthropic API expects messages format and model parameter
             # Use await with async client
             result = await self.client.messages.count_tokens(
-                messages=[{"role": "user", "content": text}], model=self.model_name
+                messages=[{"role": "user", "content": cleaned_text}],
+                model=self.model_name,
             )
             return result.input_tokens
         except Exception as e:
@@ -91,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
@@ -107,5 +127,9 @@ class AnthropicTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

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

@@ -56,12 +56,23 @@ def extract_text_from_messages(messages: list[ModelMessage]) -> str:
         if hasattr(message, "parts"):
             for part in message.parts:
                 if hasattr(part, "content") and isinstance(part.content, str):
-                    text_parts.append(part.content)
+                    # Only add non-empty content
+                    if part.content.strip():
+                        text_parts.append(part.content)
                 else:
                     # Handle non-text parts (tool calls, etc.)
-                    text_parts.append(str(part))
+                    part_str = str(part)
+                    if part_str.strip():
+                        text_parts.append(part_str)
         else:
             # Handle messages without parts
-            text_parts.append(str(message))
+            msg_str = str(message)
+            if msg_str.strip():
+                text_parts.append(msg_str)
+
+    # If no valid text parts found, return a minimal placeholder
+    # This ensures we never send completely empty content to APIs
+    if not text_parts:
+        return "."
 
     return "\n".join(text_parts)

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

@@ -57,9 +57,15 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary encoding
+        if not text or not text.strip():
+            return 0
+
         try:
             return len(self.encoding.encode(text))
-        except Exception as e:
+        except BaseException as e:
+            # Must catch BaseException to handle PanicException from tiktoken's Rust layer
+            # which can occur with extremely long texts. Regular Exception won't catch it.
             raise RuntimeError(
                 f"Failed to count tokens for OpenAI model {self.model_name}"
             ) from e
@@ -76,5 +82,9 @@ class OpenAITokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

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

@@ -88,6 +88,10 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty text to avoid unnecessary tokenization
+        if not text or not text.strip():
+            return 0
+
         await self._ensure_tokenizer()
 
         if self.sp is None:
@@ -115,5 +119,9 @@ class SentencePieceTokenCounter(TokenCounter):
         Raises:
             RuntimeError: If token counting fails
         """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
         total_text = extract_text_from_messages(messages)
         return await self.count_tokens(total_text)

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

@@ -3,6 +3,7 @@
 import hashlib
 from pathlib import Path
 
+import aiofiles
 import httpx
 
 from shotgun.logging_config import get_logger
@@ -78,7 +79,8 @@ async def download_gemini_tokenizer() -> Path:
 
         # Atomic write: write to temp file first, then rename
         temp_path = cache_path.with_suffix(".tmp")
-        temp_path.write_bytes(content)
+        async with aiofiles.open(temp_path, "wb") as f:
+            await f.write(content)
         temp_path.rename(cache_path)
 
         logger.info(f"Gemini tokenizer downloaded and cached at {cache_path}")

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

@@ -44,9 +44,6 @@ def get_token_counter(model_config: ModelConfig) -> TokenCounter:
 
     # Return cached instance if available
     if cache_key in _token_counter_cache:
-        logger.debug(
-            f"Reusing cached token counter for {model_config.provider.value}:{model_config.name}"
-        )
         return _token_counter_cache[cache_key]
 
     # Create new instance and cache it