shotgun-sh 0.1.0.dev25__py3-none-any.whl → 0.1.0.dev27__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -37,9 +37,11 @@ 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.tui.screens.chat_screen.hint_message import HintMessage
 
 from .export import create_export_agent
 from .history.compaction import apply_persistent_compaction
+from .messages import AgentSystemPrompt
 from .models import AgentDeps, AgentRuntimeOptions
 from .plan import create_plan_agent
 from .research import create_research_agent
@@ -54,7 +56,7 @@ class MessageHistoryUpdated(Message):
 
     def __init__(
         self,
-        messages: list[ModelMessage],
+        messages: list[ModelMessage | HintMessage],
         agent_type: AgentType,
         file_operations: list[FileOperation] | None = None,
     ) -> None:
@@ -143,7 +145,7 @@ class AgentManager(Widget):
         self._current_agent_type: AgentType = initial_type
 
         # Maintain shared message history
-        self.ui_message_history: list[ModelMessage] = []
+        self.ui_message_history: list[ModelMessage | HintMessage] = []
         self.message_history: list[ModelMessage] = []
         self.recently_change_files: list[FileOperation] = []
         self._stream_state: _PartialStreamState | None = None
@@ -277,15 +279,51 @@ class AgentManager(Widget):
 
         deps.agent_mode = self._current_agent_type
 
+        # Filter out system prompts from other agent types
+        from pydantic_ai.messages import ModelRequestPart
+
+        filtered_history: list[ModelMessage] = []
+        for message in message_history:
+            # Keep all non-ModelRequest messages as-is
+            if not isinstance(message, ModelRequest):
+                filtered_history.append(message)
+                continue
+
+            # Filter out AgentSystemPrompts from other agent types
+            filtered_parts: list[ModelRequestPart] = []
+            for part in message.parts:
+                # Keep non-AgentSystemPrompt parts
+                if not isinstance(part, AgentSystemPrompt):
+                    filtered_parts.append(part)
+                    continue
+
+                # Only keep system prompts from the same agent type
+                if part.agent_mode == deps.agent_mode:
+                    filtered_parts.append(part)
+
+            # Only add the message if it has parts remaining
+            if filtered_parts:
+                filtered_history.append(ModelRequest(parts=filtered_parts))
+
+        message_history = filtered_history
+
         # Add a system status message so the agent knows whats going on
         message_history = await add_system_status_message(deps, message_history)
 
-        # Check if the message history already has a system prompt
-        has_system_prompt = any(
-            hasattr(msg, "parts")
-            and any(isinstance(part, SystemPromptPart) for part in msg.parts)
-            for msg in message_history
-        )
+        # Check if the message history already has a system prompt from the same agent type
+        has_system_prompt = False
+        for message in message_history:
+            if not isinstance(message, ModelRequest):
+                continue
+
+            for part in message.parts:
+                if not isinstance(part, AgentSystemPrompt):
+                    continue
+
+                # Check if it's from the same agent type
+                if part.agent_mode == deps.agent_mode:
+                    has_system_prompt = True
+                    break
 
         # Always ensure we have a system prompt for the agent
         # (compaction may remove it from persistent history, but agent needs it)
@@ -461,8 +499,8 @@ class AgentManager(Widget):
         )
 
     def _filter_system_prompts(
-        self, messages: list[ModelMessage]
-    ) -> list[ModelMessage]:
+        self, messages: list[ModelMessage | HintMessage]
+    ) -> list[ModelMessage | HintMessage]:
         """Filter out system prompts from messages for UI display.
 
         Args:
@@ -471,10 +509,12 @@ class AgentManager(Widget):
         Returns:
             List of messages without system prompt parts
         """
-        from pydantic_ai.messages import SystemPromptPart
-
-        filtered_messages: list[ModelMessage] = []
+        filtered_messages: list[ModelMessage | HintMessage] = []
         for msg in messages:
+            if isinstance(msg, HintMessage):
+                filtered_messages.append(msg)
+                continue
+
             parts: Sequence[ModelRequestPart] | Sequence[ModelResponsePart] | None = (
                 msg.parts if hasattr(msg, "parts") else None
             )
@@ -514,6 +554,7 @@ class AgentManager(Widget):
 
         return ConversationState(
             agent_messages=self.message_history.copy(),
+            ui_messages=self.ui_message_history.copy(),
             agent_type=self._current_agent_type.value,
         )
 
@@ -524,10 +565,16 @@ class AgentManager(Widget):
             state: ConversationState object to restore
         """
         # Restore message history for agents (includes system prompts)
-        self.message_history = state.agent_messages.copy()
+        non_hint_messages = [
+            msg for msg in state.agent_messages if not isinstance(msg, HintMessage)
+        ]
+        self.message_history = non_hint_messages
 
-        # Filter out system prompts for UI display
-        self.ui_message_history = self._filter_system_prompts(state.agent_messages)
+        # Filter out system prompts for UI display while keeping hints
+        ui_source = state.ui_messages or cast(
+            list[ModelMessage | HintMessage], state.agent_messages
+        )
+        self.ui_message_history = self._filter_system_prompts(ui_source)
 
         # Restore agent type
         self._current_agent_type = AgentType(state.agent_type)
@@ -535,6 +582,10 @@ class AgentManager(Widget):
         # Notify listeners about the restored messages
         self._post_messages_updated()
 
+    def add_hint_message(self, message: HintMessage) -> None:
+        self.ui_message_history.append(message)
+        self._post_messages_updated()
+
 
 # Re-export AgentType for backward compatibility
 __all__ = [

shotgun/agents/common.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from collections.abc import Callable
+from pathlib import Path
 from typing import Any
 
 from pydantic_ai import (
@@ -15,7 +16,6 @@ from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
-    SystemPromptPart,
 )
 
 from shotgun.agents.config import ProviderType, get_config_manager, get_provider_model
@@ -28,7 +28,8 @@ from shotgun.utils.file_system_utils import get_shotgun_base_path
 
 from .history import token_limit_compactor
 from .history.compaction import apply_persistent_compaction
-from .models import AgentDeps, AgentRuntimeOptions
+from .messages import AgentSystemPrompt, SystemStatusPrompt
+from .models import AgentDeps, AgentRuntimeOptions, PipelineConfigEntry
 from .tools import (
     append_file,
     ask_user,
@@ -84,7 +85,7 @@ async def add_system_status_message(
     message_history.append(
         ModelRequest(
             parts=[
-                SystemPromptPart(content=system_state),
+                SystemStatusPrompt(content=system_state),
             ]
         )
     )
@@ -197,33 +198,24 @@ def create_base_agent(
     return agent, deps
 
 
-def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
-    """Extract table of contents from agent's markdown file.
+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.
 
     Args:
-        agent_mode: The agent mode to extract TOC for
+        file_path: Path to the markdown file
+        max_depth: Maximum heading depth (1=#, 2=##, None=all)
+        max_chars: Maximum characters for the TOC
 
     Returns:
-        Formatted TOC string (up to 2000 chars) or None if not applicable
+        Formatted TOC string or None if file doesn't exist
     """
-    # Skip for EXPORT mode or no mode
-    if (
-        not agent_mode
-        or agent_mode == AgentType.EXPORT
-        or agent_mode not in AGENT_DIRECTORIES
-    ):
-        return None
-
-    base_path = get_shotgun_base_path()
-    md_file = AGENT_DIRECTORIES[agent_mode]
-    md_path = base_path / md_file
-
-    # Check if the markdown file exists
-    if not md_path.exists():
+    if not file_path.exists():
         return None
 
     try:
-        content = md_path.read_text(encoding="utf-8")
+        content = file_path.read_text(encoding="utf-8")
         lines = content.split("\n")
 
         # Extract headings
@@ -239,6 +231,10 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
                     else:
                         break
 
+                # Skip if exceeds max_depth
+                if max_depth and level > max_depth:
+                    continue
+
                 # Get the heading text (remove the # symbols and clean up)
                 heading_text = stripped[level:].strip()
                 if heading_text:
@@ -246,21 +242,109 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
                     indent = "  " * (level - 1)
                     toc_lines.append(f"{indent}{'#' * level} {heading_text}")
 
+                    # Check if we're approaching the character limit
+                    current_length = sum(len(line) + 1 for line in toc_lines)
+                    if current_length > max_chars:
+                        # Remove the last line and add ellipsis
+                        toc_lines.pop()
+                        if toc_lines:
+                            toc_lines.append("  ...")
+                        break
+
         if not toc_lines:
             return None
 
-        # Join and truncate to 2000 characters
-        toc = "\n".join(toc_lines)
-        if len(toc) > 2000:
-            toc = toc[:1997] + "..."
-
-        return toc
+        return "\n".join(toc_lines)
 
     except Exception as e:
-        logger.debug(f"Failed to extract TOC from {md_file}: {e}")
+        logger.debug(f"Failed to extract TOC from {file_path}: {e}")
         return None
 
 
+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'
+    files to maintain context awareness while keeping context window tight.
+
+    Args:
+        agent_mode: The agent mode to extract TOC for
+
+    Returns:
+        Formatted multi-file TOC string or None if not applicable
+    """
+    # Skip if no mode
+    if not agent_mode:
+        return None
+
+    # Define pipeline order and dependencies
+    pipeline_config: dict[AgentType, PipelineConfigEntry] = {
+        AgentType.RESEARCH: PipelineConfigEntry(
+            own_file="research.md",
+            prior_files=[],  # First in pipeline
+        ),
+        AgentType.SPECIFY: PipelineConfigEntry(
+            own_file="specification.md",
+            prior_files=["research.md"],
+        ),
+        AgentType.PLAN: PipelineConfigEntry(
+            own_file="plan.md",
+            prior_files=["research.md", "specification.md"],
+        ),
+        AgentType.TASKS: PipelineConfigEntry(
+            own_file="tasks.md",
+            prior_files=["research.md", "specification.md", "plan.md"],
+        ),
+        AgentType.EXPORT: PipelineConfigEntry(
+            own_file=None,  # Export uses directory
+            prior_files=["research.md", "specification.md", "plan.md", "tasks.md"],
+        ),
+    }
+
+    # Get configuration for current agent
+    if agent_mode not in pipeline_config:
+        return None
+
+    config = pipeline_config[agent_mode]
+    base_path = get_shotgun_base_path()
+    toc_sections: list[str] = []
+
+    # Extract TOCs from prior files (high-level only)
+    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)
+        if prior_toc:
+            # Add section with XML tags
+            toc_sections.append(
+                f'<TABLE_OF_CONTENTS file_name="{prior_file}">\n{prior_toc}\n</TABLE_OF_CONTENTS>'
+            )
+
+    # 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)
+        if own_toc:
+            # Put own file TOC at the beginning with XML tags
+            toc_sections.insert(
+                0,
+                f'<TABLE_OF_CONTENTS file_name="{config.own_file}">\n{own_toc}\n</TABLE_OF_CONTENTS>',
+            )
+
+    # Combine all sections
+    if not toc_sections:
+        return None
+
+    combined_toc = "\n\n".join(toc_sections)
+
+    # Final truncation if needed (should rarely happen with our limits)
+    max_total = 3500  # Conservative total limit
+    if len(combined_toc) > max_total:
+        combined_toc = combined_toc[: max_total - 3] + "..."
+
+    return combined_toc
+
+
 def get_agent_existing_files(agent_mode: AgentType | None = None) -> list[str]:
     """Get list of existing files for the given agent mode.
 
@@ -399,7 +483,9 @@ async def add_system_prompt_message(
 
     # Create system message and prepend to message history
     system_message = ModelRequest(
-        parts=[SystemPromptPart(content=system_prompt_content)]
+        parts=[
+            AgentSystemPrompt(content=system_prompt_content, agent_mode=deps.agent_mode)
+        ]
     )
     message_history.insert(0, system_message)
     logger.debug("✅ System prompt prepended as first message")

shotgun/agents/conversation_history.py

@@ -1,7 +1,7 @@
 """Models and utilities for persisting TUI conversation history."""
 
 from datetime import datetime
-from typing import Any
+from typing import Any, cast
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai.messages import (
@@ -10,11 +10,16 @@ from pydantic_ai.messages import (
 )
 from pydantic_core import to_jsonable_python
 
+from shotgun.tui.screens.chat_screen.hint_message import HintMessage
+
+SerializedMessage = dict[str, Any]
+
 
 class ConversationState(BaseModel):
     """Represents the complete state of a conversation in memory."""
 
     agent_messages: list[ModelMessage]
+    ui_messages: list[ModelMessage | HintMessage] = Field(default_factory=list)
     agent_type: str  # Will store AgentType.value
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -24,9 +29,12 @@ class ConversationHistory(BaseModel):
     """Persistent conversation history for TUI sessions."""
 
     version: int = 1
-    agent_history: list[dict[str, Any]] = Field(
+    agent_history: list[SerializedMessage] = Field(
+        default_factory=list
+    )  # Stores serialized ModelMessage objects
+    ui_history: list[SerializedMessage] = Field(
         default_factory=list
-    )  # Will store serialized ModelMessage objects
+    )  # Stores serialized ModelMessage and HintMessage objects
     last_agent_model: str = "research"
     updated_at: datetime = Field(default_factory=datetime.now)
 
@@ -43,6 +51,25 @@ class ConversationHistory(BaseModel):
             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."""
+
+        def _serialize_message(
+            message: ModelMessage | HintMessage,
+        ) -> Any:
+            if isinstance(message, HintMessage):
+                data = message.model_dump()
+                data["message_type"] = "hint"
+                return data
+            payload = to_jsonable_python(
+                message, fallback=lambda x: str(x), exclude_none=True
+            )
+            if isinstance(payload, dict):
+                payload.setdefault("message_type", "model")
+            return payload
+
+        self.ui_history = [_serialize_message(msg) for msg in messages]
+
     def get_agent_messages(self) -> list[ModelMessage]:
         """Get agent_history as a list of ModelMessage objects.
 
@@ -54,3 +81,26 @@ class ConversationHistory(BaseModel):
 
         # Deserialize from JSON format back to ModelMessage objects
         return ModelMessagesTypeAdapter.validate_python(self.agent_history)
+
+    def get_ui_messages(self) -> list[ModelMessage | HintMessage]:
+        """Get ui_history as a list of Model or hint messages."""
+
+        if not self.ui_history:
+            # Fallback for older conversation files without UI history
+            return cast(list[ModelMessage | HintMessage], self.get_agent_messages())
+
+        messages: list[ModelMessage | HintMessage] = []
+        for item in self.ui_history:
+            message_type = item.get("message_type") if isinstance(item, dict) else None
+            if message_type == "hint":
+                messages.append(HintMessage.model_validate(item))
+                continue
+
+            # Backwards compatibility: data may not include the type marker
+            payload = item
+            if isinstance(payload, dict):
+                payload = {k: v for k, v in payload.items() if k != "message_type"}
+            deserialized = ModelMessagesTypeAdapter.validate_python([payload])
+            messages.append(deserialized[0])
+
+        return messages

shotgun/agents/history/history_processors.py

@@ -6,12 +6,12 @@ from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
     ModelResponse,
-    SystemPromptPart,
     TextPart,
     UserPromptPart,
 )
 
 from shotgun.agents.config.models import shotgun_model_request
+from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.agents.models import AgentDeps
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
@@ -20,8 +20,9 @@ from .constants import 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 (
+    get_agent_system_prompt,
     get_first_user_request,
-    get_system_prompt,
+    get_latest_system_status,
 )
 from .token_estimation import (
     calculate_max_summarization_tokens as _calculate_max_summarization_tokens,
@@ -274,31 +275,39 @@ async def token_limit_compactor(
         new_summary_part = create_marked_summary_part(summary_response)
 
         # Extract essential context from messages before the last summary (if any)
-        system_prompt = ""
+        agent_prompt = ""
+        system_status = ""
         first_user_prompt = ""
         if last_summary_index > 0:
-            # Get system and first user from original conversation
-            system_prompt = get_system_prompt(messages[:last_summary_index]) or ""
+            # Get agent system prompt and first user from original conversation
+            agent_prompt = get_agent_system_prompt(messages[:last_summary_index]) or ""
             first_user_prompt = (
                 get_first_user_request(messages[:last_summary_index]) or ""
             )
 
+        # Get the latest system status from all messages
+        system_status = get_latest_system_status(messages) or ""
+
         # Create the updated summary message
         updated_summary_message = ModelResponse(parts=[new_summary_part])
 
         # Build final compacted history with CLEAN structure
         compacted_messages: list[ModelMessage] = []
 
-        # Only add system/user context if it exists and is meaningful
-        if system_prompt or first_user_prompt:
-            compacted_messages.append(
-                ModelRequest(
-                    parts=[
-                        SystemPromptPart(content=system_prompt),
-                        UserPromptPart(content=first_user_prompt),
-                    ]
-                )
-            )
+        # Build parts for the initial request
+        from pydantic_ai.messages import ModelRequestPart
+
+        parts: list[ModelRequestPart] = []
+        if agent_prompt:
+            parts.append(AgentSystemPrompt(content=agent_prompt))
+        if system_status:
+            parts.append(SystemStatusPrompt(content=system_status))
+        if first_user_prompt:
+            parts.append(UserPromptPart(content=first_user_prompt))
+
+        # Only add if we have at least one part
+        if parts:
+            compacted_messages.append(ModelRequest(parts=parts))
 
         # Add the summary
         compacted_messages.append(updated_summary_message)
@@ -390,19 +399,26 @@ async def _full_compaction(
     marked_summary_part = create_marked_summary_part(summary_response)
 
     # Build compacted history structure
-    system_prompt = get_system_prompt(messages) or ""
+    agent_prompt = get_agent_system_prompt(messages) or ""
+    system_status = get_latest_system_status(messages) or ""
     user_prompt = get_first_user_request(messages) or ""
 
+    # Build parts for the initial request
+    from pydantic_ai.messages import ModelRequestPart
+
+    parts: list[ModelRequestPart] = []
+    if agent_prompt:
+        parts.append(AgentSystemPrompt(content=agent_prompt))
+    if system_status:
+        parts.append(SystemStatusPrompt(content=system_status))
+    if user_prompt:
+        parts.append(UserPromptPart(content=user_prompt))
+
     # Create base structure
-    compacted_messages: list[ModelMessage] = [
-        ModelRequest(
-            parts=[
-                SystemPromptPart(content=system_prompt),
-                UserPromptPart(content=user_prompt),
-            ]
-        ),
-        ModelResponse(parts=[marked_summary_part]),
-    ]
+    compacted_messages: list[ModelMessage] = []
+    if parts:
+        compacted_messages.append(ModelRequest(parts=parts))
+    compacted_messages.append(ModelResponse(parts=[marked_summary_part]))
 
     # Ensure history ends with ModelRequest for PydanticAI compatibility
     compacted_messages = ensure_ends_with_model_request(compacted_messages, messages)

shotgun/agents/history/message_utils.py

@@ -7,6 +7,8 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
+from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
+
 
 def get_first_user_request(messages: list[ModelMessage]) -> str | None:
     """Extract first user request content from messages."""
@@ -37,10 +39,46 @@ def get_user_content_from_request(request: ModelRequest) -> str | None:
 
 
 def get_system_prompt(messages: list[ModelMessage]) -> str | None:
-    """Extract system prompt from messages."""
+    """Extract system prompt from messages (any SystemPromptPart)."""
     for msg in messages:
         if isinstance(msg, ModelRequest):
             for part in msg.parts:
                 if isinstance(part, SystemPromptPart):
                     return part.content
     return None
+
+
+def get_agent_system_prompt(messages: list[ModelMessage]) -> str | None:
+    """Extract the main agent system prompt from messages.
+
+    Prioritizes AgentSystemPrompt but falls back to generic SystemPromptPart
+    if no AgentSystemPrompt is found.
+    """
+    # First try to find AgentSystemPrompt
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, AgentSystemPrompt):
+                    return part.content
+
+    # Fall back to any SystemPromptPart (excluding SystemStatusPrompt)
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, SystemPromptPart) and not isinstance(
+                    part, SystemStatusPrompt
+                ):
+                    return part.content
+
+    return None
+
+
+def get_latest_system_status(messages: list[ModelMessage]) -> str | None:
+    """Extract the most recent system status prompt from messages."""
+    # Iterate in reverse to find the most recent status
+    for msg in reversed(messages):
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, SystemStatusPrompt):
+                    return part.content
+    return None

shotgun/agents/messages.py

@@ -0,0 +1,35 @@
+"""Custom message types for Shotgun agents.
+
+This module defines specialized SystemPromptPart subclasses to distinguish
+between different types of system prompts in the agent pipeline.
+"""
+
+from dataclasses import dataclass, field
+
+from pydantic_ai.messages import SystemPromptPart
+
+from shotgun.agents.models import AgentType
+
+
+@dataclass
+class AgentSystemPrompt(SystemPromptPart):
+    """System prompt containing the main agent instructions.
+
+    This is the primary system prompt that defines the agent's role,
+    capabilities, and behavior. It should be preserved during compaction.
+    """
+
+    prompt_type: str = "agent"
+    agent_mode: AgentType | None = field(default=None)
+
+
+@dataclass
+class SystemStatusPrompt(SystemPromptPart):
+    """System prompt containing current system status information.
+
+    This includes table of contents, available files, and other contextual
+    information about the current state. Only the most recent status should
+    be preserved during compaction.
+    """
+
+    prompt_type: str = "status"