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

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,21 @@ 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"
+    )
+
+    @property
+    def has_valid_account(self) -> bool:
+        """Check if the user has a valid Shotgun Account configured.
+
+        Returns:
+            True if api_key is set and non-empty, False otherwise
+        """
+        if self.api_key is None:
+            return False
+        value = self.api_key.get_secret_value()
+        return bool(value and value.strip())
 
 
 class MarketingMessageRecord(BaseModel):
@@ -200,7 +268,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 +281,15 @@ 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",
+    )
+    router_mode: str = Field(
+        default="planning",
+        description="Router execution mode: 'planning' or 'drafting'",
+    )

shotgun/agents/config/provider.py

@@ -25,6 +25,7 @@ from .models import (
     ProviderType,
     ShotgunConfig,
 )
+from .streaming_test import check_streaming_capability
 
 logger = get_logger(__name__)
 
@@ -46,13 +47,13 @@ def get_default_model_for_provider(config: ShotgunConfig) -> ModelName:
     """
     # Priority 1: Shotgun Account
     if _get_api_key(config.shotgun.api_key):
-        return ModelName.GPT_5
+        return ModelName.GPT_5_1
 
     # Priority 2: Individual provider keys
     if _get_api_key(config.anthropic.api_key):
         return ModelName.CLAUDE_HAIKU_4_5
     if _get_api_key(config.openai.api_key):
-        return ModelName.GPT_5
+        return ModelName.GPT_5_1
     if _get_api_key(config.google.api_key):
         return ModelName.GEMINI_2_5_PRO
 
@@ -200,13 +201,16 @@ async def get_provider_model(
             model_name = provider_or_model
         else:
             # No specific model requested - use selected or default
-            model_name = config.selected_model or ModelName.GPT_5
+            model_name = config.selected_model or get_default_model_for_provider(config)
 
+        # Gracefully fall back if the selected model doesn't exist (backwards compatibility)
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = get_default_model_for_provider(config)
+
         spec = MODEL_SPECS[model_name]
 
         # Use Shotgun Account with determined model (provider = actual LLM provider)
+        # Shotgun accounts always support streaming (via LiteLLM proxy)
         return ModelConfig(
             name=spec.name,
             provider=spec.provider,  # Actual LLM provider (OPENAI/ANTHROPIC/GOOGLE)
@@ -214,6 +218,7 @@ async def get_provider_model(
             max_input_tokens=spec.max_input_tokens,
             max_output_tokens=spec.max_output_tokens,
             api_key=shotgun_api_key,
+            supports_streaming=True,  # Shotgun accounts always support streaming
         )
 
     # Priority 2: Fall back to individual provider keys
@@ -222,10 +227,12 @@ async def get_provider_model(
     if isinstance(provider_or_model, ModelName):
         # Look up the model spec
         if provider_or_model not in MODEL_SPECS:
-            raise ValueError(f"Model '{provider_or_model.value}' not found")
-        spec = MODEL_SPECS[provider_or_model]
-        provider_enum = spec.provider
-        requested_model = provider_or_model
+            requested_model = None  # Fall back to provider default
+            provider_enum = None  # Will be determined below
+        else:
+            spec = MODEL_SPECS[provider_or_model]
+            provider_enum = spec.provider
+            requested_model = provider_or_model
     else:
         # Convert string to ProviderType enum if needed (backward compatible)
         if provider_or_model:
@@ -254,12 +261,40 @@ async def get_provider_model(
         if not api_key:
             raise ValueError("OpenAI API key not configured. Set via config.")
 
-        # Use requested model or default to gpt-5
-        model_name = requested_model if requested_model else ModelName.GPT_5
+        # Use requested model or default to gpt-5.1
+        model_name = requested_model if requested_model else ModelName.GPT_5_1
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.GPT_5_1
         spec = MODEL_SPECS[model_name]
 
+        # Check and test streaming capability for GPT-5 family models
+        supports_streaming = True  # Default to True for all models
+        if model_name in (
+            ModelName.GPT_5_1,
+            ModelName.GPT_5_1_CODEX,
+            ModelName.GPT_5_1_CODEX_MINI,
+        ):
+            # Check if streaming capability has been tested
+            streaming_capability = config.openai.supports_streaming
+
+            if streaming_capability is None:
+                # Not tested yet - run streaming test (test once for all GPT-5 models)
+                logger.info("Testing streaming capability for OpenAI GPT-5 family...")
+                streaming_capability = await check_streaming_capability(
+                    api_key, model_name.value
+                )
+
+                # Save result to config (applies to all OpenAI models)
+                config.openai.supports_streaming = streaming_capability
+                await config_manager.save(config)
+                logger.info(
+                    f"Streaming test result: "
+                    f"{'enabled' if streaming_capability else 'disabled'}"
+                )
+
+            supports_streaming = streaming_capability
+
         # Create fully configured ModelConfig
         return ModelConfig(
             name=spec.name,
@@ -268,6 +303,7 @@ async def get_provider_model(
             max_input_tokens=spec.max_input_tokens,
             max_output_tokens=spec.max_output_tokens,
             api_key=api_key,
+            supports_streaming=supports_streaming,
         )
 
     elif provider_enum == ProviderType.ANTHROPIC:
@@ -277,8 +313,9 @@ async def get_provider_model(
 
         # Use requested model or default to claude-haiku-4-5
         model_name = requested_model if requested_model else ModelName.CLAUDE_HAIKU_4_5
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.CLAUDE_HAIKU_4_5
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig
@@ -298,8 +335,9 @@ async def get_provider_model(
 
         # Use requested model or default to gemini-2.5-pro
         model_name = requested_model if requested_model else ModelName.GEMINI_2_5_PRO
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.GEMINI_2_5_PRO
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig

shotgun/agents/config/streaming_test.py

@@ -0,0 +1,119 @@
+"""Utility for testing streaming capability of OpenAI models."""
+
+import logging
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Maximum number of attempts to test streaming capability
+MAX_STREAMING_TEST_ATTEMPTS = 3
+
+# Timeout for each streaming test attempt (in seconds)
+STREAMING_TEST_TIMEOUT = 10.0
+
+
+async def check_streaming_capability(
+    api_key: str, model_name: str, max_attempts: int = MAX_STREAMING_TEST_ATTEMPTS
+) -> bool:
+    """Check if the given OpenAI model supports streaming with this API key.
+
+    Retries multiple times to handle transient network issues. Only returns False
+    if streaming definitively fails after all retry attempts.
+
+    Args:
+        api_key: The OpenAI API key to test
+        model_name: The model name (e.g., "gpt-5", "gpt-5-mini")
+        max_attempts: Maximum number of attempts (default: 3)
+
+    Returns:
+        True if streaming is supported, False if it definitively fails
+    """
+    url = "https://api.openai.com/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+    # GPT-5 family uses max_completion_tokens instead of max_tokens
+    payload = {
+        "model": model_name,
+        "messages": [{"role": "user", "content": "test"}],
+        "stream": True,
+        "max_completion_tokens": 10,
+    }
+
+    last_error = None
+
+    for attempt in range(1, max_attempts + 1):
+        logger.debug(
+            f"Streaming test attempt {attempt}/{max_attempts} for {model_name}"
+        )
+
+        try:
+            async with httpx.AsyncClient(timeout=STREAMING_TEST_TIMEOUT) as client:
+                async with client.stream(
+                    "POST", url, json=payload, headers=headers
+                ) as response:
+                    # Check if we get a successful response
+                    if response.status_code != 200:
+                        last_error = f"HTTP {response.status_code}"
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name}: {last_error}"
+                        )
+
+                        # For definitive errors (403 Forbidden, 404 Not Found), don't retry
+                        if response.status_code in (403, 404):
+                            logger.info(
+                                f"Streaming definitively unsupported for {model_name} (HTTP {response.status_code})"
+                            )
+                            return False
+
+                        # For other errors, retry
+                        continue
+
+                    # Try to read at least one chunk from the stream
+                    try:
+                        async for _ in response.aiter_bytes():
+                            # Successfully received streaming data
+                            logger.info(
+                                f"Streaming test passed for {model_name} (attempt {attempt})"
+                            )
+                            return True
+                    except Exception as e:
+                        last_error = str(e)
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name} while reading stream: {e}"
+                        )
+                        continue
+
+        except httpx.TimeoutException:
+            last_error = "timeout"
+            logger.warning(
+                f"Streaming test attempt {attempt} timed out for {model_name}"
+            )
+            continue
+        except httpx.HTTPStatusError as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name}: {e}"
+            )
+            continue
+        except Exception as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name} with unexpected error: {e}"
+            )
+            continue
+
+        # If we got here without reading any chunks, streaming didn't work
+        last_error = "no data received"
+        logger.warning(
+            f"Streaming test attempt {attempt} failed for {model_name}: no data received"
+        )
+
+    # All attempts exhausted
+    logger.error(
+        f"Streaming test failed for {model_name} after {max_attempts} attempts. "
+        f"Last error: {last_error}. Assuming streaming is NOT supported."
+    )
+    return False

shotgun/agents/context_analyzer/analyzer.py

@@ -15,8 +15,12 @@ from pydantic_ai.messages import (
 )
 
 from shotgun.agents.config.models import ModelConfig
-from shotgun.agents.history.token_counting.utils import count_tokens_from_messages
-from shotgun.agents.history.token_estimation import estimate_tokens_from_messages
+from shotgun.agents.conversation.history.token_counting.utils import (
+    count_tokens_from_messages,
+)
+from shotgun.agents.conversation.history.token_estimation import (
+    estimate_tokens_from_messages,
+)
 from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.logging_config import get_logger
 from shotgun.tui.screens.chat_screen.hint_message import HintMessage

shotgun/agents/conversation/__init__.py

@@ -0,0 +1,18 @@
+"""Conversation module for managing conversation history and persistence."""
+
+from .filters import (
+    filter_incomplete_messages,
+    filter_orphaned_tool_responses,
+    is_tool_call_complete,
+)
+from .manager import ConversationManager
+from .models import ConversationHistory, ConversationState
+
+__all__ = [
+    "ConversationHistory",
+    "ConversationManager",
+    "ConversationState",
+    "filter_incomplete_messages",
+    "filter_orphaned_tool_responses",
+    "is_tool_call_complete",
+]

shotgun/agents/conversation/filters.py

@@ -0,0 +1,164 @@
+"""Filter functions for conversation message validation."""
+
+import json
+import logging
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    ToolCallPart,
+    ToolReturnPart,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def is_tool_call_complete(tool_call: ToolCallPart) -> bool:
+    """Check if a tool call has valid, complete JSON arguments.
+
+    Args:
+        tool_call: The tool call part to validate
+
+    Returns:
+        True if the tool call args are valid JSON, False otherwise
+    """
+    if tool_call.args is None:
+        return True  # No args is valid
+
+    if isinstance(tool_call.args, dict):
+        return True  # Already parsed dict is valid
+
+    if not isinstance(tool_call.args, str):
+        return False
+
+    # Try to parse the JSON string
+    try:
+        json.loads(tool_call.args)
+        return True
+    except (json.JSONDecodeError, ValueError) as e:
+        # Log incomplete tool call detection
+        args_preview = (
+            tool_call.args[:100] + "..."
+            if len(tool_call.args) > 100
+            else tool_call.args
+        )
+        logger.info(
+            "Detected incomplete tool call in validation",
+            extra={
+                "tool_name": tool_call.tool_name,
+                "tool_call_id": tool_call.tool_call_id,
+                "args_preview": args_preview,
+                "error": str(e),
+            },
+        )
+        return False
+
+
+def filter_incomplete_messages(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Filter out messages with incomplete tool calls.
+
+    Args:
+        messages: List of messages to filter
+
+    Returns:
+        List of messages with only complete tool calls
+    """
+    filtered: list[ModelMessage] = []
+    filtered_count = 0
+    filtered_tool_names: list[str] = []
+
+    for message in messages:
+        # Only check ModelResponse messages for tool calls
+        if not isinstance(message, ModelResponse):
+            filtered.append(message)
+            continue
+
+        # Check if any tool calls are incomplete
+        has_incomplete_tool_call = False
+        for part in message.parts:
+            if isinstance(part, ToolCallPart) and not is_tool_call_complete(part):
+                has_incomplete_tool_call = True
+                filtered_tool_names.append(part.tool_name)
+                break
+
+        # Only include messages without incomplete tool calls
+        if not has_incomplete_tool_call:
+            filtered.append(message)
+        else:
+            filtered_count += 1
+
+    # Log if any messages were filtered
+    if filtered_count > 0:
+        logger.info(
+            "Filtered incomplete messages before saving",
+            extra={
+                "filtered_count": filtered_count,
+                "total_messages": len(messages),
+                "filtered_tool_names": filtered_tool_names,
+            },
+        )
+
+    return filtered
+
+
+def filter_orphaned_tool_responses(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Filter out tool responses without corresponding tool calls.
+
+    This ensures message history is valid for OpenAI API which requires
+    tool responses to follow their corresponding tool calls.
+
+    Args:
+        messages: List of messages to filter
+
+    Returns:
+        List of messages with orphaned tool responses removed
+    """
+    # Collect all tool_call_ids from ToolCallPart in ModelResponse
+    valid_tool_call_ids: set[str] = set()
+    for msg in messages:
+        if isinstance(msg, ModelResponse):
+            for part in msg.parts:
+                if isinstance(part, ToolCallPart) and part.tool_call_id:
+                    valid_tool_call_ids.add(part.tool_call_id)
+
+    # Filter out orphaned ToolReturnPart from ModelRequest
+    filtered: list[ModelMessage] = []
+    orphaned_count = 0
+    orphaned_tool_names: list[str] = []
+
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            # Filter parts, removing orphaned ToolReturnPart
+            filtered_parts: list[ModelRequestPart] = []
+            request_part: ModelRequestPart
+            for request_part in msg.parts:
+                if isinstance(request_part, ToolReturnPart):
+                    if request_part.tool_call_id in valid_tool_call_ids:
+                        filtered_parts.append(request_part)
+                    else:
+                        # Skip orphaned tool response
+                        orphaned_count += 1
+                        orphaned_tool_names.append(request_part.tool_name or "unknown")
+                else:
+                    filtered_parts.append(request_part)
+
+            # Only add if there are remaining parts
+            if filtered_parts:
+                filtered.append(ModelRequest(parts=filtered_parts))
+        else:
+            filtered.append(msg)
+
+    # Log if any tool responses were filtered
+    if orphaned_count > 0:
+        logger.info(
+            "Filtered orphaned tool responses",
+            extra={
+                "orphaned_count": orphaned_count,
+                "total_messages": len(messages),
+                "orphaned_tool_names": orphaned_tool_names,
+            },
+        )
+
+    return filtered