shotgun-sh 0.2.11__py3-none-any.whl → 0.2.11.dev2__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -58,12 +58,7 @@ from shotgun.agents.context_analyzer import (
     ContextCompositionTelemetry,
     ContextFormatter,
 )
-from shotgun.agents.models import (
-    AgentResponse,
-    AgentType,
-    FileOperation,
-    FileOperationTracker,
-)
+from shotgun.agents.models import AgentResponse, AgentType, FileOperation
 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
@@ -174,14 +169,6 @@ 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.
@@ -235,7 +222,7 @@ class AgentManager(Widget):
         self.deps = deps
 
         # Create AgentRuntimeOptions from deps for agent creation
-        self._agent_runtime_options = AgentRuntimeOptions(
+        agent_runtime_options = AgentRuntimeOptions(
             interactive_mode=self.deps.interactive_mode,
             working_directory=self.deps.working_directory,
             is_tui_context=self.deps.is_tui_context,
@@ -244,18 +231,22 @@ class AgentManager(Widget):
             tasks=self.deps.tasks,
         )
 
-        # 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
+        # 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
+        )
 
         # Track current active agent
         self._current_agent_type: AgentType = initial_type
@@ -270,119 +261,6 @@ 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.
@@ -534,9 +412,6 @@ 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:
@@ -774,12 +649,6 @@ 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(
@@ -815,9 +684,11 @@ class AgentManager(Widget):
                     )
                 )
 
-            # 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([])
+            # 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)
         else:
             # No clarifying questions - show the response or a default success message
             if agent_response.response and agent_response.response.strip():
@@ -852,9 +723,10 @@ class AgentManager(Widget):
                     )
 
             # Post UI update immediately so user sees the response without delay
-            # (file operations will be posted after compaction to avoid duplicates)
-            logger.debug("Posting immediate UI update with hint message")
-            self._post_messages_updated([])
+            logger.debug(
+                "Posting immediate UI update with hint message and file operations"
+            )
+            self._post_messages_updated(file_operations)
 
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
@@ -908,7 +780,7 @@ class AgentManager(Widget):
 
         usage = result.usage()
         if hasattr(deps, "llm_model") and deps.llm_model is not None:
-            await deps.usage_manager.add_usage(
+            deps.usage_manager.add_usage(
                 usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
             )
         else:
@@ -934,9 +806,6 @@ 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()
@@ -1115,9 +984,6 @@ 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:
@@ -1145,38 +1011,6 @@ 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,7 +4,6 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
-import aiofiles
 from pydantic_ai import (
     Agent,
     RunContext,
@@ -69,7 +68,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 = await extract_markdown_toc(deps.agent_mode)
+    markdown_toc = extract_markdown_toc(deps.agent_mode)
 
     # Get current datetime with timezone information
     dt_context = get_datetime_context()
@@ -95,7 +94,7 @@ async def add_system_status_message(
     return message_history
 
 
-async def create_base_agent(
+def create_base_agent(
     system_prompt_fn: Callable[[RunContext[AgentDeps]], str],
     agent_runtime_options: AgentRuntimeOptions,
     load_codebase_understanding_tools: bool = True,
@@ -120,7 +119,7 @@ async def create_base_agent(
 
     # Get configured model or fall back to first available provider
     try:
-        model_config = await get_provider_model(provider)
+        model_config = get_provider_model(provider)
         provider_name = model_config.provider
         logger.debug(
             "🤖 Creating agent with configured %s model: %s",
@@ -195,7 +194,7 @@ async def create_base_agent(
     return agent, deps
 
 
-async def _extract_file_toc_content(
+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.
@@ -212,8 +211,7 @@ async def _extract_file_toc_content(
         return None
 
     try:
-        async with aiofiles.open(file_path, encoding="utf-8") as f:
-            content = await f.read()
+        content = file_path.read_text(encoding="utf-8")
         lines = content.split("\n")
 
         # Extract headings
@@ -259,7 +257,7 @@ async def _extract_file_toc_content(
         return None
 
 
-async def extract_markdown_toc(agent_mode: AgentType | None) -> str | 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'
@@ -311,9 +309,7 @@ async 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 = await _extract_file_toc_content(
-            file_path, max_depth=2, max_chars=500
-        )
+        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(
@@ -325,9 +321,7 @@ async 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 = await _extract_file_toc_content(
-            own_path, max_depth=None, max_chars=2000
-        )
+        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(

shotgun/agents/config/manager.py

@@ -5,8 +5,6 @@ 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
@@ -50,7 +48,7 @@ class ConfigManager:
 
         self._config: ShotgunConfig | None = None
 
-    async def load(self, force_reload: bool = True) -> ShotgunConfig:
+    def load(self, force_reload: bool = True) -> ShotgunConfig:
         """Load configuration from file.
 
         Args:
@@ -62,19 +60,18 @@ class ConfigManager:
         if self._config is not None and not force_reload:
             return self._config
 
-        if not await aiofiles.os.path.exists(self.config_path):
+        if not self.config_path.exists():
             logger.info(
                 "Configuration file not found, creating new config at: %s",
                 self.config_path,
             )
             # Create new config with generated shotgun_instance_id
-            self._config = await self.initialize()
+            self._config = self.initialize()
             return self._config
 
         try:
-            async with aiofiles.open(self.config_path, encoding="utf-8") as f:
-                content = await f.read()
-                data = json.loads(content)
+            with open(self.config_path, encoding="utf-8") as f:
+                data = json.load(f)
 
             # Migration: Rename user_id to shotgun_instance_id (config v2 -> v3)
             if "user_id" in data and SHOTGUN_INSTANCE_ID_FIELD not in data:
@@ -104,12 +101,6 @@ 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)
 
@@ -126,7 +117,7 @@ class ConfigManager:
 
                     if self._config.selected_model in MODEL_SPECS:
                         spec = MODEL_SPECS[self._config.selected_model]
-                        if not await self.has_provider_key(spec.provider):
+                        if not self.has_provider_key(spec.provider):
                             logger.info(
                                 "Selected model %s provider has no API key, finding available model",
                                 self._config.selected_model.value,
@@ -144,7 +135,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 await self.has_provider_key(provider):
+                        if self.has_provider_key(provider):
                             # Set to that provider's default model
                             from .models import MODEL_SPECS, ModelName
 
@@ -165,7 +156,7 @@ class ConfigManager:
                                 break
 
                 if should_save:
-                    await self.save(self._config)
+                    self.save(self._config)
 
             return self._config
 
@@ -174,10 +165,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 = await self.initialize()
+            self._config = self.initialize()
             return self._config
 
-    async def save(self, config: ShotgunConfig | None = None) -> None:
+    def save(self, config: ShotgunConfig | None = None) -> None:
         """Save configuration to file.
 
         Args:
@@ -193,17 +184,15 @@ class ConfigManager:
                 )
 
         # Ensure directory exists
-        await aiofiles.os.makedirs(self.config_path.parent, exist_ok=True)
+        self.config_path.parent.mkdir(parents=True, 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)
 
-            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)
+            with open(self.config_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
 
             logger.debug("Configuration saved to %s", self.config_path)
             self._config = config
@@ -212,16 +201,14 @@ class ConfigManager:
             logger.error("Failed to save configuration to %s: %s", self.config_path, e)
             raise
 
-    async def update_provider(
-        self, provider: ProviderType | str, **kwargs: Any
-    ) -> None:
+    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 = await self.load()
+        config = self.load()
 
         # Get provider config and check if it's shotgun
         provider_config, is_shotgun = self._get_provider_config_and_type(
@@ -266,11 +253,11 @@ class ConfigManager:
             # This prevents the welcome screen from showing again after user has made their choice
             config.shown_welcome_screen = True
 
-        await self.save(config)
+        self.save(config)
 
-    async def clear_provider_key(self, provider: ProviderType | str) -> None:
+    def clear_provider_key(self, provider: ProviderType | str) -> None:
         """Remove the API key for the given provider (LLM provider or shotgun)."""
-        config = await self.load()
+        config = self.load()
 
         # Get provider config (shotgun or LLM provider)
         provider_config, is_shotgun = self._get_provider_config_and_type(
@@ -283,34 +270,34 @@ class ConfigManager:
         if is_shotgun and isinstance(provider_config, ShotgunAccountConfig):
             provider_config.supabase_jwt = None
 
-        await self.save(config)
+        self.save(config)
 
-    async def update_selected_model(self, model_name: "ModelName") -> None:
+    def update_selected_model(self, model_name: "ModelName") -> None:
         """Update the selected model.
 
         Args:
             model_name: Model to select
         """
-        config = await self.load()
+        config = self.load()
         config.selected_model = model_name
-        await self.save(config)
+        self.save(config)
 
-    async def has_provider_key(self, provider: ProviderType | str) -> bool:
+    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 = await self.load(force_reload=False)
+        config = 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)
 
-    async def has_any_provider_key(self) -> bool:
+    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 = await self.load(force_reload=False)
+        config = 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))
@@ -324,7 +311,7 @@ class ConfigManager:
         has_shotgun_key = self._provider_has_api_key(config.shotgun)
         return has_llm_key or has_shotgun_key
 
-    async def initialize(self) -> ShotgunConfig:
+    def initialize(self) -> ShotgunConfig:
         """Initialize configuration with defaults and save to file.
 
         Returns:
@@ -334,7 +321,7 @@ class ConfigManager:
         config = ShotgunConfig(
             shotgun_instance_id=str(uuid.uuid4()),
         )
-        await self.save(config)
+        self.save(config)
         logger.info(
             "Configuration initialized at %s with shotgun_instance_id: %s",
             self.config_path,
@@ -390,24 +377,6 @@ 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 (
@@ -471,16 +440,16 @@ class ConfigManager:
         provider_enum = self._ensure_provider_enum(provider)
         return (self._get_provider_config(config, provider_enum), False)
 
-    async def get_shotgun_instance_id(self) -> str:
+    def get_shotgun_instance_id(self) -> str:
         """Get the shotgun instance ID from configuration.
 
         Returns:
             The unique shotgun instance ID string
         """
-        config = await self.load()
+        config = self.load()
         return config.shotgun_instance_id
 
-    async def update_shotgun_account(
+    def update_shotgun_account(
         self, api_key: str | None = None, supabase_jwt: str | None = None
     ) -> None:
         """Update Shotgun Account configuration.
@@ -489,7 +458,7 @@ class ConfigManager:
             api_key: LiteLLM proxy API key (optional)
             supabase_jwt: Supabase authentication JWT (optional)
         """
-        config = await self.load()
+        config = self.load()
 
         if api_key is not None:
             config.shotgun.api_key = SecretStr(api_key) if api_key else None
@@ -499,7 +468,7 @@ class ConfigManager:
                 SecretStr(supabase_jwt) if supabase_jwt else None
             )
 
-        await self.save(config)
+        self.save(config)
         logger.info("Updated Shotgun Account configuration")