shotgun-sh 0.2.3.dev2__py3-none-any.whl → 0.2.11.dev5__py3-none-any.whl

shotgun/agents/common.py

@@ -1,14 +1,12 @@
 """Common utilities for agent creation and management."""
 
-import asyncio
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
+import aiofiles
 from pydantic_ai import (
     Agent,
-    DeferredToolRequests,
-    DeferredToolResults,
     RunContext,
     UsageLimits,
 )
@@ -19,20 +17,19 @@ from pydantic_ai.messages import (
 )
 
 from shotgun.agents.config import ProviderType, get_provider_model
-from shotgun.agents.models import AgentType
+from shotgun.agents.models import AgentResponse, AgentType
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
 from shotgun.sdk.services import get_codebase_service
 from shotgun.utils import ensure_shotgun_directory_exists
+from shotgun.utils.datetime_utils import get_datetime_context
 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,
-    ask_user,
     codebase_shell,
     directory_lister,
     file_read,
@@ -72,7 +69,10 @@ 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()
 
     system_state = prompt_loader.render(
         "agents/state/system_state.j2",
@@ -80,6 +80,9 @@ async def add_system_status_message(
         is_tui_context=deps.is_tui_context,
         existing_files=existing_files,
         markdown_toc=markdown_toc,
+        current_datetime=dt_context.datetime_formatted,
+        timezone_name=dt_context.timezone_name,
+        utc_offset=dt_context.utc_offset,
     )
 
     message_history.append(
@@ -92,14 +95,14 @@ 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,
     additional_tools: list[Any] | None = None,
     provider: ProviderType | None = None,
     agent_mode: AgentType | None = None,
-) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
     """Create a base agent with common configuration.
 
     Args:
@@ -117,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",
@@ -157,7 +160,7 @@ def create_base_agent(
 
     agent = Agent(
         model,
-        output_type=[str, DeferredToolRequests],
+        output_type=AgentResponse,
         deps_type=AgentDeps,
         instrument=True,
         history_processors=[history_processor],
@@ -172,11 +175,6 @@ def create_base_agent(
     for tool in additional_tools or []:
         agent.tool_plain(tool)
 
-    # Register interactive tool conditionally based on deps
-    if deps.interactive_mode:
-        agent.tool(ask_user)
-        logger.debug("📞 Interactive mode enabled - ask_user tool registered")
-
     # Register common file management tools (always available)
     agent.tool(write_file)
     agent.tool(append_file)
@@ -197,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.
@@ -214,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
@@ -260,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'
@@ -312,22 +311,30 @@ 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(
-                f'<TABLE_OF_CONTENTS file_name="{prior_file}">\n{prior_toc}\n</TABLE_OF_CONTENTS>'
+                f'<TABLE_OF_CONTENTS file_name="{prior_file}">\n'
+                f"{prior_toc}\n"
+                f"</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)
+        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(
                 0,
-                f'<TABLE_OF_CONTENTS file_name="{config.own_file}">\n{own_toc}\n</TABLE_OF_CONTENTS>',
+                f'<TABLE_OF_CONTENTS file_name="{config.own_file}">\n'
+                f"{own_toc}\n"
+                f"</TABLE_OF_CONTENTS>",
             )
 
     # Combine all sections
@@ -383,23 +390,48 @@ def get_agent_existing_files(agent_mode: AgentType | None = None) -> list[str]:
                     relative_path = file_path.relative_to(base_path)
                     existing_files.append(str(relative_path))
     else:
-        # For other agents, check both .md file and directory with same name
-        allowed_file = AGENT_DIRECTORIES[agent_mode]
-
-        # Check for the .md file
-        md_file_path = base_path / allowed_file
-        if md_file_path.exists():
-            existing_files.append(allowed_file)
-
-        # Check for directory with same base name (e.g., research/ for research.md)
-        base_name = allowed_file.replace(".md", "")
-        dir_path = base_path / base_name
-        if dir_path.exists() and dir_path.is_dir():
-            # List all files in the directory
-            for file_path in dir_path.rglob("*"):
-                if file_path.is_file():
-                    relative_path = file_path.relative_to(base_path)
-                    existing_files.append(str(relative_path))
+        # For other agents, check files/directories they have access to
+        allowed_paths_raw = AGENT_DIRECTORIES[agent_mode]
+
+        # Convert single Path/string to list of Paths for uniform handling
+        if isinstance(allowed_paths_raw, str):
+            # Special case: "*" means export agent (shouldn't reach here but handle it)
+            allowed_paths = (
+                [Path(allowed_paths_raw)] if allowed_paths_raw != "*" else []
+            )
+        elif isinstance(allowed_paths_raw, Path):
+            allowed_paths = [allowed_paths_raw]
+        else:
+            # Already a list
+            allowed_paths = allowed_paths_raw
+
+        # Check each allowed path
+        for allowed_path in allowed_paths:
+            allowed_str = str(allowed_path)
+
+            # Check if it's a directory (no .md suffix)
+            if not allowed_path.suffix or not allowed_str.endswith(".md"):
+                # It's a directory - list all files within it
+                dir_path = base_path / allowed_str
+                if dir_path.exists() and dir_path.is_dir():
+                    for file_path in dir_path.rglob("*"):
+                        if file_path.is_file():
+                            relative_path = file_path.relative_to(base_path)
+                            existing_files.append(str(relative_path))
+            else:
+                # It's a file - check if it exists
+                file_path = base_path / allowed_str
+                if file_path.exists():
+                    existing_files.append(allowed_str)
+
+                # Also check for associated directory (e.g., research/ for research.md)
+                base_name = allowed_str.replace(".md", "")
+                dir_path = base_path / base_name
+                if dir_path.exists() and dir_path.is_dir():
+                    for file_path in dir_path.rglob("*"):
+                        if file_path.is_file():
+                            relative_path = file_path.relative_to(base_path)
+                            existing_files.append(str(relative_path))
 
     return existing_files
 
@@ -469,7 +501,8 @@ async def add_system_prompt_message(
     message_history = message_history or []
 
     # Create a minimal RunContext to call the system prompt function
-    # We'll pass None for model and usage since they're not used by our system prompt functions
+    # We'll pass None for model and usage since they're not used
+    # by our system prompt functions
     context = type(
         "RunContext", (), {"deps": deps, "retry": 0, "model": None, "usage": None}
     )()
@@ -493,12 +526,12 @@ async def add_system_prompt_message(
 
 
 async def run_agent(
-    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    agent: Agent[AgentDeps, AgentResponse],
     prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
     usage_limits: UsageLimits | None = None,
-) -> AgentRunResult[str | DeferredToolRequests]:
+) -> AgentRunResult[AgentResponse]:
     # Clear file tracker for new run
     deps.file_tracker.clear()
     logger.debug("🔧 Cleared file tracker for new agent run")
@@ -513,33 +546,6 @@ async def run_agent(
         message_history=message_history,
     )
 
-    # Apply persistent compaction to prevent cascading token growth across CLI commands
-    messages = await apply_persistent_compaction(result.all_messages(), deps)
-    while isinstance(result.output, DeferredToolRequests):
-        logger.info("got deferred tool requests")
-        await deps.queue.join()
-        requests = result.output
-        done, _ = await asyncio.wait(deps.tasks)
-
-        task_results = [task.result() for task in done]
-        task_results_by_tool_call_id = {
-            result.tool_call_id: result.answer for result in task_results
-        }
-        logger.info("got task results", task_results_by_tool_call_id)
-        results = DeferredToolResults()
-        for call in requests.calls:
-            results.calls[call.tool_call_id] = task_results_by_tool_call_id[
-                call.tool_call_id
-            ]
-        result = await agent.run(
-            deps=deps,
-            usage_limits=usage_limits,
-            message_history=messages,
-            deferred_tool_results=results,
-        )
-        # Apply persistent compaction to prevent cascading token growth in multi-turn loops
-        messages = await apply_persistent_compaction(result.all_messages(), deps)
-
     # Log file operations summary if any files were modified
     if deps.file_tracker.operations:
         summary = deps.file_tracker.format_summary()

shotgun/agents/config/constants.py

@@ -24,11 +24,5 @@ ANTHROPIC_PROVIDER = ConfigSection.ANTHROPIC.value
 GOOGLE_PROVIDER = ConfigSection.GOOGLE.value
 SHOTGUN_PROVIDER = ConfigSection.SHOTGUN.value
 
-# Environment variable names
-OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
-ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
-GEMINI_API_KEY_ENV = "GEMINI_API_KEY"
-SHOTGUN_API_KEY_ENV = "SHOTGUN_API_KEY"
-
 # Token limits
 MEDIUM_TEXT_8K_TOKENS = 8192  # Default max_tokens for web search requests

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,14 +144,14 @@ 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
 
                             # Find default model for this provider
                             provider_models = {
                                 ProviderType.OPENAI: ModelName.GPT_5,
-                                ProviderType.ANTHROPIC: ModelName.CLAUDE_SONNET_4_5,
+                                ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
                                 ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
                             }
 
@@ -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(
@@ -243,50 +256,61 @@ class ConfigManager:
 
                 provider_models = {
                     ProviderType.OPENAI: ModelName.GPT_5,
-                    ProviderType.ANTHROPIC: ModelName.CLAUDE_SONNET_4_5,
+                    ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
                     ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
                 }
                 if provider_enum in provider_models:
                     config.selected_model = provider_models[provider_enum]
 
-        self.save(config)
+            # Mark welcome screen as shown when BYOK provider is configured
+            # This prevents the welcome screen from showing again after user has made their choice
+            config.shown_welcome_screen = True
+
+        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, _ = self._get_provider_config_and_type(config, provider)
+        provider_config, is_shotgun = self._get_provider_config_and_type(
+            config, provider
+        )
 
         provider_config.api_key = None
-        self.save(config)
 
-    def update_selected_model(self, model_name: "ModelName") -> None:
+        # For Shotgun Account, also clear the JWT
+        if is_shotgun and isinstance(provider_config, ShotgunAccountConfig):
+            provider_config.supabase_jwt = None
+
+        await self.save(config)
+
+    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))
@@ -300,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:
@@ -310,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,
@@ -366,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 (
@@ -429,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.
@@ -447,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
@@ -457,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")