shotgun-sh 0.3.3.dev1__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/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")

shotgun/agents/tools/file_management.py

@@ -23,7 +23,10 @@ logger = get_logger(__name__)
 # - A list of Paths: multiple allowed files/directories (e.g., [Path("specification.md"), Path("contracts")])
 # - "*": any file except protected files (for export agent)
 AGENT_DIRECTORIES: dict[AgentType, str | Path | list[Path]] = {
-    AgentType.RESEARCH: Path("research.md"),
+    AgentType.RESEARCH: [
+        Path("research.md"),
+        Path("research"),
+    ],  # Research can write main file and research folder
     AgentType.SPECIFY: [
         Path("specification.md"),
         Path("contracts"),
@@ -282,3 +285,48 @@ async def append_file(ctx: RunContext[AgentDeps], filename: str, content: str) -
         Success message or error message
     """
     return await write_file(ctx, filename, content, mode="a")
+
+
+@register_tool(
+    category=ToolCategory.ARTIFACT_MANAGEMENT,
+    display_text="Deleting file",
+    key_arg="filename",
+)
+async def delete_file(ctx: RunContext[AgentDeps], filename: str) -> str:
+    """Delete a file from the .shotgun directory.
+
+    Uses the same permission model as write_file - agents can only delete
+    files they have permission to write to.
+
+    Args:
+        filename: Relative path to file within .shotgun directory
+
+    Returns:
+        Success message or error message
+
+    Raises:
+        ValueError: If path is outside .shotgun directory or agent lacks permission
+        FileNotFoundError: If file does not exist
+    """
+    logger.debug("🔧 Deleting file: %s", filename)
+
+    try:
+        # Use agent-scoped validation (same as write_file)
+        file_path = _validate_agent_scoped_path(filename, ctx.deps.agent_mode)
+
+        if not await aiofiles.os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {filename}")
+
+        # Delete the file
+        await aiofiles.os.remove(file_path)
+        logger.debug("🗑️ Deleted file: %s", filename)
+
+        # Track the file operation
+        ctx.deps.file_tracker.add_operation(file_path, FileOperationType.DELETED)
+
+        return f"Successfully deleted {filename}"
+
+    except Exception as e:
+        error_msg = f"Error deleting file '{filename}': {str(e)}"
+        logger.error("❌ File delete failed: %s", error_msg)
+        return error_msg

shotgun/agents/tools/registry.py

@@ -30,6 +30,8 @@ class ToolCategory(StrEnum):
     ARTIFACT_MANAGEMENT = "artifact_management"
     WEB_RESEARCH = "web_research"
     AGENT_RESPONSE = "agent_response"
+    PLANNING = "planning"
+    DELEGATION = "delegation"
     UNKNOWN = "unknown"
 
 

shotgun/agents/tools/web_search/__init__.py

@@ -44,9 +44,8 @@ async def get_available_web_search_tools() -> list[WebSearchTool]:
     # Check if using Shotgun Account
     config_manager = get_config_manager()
     config = await config_manager.load()
-    has_shotgun_key = config.shotgun.api_key is not None
 
-    if has_shotgun_key:
+    if config.shotgun.has_valid_account:
         logger.debug("🔑 Shotgun Account - only Gemini web search available")
 
         # Gemini: Only search tool available for Shotgun Account

shotgun/agents/tools/web_search/gemini.py

@@ -1,7 +1,7 @@
 """Gemini web search tool implementation."""
 
 from opentelemetry import trace
-from pydantic_ai.messages import ModelMessage, ModelRequest
+from pydantic_ai.messages import ModelMessage, ModelRequest, TextPart
 from pydantic_ai.settings import ModelSettings
 
 from shotgun.agents.config import get_provider_model
@@ -82,8 +82,6 @@ async def gemini_web_search_tool(query: str) -> str:
         )
 
         # Extract text from response
-        from pydantic_ai.messages import TextPart
-
         result_text = "No content returned from search"
         if response.parts:
             for part in response.parts:

shotgun/codebase/core/change_detector.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Any, cast
 
 import aiofiles
-import kuzu
+import real_ladybug as kuzu
 
 from shotgun.logging_config import get_logger
 

shotgun/codebase/core/ingestor.py

@@ -11,7 +11,7 @@ from pathlib import Path
 from typing import Any
 
 import aiofiles
-import kuzu
+import real_ladybug as kuzu
 from tree_sitter import Node, Parser, QueryCursor
 
 from shotgun.codebase.core.language_config import LANGUAGE_CONFIGS, get_language_config

shotgun/codebase/core/manager.py

@@ -12,7 +12,7 @@ from pathlib import Path
 from typing import Any, ClassVar
 
 import anyio
-import kuzu
+import real_ladybug as kuzu
 from watchdog.events import FileSystemEvent, FileSystemEventHandler
 from watchdog.observers import Observer
 

shotgun/prompts/agents/export.j2

@@ -10,6 +10,8 @@ Your primary job is to generate Agents.md or CLAUDE.md files following the https
 
 {% include 'agents/partials/common_agent_system_prompt.j2' %}
 
+{% include 'agents/partials/router_delegation_mode.j2' %}
+
 ## MEMORY MANAGEMENT PROTOCOL
 
 - You can write to ANY file EXCEPT protected files

shotgun/prompts/agents/partials/common_agent_system_prompt.j2

@@ -17,17 +17,12 @@ Your extensive expertise spans, among other things:
 
 ## AGENT FILE PERMISSIONS
 
-There are four agents in the pipeline, and each agent can ONLY write to specific files. The user can switch between agents using **Shift+Tab**.
+There are four agents in the pipeline, and each agent can ONLY write to specific files:
 
-The **Research agent** can only write to `research.md`. If the user asks about specifications, plans, or tasks, tell them: "Use **Shift+Tab** to switch to the [appropriate] agent which can edit that file for you."
-
-The **Specification agent** can only write to `specification.md` and files inside the `.shotgun/contracts/` directory. If the user asks about research, plans, or tasks, tell them which agent handles that file.
-
-The **Plan agent** can only write to `plan.md`. If the user asks about research, specifications, or tasks, tell them which agent handles that file.
-
-The **Tasks agent** can only write to `tasks.md`. If the user asks about research, specifications, or plans, tell them which agent handles that file.
-
-When a user asks you to edit a file you cannot write to, you MUST tell them which agent can help and how to switch: "I can't edit [filename] - that's handled by the [agent name] agent. Use **Shift+Tab** to switch to that agent and it can edit that file for you."
+- **Research agent**: `research.md` and `.shotgun/research/`
+- **Specification agent**: `specification.md` and `.shotgun/contracts/`
+- **Plan agent**: `plan.md`
+- **Tasks agent**: `tasks.md`
 
 ## KEY RULES
 

shotgun/prompts/agents/partials/router_delegation_mode.j2

@@ -0,0 +1,36 @@
+{% if sub_agent_context and sub_agent_context.is_router_delegated %}
+## ROUTER DELEGATION MODE
+
+You are being orchestrated by the Router agent.
+
+**Current Task Context:**
+{% if sub_agent_context.plan_goal %}
+- Plan Goal: {{ sub_agent_context.plan_goal }}
+{% endif %}
+{% if sub_agent_context.current_step_title %}
+- Current Step: {{ sub_agent_context.current_step_title }}
+{% endif %}
+
+**CRITICAL BEHAVIORAL OVERRIDES** (these override ALL other instructions):
+
+1. **DO THE WORK FIRST**: Use tools (read_file, write_file, query_graph, web_search, etc.) to COMPLETE the task BEFORE calling final_result
+2. **NO ANNOUNCEMENTS**: Do NOT announce what you're going to do - just DO IT immediately
+3. **NO "PLEASE WAIT" MESSAGES**: Do NOT tell the user to "be patient" or that "this will take a few minutes"
+4. **final_result IS FOR RESULTS ONLY**: Only call final_result AFTER you have completed actual work (written files, done research, etc.)
+5. **SKIP GREETINGS**: No greetings, pleasantries, or preamble - start working immediately
+6. **BE CONCISE**: The Router handles user communication - just return structured results
+
+**WRONG behavior:**
+```
+1. Call final_result with "I'll research this, please be patient..."
+2. Exit without doing any work
+```
+
+**CORRECT behavior:**
+```
+1. Call read_file to check existing research
+2. Call query_graph or web_search to gather information
+3. Call write_file to save research findings
+4. Call final_result with "Research complete. Updated research.md with findings on X and Y."
+```
+{% endif %}