shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.2.17__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -40,13 +40,30 @@ 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.config.models import KeyProvider
-from shotgun.agents.models import AgentResponse, 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,
+    FileOperationTracker,
+)
 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
@@ -149,6 +166,44 @@ class ClarifyingQuestionsMessage(Message):
         self.response_text = response_text
 
 
+class CompactionStartedMessage(Message):
+    """Event posted when conversation compaction starts."""
+
+
+class CompactionCompletedMessage(Message):
+    """Event posted when conversation compaction completes."""
+
+
+class AgentStreamingStarted(Message):
+    """Event posted when agent starts streaming responses."""
+
+
+class AgentStreamingCompleted(Message):
+    """Event posted when agent finishes streaming responses."""
+
+
+@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."""
@@ -180,7 +235,7 @@ class AgentManager(Widget):
         self.deps = deps
 
         # Create AgentRuntimeOptions from deps for agent creation
-        agent_runtime_options = AgentRuntimeOptions(
+        self._agent_runtime_options = AgentRuntimeOptions(
             interactive_mode=self.deps.interactive_mode,
             working_directory=self.deps.working_directory,
             is_tui_context=self.deps.is_tui_context,
@@ -189,22 +244,18 @@ class AgentManager(Widget):
             tasks=self.deps.tasks,
         )
 
-        # Initialize all agents and store their specific deps
-        self.research_agent, self.research_deps = create_research_agent(
-            agent_runtime_options=agent_runtime_options
-        )
-        self.plan_agent, self.plan_deps = create_plan_agent(
-            agent_runtime_options=agent_runtime_options
-        )
-        self.tasks_agent, self.tasks_deps = create_tasks_agent(
-            agent_runtime_options=agent_runtime_options
-        )
-        self.specify_agent, self.specify_deps = create_specify_agent(
-            agent_runtime_options=agent_runtime_options
-        )
-        self.export_agent, self.export_deps = create_export_agent(
-            agent_runtime_options=agent_runtime_options
-        )
+        # Lazy initialization - agents created on first access
+        self._research_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._research_deps: AgentDeps | None = None
+        self._plan_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._plan_deps: AgentDeps | None = None
+        self._tasks_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._tasks_deps: AgentDeps | None = None
+        self._specify_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._specify_deps: AgentDeps | None = None
+        self._export_agent: Agent[AgentDeps, AgentResponse] | None = None
+        self._export_deps: AgentDeps | None = None
+        self._agents_initialized = False
 
         # Track current active agent
         self._current_agent_type: AgentType = initial_type
@@ -219,6 +270,119 @@ class AgentManager(Widget):
         self._qa_questions: list[str] | None = None
         self._qa_mode_active: bool = False
 
+    async def _ensure_agents_initialized(self) -> None:
+        """Ensure all agents are initialized (lazy initialization)."""
+        if self._agents_initialized:
+            return
+
+        # Initialize all agents asynchronously
+        self._research_agent, self._research_deps = await create_research_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
+        self._plan_agent, self._plan_deps = await create_plan_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
+        self._tasks_agent, self._tasks_deps = await create_tasks_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
+        self._specify_agent, self._specify_deps = await create_specify_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
+        self._export_agent, self._export_deps = await create_export_agent(
+            agent_runtime_options=self._agent_runtime_options
+        )
+        self._agents_initialized = True
+
+    @property
+    def research_agent(self) -> Agent[AgentDeps, AgentResponse]:
+        """Get research agent (must call _ensure_agents_initialized first)."""
+        if self._research_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._research_agent
+
+    @property
+    def research_deps(self) -> AgentDeps:
+        """Get research deps (must call _ensure_agents_initialized first)."""
+        if self._research_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._research_deps
+
+    @property
+    def plan_agent(self) -> Agent[AgentDeps, AgentResponse]:
+        """Get plan agent (must call _ensure_agents_initialized first)."""
+        if self._plan_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._plan_agent
+
+    @property
+    def plan_deps(self) -> AgentDeps:
+        """Get plan deps (must call _ensure_agents_initialized first)."""
+        if self._plan_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._plan_deps
+
+    @property
+    def tasks_agent(self) -> Agent[AgentDeps, AgentResponse]:
+        """Get tasks agent (must call _ensure_agents_initialized first)."""
+        if self._tasks_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._tasks_agent
+
+    @property
+    def tasks_deps(self) -> AgentDeps:
+        """Get tasks deps (must call _ensure_agents_initialized first)."""
+        if self._tasks_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._tasks_deps
+
+    @property
+    def specify_agent(self) -> Agent[AgentDeps, AgentResponse]:
+        """Get specify agent (must call _ensure_agents_initialized first)."""
+        if self._specify_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._specify_agent
+
+    @property
+    def specify_deps(self) -> AgentDeps:
+        """Get specify deps (must call _ensure_agents_initialized first)."""
+        if self._specify_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._specify_deps
+
+    @property
+    def export_agent(self) -> Agent[AgentDeps, AgentResponse]:
+        """Get export agent (must call _ensure_agents_initialized first)."""
+        if self._export_agent is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._export_agent
+
+    @property
+    def export_deps(self) -> AgentDeps:
+        """Get export deps (must call _ensure_agents_initialized first)."""
+        if self._export_deps is None:
+            raise RuntimeError(
+                "Agents not initialized. Call _ensure_agents_initialized() first."
+            )
+        return self._export_deps
+
     @property
     def current_agent(self) -> Agent[AgentDeps, AgentResponse]:
         """Get the currently active agent.
@@ -370,6 +534,9 @@ class AgentManager(Widget):
         Returns:
             The agent run result.
         """
+        # Ensure agents are initialized before running
+        await self._ensure_agents_initialized()
+
         logger.info(f"Running agent {self._current_agent_type.value}")
         # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
@@ -382,19 +549,11 @@ class AgentManager(Widget):
         # Clear file tracker before each run to track only this run's operations
         deps.file_tracker.clear()
 
-        # Add user prompt if present (will be shown immediately via post_messages_updated)
-        if prompt:
-            user_request = ModelRequest.user_text_prompt(prompt)
-            self.ui_message_history.append(user_request)
-
-        # Always post update before run to show user message (or current state if no prompt)
-        self._post_messages_updated()
+        # Don't manually add the user prompt - Pydantic AI will include it in result.new_messages()
+        # This prevents duplicates and confusion with incremental mounting
 
-        # Save history WITHOUT the just-added prompt to avoid duplicates
-        # (result.new_messages() will include the prompt)
-        original_messages = (
-            self.ui_message_history[:-1] if prompt else self.ui_message_history.copy()
-        )
+        # Save current message history before the run
+        original_messages = self.ui_message_history.copy()
 
         # Start with persistent message history
         message_history = self.message_history
@@ -562,11 +721,35 @@ class AgentManager(Widget):
             },
         )
 
-        # Always add the agent's response messages to maintain conversation history
-        self.ui_message_history = original_messages + cast(
+        # 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
@@ -591,6 +774,12 @@ class AgentManager(Widget):
                     HintMessage(message=agent_response.response)
                 )
 
+            # Add file operation hints before questions (so they appear first in UI)
+            if file_operations:
+                file_hint = self._create_file_operation_hint(file_operations)
+                if file_hint:
+                    self.ui_message_history.append(HintMessage(message=file_hint))
+
             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(
@@ -626,11 +815,9 @@ class AgentManager(Widget):
                     )
                 )
 
-            # 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)
+            # Post UI update with hint messages (file operations will be posted after compaction)
+            logger.debug("Posting UI update for Q&A mode with hint messages")
+            self._post_messages_updated([])
         else:
             # No clarifying questions - show the response or a default success message
             if agent_response.response and agent_response.response.strip():
@@ -665,19 +852,31 @@ class AgentManager(Widget):
                     )
 
             # 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)
+            # (file operations will be posted after compaction to avoid duplicates)
+            logger.debug("Posting immediate UI update with hint message")
+            self._post_messages_updated([])
 
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
+        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={
@@ -699,9 +898,17 @@ class AgentManager(Widget):
             # 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,
+        )
+
         usage = result.usage()
         if hasattr(deps, "llm_model") and deps.llm_model is not None:
-            deps.usage_manager.add_usage(
+            await deps.usage_manager.add_usage(
                 usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
             )
         else:
@@ -710,8 +917,13 @@ class AgentManager(Widget):
                 extra={"agent_mode": self._current_agent_type.value},
             )
 
-        # UI updates are now posted immediately in each branch (Q&A or non-Q&A)
-        # before compaction, so no duplicate posting needed here
+        # 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
 
@@ -722,6 +934,9 @@ class AgentManager(Widget):
     ) -> None:
         """Process streamed events and forward partial updates to the UI."""
 
+        # Notify UI that streaming has started
+        self.post_message(AgentStreamingStarted())
+
         state = self._stream_state
         if state is None:
             state = self._stream_state = _PartialStreamState()
@@ -900,6 +1115,9 @@ class AgentManager(Widget):
             self._post_partial_message(True)
         state.current_response = None
 
+        # Notify UI that streaming has completed
+        self.post_message(AgentStreamingCompleted())
+
     def _build_partial_response(
         self, parts: list[ModelResponsePart | ToolCallPartDelta]
     ) -> ModelResponse | None:
@@ -927,6 +1145,38 @@ class AgentManager(Widget):
             )
         )
 
+    def _create_file_operation_hint(
+        self, file_operations: list[FileOperation]
+    ) -> str | None:
+        """Create a hint message for file operations.
+
+        Args:
+            file_operations: List of file operations to create a hint for
+
+        Returns:
+            Hint message string or None if no operations
+        """
+        if not file_operations:
+            return None
+
+        tracker = FileOperationTracker(operations=file_operations)
+        display_path = tracker.get_display_path()
+
+        if not display_path:
+            return None
+
+        path_obj = Path(display_path)
+
+        if len(file_operations) == 1:
+            return f"📝 Modified: `{display_path}`"
+        else:
+            num_files = len({op.file_path for op in file_operations})
+            if path_obj.is_dir():
+                return f"📁 Modified {num_files} files in: `{display_path}`"
+            else:
+                # Common path is a file, show parent directory
+                return f"📁 Modified {num_files} files in: `{path_obj.parent}`"
+
     def _post_messages_updated(
         self, file_operations: list[FileOperation] | None = None
     ) -> None:
@@ -988,6 +1238,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.
 
@@ -1035,7 +1341,9 @@ class AgentManager(Widget):
 __all__ = [
     "AgentManager",
     "AgentType",
+    "ClarifyingQuestionsMessage",
+    "CompactionCompletedMessage",
+    "CompactionStartedMessage",
     "MessageHistoryUpdated",
     "PartialResponseMessage",
-    "ClarifyingQuestionsMessage",
 ]

shotgun/agents/common.py

@@ -4,6 +4,7 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
+import aiofiles
 from pydantic_ai import (
     Agent,
     RunContext,
@@ -68,7 +69,7 @@ async def add_system_status_message(
     existing_files = get_agent_existing_files(deps.agent_mode)
 
     # Extract table of contents from the agent's markdown file
-    markdown_toc = extract_markdown_toc(deps.agent_mode)
+    markdown_toc = await extract_markdown_toc(deps.agent_mode)
 
     # Get current datetime with timezone information
     dt_context = get_datetime_context()
@@ -94,7 +95,7 @@ async def add_system_status_message(
     return message_history
 
 
-def create_base_agent(
+async def create_base_agent(
     system_prompt_fn: Callable[[RunContext[AgentDeps]], str],
     agent_runtime_options: AgentRuntimeOptions,
     load_codebase_understanding_tools: bool = True,
@@ -119,7 +120,7 @@ def create_base_agent(
 
     # Get configured model or fall back to first available provider
     try:
-        model_config = get_provider_model(provider)
+        model_config = await get_provider_model(provider)
         provider_name = model_config.provider
         logger.debug(
             "🤖 Creating agent with configured %s model: %s",
@@ -194,7 +195,7 @@ def create_base_agent(
     return agent, deps
 
 
-def _extract_file_toc_content(
+async def _extract_file_toc_content(
     file_path: Path, max_depth: int | None = None, max_chars: int = 500
 ) -> str | None:
     """Extract TOC from a single file with depth and character limits.
@@ -211,7 +212,8 @@ def _extract_file_toc_content(
         return None
 
     try:
-        content = file_path.read_text(encoding="utf-8")
+        async with aiofiles.open(file_path, encoding="utf-8") as f:
+            content = await f.read()
         lines = content.split("\n")
 
         # Extract headings
@@ -257,7 +259,7 @@ def _extract_file_toc_content(
         return None
 
 
-def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
+async def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
     """Extract TOCs from current and prior agents' files in the pipeline.
 
     Shows full TOC of agent's own file and high-level summaries of prior agents'
@@ -309,7 +311,9 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
     for prior_file in config.prior_files:
         file_path = base_path / prior_file
         # Only show # and ## headings from prior files, max 500 chars each
-        prior_toc = _extract_file_toc_content(file_path, max_depth=2, max_chars=500)
+        prior_toc = await _extract_file_toc_content(
+            file_path, max_depth=2, max_chars=500
+        )
         if prior_toc:
             # Add section with XML tags
             toc_sections.append(
@@ -321,7 +325,9 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
     # Extract TOC from own file (full detail)
     if config.own_file:
         own_path = base_path / config.own_file
-        own_toc = _extract_file_toc_content(own_path, max_depth=None, max_chars=2000)
+        own_toc = await _extract_file_toc_content(
+            own_path, max_depth=None, max_chars=2000
+        )
         if own_toc:
             # Put own file TOC at the beginning with XML tags
             toc_sections.insert(

shotgun/agents/config/constants.py

@@ -24,11 +24,5 @@ ANTHROPIC_PROVIDER = ConfigSection.ANTHROPIC.value
 GOOGLE_PROVIDER = ConfigSection.GOOGLE.value
 SHOTGUN_PROVIDER = ConfigSection.SHOTGUN.value
 
-# Environment variable names
-OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
-ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
-GEMINI_API_KEY_ENV = "GEMINI_API_KEY"
-SHOTGUN_API_KEY_ENV = "SHOTGUN_API_KEY"
-
 # Token limits
 MEDIUM_TEXT_8K_TOKENS = 8192  # Default max_tokens for web search requests