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

shotgun/agents/router/tools/plan_tools.py

@@ -0,0 +1,322 @@
+"""Plan management tools for the Router agent.
+
+These tools allow the router to create and manage execution plans.
+All tools use Pydantic models for inputs and outputs.
+
+IMPORTANT: There is NO get_plan() tool - the plan is shown in the system
+status message every turn so the router always has visibility into it.
+"""
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.router.models import (
+    AddStepInput,
+    CreatePlanInput,
+    ExecutionPlan,
+    ExecutionStep,
+    MarkStepDoneInput,
+    PendingApproval,
+    PendingCheckpoint,
+    PlanApprovalStatus,
+    RemoveStepInput,
+    RouterDeps,
+    RouterMode,
+    ToolResult,
+)
+from shotgun.agents.tools.registry import ToolCategory, register_tool
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+def _notify_plan_changed(deps: RouterDeps) -> None:
+    """Notify TUI of plan changes via callback if registered.
+
+    This helper is called after any plan modification to update the
+    Plan Panel widget in the TUI.
+
+    Args:
+        deps: RouterDeps containing the on_plan_changed callback.
+    """
+    if deps.on_plan_changed:
+        deps.on_plan_changed(deps.current_plan)
+
+
+@register_tool(
+    category=ToolCategory.PLANNING,
+    display_text="Creating execution plan",
+    key_arg="input",
+)
+async def create_plan(
+    ctx: RunContext[RouterDeps], input: CreatePlanInput
+) -> ToolResult:
+    """Create a new execution plan for the current task.
+
+    This replaces any existing plan. The plan is stored in-memory in RouterDeps,
+    NOT in a file. It will be shown in the system status message.
+
+    Args:
+        ctx: RunContext with RouterDeps
+        input: CreatePlanInput with goal and steps
+
+    Returns:
+        ToolResult indicating success or failure
+    """
+    logger.debug("Creating execution plan with goal: %s", input.goal)
+
+    # Convert step inputs to ExecutionStep objects
+    steps = [
+        ExecutionStep(
+            id=step_input.id,
+            title=step_input.title,
+            objective=step_input.objective,
+            done=False,
+        )
+        for step_input in input.steps
+    ]
+
+    # Create and store the plan
+    plan = ExecutionPlan(
+        goal=input.goal,
+        steps=steps,
+        current_step_index=0,
+    )
+
+    ctx.deps.current_plan = plan
+
+    # Set pending approval for multi-step plans in Planning mode
+    # The TUI will detect this and show the PlanApprovalWidget
+    if ctx.deps.router_mode == RouterMode.PLANNING and plan.needs_approval():
+        ctx.deps.pending_approval = PendingApproval(plan=plan)
+        ctx.deps.approval_status = PlanApprovalStatus.PENDING
+        # Plan is NOT executing yet - user must approve first
+        ctx.deps.is_executing = False
+        logger.debug(
+            "Set pending approval for plan with %d steps",
+            len(steps),
+        )
+    else:
+        # Single-step plans or Drafting mode - skip approval and start executing
+        ctx.deps.approval_status = PlanApprovalStatus.SKIPPED
+        ctx.deps.is_executing = True
+        logger.debug("Plan approved automatically, is_executing=True")
+
+    logger.info(
+        "Created execution plan with %d steps: %s",
+        len(steps),
+        input.goal,
+    )
+
+    _notify_plan_changed(ctx.deps)
+
+    # Return different message based on whether approval is needed
+    if ctx.deps.pending_approval is not None:
+        return ToolResult(
+            success=True,
+            message=f"Created plan with {len(steps)} steps. Goal: {input.goal}\n\n"
+            "IMPORTANT: This plan requires user approval before execution. "
+            "You MUST call final_result NOW to present this plan to the user. "
+            "Do NOT attempt to delegate or execute any steps yet.",
+        )
+
+    return ToolResult(
+        success=True,
+        message=f"Created plan with {len(steps)} steps. Goal: {input.goal}",
+    )
+
+
+@register_tool(
+    category=ToolCategory.PLANNING,
+    display_text="Marking step complete",
+    key_arg="input",
+)
+async def mark_step_done(
+    ctx: RunContext[RouterDeps], input: MarkStepDoneInput
+) -> ToolResult:
+    """Mark a step in the execution plan as complete.
+
+    Args:
+        ctx: RunContext with RouterDeps
+        input: MarkStepDoneInput with step_id
+
+    Returns:
+        ToolResult indicating success or failure
+    """
+    plan = ctx.deps.current_plan
+
+    if plan is None:
+        return ToolResult(
+            success=False,
+            message="No execution plan exists. Create a plan first.",
+        )
+
+    # Find the step by ID
+    for _i, step in enumerate(plan.steps):
+        if step.id == input.step_id:
+            step.done = True
+            logger.info("Marked step '%s' as done", input.step_id)
+
+            # Advance current_step_index to next incomplete step
+            while (
+                plan.current_step_index < len(plan.steps)
+                and plan.steps[plan.current_step_index].done
+            ):
+                plan.current_step_index += 1
+
+            # Check if plan is complete
+            if plan.is_complete():
+                ctx.deps.is_executing = False
+                logger.debug("Plan complete, is_executing=False")
+            # Set pending checkpoint for Planning mode
+            # The TUI will detect this and show the StepCheckpointWidget
+            elif ctx.deps.router_mode == RouterMode.PLANNING:
+                # Use current_step() since the while loop above already advanced
+                # current_step_index to the next incomplete step
+                next_step = plan.current_step()
+                ctx.deps.pending_checkpoint = PendingCheckpoint(
+                    completed_step=step, next_step=next_step
+                )
+                logger.debug(
+                    "Set pending checkpoint: completed='%s', next='%s'",
+                    step.title,
+                    next_step.title if next_step else None,
+                )
+
+            _notify_plan_changed(ctx.deps)
+
+            return ToolResult(
+                success=True,
+                message=f"Marked step '{step.title}' as complete.",
+            )
+
+    return ToolResult(
+        success=False,
+        message=f"Step with ID '{input.step_id}' not found in plan.",
+    )
+
+
+@register_tool(
+    category=ToolCategory.PLANNING,
+    display_text="Adding step to plan",
+    key_arg="input",
+)
+async def add_step(ctx: RunContext[RouterDeps], input: AddStepInput) -> ToolResult:
+    """Add a new step to the execution plan.
+
+    The step can be inserted after a specific step (by ID) or appended to the end.
+
+    Args:
+        ctx: RunContext with RouterDeps
+        input: AddStepInput with step details and optional after_step_id
+
+    Returns:
+        ToolResult indicating success or failure
+    """
+    plan = ctx.deps.current_plan
+
+    if plan is None:
+        return ToolResult(
+            success=False,
+            message="No execution plan exists. Create a plan first.",
+        )
+
+    # Check for duplicate ID
+    existing_ids = {step.id for step in plan.steps}
+    if input.step.id in existing_ids:
+        return ToolResult(
+            success=False,
+            message=f"Step with ID '{input.step.id}' already exists in plan.",
+        )
+
+    # Create the new step
+    new_step = ExecutionStep(
+        id=input.step.id,
+        title=input.step.title,
+        objective=input.step.objective,
+        done=False,
+    )
+
+    # Insert at the specified position
+    if input.after_step_id is None:
+        # Append to end
+        plan.steps.append(new_step)
+        logger.info("Appended step '%s' to end of plan", input.step.id)
+
+        _notify_plan_changed(ctx.deps)
+
+        return ToolResult(
+            success=True,
+            message=f"Added step '{new_step.title}' at end of plan.",
+        )
+
+    # Find the position to insert after
+    for i, step in enumerate(plan.steps):
+        if step.id == input.after_step_id:
+            plan.steps.insert(i + 1, new_step)
+            logger.info(
+                "Inserted step '%s' after '%s'",
+                input.step.id,
+                input.after_step_id,
+            )
+
+            _notify_plan_changed(ctx.deps)
+
+            return ToolResult(
+                success=True,
+                message=f"Added step '{new_step.title}' after '{step.title}'.",
+            )
+
+    return ToolResult(
+        success=False,
+        message=f"Step with ID '{input.after_step_id}' not found in plan.",
+    )
+
+
+@register_tool(
+    category=ToolCategory.PLANNING,
+    display_text="Removing step from plan",
+    key_arg="input",
+)
+async def remove_step(
+    ctx: RunContext[RouterDeps], input: RemoveStepInput
+) -> ToolResult:
+    """Remove a step from the execution plan.
+
+    Args:
+        ctx: RunContext with RouterDeps
+        input: RemoveStepInput with step_id
+
+    Returns:
+        ToolResult indicating success or failure
+    """
+    plan = ctx.deps.current_plan
+
+    if plan is None:
+        return ToolResult(
+            success=False,
+            message="No execution plan exists. Create a plan first.",
+        )
+
+    # Find and remove the step
+    for i, step in enumerate(plan.steps):
+        if step.id == input.step_id:
+            removed_step = plan.steps.pop(i)
+            logger.info("Removed step '%s' from plan", input.step_id)
+
+            # Adjust current_step_index if needed
+            if plan.current_step_index > i:
+                plan.current_step_index -= 1
+            elif plan.current_step_index >= len(plan.steps):
+                plan.current_step_index = max(0, len(plan.steps) - 1)
+
+            _notify_plan_changed(ctx.deps)
+
+            return ToolResult(
+                success=True,
+                message=f"Removed step '{removed_step.title}' from plan.",
+            )
+
+    return ToolResult(
+        success=False,
+        message=f"Step with ID '{input.step_id}' not found in plan.",
+    )

shotgun/agents/runner.py

@@ -0,0 +1,230 @@
+"""Unified agent execution with consistent error handling.
+
+This module provides a reusable agent runner that wraps agent execution exceptions
+in user-friendly custom exceptions that can be caught and displayed by TUI or CLI.
+"""
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, NoReturn
+
+from anthropic import APIStatusError as AnthropicAPIStatusError
+from openai import APIStatusError as OpenAIAPIStatusError
+from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior
+
+from shotgun.agents.error.models import AgentErrorContext
+from shotgun.exceptions import (
+    AgentCancelledException,
+    BudgetExceededException,
+    BYOKAuthenticationException,
+    BYOKGenericAPIException,
+    BYOKQuotaBillingException,
+    BYOKRateLimitException,
+    BYOKServiceOverloadException,
+    ContextSizeLimitExceeded,
+    GenericAPIStatusException,
+    ShotgunRateLimitException,
+    ShotgunServiceOverloadException,
+    UnknownAgentException,
+)
+
+if TYPE_CHECKING:
+    from shotgun.agents.agent_manager import AgentManager
+
+logger = logging.getLogger(__name__)
+
+
+class AgentRunner:
+    """Unified agent execution wrapper with consistent error handling.
+
+    This class wraps agent execution and converts any exceptions into
+    user-friendly custom exceptions that can be caught and displayed by the
+    calling interface (TUI or CLI).
+
+    The runner:
+    - Executes the agent
+    - Logs errors for debugging
+    - Wraps exceptions in custom exception types (AgentCancelledException,
+      BYOKRateLimitException, etc.)
+    - Lets exceptions propagate to caller for display
+
+    Example:
+        >>> runner = AgentRunner(agent_manager)
+        >>> try:
+        >>>     await runner.run("Write a hello world function")
+        >>> except ContextSizeLimitExceeded as e:
+        >>>     print(e.to_markdown())
+        >>> except BYOKRateLimitException as e:
+        >>>     print(e.to_plain_text())
+    """
+
+    def __init__(self, agent_manager: "AgentManager"):
+        """Initialize the agent runner.
+
+        Args:
+            agent_manager: The agent manager to execute
+        """
+        self.agent_manager = agent_manager
+
+    async def run(self, prompt: str) -> None:
+        """Run the agent with the given prompt.
+
+        Args:
+            prompt: The user's prompt/query
+
+        Raises:
+            Custom exceptions for different error types:
+            - AgentCancelledException: User cancelled the operation
+            - ContextSizeLimitExceeded: Context too large for model
+            - BudgetExceededException: Shotgun Account budget exceeded
+            - BYOKRateLimitException: BYOK rate limit hit
+            - BYOKQuotaBillingException: BYOK quota/billing issue
+            - BYOKAuthenticationException: BYOK authentication failed
+            - BYOKServiceOverloadException: BYOK service overloaded
+            - BYOKGenericAPIException: Generic BYOK API error
+            - ShotgunServiceOverloadException: Shotgun service overloaded
+            - ShotgunRateLimitException: Shotgun rate limit hit
+            - GenericAPIStatusException: Generic API error
+            - UnknownAgentException: Unknown/unclassified error
+        """
+        try:
+            await self.agent_manager.run(prompt=prompt)
+
+        except asyncio.CancelledError as e:
+            # User cancelled - wrap and re-raise as our custom exception
+            context = self._create_error_context(e)
+            self._classify_and_raise(context)
+
+        except ContextSizeLimitExceeded as e:
+            # Already a custom exception - log and re-raise
+            logger.info(
+                "Context size limit exceeded",
+                extra={
+                    "max_tokens": e.max_tokens,
+                    "model_name": e.model_name,
+                },
+            )
+            raise
+
+        except Exception as e:
+            # Log with full stack trace to shotgun.log
+            logger.exception(
+                "Agent run failed",
+                extra={
+                    "agent_mode": self.agent_manager._current_agent_type.value,
+                    "error_type": type(e).__name__,
+                },
+            )
+
+            # Create error context and wrap/raise custom exception
+            context = self._create_error_context(e)
+            self._classify_and_raise(context)
+
+    def _create_error_context(self, exception: BaseException) -> AgentErrorContext:
+        """Create error context from exception and agent state.
+
+        Args:
+            exception: The exception that was raised
+
+        Returns:
+            AgentErrorContext with all necessary information for classification
+        """
+        return AgentErrorContext(
+            exception=exception,
+            is_shotgun_account=self.agent_manager.deps.llm_model.is_shotgun_account,
+        )
+
+    def _classify_and_raise(self, context: AgentErrorContext) -> NoReturn:
+        """Classify an exception and raise the appropriate custom exception.
+
+        Args:
+            context: Context information about the error
+
+        Raises:
+            Custom exception based on the error type
+        """
+        exception = context.exception
+        error_name = type(exception).__name__
+        error_message = str(exception)
+
+        # Check for cancellation
+        if isinstance(exception, asyncio.CancelledError):
+            raise AgentCancelledException() from exception
+
+        # Check for context size limit exceeded
+        if isinstance(exception, ContextSizeLimitExceeded):
+            # Already the right exception type, re-raise it
+            raise exception
+
+        # Check for budget exceeded (Shotgun Account only)
+        if (
+            context.is_shotgun_account
+            and "apistatuserror" in error_name.lower()
+            and "budget" in error_message.lower()
+            and "exceeded" in error_message.lower()
+        ):
+            raise BudgetExceededException(message=error_message) from exception
+
+        # Check for empty model response (e.g., model unavailable or misconfigured)
+        if isinstance(exception, UnexpectedModelBehavior):
+            raise GenericAPIStatusException(
+                "The model returned an empty response. This may indicate:\n"
+                "- The model is unavailable or misconfigured\n"
+                "- A temporary service issue\n\n"
+                "Try switching to a different model or try again later."
+            ) from exception
+
+        # Detect API errors
+        is_api_error = False
+        if isinstance(exception, OpenAIAPIStatusError):
+            is_api_error = True
+        elif isinstance(exception, AnthropicAPIStatusError):
+            is_api_error = True
+        elif isinstance(exception, ModelHTTPError):
+            # pydantic_ai wraps API errors in ModelHTTPError
+            # Check for HTTP error status codes (4xx client errors)
+            if 400 <= exception.status_code < 500:
+                is_api_error = True
+
+        # BYOK user API errors
+        if not context.is_shotgun_account and is_api_error:
+            self._raise_byok_api_error(error_message, exception)
+
+        # Shotgun Account specific errors
+        if "APIStatusError" in error_name:
+            if "overload" in error_message.lower():
+                raise ShotgunServiceOverloadException(error_message) from exception
+            elif "rate" in error_message.lower():
+                raise ShotgunRateLimitException(error_message) from exception
+            else:
+                raise GenericAPIStatusException(error_message) from exception
+
+        # Unknown error - wrap in our custom exception
+        raise UnknownAgentException(exception) from exception
+
+    def _raise_byok_api_error(
+        self, error_message: str, original_exception: Exception
+    ) -> NoReturn:
+        """Classify and raise API errors for BYOK users into specific types.
+
+        Args:
+            error_message: The error message from the API
+            original_exception: The original exception
+
+        Raises:
+            Specific BYOK exception type
+        """
+        error_lower = error_message.lower()
+
+        if "rate" in error_lower:
+            raise BYOKRateLimitException(error_message) from original_exception
+        elif "quota" in error_lower or "billing" in error_lower:
+            raise BYOKQuotaBillingException(error_message) from original_exception
+        elif "authentication" in error_lower or (
+            "invalid" in error_lower and "key" in error_lower
+        ):
+            raise BYOKAuthenticationException(error_message) from original_exception
+        elif "overload" in error_lower:
+            raise BYOKServiceOverloadException(error_message) from original_exception
+        else:
+            raise BYOKGenericAPIException(error_message) from original_exception

shotgun/agents/specify.py

@@ -2,16 +2,15 @@
 
 from functools import partial
 
-from pydantic_ai import (
-    Agent,
-)
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config import ProviderType
+from shotgun.agents.models import ShotgunAgent
 from shotgun.logging_config import get_logger
 
 from .common import (
+    EventStreamHandler,
     add_system_status_message,
     build_agent_system_prompt,
     create_base_agent,
@@ -25,7 +24,7 @@ logger = get_logger(__name__)
 
 async def create_specify_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create a specify agent with artifact management capabilities.
 
     Args:
@@ -51,26 +50,25 @@ async def create_specify_agent(
 
 
 async def run_specify_agent(
-    agent: Agent[AgentDeps, AgentResponse],
-    requirement: str,
+    agent: ShotgunAgent,
+    prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
-    """Create or update specifications based on the given requirement.
+    """Create or update specifications based on the given prompt.
 
     Args:
         agent: The configured specify agent
-        requirement: The specification requirement or instruction
+        prompt: The specification prompt or instruction
         deps: Agent dependencies
         message_history: Optional message history for conversation continuity
+        event_stream_handler: Optional callback for streaming events
 
     Returns:
         AgentRunResult containing the specification process output
     """
-    logger.debug("📋 Starting specification for requirement: %s", requirement)
-
-    # Simple prompt - the agent system prompt has all the artifact instructions
-    full_prompt = f"Create a comprehensive specification for: {requirement}"
+    logger.debug("📋 Starting specification for prompt: %s", prompt)
 
     try:
         # Create usage limits for responsible API usage
@@ -80,10 +78,11 @@ async def run_specify_agent(
 
         result = await run_agent(
             agent=agent,
-            prompt=full_prompt,
+            prompt=prompt,
             deps=deps,
             message_history=message_history,
             usage_limits=usage_limits,
+            event_stream_handler=event_stream_handler,
         )
 
         logger.debug("✅ Specification completed successfully")

shotgun/agents/tasks.py

@@ -2,16 +2,15 @@
 
 from functools import partial
 
-from pydantic_ai import (
-    Agent,
-)
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config import ProviderType
+from shotgun.agents.models import ShotgunAgent
 from shotgun.logging_config import get_logger
 
 from .common import (
+    EventStreamHandler,
     add_system_status_message,
     build_agent_system_prompt,
     create_base_agent,
@@ -25,7 +24,7 @@ logger = get_logger(__name__)
 
 async def create_tasks_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create a tasks agent with file management capabilities.
 
     Args:
@@ -49,39 +48,39 @@ async def create_tasks_agent(
 
 
 async def run_tasks_agent(
-    agent: Agent[AgentDeps, AgentResponse],
-    instruction: str,
+    agent: ShotgunAgent,
+    prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
-    """Create or update tasks based on the given instruction.
+    """Create or update tasks based on the given prompt.
 
     Args:
         agent: The configured tasks agent
-        instruction: The task creation/update instruction
+        prompt: The task creation/update prompt
         deps: Agent dependencies
         message_history: Optional message history for conversation continuity
+        event_stream_handler: Optional callback for streaming events
 
     Returns:
         AgentRunResult containing the task creation process output
     """
-    logger.debug("📋 Starting task creation for instruction: %s", instruction)
+    logger.debug("📋 Starting task creation for prompt: %s", prompt)
 
     message_history = await add_system_status_message(deps, message_history)
 
-    # Let the agent use its tools to read existing tasks, plan, and research
-    full_prompt = f"Create or update tasks based on: {instruction}"
-
     try:
         # Create usage limits for responsible API usage
         usage_limits = create_usage_limits()
 
         result = await run_agent(
             agent=agent,
-            prompt=full_prompt,
+            prompt=prompt,
             deps=deps,
             message_history=message_history,
             usage_limits=usage_limits,
+            event_stream_handler=event_stream_handler,
         )
 
         logger.debug("✅ Task creation completed successfully")