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

shotgun/agents/context_analyzer/models.py

@@ -0,0 +1,212 @@
+"""Pydantic models for context analysis."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TokenAllocation(BaseModel):
+    """Token counts allocated from API usage data by message/tool type.
+
+    Used internally by ContextAnalyzer to track token distribution across
+    different message types and tool categories.
+    """
+
+    user: int = Field(ge=0, default=0, description="Tokens from user prompts")
+    agent_responses: int = Field(
+        ge=0, default=0, description="Tokens from agent text responses"
+    )
+    system_prompts: int = Field(
+        ge=0, default=0, description="Tokens from system prompts"
+    )
+    system_status: int = Field(
+        ge=0, default=0, description="Tokens from system status messages"
+    )
+    codebase_understanding: int = Field(
+        ge=0, default=0, description="Tokens from codebase understanding tools"
+    )
+    artifact_management: int = Field(
+        ge=0, default=0, description="Tokens from artifact management tools"
+    )
+    web_research: int = Field(
+        ge=0, default=0, description="Tokens from web research tools"
+    )
+    unknown: int = Field(ge=0, default=0, description="Tokens from uncategorized tools")
+
+
+class MessageTypeStats(BaseModel):
+    """Statistics for a specific message type."""
+
+    count: int = Field(ge=0, description="Number of messages of this type")
+    tokens: int = Field(ge=0, description="Total tokens consumed by this type")
+
+    @property
+    def avg_tokens(self) -> float:
+        """Calculate average tokens per message."""
+        return self.tokens / self.count if self.count > 0 else 0.0
+
+
+class ContextAnalysis(BaseModel):
+    """Complete analysis of conversation context composition."""
+
+    user_messages: MessageTypeStats
+    agent_responses: MessageTypeStats
+    system_prompts: MessageTypeStats
+    system_status: MessageTypeStats
+    codebase_understanding: MessageTypeStats
+    artifact_management: MessageTypeStats
+    web_research: MessageTypeStats
+    unknown: MessageTypeStats
+    hint_messages: MessageTypeStats
+    total_tokens: int = Field(ge=0, description="Total tokens including hints")
+    total_messages: int = Field(ge=0, description="Total message count including hints")
+    context_window: int = Field(ge=0, description="Model's maximum input tokens")
+    agent_context_tokens: int = Field(
+        ge=0,
+        description="Tokens that actually consume agent context (excluding UI-only)",
+    )
+    model_name: str = Field(description="Name of the model being used")
+    max_usable_tokens: int = Field(
+        ge=0, description="80% of max_input_tokens (usable limit)"
+    )
+    free_space_tokens: int = Field(
+        description="Remaining tokens available (negative if over capacity)"
+    )
+
+    def get_percentage(self, stats: MessageTypeStats) -> float:
+        """Calculate percentage of agent context tokens for a message type.
+
+        Args:
+            stats: Message type statistics to calculate percentage for
+
+        Returns:
+            Percentage of total agent context tokens (0-100)
+        """
+        return (
+            (stats.tokens / self.agent_context_tokens * 100)
+            if self.agent_context_tokens > 0
+            else 0.0
+        )
+
+
+class ContextCompositionTelemetry(BaseModel):
+    """Telemetry data for context composition tracking to PostHog."""
+
+    # Context usage
+    total_messages: int = Field(ge=0)
+    agent_context_tokens: int = Field(ge=0)
+    context_window: int = Field(ge=0)
+    max_usable_tokens: int = Field(ge=0)
+    free_space_tokens: int = Field(ge=0)
+    usage_percentage: float = Field(ge=0, le=100)
+
+    # Message type counts
+    user_messages_count: int = Field(ge=0)
+    agent_responses_count: int = Field(ge=0)
+    system_prompts_count: int = Field(ge=0)
+    system_status_count: int = Field(ge=0)
+    codebase_understanding_count: int = Field(ge=0)
+    artifact_management_count: int = Field(ge=0)
+    web_research_count: int = Field(ge=0)
+    unknown_tools_count: int = Field(ge=0)
+
+    # Token distribution percentages
+    user_messages_pct: float = Field(ge=0, le=100)
+    agent_responses_pct: float = Field(ge=0, le=100)
+    system_prompts_pct: float = Field(ge=0, le=100)
+    system_status_pct: float = Field(ge=0, le=100)
+    codebase_understanding_pct: float = Field(ge=0, le=100)
+    artifact_management_pct: float = Field(ge=0, le=100)
+    web_research_pct: float = Field(ge=0, le=100)
+    unknown_tools_pct: float = Field(ge=0, le=100)
+
+    # Compaction info
+    compaction_occurred: bool
+    messages_before_compaction: int | None = None
+    messages_after_compaction: int | None = None
+    compaction_reduction_pct: float | None = None
+
+    @classmethod
+    def from_analysis(
+        cls,
+        analysis: "ContextAnalysis",
+        compaction_occurred: bool = False,
+        messages_before_compaction: int | None = None,
+    ) -> "ContextCompositionTelemetry":
+        """Create telemetry from context analysis.
+
+        Args:
+            analysis: The context analysis to convert
+            compaction_occurred: Whether message compaction occurred
+            messages_before_compaction: Number of messages before compaction
+
+        Returns:
+            ContextCompositionTelemetry instance
+        """
+        total_messages = analysis.total_messages - analysis.hint_messages.count
+        usage_pct = (
+            round((analysis.agent_context_tokens / analysis.max_usable_tokens * 100), 1)
+            if analysis.max_usable_tokens > 0
+            else 0
+        )
+
+        # Calculate compaction metrics
+        messages_after: int | None = None
+        compaction_reduction_pct: float | None = None
+
+        if compaction_occurred and messages_before_compaction is not None:
+            messages_after = total_messages
+            if messages_before_compaction > 0:
+                compaction_reduction_pct = round(
+                    (1 - (total_messages / messages_before_compaction)) * 100, 1
+                )
+
+        return cls(
+            # Context usage
+            total_messages=total_messages,
+            agent_context_tokens=analysis.agent_context_tokens,
+            context_window=analysis.context_window,
+            max_usable_tokens=analysis.max_usable_tokens,
+            free_space_tokens=analysis.free_space_tokens,
+            usage_percentage=usage_pct,
+            # Message type counts
+            user_messages_count=analysis.user_messages.count,
+            agent_responses_count=analysis.agent_responses.count,
+            system_prompts_count=analysis.system_prompts.count,
+            system_status_count=analysis.system_status.count,
+            codebase_understanding_count=analysis.codebase_understanding.count,
+            artifact_management_count=analysis.artifact_management.count,
+            web_research_count=analysis.web_research.count,
+            unknown_tools_count=analysis.unknown.count,
+            # Token distribution percentages
+            user_messages_pct=round(analysis.get_percentage(analysis.user_messages), 1),
+            agent_responses_pct=round(
+                analysis.get_percentage(analysis.agent_responses), 1
+            ),
+            system_prompts_pct=round(
+                analysis.get_percentage(analysis.system_prompts), 1
+            ),
+            system_status_pct=round(analysis.get_percentage(analysis.system_status), 1),
+            codebase_understanding_pct=round(
+                analysis.get_percentage(analysis.codebase_understanding), 1
+            ),
+            artifact_management_pct=round(
+                analysis.get_percentage(analysis.artifact_management), 1
+            ),
+            web_research_pct=round(analysis.get_percentage(analysis.web_research), 1),
+            unknown_tools_pct=round(analysis.get_percentage(analysis.unknown), 1),
+            # Compaction info
+            compaction_occurred=compaction_occurred,
+            messages_before_compaction=messages_before_compaction,
+            messages_after_compaction=messages_after,
+            compaction_reduction_pct=compaction_reduction_pct,
+        )
+
+
+class ContextAnalysisOutput(BaseModel):
+    """Output format for context analysis with multiple representations."""
+
+    markdown: str = Field(description="Markdown-formatted analysis for display")
+    json_data: dict[str, Any] = Field(
+        description="JSON representation of analysis data"
+    )

shotgun/agents/conversation_history.py

@@ -1,5 +1,7 @@
 """Models and utilities for persisting TUI conversation history."""
 
+import json
+import logging
 from datetime import datetime
 from typing import Any, cast
 
@@ -7,14 +9,108 @@ from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai.messages import (
     ModelMessage,
     ModelMessagesTypeAdapter,
+    ModelResponse,
+    ToolCallPart,
 )
 from pydantic_core import to_jsonable_python
 
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage
 
+__all__ = ["HintMessage", "ConversationHistory"]
+
+logger = logging.getLogger(__name__)
+
 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."""
 
@@ -46,14 +142,41 @@ class ConversationHistory(BaseModel):
         Args:
             messages: List of ModelMessage objects to serialize and store
         """
+        # Filter out messages with incomplete tool calls to prevent corruption
+        filtered_messages = filter_incomplete_messages(messages)
+
         # Serialize ModelMessage list to JSON-serializable format
         self.agent_history = to_jsonable_python(
-            messages, fallback=lambda x: str(x), exclude_none=True
+            filtered_messages, fallback=lambda x: str(x), exclude_none=True
         )
 
     def set_ui_messages(self, messages: list[ModelMessage | HintMessage]) -> None:
         """Set ui_history from a list of UI messages."""
 
+        # Filter out ModelMessages with incomplete tool calls (keep all HintMessages)
+        # We need to maintain message order, so we'll check each message individually
+        filtered_messages: list[ModelMessage | HintMessage] = []
+
+        for msg in messages:
+            if isinstance(msg, HintMessage):
+                # Always keep hint messages
+                filtered_messages.append(msg)
+            elif isinstance(msg, ModelResponse):
+                # Check if this ModelResponse has incomplete tool calls
+                has_incomplete = False
+                for part in msg.parts:
+                    if isinstance(part, ToolCallPart) and not is_tool_call_complete(
+                        part
+                    ):
+                        has_incomplete = True
+                        break
+
+                if not has_incomplete:
+                    filtered_messages.append(msg)
+            else:
+                # Keep all other ModelMessage types (ModelRequest, etc.)
+                filtered_messages.append(msg)
+
         def _serialize_message(
             message: ModelMessage | HintMessage,
         ) -> Any:
@@ -68,7 +191,7 @@ class ConversationHistory(BaseModel):
                 payload.setdefault("message_type", "model")
             return payload
 
-        self.ui_history = [_serialize_message(msg) for msg in messages]
+        self.ui_history = [_serialize_message(msg) for msg in filtered_messages]
 
     def get_agent_messages(self) -> list[ModelMessage]:
         """Get agent_history as a list of ModelMessage objects.

shotgun/agents/conversation_manager.py

@@ -1,6 +1,7 @@
 """Manager for handling conversation persistence operations."""
 
 import json
+import shutil
 from pathlib import Path
 
 from shotgun.logging_config import get_logger
@@ -77,9 +78,30 @@ class ConversationManager:
             )
             return conversation
 
-        except Exception as e:
+        except (json.JSONDecodeError, ValueError) as e:
+            # Handle corrupted JSON or validation errors
+            logger.error(
+                "Corrupted conversation file at %s: %s. Creating backup and starting fresh.",
+                self.conversation_path,
+                e,
+            )
+
+            # Create a backup of the corrupted file for debugging
+            backup_path = self.conversation_path.with_suffix(".json.backup")
+            try:
+                shutil.copy2(self.conversation_path, backup_path)
+                logger.info("Backed up corrupted conversation to %s", backup_path)
+            except Exception as backup_error:  # pragma: no cover
+                logger.warning("Failed to backup corrupted file: %s", backup_error)
+
+            return None
+
+        except Exception as e:  # pragma: no cover
+            # Catch-all for unexpected errors
             logger.error(
-                "Failed to load conversation from %s: %s", self.conversation_path, e
+                "Unexpected error loading conversation from %s: %s",
+                self.conversation_path,
+                e,
             )
             return None
 

shotgun/agents/export.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_export_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create an export agent with file management capabilities.
 
     Args:
@@ -50,11 +49,11 @@ def create_export_agent(
 
 
 async def run_export_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]:
     """Export artifacts based on the given instruction.
 
     Args:

shotgun/agents/history/compaction.py

@@ -13,7 +13,7 @@ logger = get_logger(__name__)
 
 
 async def apply_persistent_compaction(
-    messages: list[ModelMessage], deps: AgentDeps
+    messages: list[ModelMessage], deps: AgentDeps, force: bool = False
 ) -> list[ModelMessage]:
     """Apply compaction to message history for persistent storage.
 
@@ -23,6 +23,7 @@ async def apply_persistent_compaction(
     Args:
         messages: Full message history from agent run
         deps: Agent dependencies containing model config
+        force: If True, force compaction even if below token threshold
 
     Returns:
         Compacted message history that should be stored as conversation state
@@ -46,7 +47,7 @@ async def apply_persistent_compaction(
                 self.usage = usage
 
         ctx = MockContext(deps, usage)
-        compacted_messages = await token_limit_compactor(ctx, messages)
+        compacted_messages = await token_limit_compactor(ctx, messages, force=force)
 
         # Log the result for monitoring
         original_size = len(messages)
@@ -59,17 +60,21 @@ async def apply_persistent_compaction(
                 f"({reduction_pct:.1f}% reduction)"
             )
 
-            # Track persistent compaction event
+            # Track persistent compaction event with simple metrics (fast, no token counting)
             track_event(
                 "persistent_compaction_applied",
                 {
+                    # Basic compaction metrics
                     "messages_before": original_size,
                     "messages_after": compacted_size,
-                    "tokens_before": estimated_tokens,
                     "reduction_percentage": round(reduction_pct, 2),
                     "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,
                 },
             )
         else:

shotgun/agents/history/context_extraction.py

@@ -1,5 +1,9 @@
 """Context extraction utilities for history processing."""
 
+import json
+import logging
+import traceback
+
 from pydantic_ai.messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
@@ -16,6 +20,46 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
+logger = logging.getLogger(__name__)
+
+
+def _safely_parse_tool_args(args: dict[str, object] | str | None) -> dict[str, object]:
+    """Safely parse tool call arguments, handling incomplete/invalid JSON.
+
+    Args:
+        args: Tool call arguments (dict, JSON string, or None)
+
+    Returns:
+        Parsed args dict, or empty dict if parsing fails
+    """
+    if args is None:
+        return {}
+
+    if isinstance(args, dict):
+        return args
+
+    if not isinstance(args, str):
+        return {}
+
+    try:
+        parsed = json.loads(args)
+        return parsed if isinstance(parsed, dict) else {}
+    except (json.JSONDecodeError, ValueError) as e:
+        # Only log warning if it looks like JSON (starts with { or [) - incomplete JSON
+        # Plain strings are valid args and shouldn't trigger warnings
+        stripped_args = args.strip()
+        if stripped_args.startswith(("{", "[")):
+            args_preview = args[:100] + "..." if len(args) > 100 else args
+            logger.warning(
+                "Detected incomplete/invalid JSON in tool call args during parsing",
+                extra={
+                    "args_preview": args_preview,
+                    "error": str(e),
+                    "args_length": len(args),
+                },
+            )
+        return {}
+
 
 def extract_context_from_messages(messages: list[ModelMessage]) -> str:
     """Extract context from a list of messages for summarization."""
@@ -87,12 +131,55 @@ def extract_context_from_part(
         return f"<ASSISTANT_TEXT>\n{message_part.content}\n</ASSISTANT_TEXT>"
 
     elif isinstance(message_part, ToolCallPart):
-        if isinstance(message_part.args, dict):
-            args_str = ", ".join(f"{k}={repr(v)}" for k, v in message_part.args.items())
-            tool_call_str = f"{message_part.tool_name}({args_str})"
-        else:
-            tool_call_str = f"{message_part.tool_name}({message_part.args})"
-        return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        # Safely parse args to avoid crashes from incomplete JSON during streaming
+        try:
+            parsed_args = _safely_parse_tool_args(message_part.args)
+            if parsed_args:
+                # Successfully parsed as dict - format nicely
+                args_str = ", ".join(f"{k}={repr(v)}" for k, v in parsed_args.items())
+                tool_call_str = f"{message_part.tool_name}({args_str})"
+            elif isinstance(message_part.args, str) and message_part.args:
+                # Non-empty string that didn't parse as JSON
+                # Check if it looks like JSON (starts with { or [) - if so, it's incomplete
+                stripped_args = message_part.args.strip()
+                if stripped_args.startswith(("{", "[")):
+                    # Looks like incomplete JSON - log warning and show empty parens
+                    args_preview = (
+                        stripped_args[:100] + "..."
+                        if len(stripped_args) > 100
+                        else stripped_args
+                    )
+                    stack_trace = "".join(traceback.format_stack())
+                    logger.warning(
+                        "ToolCallPart with unparseable args encountered during context extraction",
+                        extra={
+                            "tool_name": message_part.tool_name,
+                            "tool_call_id": message_part.tool_call_id,
+                            "args_preview": args_preview,
+                            "args_type": type(message_part.args).__name__,
+                            "stack_trace": stack_trace,
+                        },
+                    )
+                    tool_call_str = f"{message_part.tool_name}()"
+                else:
+                    # Plain string arg - display as-is
+                    tool_call_str = f"{message_part.tool_name}({message_part.args})"
+            else:
+                # No args
+                tool_call_str = f"{message_part.tool_name}()"
+            return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        except Exception as e:  # pragma: no cover - defensive catch-all
+            # If anything goes wrong, log full exception with stack trace
+            logger.error(
+                "Unexpected error processing ToolCallPart",
+                exc_info=True,
+                extra={
+                    "tool_name": message_part.tool_name,
+                    "tool_call_id": message_part.tool_call_id,
+                    "error": str(e),
+                },
+            )
+            return f"<TOOL_CALL>\n{message_part.tool_name}()\n</TOOL_CALL>"
 
     elif isinstance(message_part, BuiltinToolCallPart):
         return f"<BUILTIN_TOOL_CALL>\n{message_part.tool_name}\n</BUILTIN_TOOL_CALL>"