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

shotgun/agents/router/router.py

@@ -0,0 +1,185 @@
+"""Router agent factory for the intelligent orchestrator.
+
+The Router agent is the single user-facing interface that orchestrates
+sub-agents (Research, Specify, Plan, Tasks, Export) based on user intent.
+"""
+
+import traceback
+from functools import partial
+
+from pydantic_ai import Agent, Tool
+from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.agents.common import (
+    add_system_status_message,
+    build_agent_system_prompt,
+    create_usage_limits,
+    run_agent,
+)
+from shotgun.agents.config import ProviderType, get_provider_model
+from shotgun.agents.conversation.history import token_limit_compactor
+from shotgun.agents.models import AgentResponse, AgentRuntimeOptions, AgentType
+from shotgun.agents.router.models import RouterDeps
+from shotgun.agents.router.tools import (
+    add_step,
+    create_plan,
+    mark_step_done,
+    remove_step,
+)
+from shotgun.agents.router.tools.delegation_tools import (
+    delegate_to_export,
+    delegate_to_plan,
+    delegate_to_research,
+    delegate_to_specification,
+    delegate_to_tasks,
+    prepare_delegation_tool,
+)
+from shotgun.agents.tools import read_file
+from shotgun.logging_config import get_logger
+from shotgun.sdk.services import get_codebase_service
+from shotgun.utils import ensure_shotgun_directory_exists
+
+logger = get_logger(__name__)
+
+
+async def create_router_agent(
+    agent_runtime_options: AgentRuntimeOptions,
+    provider: ProviderType | None = None,
+) -> tuple[Agent[RouterDeps, AgentResponse], RouterDeps]:
+    """Create the Router agent with plan management and delegation capabilities.
+
+    The Router is the intelligent orchestrator that:
+    - Understands user intent
+    - Creates and manages execution plans
+    - Delegates work to specialized sub-agents
+    - Operates in Planning (incremental) or Drafting (auto-execute) mode
+
+    Args:
+        agent_runtime_options: Runtime options for the agent
+        provider: Optional provider override. If None, uses configured default
+
+    Returns:
+        Tuple of (Configured Router agent, RouterDeps with plan management state)
+    """
+    logger.debug("Initializing router agent")
+    ensure_shotgun_directory_exists()
+
+    # Get configured model
+    try:
+        model_config = await get_provider_model(provider)
+        logger.debug(
+            "Router agent using %s model: %s",
+            model_config.provider.value.upper(),
+            model_config.name,
+        )
+        model = model_config.model_instance
+    except Exception as e:
+        logger.error("Failed to load configured model for router: %s", e)
+        raise ValueError("Configured model is required for router agent") from e
+
+    # Create RouterDeps (extends AgentDeps with router-specific state)
+    codebase_service = get_codebase_service()
+    system_prompt_fn = partial(build_agent_system_prompt, "router")
+
+    deps = RouterDeps(
+        **agent_runtime_options.model_dump(),
+        llm_model=model_config,
+        codebase_service=codebase_service,
+        system_prompt_fn=system_prompt_fn,
+        agent_mode=AgentType.ROUTER,
+    )
+
+    # Create history processor with access to deps via closure
+    async def history_processor(messages: list[ModelMessage]) -> list[ModelMessage]:
+        """History processor with access to deps via closure."""
+
+        class ProcessorContext:
+            def __init__(self, deps: RouterDeps):
+                self.deps = deps
+                self.usage = None
+
+        ctx = ProcessorContext(deps)
+        return await token_limit_compactor(ctx, messages)  # type: ignore[arg-type]
+
+    # Delegation tools with prepare function - only visible after plan is approved
+    # in Planning mode, always available in Drafting mode
+    delegation_tools = [
+        Tool(delegate_to_research, prepare=prepare_delegation_tool),
+        Tool(delegate_to_specification, prepare=prepare_delegation_tool),
+        Tool(delegate_to_plan, prepare=prepare_delegation_tool),
+        Tool(delegate_to_tasks, prepare=prepare_delegation_tool),
+        Tool(delegate_to_export, prepare=prepare_delegation_tool),
+    ]
+
+    # Create the agent with delegation tools that have prepare functions
+    agent: Agent[RouterDeps, AgentResponse] = Agent(
+        model,
+        output_type=AgentResponse,
+        deps_type=RouterDeps,
+        instrument=True,
+        history_processors=[history_processor],
+        retries=3,
+        tools=delegation_tools,
+    )
+
+    # Register plan management tools (router-specific, always available)
+    agent.tool(create_plan)
+    agent.tool(mark_step_done)
+    agent.tool(add_step)
+    agent.tool(remove_step)
+
+    # Register read-only file access for .shotgun/ directory
+    agent.tool(read_file)
+
+    # Note: The Router does NOT have write_file, append_file, or codebase tools.
+    # All file modifications and codebase understanding must be delegated to
+    # the appropriate sub-agent (Research, Specify, Plan, Tasks, Export).
+
+    logger.debug("Router agent tools registered")
+    logger.info(
+        "Router agent created in %s mode",
+        deps.router_mode.value.upper(),
+    )
+
+    return agent, deps
+
+
+async def run_router_agent(
+    agent: Agent[RouterDeps, AgentResponse],
+    prompt: str,
+    deps: RouterDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[AgentResponse]:
+    """Run the router agent with a user prompt.
+
+    Args:
+        agent: The configured router agent
+        prompt: User's request
+        deps: RouterDeps with plan management state
+        message_history: Optional existing message history
+
+    Returns:
+        Agent run result with response and any clarifying questions
+    """
+    logger.debug("Running router agent with prompt: %s", prompt[:100])
+
+    message_history = await add_system_status_message(deps, message_history)
+
+    try:
+        usage_limits = create_usage_limits()
+
+        result = await run_agent(
+            agent=agent,  # type: ignore[arg-type]
+            prompt=prompt,
+            deps=deps,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        logger.debug("Router agent completed successfully")
+        return result
+
+    except Exception:
+        logger.error("Router agent error:\n%s", traceback.format_exc())
+        raise

shotgun/agents/router/tools/__init__.py

@@ -0,0 +1,18 @@
+"""Router tools package."""
+
+from shotgun.agents.router.tools.plan_tools import (
+    add_step,
+    create_plan,
+    mark_step_done,
+    remove_step,
+)
+
+# Note: Delegation tools are imported directly in router.py to use
+# the prepare_delegation_tool function for conditional visibility.
+
+__all__ = [
+    "add_step",
+    "create_plan",
+    "mark_step_done",
+    "remove_step",
+]

shotgun/agents/router/tools/delegation_tools.py

@@ -0,0 +1,503 @@
+"""Delegation tools for the Router agent.
+
+These tools allow the Router to delegate work to specialized sub-agents
+(Research, Specify, Plan, Tasks, Export) for specific tasks.
+
+Sub-agents run with isolated message histories to prevent context window bloat.
+"""
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from pydantic_ai import RunContext
+from pydantic_ai.tools import ToolDefinition
+
+from shotgun.agents.export import create_export_agent, run_export_agent
+from shotgun.agents.models import (
+    AgentDeps,
+    AgentRuntimeOptions,
+    AgentType,
+    ShotgunAgent,
+    SubAgentContext,
+)
+from shotgun.agents.plan import create_plan_agent, run_plan_agent
+from shotgun.agents.research import create_research_agent, run_research_agent
+from shotgun.agents.router.models import (
+    DelegationInput,
+    DelegationResult,
+    RouterDeps,
+    RouterMode,
+    SubAgentCacheEntry,
+)
+from shotgun.agents.specify import create_specify_agent, run_specify_agent
+from shotgun.agents.tasks import create_tasks_agent, run_tasks_agent
+from shotgun.agents.tools.registry import ToolCategory, register_tool
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+# =============================================================================
+# Tool Preparation (Conditional Availability)
+# =============================================================================
+
+
+async def prepare_delegation_tool(
+    ctx: RunContext[RouterDeps], tool_def: ToolDefinition
+) -> ToolDefinition | None:
+    """Prepare function to conditionally show delegation tools.
+
+    In Planning mode, delegation tools are ONLY available when:
+    1. A plan exists (current_plan is not None)
+    2. The plan has been approved (pending_approval is None)
+
+    In Drafting mode, delegation tools are always available.
+
+    Args:
+        ctx: RunContext with RouterDeps containing plan state.
+        tool_def: The tool definition to conditionally return.
+
+    Returns:
+        The tool_def if delegation is allowed, None to hide the tool.
+    """
+    deps = ctx.deps
+
+    # Drafting mode - tools always available
+    if deps.router_mode == RouterMode.DRAFTING:
+        return tool_def
+
+    # Planning mode - check plan state
+    if deps.current_plan is None:
+        logger.debug("Hiding %s: no plan exists in Planning mode", tool_def.name)
+        return None
+
+    if deps.pending_approval is not None:
+        logger.debug("Hiding %s: plan pending user approval", tool_def.name)
+        return None
+
+    # Plan exists and is approved - allow delegation
+    return tool_def
+
+
+# Type aliases for factory functions
+CreateAgentFn = Callable[
+    [AgentRuntimeOptions], Awaitable[tuple[ShotgunAgent, AgentDeps]]
+]
+RunAgentFn = Callable[..., Awaitable[Any]]
+
+# Maximum retries for transient errors
+MAX_RETRIES = 2
+
+# Map agent types to their factory and run functions
+AGENT_FACTORIES: dict[AgentType, tuple[CreateAgentFn, RunAgentFn]] = {
+    AgentType.RESEARCH: (create_research_agent, run_research_agent),
+    AgentType.SPECIFY: (create_specify_agent, run_specify_agent),
+    AgentType.PLAN: (create_plan_agent, run_plan_agent),
+    AgentType.TASKS: (create_tasks_agent, run_tasks_agent),
+    AgentType.EXPORT: (create_export_agent, run_export_agent),
+}
+
+
+def _is_retryable_error(exception: BaseException) -> bool:
+    """Check if exception should trigger a retry.
+
+    Args:
+        exception: The exception to check.
+
+    Returns:
+        True if the exception is a transient error that should be retried.
+    """
+    # ValueError for truncated/incomplete JSON
+    if isinstance(exception, ValueError):
+        error_str = str(exception)
+        return "EOF while parsing" in error_str or (
+            "JSON" in error_str and "parsing" in error_str
+        )
+
+    # API errors (overload, rate limits)
+    exception_name = type(exception).__name__
+    if "APIStatusError" in exception_name:
+        error_str = str(exception)
+        return "overload" in error_str.lower() or "rate" in error_str.lower()
+
+    # Network errors
+    if "ConnectionError" in exception_name or "TimeoutError" in exception_name:
+        return True
+
+    return False
+
+
+def _create_agent_runtime_options(deps: RouterDeps) -> AgentRuntimeOptions:
+    """Create AgentRuntimeOptions from RouterDeps for sub-agent creation.
+
+    Args:
+        deps: RouterDeps containing shared runtime configuration.
+
+    Returns:
+        AgentRuntimeOptions configured for sub-agent creation.
+    """
+    return AgentRuntimeOptions(
+        interactive_mode=deps.interactive_mode,
+        working_directory=deps.working_directory,
+        is_tui_context=deps.is_tui_context,
+        max_iterations=deps.max_iterations,
+        queue=deps.queue,
+        tasks=deps.tasks,
+    )
+
+
+async def _get_or_create_sub_agent(
+    deps: RouterDeps,
+    agent_type: AgentType,
+) -> SubAgentCacheEntry:
+    """Get a cached sub-agent or create a new one.
+
+    Args:
+        deps: RouterDeps with sub_agent_cache.
+        agent_type: The type of agent to get or create.
+
+    Returns:
+        Tuple of (agent, agent_deps) for the requested agent type.
+
+    Raises:
+        ValueError: If agent_type is not supported for delegation.
+    """
+    # Check cache first
+    if agent_type in deps.sub_agent_cache:
+        logger.debug("Using cached %s agent", agent_type.value)
+        return deps.sub_agent_cache[agent_type]
+
+    # Get factory functions
+    if agent_type not in AGENT_FACTORIES:
+        raise ValueError(f"Agent type {agent_type} is not supported for delegation")
+
+    create_fn, _ = AGENT_FACTORIES[agent_type]
+    runtime_options = _create_agent_runtime_options(deps)
+
+    logger.debug("Creating new %s agent for delegation", agent_type.value)
+    agent, agent_deps = await create_fn(runtime_options)
+
+    # Cache for reuse
+    cache_entry: SubAgentCacheEntry = (agent, agent_deps)
+    deps.sub_agent_cache[agent_type] = cache_entry
+
+    return cache_entry
+
+
+def _build_sub_agent_context(deps: RouterDeps) -> SubAgentContext:
+    """Build SubAgentContext with plan information.
+
+    Args:
+        deps: RouterDeps with current plan.
+
+    Returns:
+        SubAgentContext configured for delegation.
+    """
+    current_step = deps.current_plan.current_step() if deps.current_plan else None
+
+    return SubAgentContext(
+        is_router_delegated=True,
+        plan_goal=deps.current_plan.goal if deps.current_plan else "",
+        current_step_id=current_step.id if current_step else "",
+        current_step_title=current_step.title if current_step else "",
+    )
+
+
+async def _run_sub_agent(
+    ctx: RunContext[RouterDeps],
+    agent_type: AgentType,
+    task: str,
+    context_hint: str | None = None,
+) -> DelegationResult:
+    """Run a sub-agent with the given task.
+
+    This helper function handles:
+    - Checking for pending approval (blocks delegation until user approves)
+    - Getting or creating the sub-agent from cache
+    - Setting up SubAgentContext
+    - Managing active_sub_agent state for UI updates
+    - Running the sub-agent with isolated message history
+    - Extracting files_modified and handling errors with retries
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        agent_type: The type of sub-agent to run.
+        task: The task to delegate to the sub-agent.
+        context_hint: Optional context to help the sub-agent.
+
+    Returns:
+        DelegationResult with success/failure status, response, and files_modified.
+    """
+    deps = ctx.deps
+
+    # Note: Delegation checks are now handled by prepare_delegation_tool which
+    # hides delegation tools entirely when delegation isn't allowed. This is a
+    # cleaner approach than returning errors - the LLM simply won't see the tools.
+
+    # Build the prompt with context hint if provided
+    prompt = task
+    if context_hint:
+        prompt = f"{task}\n\nContext: {context_hint}"
+
+    # Get or create the sub-agent
+    try:
+        agent, sub_agent_deps = await _get_or_create_sub_agent(deps, agent_type)
+    except ValueError as e:
+        return DelegationResult(
+            success=False,
+            error=str(e),
+            response="",
+            files_modified=[],
+        )
+
+    # Set up SubAgentContext so sub-agent knows it's being orchestrated
+    sub_agent_deps.sub_agent_context = _build_sub_agent_context(deps)
+
+    # Clear sub-agent's file tracker for fresh tracking
+    sub_agent_deps.file_tracker.clear()
+
+    # Set active_sub_agent for UI mode indicator
+    deps.active_sub_agent = agent_type
+    logger.info("Delegating to %s agent: %s", agent_type.value, task[:100])
+
+    # Get the run function for this agent type
+    _, run_fn = AGENT_FACTORIES[agent_type]
+
+    # Retry loop for transient errors
+    last_error: BaseException | None = None
+    for attempt in range(MAX_RETRIES + 1):
+        try:
+            # Run sub-agent with isolated message history and streaming support
+            result = await run_fn(
+                agent=agent,
+                prompt=prompt,
+                deps=sub_agent_deps,
+                message_history=[],  # Isolated context
+                event_stream_handler=deps.parent_stream_handler,  # Forward streaming
+            )
+
+            # Extract response text
+            response_text = ""
+            if result and result.output:
+                response_text = result.output.response
+
+            # Extract files modified
+            files_modified = [
+                op.file_path for op in sub_agent_deps.file_tracker.operations
+            ]
+
+            # Check for clarifying questions
+            has_questions = False
+            questions: list[str] = []
+            if result and result.output and result.output.clarifying_questions:
+                has_questions = True
+                questions = result.output.clarifying_questions
+
+            logger.info(
+                "Sub-agent %s completed. Files modified: %s",
+                agent_type.value,
+                files_modified,
+            )
+
+            # Clear active_sub_agent
+            deps.active_sub_agent = None
+
+            return DelegationResult(
+                success=True,
+                response=response_text,
+                files_modified=files_modified,
+                has_questions=has_questions,
+                questions=questions,
+            )
+
+        except Exception as e:
+            last_error = e
+            if _is_retryable_error(e) and attempt < MAX_RETRIES:
+                logger.warning(
+                    "Sub-agent %s failed (attempt %d/%d), retrying: %s",
+                    agent_type.value,
+                    attempt + 1,
+                    MAX_RETRIES + 1,
+                    str(e),
+                )
+                continue
+
+            # Non-retryable error or max retries exceeded
+            logger.error(
+                "Sub-agent %s failed after %d attempts: %s",
+                agent_type.value,
+                attempt + 1,
+                str(e),
+            )
+            break
+
+    # Clear active_sub_agent on failure
+    deps.active_sub_agent = None
+
+    return DelegationResult(
+        success=False,
+        error=str(last_error) if last_error else "Unknown error",
+        response="",
+        files_modified=[],
+    )
+
+
+# =============================================================================
+# Delegation Tools
+# =============================================================================
+
+
+@register_tool(
+    category=ToolCategory.DELEGATION,
+    display_text="Delegating to Research agent",
+    key_arg="task",
+)
+async def delegate_to_research(
+    ctx: RunContext[RouterDeps],
+    input: DelegationInput,
+) -> DelegationResult:
+    """Delegate a task to the Research agent.
+
+    The Research agent specializes in:
+    - Finding information via web search
+    - Analyzing code and documentation
+    - Gathering background research
+    - Saving research findings to .shotgun/research.md
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        input: DelegationInput with task and optional context_hint.
+
+    Returns:
+        DelegationResult with the research findings.
+    """
+    return await _run_sub_agent(
+        ctx,
+        AgentType.RESEARCH,
+        input.task,
+        input.context_hint,
+    )
+
+
+@register_tool(
+    category=ToolCategory.DELEGATION,
+    display_text="Delegating to Specification agent",
+    key_arg="task",
+)
+async def delegate_to_specification(
+    ctx: RunContext[RouterDeps],
+    input: DelegationInput,
+) -> DelegationResult:
+    """Delegate a task to the Specification agent.
+
+    The Specification agent specializes in:
+    - Writing and updating .shotgun/specification.md
+    - Creating Pydantic contracts in .shotgun/contracts/
+    - Defining technical requirements and interfaces
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        input: DelegationInput with task and optional context_hint.
+
+    Returns:
+        DelegationResult with the specification updates.
+    """
+    return await _run_sub_agent(
+        ctx,
+        AgentType.SPECIFY,
+        input.task,
+        input.context_hint,
+    )
+
+
+@register_tool(
+    category=ToolCategory.DELEGATION,
+    display_text="Delegating to Plan agent",
+    key_arg="task",
+)
+async def delegate_to_plan(
+    ctx: RunContext[RouterDeps],
+    input: DelegationInput,
+) -> DelegationResult:
+    """Delegate a task to the Plan agent.
+
+    The Plan agent specializes in:
+    - Writing and updating .shotgun/plan.md
+    - Creating implementation plans with stages
+    - Defining technical approach and architecture
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        input: DelegationInput with task and optional context_hint.
+
+    Returns:
+        DelegationResult with the plan updates.
+    """
+    return await _run_sub_agent(
+        ctx,
+        AgentType.PLAN,
+        input.task,
+        input.context_hint,
+    )
+
+
+@register_tool(
+    category=ToolCategory.DELEGATION,
+    display_text="Delegating to Tasks agent",
+    key_arg="task",
+)
+async def delegate_to_tasks(
+    ctx: RunContext[RouterDeps],
+    input: DelegationInput,
+) -> DelegationResult:
+    """Delegate a task to the Tasks agent.
+
+    The Tasks agent specializes in:
+    - Writing and updating .shotgun/tasks.md
+    - Creating actionable implementation tasks
+    - Breaking down work into manageable items
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        input: DelegationInput with task and optional context_hint.
+
+    Returns:
+        DelegationResult with the task list updates.
+    """
+    return await _run_sub_agent(
+        ctx,
+        AgentType.TASKS,
+        input.task,
+        input.context_hint,
+    )
+
+
+@register_tool(
+    category=ToolCategory.DELEGATION,
+    display_text="Delegating to Export agent",
+    key_arg="task",
+)
+async def delegate_to_export(
+    ctx: RunContext[RouterDeps],
+    input: DelegationInput,
+) -> DelegationResult:
+    """Delegate a task to the Export agent.
+
+    The Export agent specializes in:
+    - Exporting artifacts and deliverables
+    - Generating outputs to .shotgun/export/
+    - Creating documentation exports
+
+    Args:
+        ctx: RunContext with RouterDeps.
+        input: DelegationInput with task and optional context_hint.
+
+    Returns:
+        DelegationResult with the export results.
+    """
+    return await _run_sub_agent(
+        ctx,
+        AgentType.EXPORT,
+        input.task,
+        input.context_hint,
+    )