shotgun-sh 0.1.0.dev26__tar.gz → 0.1.0.dev28__tar.gz

{shotgun_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.1.0.dev26
+Version: 0.1.0.dev28
 Summary: AI-powered research, planning, and task management CLI tool
 Project-URL: Homepage, https://shotgun.sh/
 Project-URL: Repository, https://github.com/shotgun-sh/shotgun

{shotgun_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shotgun-sh"
-version = "0.1.0.dev26"
+version = "0.1.0.dev28"
 description = "AI-powered research, planning, and task management CLI tool"
 readme = "README.md"
 license = { text = "MIT" }

{shotgun_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/src/shotgun/agents/agent_manager.py

@@ -41,6 +41,7 @@ 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
@@ -269,6 +270,9 @@ class AgentManager(Widget):
         if deps is None:
             raise ValueError("AgentDeps must be provided")
 
+        # Clear file tracker before each run to track only this run's operations
+        deps.file_tracker.clear()
+
         if prompt:
             self.ui_message_history.append(ModelRequest.user_text_prompt(prompt))
         self._post_messages_updated()
@@ -278,15 +282,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)
@@ -472,8 +512,6 @@ class AgentManager(Widget):
         Returns:
             List of messages without system prompt parts
         """
-        from pydantic_ai.messages import SystemPromptPart
-
         filtered_messages: list[ModelMessage | HintMessage] = []
         for msg in messages:
             if isinstance(msg, HintMessage):

{shotgun_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/src/shotgun/agents/common.py

@@ -16,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
@@ -29,6 +28,7 @@ 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 .messages import AgentSystemPrompt, SystemStatusPrompt
 from .models import AgentDeps, AgentRuntimeOptions, PipelineConfigEntry
 from .tools import (
     append_file,
@@ -85,7 +85,7 @@ async def add_system_status_message(
     message_history.append(
         ModelRequest(
             parts=[
-                SystemPromptPart(content=system_state),
+                SystemStatusPrompt(content=system_state),
             ]
         )
     )
@@ -315,10 +315,9 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
         # 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 header
-            file_label = prior_file.replace(".md", "").replace("_", " ").title()
+            # Add section with XML tags
             toc_sections.append(
-                f"=== Prior Context: {file_label} (summary) ===\n{prior_toc}"
+                f'<TABLE_OF_CONTENTS file_name="{prior_file}">\n{prior_toc}\n</TABLE_OF_CONTENTS>'
             )
 
     # Extract TOC from own file (full detail)
@@ -326,10 +325,10 @@ def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
         own_path = base_path / config.own_file
         own_toc = _extract_file_toc_content(own_path, max_depth=None, max_chars=2000)
         if own_toc:
-            file_label = config.own_file.replace(".md", "").replace("_", " ").title()
-            # Put own file TOC at the beginning
+            # Put own file TOC at the beginning with XML tags
             toc_sections.insert(
-                0, f"=== Your Current Document: {file_label} ===\n{own_toc}"
+                0,
+                f'<TABLE_OF_CONTENTS file_name="{config.own_file}">\n{own_toc}\n</TABLE_OF_CONTENTS>',
             )
 
     # Combine all sections
@@ -484,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_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/src/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_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/src/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_sh-0.1.0.dev28/src/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"

{shotgun_sh-0.1.0.dev26 → shotgun_sh-0.1.0.dev28}/src/shotgun/agents/tools/file_management.py

@@ -20,7 +20,15 @@ AGENT_DIRECTORIES = {
     AgentType.SPECIFY: "specification.md",
     AgentType.PLAN: "plan.md",
     AgentType.TASKS: "tasks.md",
-    AgentType.EXPORT: "exports/",
+    AgentType.EXPORT: "*",  # Export agent can write anywhere except protected files
+}
+
+# Files protected from export agent modifications
+PROTECTED_AGENT_FILES = {
+    "research.md",
+    "specification.md",
+    "plan.md",
+    "tasks.md",
 }
 
 
@@ -40,21 +48,17 @@ def _validate_agent_scoped_path(filename: str, agent_mode: AgentType | None) ->
     base_path = get_shotgun_base_path()
 
     if agent_mode and agent_mode in AGENT_DIRECTORIES:
-        # For export mode, allow writing to any file in exports/ directory
+        # For export mode, allow writing to any file except protected agent files
         if agent_mode == AgentType.EXPORT:
-            # Ensure the filename starts with exports/ or is being written to exports/
-            if not filename.startswith("exports/"):
-                filename = f"exports/{filename}"
-            full_path = (base_path / filename).resolve()
-
-            # Ensure it's within .shotgun/exports/
-            exports_dir = base_path / "exports"
-            try:
-                full_path.relative_to(exports_dir.resolve())
-            except ValueError as e:
+            # Check if trying to write to a protected file
+            if filename in PROTECTED_AGENT_FILES:
                 raise ValueError(
-                    f"Export agent can only write to exports/ directory. Path '{filename}' is not allowed"
-                ) from e
+                    f"Export agent cannot write to protected file '{filename}'. "
+                    f"Protected files are: {', '.join(sorted(PROTECTED_AGENT_FILES))}"
+                )
+
+            # Allow writing anywhere else in .shotgun directory
+            full_path = (base_path / filename).resolve()
         else:
             # For other agents, only allow writing to their specific file
             allowed_file = AGENT_DIRECTORIES[agent_mode]