shotgun-sh 0.2.17__py3-none-any.whl → 0.3.3.dev1__py3-none-any.whl

shotgun/agents/config/manager.py

@@ -1,7 +1,9 @@
 """Configuration manager for Shotgun CLI."""
 
 import json
+import shutil
 import uuid
+from datetime import datetime
 from pathlib import Path
 from typing import Any
 
@@ -30,9 +32,196 @@ from .models import (
 
 logger = get_logger(__name__)
 
+
+class ConfigMigrationError(Exception):
+    """Exception raised when config migration fails."""
+
+    def __init__(self, message: str, backup_path: Path | None = None):
+        """Initialize with error message and optional backup path.
+
+        Args:
+            message: Error message describing what went wrong
+            backup_path: Path to backup file if one was created
+        """
+        self.backup_path = backup_path
+        super().__init__(message)
+
+
 # Type alias for provider configuration objects
 ProviderConfig = OpenAIConfig | AnthropicConfig | GoogleConfig | ShotgunAccountConfig
 
+# Current config version
+CURRENT_CONFIG_VERSION = 5
+
+# Backup directory name
+BACKUP_DIR_NAME = "backup"
+
+
+def get_backup_dir(config_path: Path) -> Path:
+    """Get the backup directory path for a given config file.
+
+    Args:
+        config_path: Path to the config file
+
+    Returns:
+        Path to the backup directory (e.g., ~/.shotgun-sh/backup/)
+    """
+    return config_path.parent / BACKUP_DIR_NAME
+
+
+def _create_backup(config_path: Path) -> Path:
+    """Create a timestamped backup of the config file before migration.
+
+    Backups are saved to ~/.shotgun-sh/backup/ directory.
+
+    Args:
+        config_path: Path to the config file to backup
+
+    Returns:
+        Path to the backup file in the backup directory
+
+    Raises:
+        OSError: If backup creation fails
+    """
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_dir = get_backup_dir(config_path)
+    backup_path = backup_dir / f"config.backup.{timestamp}.json"
+
+    try:
+        # Create backup directory if it doesn't exist
+        backup_dir.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(config_path, backup_path)
+        logger.info(f"Created config backup at {backup_path}")
+        return backup_path
+    except Exception as e:
+        logger.error(f"Failed to create config backup: {e}")
+        raise OSError(f"Failed to create config backup: {e}") from e
+
+
+def _migrate_v2_to_v3(data: dict[str, Any]) -> dict[str, Any]:
+    """Migrate config from version 2 to version 3.
+
+    Changes:
+    - Rename 'user_id' field to 'shotgun_instance_id'
+
+    Args:
+        data: Config data dict at version 2
+
+    Returns:
+        Modified config data dict at version 3
+    """
+    if "user_id" in data and SHOTGUN_INSTANCE_ID_FIELD not in data:
+        data[SHOTGUN_INSTANCE_ID_FIELD] = data.pop("user_id")
+        data["config_version"] = 3
+        logger.info("Migrated config v2->v3: renamed user_id to shotgun_instance_id")
+
+    return data
+
+
+def _migrate_v3_to_v4(data: dict[str, Any]) -> dict[str, Any]:
+    """Migrate config from version 3 to version 4.
+
+    Changes:
+    - Add 'marketing' field with empty messages dict
+    - Set 'shown_welcome_screen' to False for existing BYOK users
+
+    Args:
+        data: Config data dict at version 3
+
+    Returns:
+        Modified config data dict at version 4
+    """
+    # Add marketing config
+    if "marketing" not in data:
+        data["marketing"] = {"messages": {}}
+        logger.info("Migrated config v3->v4: added marketing configuration")
+
+    # Set shown_welcome_screen for existing BYOK users
+    # If shown_welcome_screen doesn't exist AND any BYOK provider has a key,
+    # set it to False so they see the welcome screen once
+    if "shown_welcome_screen" not in data:
+        has_byok_key = False
+        for section in ["openai", "anthropic", "google"]:
+            if (
+                section in data
+                and isinstance(data[section], dict)
+                and data[section].get("api_key")
+            ):
+                has_byok_key = True
+                break
+
+        if has_byok_key:
+            data["shown_welcome_screen"] = False
+            logger.info(
+                "Existing BYOK user detected: set shown_welcome_screen=False to show welcome screen"
+            )
+
+    data["config_version"] = 4
+    return data
+
+
+def _migrate_v4_to_v5(data: dict[str, Any]) -> dict[str, Any]:
+    """Migrate config from version 4 to version 5.
+
+    Changes:
+    - Add 'supports_streaming' field to OpenAI config (initially None for auto-detection)
+
+    Args:
+        data: Config data dict at version 4
+
+    Returns:
+        Modified config data dict at version 5
+    """
+    if "openai" in data and isinstance(data["openai"], dict):
+        if "supports_streaming" not in data["openai"]:
+            data["openai"]["supports_streaming"] = None
+            logger.info(
+                "Migrated config v4->v5: added streaming capability detection for OpenAI"
+            )
+
+    data["config_version"] = 5
+    return data
+
+
+def _apply_migrations(data: dict[str, Any]) -> dict[str, Any]:
+    """Apply all necessary migrations to bring config to current version.
+
+    Migrations are applied sequentially from the config's current version
+    to CURRENT_CONFIG_VERSION.
+
+    Args:
+        data: Config data dict at any version
+
+    Returns:
+        Config data dict at CURRENT_CONFIG_VERSION
+    """
+    # Get current version (default to 2 for very old configs)
+    current_version = data.get("config_version", 2)
+
+    # Define migrations in order
+    migrations = {
+        2: _migrate_v2_to_v3,
+        3: _migrate_v3_to_v4,
+        4: _migrate_v4_to_v5,
+    }
+
+    # Apply migrations sequentially
+    while current_version < CURRENT_CONFIG_VERSION:
+        if current_version in migrations:
+            logger.info(
+                f"Applying migration from v{current_version} to v{current_version + 1}"
+            )
+            data = migrations[current_version](data)
+            current_version = data.get("config_version", current_version + 1)
+        else:
+            logger.warning(
+                f"No migration defined for v{current_version}, skipping to v{current_version + 1}"
+            )
+            current_version += 1
+            data["config_version"] = current_version
+
+    return data
+
 
 class ConfigManager:
     """Manager for Shotgun configuration."""
@@ -71,71 +260,88 @@ class ConfigManager:
             self._config = await self.initialize()
             return self._config
 
+        backup_path: Path | None = None
         try:
             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:
-                data[SHOTGUN_INSTANCE_ID_FIELD] = data.pop("user_id")
-                data["config_version"] = 3
+            # Get current version to determine if migration is needed
+            current_version = data.get("config_version", 2)
+
+            # Create backup before migration if config needs upgrading
+            if current_version < CURRENT_CONFIG_VERSION:
                 logger.info(
-                    "Migrated config v2->v3: renamed user_id to shotgun_instance_id"
+                    f"Config needs migration from v{current_version} to v{CURRENT_CONFIG_VERSION}"
                 )
-
-            # Migration: Set shown_welcome_screen for existing BYOK users
-            # If shown_welcome_screen doesn't exist AND any BYOK provider has a key,
-            # set it to False so they see the welcome screen once
-            if "shown_welcome_screen" not in data:
-                has_byok_key = False
-                for section in ["openai", "anthropic", "google"]:
-                    if (
-                        section in data
-                        and isinstance(data[section], dict)
-                        and data[section].get("api_key")
-                    ):
-                        has_byok_key = True
-                        break
-
-                if has_byok_key:
-                    data["shown_welcome_screen"] = False
-                    logger.info(
-                        "Existing BYOK user detected: set shown_welcome_screen=False to show welcome screen"
+                try:
+                    backup_path = _create_backup(self.config_path)
+                except OSError as backup_error:
+                    logger.warning(
+                        f"Could not create backup before migration: {backup_error}"
+                    )
+                    # Continue without backup - better than failing completely
+
+            # Apply all necessary migrations to bring config to current version
+            try:
+                data = _apply_migrations(data)
+            except Exception as migration_error:
+                error_msg = (
+                    f"Failed to migrate configuration from v{current_version} to v{CURRENT_CONFIG_VERSION}. "
+                    f"Error: {migration_error}"
+                )
+                if backup_path:
+                    error_msg += f"\n\nYour original config has been backed up to:\n{backup_path}"
+                    error_msg += (
+                        "\n\nTo start fresh, delete or rename your config file:\n"
+                        f"  rm {self.config_path}\n"
+                        f"  shotgun config init\n\n"
+                        "To restore your backup:\n"
+                        f"  cp {backup_path} {self.config_path}"
                     )
+                else:
+                    error_msg += "\n\nTo start fresh, run: shotgun config init"
 
-            # 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")
+                raise ConfigMigrationError(error_msg, backup_path) from migration_error
 
             # Convert plain text secrets to SecretStr objects
             self._convert_secrets_to_secretstr(data)
 
+            # Clean up invalid selected_model before Pydantic validation
+            if "selected_model" in data and data["selected_model"] is not None:
+                from .models import MODEL_SPECS, ModelName
+
+                try:
+                    # Try to convert to ModelName enum
+                    model_name = ModelName(data["selected_model"])
+                    # Check if it exists in MODEL_SPECS
+                    if model_name not in MODEL_SPECS:
+                        data["selected_model"] = None
+                except (ValueError, KeyError):
+                    # Invalid model name - reset to None
+                    data["selected_model"] = None
+
             self._config = ShotgunConfig.model_validate(data)
             logger.debug("Configuration loaded successfully from %s", self.config_path)
 
-            # Validate selected_model if in BYOK mode (no Shotgun key)
-            if not self._provider_has_api_key(self._config.shotgun):
-                should_save = False
+            # Clear migration_failed flag if config loaded successfully
+            should_save = False
+            if self._config.migration_failed:
+                self._config.migration_failed = False
+                self._config.migration_backup_path = None
+                should_save = True
 
+            # Validate selected_model for BYOK mode - verify provider has a key
+            if not self._provider_has_api_key(self._config.shotgun):
                 # If selected_model is set, verify its provider has a key
                 if self._config.selected_model:
                     from .models import MODEL_SPECS
 
-                    if self._config.selected_model in MODEL_SPECS:
-                        spec = MODEL_SPECS[self._config.selected_model]
-                        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,
-                            )
-                            self._config.selected_model = None
-                            should_save = True
-                    else:
+                    spec = MODEL_SPECS[self._config.selected_model]
+                    if not await self.has_provider_key(spec.provider):
+                        # Provider has no key - reset to None
                         logger.info(
-                            "Selected model %s not found in MODEL_SPECS, resetting",
+                            "Selected model %s provider has no API key, finding available model",
                             self._config.selected_model.value,
                         )
                         self._config.selected_model = None
@@ -150,17 +356,13 @@ class ConfigManager:
 
                             # Find default model for this provider
                             provider_models = {
-                                ProviderType.OPENAI: ModelName.GPT_5,
+                                ProviderType.OPENAI: ModelName.GPT_5_1,
                                 ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
                                 ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
                             }
 
                             if provider in provider_models:
                                 self._config.selected_model = provider_models[provider]
-                                logger.info(
-                                    "Set selected_model to %s (first available provider)",
-                                    self._config.selected_model.value,
-                                )
                                 should_save = True
                                 break
 
@@ -169,12 +371,56 @@ class ConfigManager:
 
             return self._config
 
-        except Exception as e:
+        except ConfigMigrationError as migration_error:
+            # Migration failed - automatically create fresh config with migration info
             logger.error(
-                "Failed to load configuration from %s: %s", self.config_path, e
+                "Config migration failed, creating fresh config: %s", migration_error
             )
-            logger.info("Creating new configuration with generated shotgun_instance_id")
+            backup_path = migration_error.backup_path
+
+            # Create fresh config with migration failure info
+            self._config = await self.initialize()
+            self._config.migration_failed = True
+            if backup_path:
+                self._config.migration_backup_path = str(backup_path)
+
+            # Save the fresh config
+            await self.save(self._config)
+            logger.info("Created fresh config after migration failure")
+
+            return self._config
+
+        except json.JSONDecodeError as json_error:
+            # Invalid JSON - create backup and fresh config
+            logger.error("Config file has invalid JSON: %s", json_error)
+
+            try:
+                backup_path = _create_backup(self.config_path)
+            except OSError:
+                backup_path = None
+
+            self._config = await self.initialize()
+            self._config.migration_failed = True
+            if backup_path:
+                self._config.migration_backup_path = str(backup_path)
+
+            await self.save(self._config)
+            logger.info("Created fresh config after JSON parse error")
+
+            return self._config
+
+        except Exception as e:
+            # Generic error - create fresh config
+            logger.error("Failed to load config: %s", e)
+
             self._config = await self.initialize()
+            self._config.migration_failed = True
+            if backup_path:
+                self._config.migration_backup_path = str(backup_path)
+
+            await self.save(self._config)
+            logger.info("Created fresh config after load error")
+
             return self._config
 
     async def save(self, config: ShotgunConfig | None = None) -> None:
@@ -237,6 +483,11 @@ class ConfigManager:
                 SecretStr(api_key_value) if api_key_value is not None else None
             )
 
+            # Reset streaming capabilities when OpenAI API key is changed
+            if not is_shotgun and provider_enum == ProviderType.OPENAI:
+                if isinstance(provider_config, OpenAIConfig):
+                    provider_config.supports_streaming = None
+
         # Reject other fields
         unsupported_fields = set(kwargs.keys()) - {API_KEY_FIELD}
         if unsupported_fields:
@@ -255,7 +506,7 @@ class ConfigManager:
                 from .models import ModelName
 
                 provider_models = {
-                    ProviderType.OPENAI: ModelName.GPT_5,
+                    ProviderType.OPENAI: ModelName.GPT_5_1,
                     ProviderType.ANTHROPIC: ModelName.CLAUDE_HAIKU_4_5,
                     ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
                 }
@@ -266,6 +517,11 @@ class ConfigManager:
             # This prevents the welcome screen from showing again after user has made their choice
             config.shown_welcome_screen = True
 
+        # Clear migration failure flag when user successfully configures a provider
+        if API_KEY_FIELD in kwargs and api_key_value is not None:
+            config.migration_failed = False
+            config.migration_backup_path = None
+
         await self.save(config)
 
     async def clear_provider_key(self, provider: ProviderType | str) -> None:
@@ -283,6 +539,13 @@ class ConfigManager:
         if is_shotgun and isinstance(provider_config, ShotgunAccountConfig):
             provider_config.supabase_jwt = None
 
+        # Reset streaming capabilities when OpenAI API key is cleared
+        if not is_shotgun:
+            provider_enum = self._ensure_provider_enum(provider)
+            if provider_enum == ProviderType.OPENAI:
+                if isinstance(provider_config, OpenAIConfig):
+                    provider_config.supports_streaming = None
+
         await self.save(config)
 
     async def update_selected_model(self, model_name: "ModelName") -> None:
@@ -481,13 +744,17 @@ class ConfigManager:
         return config.shotgun_instance_id
 
     async def update_shotgun_account(
-        self, api_key: str | None = None, supabase_jwt: str | None = None
+        self,
+        api_key: str | None = None,
+        supabase_jwt: str | None = None,
+        workspace_id: str | None = None,
     ) -> None:
         """Update Shotgun Account configuration.
 
         Args:
             api_key: LiteLLM proxy API key (optional)
             supabase_jwt: Supabase authentication JWT (optional)
+            workspace_id: Default workspace ID for shared specs (optional)
         """
         config = await self.load()
 
@@ -499,6 +766,9 @@ class ConfigManager:
                 SecretStr(supabase_jwt) if supabase_jwt else None
             )
 
+        if workspace_id is not None:
+            config.shotgun.workspace_id = workspace_id
+
         await self.save(config)
         logger.info("Updated Shotgun Account configuration")
 

shotgun/agents/config/models.py

@@ -25,13 +25,17 @@ class KeyProvider(StrEnum):
 class ModelName(StrEnum):
     """Available AI model names."""
 
-    GPT_5 = "gpt-5"
-    GPT_5_MINI = "gpt-5-mini"
-    CLAUDE_OPUS_4_1 = "claude-opus-4-1"
+    GPT_5_1 = "gpt-5.1"
+    GPT_5_1_CODEX = "gpt-5.1-codex"
+    GPT_5_1_CODEX_MINI = "gpt-5.1-codex-mini"
+    CLAUDE_OPUS_4_5 = "claude-opus-4-5"
+    CLAUDE_SONNET_4 = "claude-sonnet-4"
     CLAUDE_SONNET_4_5 = "claude-sonnet-4-5"
     CLAUDE_HAIKU_4_5 = "claude-haiku-4-5"
     GEMINI_2_5_PRO = "gemini-2.5-pro"
     GEMINI_2_5_FLASH = "gemini-2.5-flash"
+    GEMINI_2_5_FLASH_LITE = "gemini-2.5-flash-lite"
+    GEMINI_3_PRO_PREVIEW = "gemini-3-pro-preview"
 
 
 class ModelSpec(BaseModel):
@@ -56,6 +60,10 @@ class ModelConfig(BaseModel):
     max_input_tokens: int
     max_output_tokens: int
     api_key: str
+    supports_streaming: bool = Field(
+        default=True,
+        description="Whether this model configuration supports streaming. False only for BYOK GPT-5 models without streaming enabled.",
+    )
     _model_instance: Model | None = PrivateAttr(default=None)
 
     class Config:
@@ -82,32 +90,41 @@ class ModelConfig(BaseModel):
         }
         return f"{provider_prefix[self.provider]}:{self.name}"
 
+    @property
+    def is_shotgun_account(self) -> bool:
+        """Check if this model is using Shotgun Account authentication.
+
+        Returns:
+            True if using Shotgun Account, False if BYOK
+        """
+        return self.key_provider == KeyProvider.SHOTGUN
+
 
 # Model specifications registry (static metadata)
 MODEL_SPECS: dict[ModelName, ModelSpec] = {
-    ModelName.GPT_5: ModelSpec(
-        name=ModelName.GPT_5,
+    ModelName.GPT_5_1: ModelSpec(
+        name=ModelName.GPT_5_1,
         provider=ProviderType.OPENAI,
-        max_input_tokens=400_000,
+        max_input_tokens=272_000,
         max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5",
-        short_name="GPT-5",
+        litellm_proxy_model_name="openai/gpt-5.1",
+        short_name="GPT-5.1",
     ),
-    ModelName.GPT_5_MINI: ModelSpec(
-        name=ModelName.GPT_5_MINI,
+    ModelName.GPT_5_1_CODEX: ModelSpec(
+        name=ModelName.GPT_5_1_CODEX,
         provider=ProviderType.OPENAI,
-        max_input_tokens=400_000,
+        max_input_tokens=272_000,
         max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5-mini",
-        short_name="GPT-5 Mini",
+        litellm_proxy_model_name="openai/gpt-5.1-codex",
+        short_name="GPT-5.1 Codex",
     ),
-    ModelName.CLAUDE_OPUS_4_1: ModelSpec(
-        name=ModelName.CLAUDE_OPUS_4_1,
-        provider=ProviderType.ANTHROPIC,
-        max_input_tokens=200_000,
-        max_output_tokens=32_000,
-        litellm_proxy_model_name="anthropic/claude-opus-4-1",
-        short_name="Opus 4.1",
+    ModelName.GPT_5_1_CODEX_MINI: ModelSpec(
+        name=ModelName.GPT_5_1_CODEX_MINI,
+        provider=ProviderType.OPENAI,
+        max_input_tokens=272_000,
+        max_output_tokens=128_000,
+        litellm_proxy_model_name="openai/gpt-5.1-codex-mini",
+        short_name="GPT-5.1 Codex Mini",
     ),
     ModelName.CLAUDE_SONNET_4_5: ModelSpec(
         name=ModelName.CLAUDE_SONNET_4_5,
@@ -141,6 +158,38 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         litellm_proxy_model_name="gemini/gemini-2.5-flash",
         short_name="Gemini 2.5 Flash",
     ),
+    ModelName.CLAUDE_OPUS_4_5: ModelSpec(
+        name=ModelName.CLAUDE_OPUS_4_5,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="anthropic/claude-opus-4-5",
+        short_name="Opus 4.5",
+    ),
+    ModelName.CLAUDE_SONNET_4: ModelSpec(
+        name=ModelName.CLAUDE_SONNET_4,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="anthropic/claude-sonnet-4",
+        short_name="Sonnet 4",
+    ),
+    ModelName.GEMINI_2_5_FLASH_LITE: ModelSpec(
+        name=ModelName.GEMINI_2_5_FLASH_LITE,
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_048_576,
+        max_output_tokens=65_536,
+        litellm_proxy_model_name="gemini/gemini-2.5-flash-lite",
+        short_name="Gemini 2.5 Flash Lite",
+    ),
+    ModelName.GEMINI_3_PRO_PREVIEW: ModelSpec(
+        name=ModelName.GEMINI_3_PRO_PREVIEW,
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_048_576,
+        max_output_tokens=65_536,
+        litellm_proxy_model_name="gemini/gemini-3-pro-preview",
+        short_name="Gemini 3 Pro",
+    ),
 }
 
 
@@ -148,6 +197,10 @@ class OpenAIConfig(BaseModel):
     """Configuration for OpenAI provider."""
 
     api_key: SecretStr | None = None
+    supports_streaming: bool | None = Field(
+        default=None,
+        description="Whether streaming is supported for this API key. None = not tested yet",
+    )
 
 
 class AnthropicConfig(BaseModel):
@@ -169,6 +222,9 @@ class ShotgunAccountConfig(BaseModel):
     supabase_jwt: SecretStr | None = Field(
         default=None, description="Supabase authentication JWT"
     )
+    workspace_id: str | None = Field(
+        default=None, description="Default workspace ID for shared specs"
+    )
 
 
 class MarketingMessageRecord(BaseModel):
@@ -200,7 +256,7 @@ class ShotgunConfig(BaseModel):
     shotgun_instance_id: str = Field(
         description="Unique shotgun instance identifier (also used for anonymous telemetry)",
     )
-    config_version: int = Field(default=4, description="Configuration schema version")
+    config_version: int = Field(default=5, description="Configuration schema version")
     shown_welcome_screen: bool = Field(
         default=False,
         description="Whether the welcome screen has been shown to the user",
@@ -213,3 +269,11 @@ class ShotgunConfig(BaseModel):
         default_factory=MarketingConfig,
         description="Marketing messages configuration and tracking",
     )
+    migration_failed: bool = Field(
+        default=False,
+        description="Whether the last config migration failed (cleared after user configures a provider)",
+    )
+    migration_backup_path: str | None = Field(
+        default=None,
+        description="Path to the backup file created when migration failed",
+    )