shotgun-sh 0.2.11.dev1__py3-none-any.whl → 0.2.17.dev1__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -58,7 +58,12 @@ from shotgun.agents.context_analyzer import (
     ContextCompositionTelemetry,
     ContextFormatter,
 )
-from shotgun.agents.models import AgentResponse, AgentType, FileOperation
+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
@@ -169,6 +174,14 @@ 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.
@@ -222,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,
@@ -231,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
@@ -261,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.
@@ -412,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:
@@ -649,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(
@@ -684,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():
@@ -723,10 +852,9 @@ 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()
@@ -780,7 +908,7 @@ class AgentManager(Widget):
 
         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:
@@ -806,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()
@@ -984,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:
@@ -1011,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:

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/manager.py

@@ -5,6 +5,8 @@ import uuid
 from pathlib import Path
 from typing import Any
 
+import aiofiles
+import aiofiles.os
 from pydantic import SecretStr
 
 from shotgun.logging_config import get_logger
@@ -48,7 +50,7 @@ class ConfigManager:
 
         self._config: ShotgunConfig | None = None
 
-    def load(self, force_reload: bool = True) -> ShotgunConfig:
+    async def load(self, force_reload: bool = True) -> ShotgunConfig:
         """Load configuration from file.
 
         Args:
@@ -60,18 +62,19 @@ class ConfigManager:
         if self._config is not None and not force_reload:
             return self._config
 
-        if not self.config_path.exists():
+        if not await aiofiles.os.path.exists(self.config_path):
             logger.info(
                 "Configuration file not found, creating new config at: %s",
                 self.config_path,
             )
             # Create new config with generated shotgun_instance_id
-            self._config = self.initialize()
+            self._config = await self.initialize()
             return self._config
 
         try:
-            with open(self.config_path, encoding="utf-8") as f:
-                data = json.load(f)
+            async with aiofiles.open(self.config_path, encoding="utf-8") as f:
+                content = await f.read()
+                data = json.loads(content)
 
             # Migration: Rename user_id to shotgun_instance_id (config v2 -> v3)
             if "user_id" in data and SHOTGUN_INSTANCE_ID_FIELD not in data:
@@ -101,6 +104,12 @@ class ConfigManager:
                         "Existing BYOK user detected: set shown_welcome_screen=False to show welcome screen"
                     )
 
+            # Migration: Add marketing config for v3 -> v4
+            if "marketing" not in data:
+                data["marketing"] = {"messages": {}}
+                data["config_version"] = 4
+                logger.info("Migrated config v3->v4: added marketing configuration")
+
             # Convert plain text secrets to SecretStr objects
             self._convert_secrets_to_secretstr(data)
 
@@ -117,7 +126,7 @@ class ConfigManager:
 
                     if self._config.selected_model in MODEL_SPECS:
                         spec = MODEL_SPECS[self._config.selected_model]
-                        if not self.has_provider_key(spec.provider):
+                        if not await self.has_provider_key(spec.provider):
                             logger.info(
                                 "Selected model %s provider has no API key, finding available model",
                                 self._config.selected_model.value,
@@ -135,7 +144,7 @@ class ConfigManager:
                 # If no selected_model or it was invalid, find first available model
                 if not self._config.selected_model:
                     for provider in ProviderType:
-                        if self.has_provider_key(provider):
+                        if await self.has_provider_key(provider):
                             # Set to that provider's default model
                             from .models import MODEL_SPECS, ModelName
 
@@ -156,7 +165,7 @@ class ConfigManager:
                                 break
 
                 if should_save:
-                    self.save(self._config)
+                    await self.save(self._config)
 
             return self._config
 
@@ -165,10 +174,10 @@ class ConfigManager:
                 "Failed to load configuration from %s: %s", self.config_path, e
             )
             logger.info("Creating new configuration with generated shotgun_instance_id")
-            self._config = self.initialize()
+            self._config = await self.initialize()
             return self._config
 
-    def save(self, config: ShotgunConfig | None = None) -> None:
+    async def save(self, config: ShotgunConfig | None = None) -> None:
         """Save configuration to file.
 
         Args:
@@ -184,15 +193,17 @@ class ConfigManager:
                 )
 
         # Ensure directory exists
-        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+        await aiofiles.os.makedirs(self.config_path.parent, exist_ok=True)
 
         try:
             # Convert SecretStr to plain text for JSON serialization
             data = config.model_dump()
             self._convert_secretstr_to_plain(data)
+            self._convert_datetime_to_isoformat(data)
 
-            with open(self.config_path, "w", encoding="utf-8") as f:
-                json.dump(data, f, indent=2, ensure_ascii=False)
+            json_content = json.dumps(data, indent=2, ensure_ascii=False)
+            async with aiofiles.open(self.config_path, "w", encoding="utf-8") as f:
+                await f.write(json_content)
 
             logger.debug("Configuration saved to %s", self.config_path)
             self._config = config
@@ -201,14 +212,16 @@ class ConfigManager:
             logger.error("Failed to save configuration to %s: %s", self.config_path, e)
             raise
 
-    def update_provider(self, provider: ProviderType | str, **kwargs: Any) -> None:
+    async def update_provider(
+        self, provider: ProviderType | str, **kwargs: Any
+    ) -> None:
         """Update provider configuration.
 
         Args:
             provider: Provider to update
             **kwargs: Configuration fields to update (only api_key supported)
         """
-        config = self.load()
+        config = await self.load()
 
         # Get provider config and check if it's shotgun
         provider_config, is_shotgun = self._get_provider_config_and_type(
@@ -253,11 +266,11 @@ class ConfigManager:
             # This prevents the welcome screen from showing again after user has made their choice
             config.shown_welcome_screen = True
 
-        self.save(config)
+        await self.save(config)
 
-    def clear_provider_key(self, provider: ProviderType | str) -> None:
+    async def clear_provider_key(self, provider: ProviderType | str) -> None:
         """Remove the API key for the given provider (LLM provider or shotgun)."""
-        config = self.load()
+        config = await self.load()
 
         # Get provider config (shotgun or LLM provider)
         provider_config, is_shotgun = self._get_provider_config_and_type(
@@ -270,34 +283,34 @@ class ConfigManager:
         if is_shotgun and isinstance(provider_config, ShotgunAccountConfig):
             provider_config.supabase_jwt = None
 
-        self.save(config)
+        await self.save(config)
 
-    def update_selected_model(self, model_name: "ModelName") -> None:
+    async def update_selected_model(self, model_name: "ModelName") -> None:
         """Update the selected model.
 
         Args:
             model_name: Model to select
         """
-        config = self.load()
+        config = await self.load()
         config.selected_model = model_name
-        self.save(config)
+        await self.save(config)
 
-    def has_provider_key(self, provider: ProviderType | str) -> bool:
+    async def has_provider_key(self, provider: ProviderType | str) -> bool:
         """Check if the given provider has a non-empty API key configured.
 
         This checks only the configuration file.
         """
         # Use force_reload=False to avoid infinite loop when called from load()
-        config = self.load(force_reload=False)
+        config = await self.load(force_reload=False)
         provider_enum = self._ensure_provider_enum(provider)
         provider_config = self._get_provider_config(config, provider_enum)
 
         return self._provider_has_api_key(provider_config)
 
-    def has_any_provider_key(self) -> bool:
+    async def has_any_provider_key(self) -> bool:
         """Determine whether any provider has a configured API key."""
         # Use force_reload=False to avoid infinite loop when called from load()
-        config = self.load(force_reload=False)
+        config = await self.load(force_reload=False)
         # Check LLM provider keys (BYOK)
         has_llm_key = any(
             self._provider_has_api_key(self._get_provider_config(config, provider))
@@ -311,7 +324,7 @@ class ConfigManager:
         has_shotgun_key = self._provider_has_api_key(config.shotgun)
         return has_llm_key or has_shotgun_key
 
-    def initialize(self) -> ShotgunConfig:
+    async def initialize(self) -> ShotgunConfig:
         """Initialize configuration with defaults and save to file.
 
         Returns:
@@ -321,7 +334,7 @@ class ConfigManager:
         config = ShotgunConfig(
             shotgun_instance_id=str(uuid.uuid4()),
         )
-        self.save(config)
+        await self.save(config)
         logger.info(
             "Configuration initialized at %s with shotgun_instance_id: %s",
             self.config_path,
@@ -377,6 +390,24 @@ class ConfigManager:
                             SUPABASE_JWT_FIELD
                         ].get_secret_value()
 
+    def _convert_datetime_to_isoformat(self, data: dict[str, Any]) -> None:
+        """Convert datetime objects in data to ISO8601 format strings for JSON serialization."""
+        from datetime import datetime
+
+        def convert_dict(d: dict[str, Any]) -> None:
+            """Recursively convert datetime objects in a dict."""
+            for key, value in d.items():
+                if isinstance(value, datetime):
+                    d[key] = value.isoformat()
+                elif isinstance(value, dict):
+                    convert_dict(value)
+                elif isinstance(value, list):
+                    for item in value:
+                        if isinstance(item, dict):
+                            convert_dict(item)
+
+        convert_dict(data)
+
     def _ensure_provider_enum(self, provider: ProviderType | str) -> ProviderType:
         """Normalize provider values to ProviderType enum."""
         return (
@@ -440,16 +471,16 @@ class ConfigManager:
         provider_enum = self._ensure_provider_enum(provider)
         return (self._get_provider_config(config, provider_enum), False)
 
-    def get_shotgun_instance_id(self) -> str:
+    async def get_shotgun_instance_id(self) -> str:
         """Get the shotgun instance ID from configuration.
 
         Returns:
             The unique shotgun instance ID string
         """
-        config = self.load()
+        config = await self.load()
         return config.shotgun_instance_id
 
-    def update_shotgun_account(
+    async def update_shotgun_account(
         self, api_key: str | None = None, supabase_jwt: str | None = None
     ) -> None:
         """Update Shotgun Account configuration.
@@ -458,7 +489,7 @@ class ConfigManager:
             api_key: LiteLLM proxy API key (optional)
             supabase_jwt: Supabase authentication JWT (optional)
         """
-        config = self.load()
+        config = await self.load()
 
         if api_key is not None:
             config.shotgun.api_key = SecretStr(api_key) if api_key else None
@@ -468,7 +499,7 @@ class ConfigManager:
                 SecretStr(supabase_jwt) if supabase_jwt else None
             )
 
-        self.save(config)
+        await self.save(config)
         logger.info("Updated Shotgun Account configuration")