shotgun-sh 0.1.0.dev14__py3-none-any.whl → 0.1.0.dev16__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -1,16 +1,30 @@
 """Agent manager for coordinating multiple AI agents with shared message history."""
 
+import logging
+from collections.abc import AsyncIterable
+from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any
+from typing import Any, cast
 
 from pydantic_ai import (
     Agent,
     DeferredToolRequests,
     DeferredToolResults,
+    RunContext,
     UsageLimits,
 )
 from pydantic_ai.agent import AgentRunResult
-from pydantic_ai.messages import ModelMessage, ModelRequest
+from pydantic_ai.messages import (
+    AgentStreamEvent,
+    FinalResultEvent,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    PartDeltaEvent,
+    PartStartEvent,
+    ToolCallPartDelta,
+)
 from textual.message import Message
 from textual.widget import Widget
 
@@ -18,8 +32,11 @@ from .history.compaction import apply_persistent_compaction
 from .models import AgentDeps, AgentRuntimeOptions, FileOperation
 from .plan import create_plan_agent
 from .research import create_research_agent
+from .specify import create_specify_agent
 from .tasks import create_tasks_agent
 
+logger = logging.getLogger(__name__)
+
 
 class AgentType(Enum):
     """Enumeration for available agent types (for Python < 3.11)."""
@@ -27,21 +44,48 @@ class AgentType(Enum):
     RESEARCH = "research"
     PLAN = "plan"
     TASKS = "tasks"
+    SPECIFY = "specify"
 
 
 class MessageHistoryUpdated(Message):
     """Event posted when the message history is updated."""
 
-    def __init__(self, messages: list[ModelMessage], agent_type: AgentType) -> None:
+    def __init__(
+        self,
+        messages: list[ModelMessage],
+        agent_type: AgentType,
+        file_operations: list[FileOperation] | None = None,
+    ) -> None:
         """Initialize the message history updated event.
 
         Args:
             messages: The updated message history.
             agent_type: The type of agent that triggered the update.
+            file_operations: List of file operations from this run.
         """
         super().__init__()
         self.messages = messages
         self.agent_type = agent_type
+        self.file_operations = file_operations or []
+
+
+class PartialResponseMessage(Message):
+    """Event posted when a partial response is received."""
+
+    def __init__(self, message: ModelResponse | None, is_last: bool) -> None:
+        """Initialize the partial response message."""
+        super().__init__()
+        self.message = message
+        self.is_last = is_last
+
+
+@dataclass(slots=True)
+class _PartialStreamState:
+    """Tracks partial response parts while streaming a single agent run."""
+
+    parts: list[ModelResponsePart | ToolCallPartDelta] = field(default_factory=list)
+    latest_partial: ModelResponse | None = None
+    final_sent: bool = False
 
 
 class AgentManager(Widget):
@@ -58,6 +102,8 @@ class AgentManager(Widget):
             deps: Optional agent dependencies. If not provided, defaults to interactive mode.
         """
         super().__init__()
+        self.display = False
+
         # Use provided deps or create default with interactive mode
         self.deps = deps
 
@@ -73,14 +119,17 @@ class AgentManager(Widget):
             tasks=self.deps.tasks,
         )
 
-        # Initialize all agents with the same deps
-        self.research_agent, _ = create_research_agent(
+        # 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, _ = create_plan_agent(
+        self.plan_agent, self.plan_deps = create_plan_agent(
             agent_runtime_options=agent_runtime_options
         )
-        self.tasks_agent, _ = create_tasks_agent(
+        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
         )
 
@@ -91,6 +140,7 @@ class AgentManager(Widget):
         self.ui_message_history: list[ModelMessage] = []
         self.message_history: list[ModelMessage] = []
         self.recently_change_files: list[FileOperation] = []
+        self._stream_state: _PartialStreamState | None = None
 
     @property
     def current_agent(self) -> Agent[AgentDeps, str | DeferredToolRequests]:
@@ -116,9 +166,52 @@ class AgentManager(Widget):
             AgentType.RESEARCH: self.research_agent,
             AgentType.PLAN: self.plan_agent,
             AgentType.TASKS: self.tasks_agent,
+            AgentType.SPECIFY: self.specify_agent,
         }
         return agent_map[agent_type]
 
+    def _get_agent_deps(self, agent_type: AgentType) -> AgentDeps:
+        """Get agent-specific deps by type.
+
+        Args:
+            agent_type: The type of agent to retrieve deps for.
+
+        Returns:
+            The agent-specific dependencies.
+        """
+        deps_map = {
+            AgentType.RESEARCH: self.research_deps,
+            AgentType.PLAN: self.plan_deps,
+            AgentType.TASKS: self.tasks_deps,
+            AgentType.SPECIFY: self.specify_deps,
+        }
+        return deps_map[agent_type]
+
+    def _create_merged_deps(self, agent_type: AgentType) -> AgentDeps:
+        """Create merged dependencies combining shared and agent-specific deps.
+
+        This preserves the agent's system_prompt_fn while using shared runtime state.
+
+        Args:
+            agent_type: The type of agent to create merged deps for.
+
+        Returns:
+            Merged AgentDeps with agent-specific system_prompt_fn.
+        """
+        agent_deps = self._get_agent_deps(agent_type)
+
+        # Ensure shared deps is not None (should be guaranteed by __init__)
+        if self.deps is None:
+            raise ValueError("Shared deps is None - this should not happen")
+
+        # Create new deps with shared runtime state but agent's system_prompt_fn
+        # Use a copy of the shared deps and update the system_prompt_fn
+        merged_deps = self.deps.model_copy(
+            update={"system_prompt_fn": agent_deps.system_prompt_fn}
+        )
+
+        return merged_deps
+
     def set_agent(self, agent_type: AgentType) -> None:
         """Set the current active agent.
 
@@ -159,9 +252,9 @@ class AgentManager(Widget):
         Returns:
             The agent run result.
         """
-        # Use manager's deps if not provided
+        # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
-            deps = self.deps
+            deps = self._create_merged_deps(self._current_agent_type)
 
         # Ensure deps is not None
         if deps is None:
@@ -171,39 +264,176 @@ class AgentManager(Widget):
             self.ui_message_history.append(ModelRequest.user_text_prompt(prompt))
         self._post_messages_updated()
 
-        # Run the agent with the shared message history
-        result: AgentRunResult[
-            str | DeferredToolRequests
-        ] = await self.current_agent.run(
-            prompt,
-            deps=deps,
-            usage_limits=usage_limits,
-            message_history=self.message_history,
-            deferred_tool_results=deferred_tool_results,
-            **kwargs,
+        # Ensure system prompt is added to message history before running agent
+        from pydantic_ai.messages import SystemPromptPart
+
+        from shotgun.agents.common import add_system_prompt_message
+
+        # Start with persistent message history
+        message_history = self.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
+        )
+
+        # Always ensure we have a system prompt for the agent
+        # (compaction may remove it from persistent history, but agent needs it)
+        if not has_system_prompt:
+            message_history = await add_system_prompt_message(deps, message_history)
+
+        # Run the agent with streaming support (from origin/main)
+        self._stream_state = _PartialStreamState()
+
+        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()
         )
 
-        # Update the shared message history with all messages from this run
+        try:
+            result: AgentRunResult[
+                str | DeferredToolRequests
+            ] = await self.current_agent.run(
+                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,
+                **kwargs,
+            )
+        finally:
+            # If the stream ended unexpectedly without a final result, clear accumulated state.
+            if self._stream_state is not None and not self._stream_state.final_sent:
+                partial_message = self._build_partial_response(self._stream_state.parts)
+                if partial_message is not None:
+                    self._post_partial_message(partial_message, True)
+            self._stream_state = None
+
         self.ui_message_history = self.ui_message_history + [
             mes for mes in result.new_messages() if not isinstance(mes, ModelRequest)
         ]
 
         # Apply compaction to persistent message history to prevent cascading growth
-        self.message_history = await apply_persistent_compaction(
-            result.all_messages(), deps
-        )
-        self._post_messages_updated()
+        all_messages = result.all_messages()
+        self.message_history = await apply_persistent_compaction(all_messages, deps)
 
         # Log file operations summary if any files were modified
-        self.recently_change_files = deps.file_tracker.operations.copy()
+        file_operations = deps.file_tracker.operations.copy()
+        self.recently_change_files = file_operations
+
+        self._post_messages_updated(file_operations)
 
         return result
 
-    def _post_messages_updated(self) -> None:
+    async def _handle_event_stream(
+        self,
+        _ctx: RunContext[AgentDeps],
+        stream: AsyncIterable[AgentStreamEvent],
+    ) -> None:
+        """Process streamed events and forward partial updates to the UI."""
+
+        state = self._stream_state
+        if state is None:
+            state = self._stream_state = _PartialStreamState()
+
+        partial_parts = state.parts
+
+        async for event in stream:
+            try:
+                if isinstance(event, PartStartEvent):
+                    index = event.index
+                    if index < len(partial_parts):
+                        partial_parts[index] = event.part
+                    elif index == len(partial_parts):
+                        partial_parts.append(event.part)
+                    else:
+                        logger.warning(
+                            "Received PartStartEvent with out-of-bounds index",
+                            extra={"index": index, "current_len": len(partial_parts)},
+                        )
+                        partial_parts.append(event.part)
+
+                    partial_message = self._build_partial_response(partial_parts)
+                    if partial_message is not None:
+                        state.latest_partial = partial_message
+                        self._post_partial_message(partial_message, False)
+
+                elif isinstance(event, PartDeltaEvent):
+                    index = event.index
+                    if index >= len(partial_parts):
+                        logger.warning(
+                            "Received PartDeltaEvent before corresponding start event",
+                            extra={"index": index, "current_len": len(partial_parts)},
+                        )
+                        continue
+
+                    try:
+                        updated_part = event.delta.apply(
+                            cast(ModelResponsePart, partial_parts[index])
+                        )
+                    except Exception:  # pragma: no cover - defensive logging
+                        logger.exception(
+                            "Failed to apply part delta", extra={"event": event}
+                        )
+                        continue
+
+                    partial_parts[index] = updated_part
+
+                    partial_message = self._build_partial_response(partial_parts)
+                    if partial_message is not None:
+                        state.latest_partial = partial_message
+                        self._post_partial_message(partial_message, False)
+
+                elif isinstance(event, FinalResultEvent):
+                    final_message = (
+                        state.latest_partial
+                        or self._build_partial_response(partial_parts)
+                    )
+                    self._post_partial_message(final_message, True)
+                    state.latest_partial = None
+                    state.final_sent = True
+                    partial_parts.clear()
+                    self._stream_state = None
+                    break
+
+                # Ignore other AgentStreamEvent variants (e.g. tool call notifications) for partial UI updates.
+
+            except Exception:  # pragma: no cover - defensive logging
+                logger.exception(
+                    "Error while handling agent stream event", extra={"event": event}
+                )
+
+    def _build_partial_response(
+        self, parts: list[ModelResponsePart | ToolCallPartDelta]
+    ) -> ModelResponse | None:
+        """Create a `ModelResponse` from the currently streamed parts."""
+
+        completed_parts = [
+            part for part in parts if not isinstance(part, ToolCallPartDelta)
+        ]
+        if not completed_parts:
+            return None
+        return ModelResponse(parts=list(completed_parts))
+
+    def _post_partial_message(
+        self, message: ModelResponse | None, is_last: bool
+    ) -> None:
+        """Post a partial message to the UI."""
+        self.post_message(PartialResponseMessage(message, is_last))
+
+    def _post_messages_updated(
+        self, file_operations: list[FileOperation] | None = None
+    ) -> None:
         # Post event to notify listeners of the message history update
         self.post_message(
             MessageHistoryUpdated(
                 messages=self.ui_message_history.copy(),
                 agent_type=self._current_agent_type,
+                file_operations=file_operations,
             )
         )

shotgun/agents/models.py

@@ -1,5 +1,6 @@
 """Pydantic models for agent dependencies and configuration."""
 
+import os
 from asyncio import Future, Queue
 from collections.abc import Callable
 from datetime import datetime
@@ -187,6 +188,25 @@ class FileOperationTracker(BaseModel):
 
         return "\n".join(lines)
 
+    def get_display_path(self) -> str | None:
+        """Get a single file path or common parent directory for display.
+
+        Returns:
+            Path string to display, or None if no files were modified
+        """
+        if not self.operations:
+            return None
+
+        unique_paths = list({op.file_path for op in self.operations})
+
+        if len(unique_paths) == 1:
+            # Single file - return its path
+            return unique_paths[0]
+
+        # Multiple files - find common parent directory
+        common_path = os.path.commonpath(unique_paths)
+        return common_path
+
 
 class AgentDeps(AgentRuntimeOptions):
     """Dependencies passed to all agents for configuration and runtime behavior."""

shotgun/agents/tools/artifact_management.py

@@ -347,14 +347,13 @@ async def read_artifact_section(
 
         section = service.get_section(artifact_id, mode, section_number)
 
-        # Return formatted content with title
-        formatted_content = f"# {section.title}\n\n{section.content}"
+        # Return section content (already contains title header from file storage)
         logger.debug(
             "📄 Read section %d with %d characters",
             section_number,
             len(section.content),
         )
-        return formatted_content
+        return section.content
 
     except Exception as e:
         error_msg = (