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

shotgun/agents/agent_manager.py

@@ -1,17 +1,26 @@
 """Agent manager for coordinating multiple AI agents with shared message history."""
 
+import json
 import logging
 from collections.abc import AsyncIterable, Sequence
 from dataclasses import dataclass, field, is_dataclass, replace
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
+import logfire
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
+
 if TYPE_CHECKING:
     from shotgun.agents.conversation_history import ConversationState
 
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
-    DeferredToolResults,
     RunContext,
     UsageLimits,
 )
@@ -31,12 +40,25 @@ from pydantic_ai.messages import (
     SystemPromptPart,
     ToolCallPart,
     ToolCallPartDelta,
+    UserPromptPart,
 )
 from textual.message import Message
 from textual.widget import Widget
 
 from shotgun.agents.common import add_system_prompt_message, add_system_status_message
-from shotgun.agents.models import AgentType, FileOperation
+from shotgun.agents.config.models import (
+    KeyProvider,
+    ModelConfig,
+    ModelName,
+    ProviderType,
+)
+from shotgun.agents.context_analyzer import (
+    ContextAnalysis,
+    ContextAnalyzer,
+    ContextCompositionTelemetry,
+    ContextFormatter,
+)
+from shotgun.agents.models import AgentResponse, AgentType, FileOperation
 from shotgun.posthog_telemetry import track_event
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage
 from shotgun.utils.source_detection import detect_source
@@ -44,7 +66,7 @@ from shotgun.utils.source_detection import detect_source
 from .export import create_export_agent
 from .history.compaction import apply_persistent_compaction
 from .messages import AgentSystemPrompt
-from .models import AgentDeps, AgentRuntimeOptions, UserAnswer
+from .models import AgentDeps, AgentRuntimeOptions
 from .plan import create_plan_agent
 from .research import create_research_agent
 from .specify import create_specify_agent
@@ -53,6 +75,35 @@ from .tasks import create_tasks_agent
 logger = logging.getLogger(__name__)
 
 
+def _is_retryable_error(exception: BaseException) -> bool:
+    """Check if exception should trigger a retry.
+
+    Args:
+        exception: The exception to check.
+
+    Returns:
+        True if the exception is a transient error that should be retried.
+    """
+    # ValueError for truncated/incomplete JSON
+    if isinstance(exception, ValueError):
+        error_str = str(exception)
+        return "EOF while parsing" in error_str or (
+            "JSON" in error_str and "parsing" in error_str
+        )
+
+    # API errors (overload, rate limits)
+    exception_name = type(exception).__name__
+    if "APIStatusError" in exception_name:
+        error_str = str(exception)
+        return "overload" in error_str.lower() or "rate" in error_str.lower()
+
+    # Network errors
+    if "ConnectionError" in exception_name or "TimeoutError" in exception_name:
+        return True
+
+    return False
+
+
 class MessageHistoryUpdated(Message):
     """Event posted when the message history is updated."""
 
@@ -91,6 +142,55 @@ class PartialResponseMessage(Message):
         self.is_last = is_last
 
 
+class ClarifyingQuestionsMessage(Message):
+    """Event posted when agent returns clarifying questions."""
+
+    def __init__(
+        self,
+        questions: list[str],
+        response_text: str,
+    ) -> None:
+        """Initialize the clarifying questions message.
+
+        Args:
+            questions: List of clarifying questions from the agent
+            response_text: The agent's response text before asking questions
+        """
+        super().__init__()
+        self.questions = questions
+        self.response_text = response_text
+
+
+class CompactionStartedMessage(Message):
+    """Event posted when conversation compaction starts."""
+
+
+class CompactionCompletedMessage(Message):
+    """Event posted when conversation compaction completes."""
+
+
+@dataclass(frozen=True)
+class ModelConfigUpdated:
+    """Data returned when AI model configuration changes.
+
+    Used as a return value from ModelPickerScreen to communicate model
+    selection back to the calling screen.
+
+    Attributes:
+        old_model: Previous model name (None if first selection)
+        new_model: New model name
+        provider: LLM provider (OpenAI, Anthropic, Google)
+        key_provider: Authentication method (BYOK or Shotgun)
+        model_config: Complete model configuration
+    """
+
+    old_model: ModelName | None
+    new_model: ModelName
+    provider: ProviderType
+    key_provider: KeyProvider
+    model_config: ModelConfig
+
+
 @dataclass(slots=True)
 class _PartialStreamState:
     """Tracks streamed messages while handling a single agent run."""
@@ -157,8 +257,12 @@ class AgentManager(Widget):
         self.recently_change_files: list[FileOperation] = []
         self._stream_state: _PartialStreamState | None = None
 
+        # Q&A mode state for structured output questions
+        self._qa_questions: list[str] | None = None
+        self._qa_mode_active: bool = False
+
     @property
-    def current_agent(self) -> Agent[AgentDeps, str | DeferredToolRequests]:
+    def current_agent(self) -> Agent[AgentDeps, AgentResponse]:
         """Get the currently active agent.
 
         Returns:
@@ -166,9 +270,7 @@ class AgentManager(Widget):
         """
         return self._get_agent(self._current_agent_type)
 
-    def _get_agent(
-        self, agent_type: AgentType
-    ) -> Agent[AgentDeps, str | DeferredToolRequests]:
+    def _get_agent(self, agent_type: AgentType) -> Agent[AgentDeps, AgentResponse]:
         """Get agent by type.
 
         Args:
@@ -245,15 +347,57 @@ class AgentManager(Widget):
                 f"Invalid agent type: {agent_type}. Must be one of: {', '.join(e.value for e in AgentType)}"
             ) from None
 
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=8),
+        retry=retry_if_exception(_is_retryable_error),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )
+    async def _run_agent_with_retry(
+        self,
+        agent: Agent[AgentDeps, AgentResponse],
+        prompt: str | None,
+        deps: AgentDeps,
+        usage_limits: UsageLimits | None,
+        message_history: list[ModelMessage],
+        event_stream_handler: Any,
+        **kwargs: Any,
+    ) -> AgentRunResult[AgentResponse]:
+        """Run agent with automatic retry on transient errors.
+
+        Args:
+            agent: The agent to run.
+            prompt: Optional prompt to send to the agent.
+            deps: Agent dependencies.
+            usage_limits: Optional usage limits.
+            message_history: Message history to provide to agent.
+            event_stream_handler: Event handler for streaming.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The agent run result.
+
+        Raises:
+            Various exceptions if all retries fail.
+        """
+        return await agent.run(
+            prompt,
+            deps=deps,
+            usage_limits=usage_limits,
+            message_history=message_history,
+            event_stream_handler=event_stream_handler,
+            **kwargs,
+        )
+
     async def run(
         self,
         prompt: str | None = None,
         *,
         deps: AgentDeps | None = None,
         usage_limits: UsageLimits | None = None,
-        deferred_tool_results: DeferredToolResults | None = None,
         **kwargs: Any,
-    ) -> AgentRunResult[str | DeferredToolRequests]:
+    ) -> AgentRunResult[AgentResponse]:
         """Run the current agent with automatic message history management.
 
         This method wraps the agent's run method, automatically injecting the
@@ -263,7 +407,6 @@ class AgentManager(Widget):
             prompt: Optional prompt to send to the agent.
             deps: Optional dependencies override (defaults to manager's deps).
             usage_limits: Optional usage limits for the agent run.
-            deferred_tool_results: Optional deferred tool results for continuing a conversation.
             **kwargs: Additional keyword arguments to pass to the agent.
 
         Returns:
@@ -273,15 +416,6 @@ class AgentManager(Widget):
         # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
             deps = self._create_merged_deps(self._current_agent_type)
-        ask_user_part = self.get_unanswered_ask_user_part()
-        if ask_user_part and prompt:
-            if not deferred_tool_results:
-                deferred_tool_results = DeferredToolResults()
-            deferred_tool_results.calls[ask_user_part.tool_call_id] = UserAnswer(
-                answer=prompt,
-                tool_call_id=ask_user_part.tool_call_id,
-            )
-            prompt = None
 
         # Ensure deps is not None
         if deps is None:
@@ -289,13 +423,12 @@ class AgentManager(Widget):
 
         # Clear file tracker before each run to track only this run's operations
         deps.file_tracker.clear()
-        # preprocess messages; maybe we need to include the user answer in the message history
 
-        original_messages = self.ui_message_history.copy()
+        # Don't manually add the user prompt - Pydantic AI will include it in result.new_messages()
+        # This prevents duplicates and confusion with incremental mounting
 
-        if prompt:
-            self.ui_message_history.append(ModelRequest.user_text_prompt(prompt))
-        self._post_messages_updated()
+        # Save current message history before the run
+        original_messages = self.ui_message_history.copy()
 
         # Start with persistent message history
         message_history = self.message_history
@@ -359,52 +492,309 @@ class AgentManager(Widget):
         model_name = ""
         if hasattr(deps, "llm_model") and deps.llm_model is not None:
             model_name = deps.llm_model.name
-        is_gpt5 = (  # streaming is likely not supported for gpt5. It varies between keys.
-            "gpt-5" in model_name.lower()
+
+        # Check if it's a Shotgun account
+        is_shotgun_account = (
+            hasattr(deps, "llm_model")
+            and deps.llm_model is not None
+            and deps.llm_model.key_provider == KeyProvider.SHOTGUN
         )
 
+        # Only disable streaming for GPT-5 if NOT a Shotgun account
+        # Shotgun accounts support streaming for GPT-5
+        is_gpt5_byok = "gpt-5" in model_name.lower() and not is_shotgun_account
+
         # Track message send event
         event_name = f"message_send_{self._current_agent_type.value}"
         track_event(
             event_name,
             {
                 "has_prompt": prompt is not None,
-                "has_deferred_results": deferred_tool_results is not None,
                 "model_name": model_name,
             },
         )
 
         try:
-            result: AgentRunResult[
-                str | DeferredToolRequests
-            ] = await self.current_agent.run(
-                prompt,
+            result: AgentRunResult[AgentResponse] = await self._run_agent_with_retry(
+                agent=self.current_agent,
+                prompt=prompt,
                 deps=deps,
                 usage_limits=usage_limits,
                 message_history=message_history,
-                deferred_tool_results=deferred_tool_results,
-                event_stream_handler=self._handle_event_stream if not is_gpt5 else None,
+                event_stream_handler=self._handle_event_stream
+                if not is_gpt5_byok
+                else None,
                 **kwargs,
             )
+        except ValueError as e:
+            # Handle truncated/incomplete JSON in tool calls specifically
+            error_str = str(e)
+            if "EOF while parsing" in error_str or (
+                "JSON" in error_str and "parsing" in error_str
+            ):
+                logger.error(
+                    "Tool call with truncated/incomplete JSON arguments detected",
+                    extra={
+                        "agent_mode": self._current_agent_type.value,
+                        "model_name": model_name,
+                        "error": error_str,
+                    },
+                )
+                logfire.error(
+                    "Tool call with truncated JSON arguments",
+                    agent_mode=self._current_agent_type.value,
+                    model_name=model_name,
+                    error=error_str,
+                )
+                # Add helpful hint message for the user
+                self.ui_message_history.append(
+                    HintMessage(
+                        message="⚠️ The agent attempted an operation with arguments that were too large (truncated JSON). "
+                        "Try breaking your request into smaller steps or more focused contracts."
+                    )
+                )
+                self._post_messages_updated()
+            # Re-raise to maintain error visibility
+            raise
+        except Exception as e:
+            # Log the error with full stack trace to shotgun.log and Logfire
+            logger.exception(
+                "Agent execution failed",
+                extra={
+                    "agent_mode": self._current_agent_type.value,
+                    "model_name": model_name,
+                    "error_type": type(e).__name__,
+                },
+            )
+            logfire.exception(
+                "Agent execution failed",
+                agent_mode=self._current_agent_type.value,
+                model_name=model_name,
+                error_type=type(e).__name__,
+            )
+            # Re-raise to let TUI handle user messaging
+            raise
         finally:
             self._stream_state = None
 
-        self.ui_message_history = original_messages + cast(
+        # Agent ALWAYS returns AgentResponse with structured output
+        agent_response = result.output
+        logger.debug(
+            "Agent returned structured AgentResponse",
+            extra={
+                "has_response": agent_response.response is not None,
+                "response_length": len(agent_response.response)
+                if agent_response.response
+                else 0,
+                "response_preview": agent_response.response[:100] + "..."
+                if agent_response.response and len(agent_response.response) > 100
+                else agent_response.response or "(empty)",
+                "has_clarifying_questions": bool(agent_response.clarifying_questions),
+                "num_clarifying_questions": len(agent_response.clarifying_questions)
+                if agent_response.clarifying_questions
+                else 0,
+            },
+        )
+
+        # Merge agent's response messages, avoiding duplicates
+        # The TUI may have already added the user prompt, so check for it
+        new_messages = cast(
             list[ModelRequest | ModelResponse | HintMessage], result.new_messages()
         )
 
+        # Deduplicate: skip user prompts that are already in original_messages
+        deduplicated_new_messages = []
+        for msg in new_messages:
+            # Check if this is a user prompt that's already in original_messages
+            if isinstance(msg, ModelRequest) and any(
+                isinstance(part, UserPromptPart) for part in msg.parts
+            ):
+                # Check if an identical user prompt is already in original_messages
+                already_exists = any(
+                    isinstance(existing, ModelRequest)
+                    and any(isinstance(p, UserPromptPart) for p in existing.parts)
+                    and existing.parts == msg.parts
+                    for existing in original_messages[
+                        -5:
+                    ]  # Check last 5 messages for efficiency
+                )
+                if already_exists:
+                    continue  # Skip this duplicate user prompt
+
+            deduplicated_new_messages.append(msg)
+
+        self.ui_message_history = original_messages + deduplicated_new_messages
+
+        # Get file operations early so we can use them for contextual messages
+        file_operations = deps.file_tracker.operations.copy()
+        self.recently_change_files = file_operations
+
+        logger.debug(
+            "File operations tracked",
+            extra={
+                "num_file_operations": len(file_operations),
+                "operation_files": [Path(op.file_path).name for op in file_operations],
+            },
+        )
+
+        # Check if there are clarifying questions
+        if agent_response.clarifying_questions:
+            logger.info(
+                f"Agent has {len(agent_response.clarifying_questions)} clarifying questions"
+            )
+
+            # Add agent's response first if present
+            if agent_response.response:
+                self.ui_message_history.append(
+                    HintMessage(message=agent_response.response)
+                )
+
+            if len(agent_response.clarifying_questions) == 1:
+                # Single question - treat as non-blocking suggestion, DON'T enter Q&A mode
+                self.ui_message_history.append(
+                    HintMessage(message=f"💡 {agent_response.clarifying_questions[0]}")
+                )
+            else:
+                # Multiple questions (2+) - enter Q&A mode
+                self._qa_questions = agent_response.clarifying_questions
+                self._qa_mode_active = True
+
+                # Show intro with list, then first question
+                questions_list_with_intro = (
+                    f"I have {len(agent_response.clarifying_questions)} questions:\n\n"
+                    + "\n".join(
+                        f"{i + 1}. {q}"
+                        for i, q in enumerate(agent_response.clarifying_questions)
+                    )
+                )
+                self.ui_message_history.append(
+                    HintMessage(message=questions_list_with_intro)
+                )
+                self.ui_message_history.append(
+                    HintMessage(
+                        message=f"**Q1:** {agent_response.clarifying_questions[0]}"
+                    )
+                )
+
+                # Post event to TUI to update Q&A mode state (only for multiple questions)
+                self.post_message(
+                    ClarifyingQuestionsMessage(
+                        questions=agent_response.clarifying_questions,
+                        response_text=agent_response.response,
+                    )
+                )
+
+            # Post UI update with hint messages and file operations
+            logger.debug(
+                "Posting UI update for Q&A mode with hint messages and file operations"
+            )
+            self._post_messages_updated(file_operations)
+        else:
+            # No clarifying questions - show the response or a default success message
+            if agent_response.response and agent_response.response.strip():
+                logger.debug(
+                    "Adding agent response as hint",
+                    extra={
+                        "response_preview": agent_response.response[:100] + "..."
+                        if len(agent_response.response) > 100
+                        else agent_response.response,
+                        "has_file_operations": len(file_operations) > 0,
+                    },
+                )
+                self.ui_message_history.append(
+                    HintMessage(message=agent_response.response)
+                )
+            else:
+                # Fallback: response is empty or whitespace
+                logger.debug(
+                    "Agent response was empty, using fallback completion message",
+                    extra={"has_file_operations": len(file_operations) > 0},
+                )
+                # Show contextual message based on whether files were modified
+                if file_operations:
+                    self.ui_message_history.append(
+                        HintMessage(
+                            message="✅ Task completed - files have been modified"
+                        )
+                    )
+                else:
+                    self.ui_message_history.append(
+                        HintMessage(message="✅ Task completed")
+                    )
+
+            # Post UI update immediately so user sees the response without delay
+            logger.debug(
+                "Posting immediate UI update with hint message and file operations"
+            )
+            self._post_messages_updated(file_operations)
+
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
-        self.message_history = await apply_persistent_compaction(all_messages, deps)
-        usage = result.usage()
-        deps.usage_manager.add_usage(
-            usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
+        messages_before_compaction = len(all_messages)
+        compaction_occurred = False
+
+        try:
+            logger.debug(
+                "Starting message history compaction",
+                extra={"message_count": len(all_messages)},
+            )
+            # Notify UI that compaction is starting
+            self.post_message(CompactionStartedMessage())
+
+            self.message_history = await apply_persistent_compaction(all_messages, deps)
+
+            # Track if compaction actually modified the history
+            compaction_occurred = len(self.message_history) != len(all_messages)
+
+            # Notify UI that compaction is complete
+            self.post_message(CompactionCompletedMessage())
+
+            logger.debug(
+                "Completed message history compaction",
+                extra={
+                    "original_count": len(all_messages),
+                    "compacted_count": len(self.message_history),
+                },
+            )
+        except Exception as e:
+            # If compaction fails, log full error with stack trace and use uncompacted messages
+            logger.error(
+                "Failed to compact message history - using uncompacted messages",
+                exc_info=True,
+                extra={
+                    "error": str(e),
+                    "message_count": len(all_messages),
+                    "agent_mode": self._current_agent_type.value,
+                },
+            )
+            # Fallback: use uncompacted messages to prevent data loss
+            self.message_history = all_messages
+
+        # Track context composition telemetry
+        await self._track_context_analysis(
+            compaction_occurred=compaction_occurred,
+            messages_before_compaction=messages_before_compaction
+            if compaction_occurred
+            else None,
         )
 
-        # Log file operations summary if any files were modified
-        file_operations = deps.file_tracker.operations.copy()
-        self.recently_change_files = file_operations
+        usage = result.usage()
+        if hasattr(deps, "llm_model") and deps.llm_model is not None:
+            deps.usage_manager.add_usage(
+                usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
+            )
+        else:
+            logger.warning(
+                "llm_model is None, skipping usage tracking",
+                extra={"agent_mode": self._current_agent_type.value},
+            )
 
+        # Post final UI update after compaction completes
+        # This ensures widgets that depend on message_history (like context indicator)
+        # receive the updated history after compaction
+        logger.debug(
+            "Posting final UI update after compaction with updated message_history"
+        )
         self._post_messages_updated(file_operations)
 
         return result
@@ -480,6 +870,39 @@ class AgentManager(Widget):
                     # Detect source from call stack
                     source = detect_source()
 
+                    # Log if tool call has incomplete args (for debugging truncated JSON)
+                    if isinstance(event.part.args, str):
+                        try:
+                            json.loads(event.part.args)
+                        except (json.JSONDecodeError, ValueError):
+                            args_preview = (
+                                event.part.args[:100] + "..."
+                                if len(event.part.args) > 100
+                                else event.part.args
+                            )
+                            logger.warning(
+                                "FunctionToolCallEvent received with incomplete JSON args",
+                                extra={
+                                    "tool_name": event.part.tool_name,
+                                    "tool_call_id": event.part.tool_call_id,
+                                    "args_preview": args_preview,
+                                    "args_length": len(event.part.args)
+                                    if event.part.args
+                                    else 0,
+                                    "agent_mode": self._current_agent_type.value,
+                                },
+                            )
+                            logfire.warn(
+                                "FunctionToolCallEvent received with incomplete JSON args",
+                                tool_name=event.part.tool_name,
+                                tool_call_id=event.part.tool_call_id,
+                                args_preview=args_preview,
+                                args_length=len(event.part.args)
+                                if event.part.args
+                                else 0,
+                                agent_mode=self._current_agent_type.value,
+                            )
+
                     track_event(
                         "tool_called",
                         {
@@ -649,6 +1072,62 @@ class AgentManager(Widget):
     def get_usage_hint(self) -> str | None:
         return self.deps.usage_manager.build_usage_hint()
 
+    async def get_context_hint(self) -> str | None:
+        """Get conversation context analysis as a formatted hint.
+
+        Returns:
+            Markdown-formatted string with context composition statistics, or None if unavailable
+        """
+        analysis = await self.get_context_analysis()
+        if analysis:
+            return ContextFormatter.format_markdown(analysis)
+        return None
+
+    async def get_context_analysis(self) -> ContextAnalysis | None:
+        """Get conversation context analysis as structured data.
+
+        Returns:
+            ContextAnalysis object with token usage data, or None if unavailable
+        """
+
+        try:
+            analyzer = ContextAnalyzer(self.deps.llm_model)
+            return await analyzer.analyze_conversation(
+                self.message_history, self.ui_message_history
+            )
+        except Exception as e:
+            logger.error(f"Failed to generate context analysis: {e}", exc_info=True)
+            return None
+
+    async def _track_context_analysis(
+        self,
+        compaction_occurred: bool = False,
+        messages_before_compaction: int | None = None,
+    ) -> None:
+        """Track context composition telemetry to PostHog.
+
+        Args:
+            compaction_occurred: Whether compaction was applied
+            messages_before_compaction: Message count before compaction, if it occurred
+        """
+        try:
+            analyzer = ContextAnalyzer(self.deps.llm_model)
+            analysis = await analyzer.analyze_conversation(
+                self.message_history, self.ui_message_history
+            )
+
+            # Create telemetry model from analysis
+            telemetry = ContextCompositionTelemetry.from_analysis(
+                analysis,
+                compaction_occurred=compaction_occurred,
+                messages_before_compaction=messages_before_compaction,
+            )
+
+            # Send to PostHog using model_dump() for dict conversion
+            track_event("agent_context_composition", telemetry.model_dump())
+        except Exception as e:
+            logger.warning(f"Failed to track context analysis: {e}")
+
     def get_conversation_state(self) -> "ConversationState":
         """Get the current conversation state.
 
@@ -691,27 +1170,14 @@ class AgentManager(Widget):
         self.ui_message_history.append(message)
         self._post_messages_updated()
 
-    def get_unanswered_ask_user_part(self) -> ToolCallPart | None:
-        if not self.message_history:
-            return None
-        self.last_response = self.message_history[-1]
-        ## we're searching for unanswered ask_user parts
-        found_tool = next(
-            (
-                part
-                for part in self.message_history[-1].parts
-                if isinstance(part, ToolCallPart) and part.tool_name == "ask_user"
-            ),
-            None,
-        )
-
-        return found_tool
-
 
 # Re-export AgentType for backward compatibility
 __all__ = [
     "AgentManager",
     "AgentType",
+    "ClarifyingQuestionsMessage",
+    "CompactionCompletedMessage",
+    "CompactionStartedMessage",
     "MessageHistoryUpdated",
     "PartialResponseMessage",
 ]