webber/webber-api/src/domains/agents/task/agent.py

"""
Task Agent implementation using PydanticAI.

Full orchestrator agent that can:
- Execute multi-step tasks autonomously
- Use all tools (read + write) based on permission mode
- Spawn sub-agents (Explore, Plan) for focused work
- Stream structured events instead of raw text
"""
import asyncio
import os
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel

from src.domains.agents.base import BaseAgent, AgentContext, register_agent
from src.domains.agents.schemas import PermissionMode, StreamEvent, StreamEventType
from src.domains.agents.task.prompts import TASK_SYSTEM_PROMPT, TASK_PLAN_MODE_PROMPT
from src.ollama.provider import get_ollama_provider
from src.shared.config import get_settings
from src.shared.logging import logged, get_logger, trace_span

logger = get_logger(__name__)


@dataclass
class TaskContext(AgentContext):
    """
    Context for task agent tools.

    Passed to all tool functions via RunContext.
    Extends base AgentContext with permission mode and event queue.
    """
    mode: PermissionMode = PermissionMode.default
    # Prep for approval flow - tools can check this
    pending_approvals: list[str] = field(default_factory=list)
    # Event queue for streaming events from tools
    event_queue: asyncio.Queue | None = field(default=None, repr=False)
    # Track tool calls for retry logic
    tools_called: int = 0


def _emit_event(ctx: AgentContext, event: StreamEvent) -> None:
    """Emit an event to the queue if available."""
    if hasattr(ctx, 'event_queue') and ctx.event_queue is not None:
        ctx.event_queue.put_nowait(event)


def _summarize_result(result: str, max_len: int = 80) -> str:
    """Create a brief summary of a tool result."""
    # Count lines if multiline
    lines = result.strip().split('\n')
    if len(lines) > 1:
        return f"{len(lines)} lines"
    # Single line - truncate if needed
    if len(result) > max_len:
        return result[:max_len] + "..."
    return result


class TaskAgentImpl(BaseAgent):
    """
    Full orchestrator agent for autonomous task execution.

    Tool access depends on permission mode:
    - plan: Read-only tools only (safe exploration)
    - default: All tools (approval required for writes - future)
    - auto_accept: All tools (no approval prompts)

    Can spawn Explore and Plan agents to offload focused tasks,
    keeping context efficient across complex multi-step work.
    """

    name = "task"
    description = "Autonomous multi-step task execution with sub-agent orchestration"

    def __init__(self):
        """Initialize the task agent."""
        # Cache agents by mode to avoid recreating
        self._agents: dict[PermissionMode, Agent[TaskContext, str]] = {}
        self._settings = get_settings()

    @property
    def agent(self) -> Agent[TaskContext, str]:
        """Default agent (full mode) for compatibility."""
        return self._get_agent_for_mode(PermissionMode.default)

    def _get_agent_for_mode(self, mode: PermissionMode) -> Agent[TaskContext, str]:
        """Get or create agent configured for the specified mode."""
        if mode not in self._agents:
            self._agents[mode] = self._create_agent(mode)
        return self._agents[mode]

    def _create_agent(self, mode: PermissionMode = PermissionMode.default) -> Agent[TaskContext, str]:
        """Create the PydanticAI agent with Ollama backend."""
        # Use sanitized Ollama provider to fix content: null issues
        model = OpenAIModel(
            model_name=self._settings.ollama_agent_model,
            provider=get_ollama_provider(),
        )

        # Select system prompt based on mode
        system_prompt = TASK_PLAN_MODE_PROMPT if mode == PermissionMode.plan else TASK_SYSTEM_PROMPT

        agent: Agent[TaskContext, str] = Agent(
            model=model,
            system_prompt=system_prompt,
            deps_type=TaskContext,
            output_type=str,
            # Mistral Nemo settings:
            # - temperature 0.3 (Nemo needs slightly higher than 0.0)
            # - tool_choice "required" forces tool use
            model_settings={
                "temperature": 0.3,
                "extra_body": {"tool_choice": "required"},
            },
        )

        # Register tools based on mode
        self._register_tools(agent, mode)

        return agent

    def _register_tools(self, agent: Agent[TaskContext, str], mode: PermissionMode) -> None:
        """Register tools with the agent based on permission mode."""
        from src.domains.agents.task.tools_streaming import (
            register_task_tools_streaming,
            register_readonly_tools_streaming,
        )

        if mode == PermissionMode.plan:
            # Plan mode: read-only tools only
            register_readonly_tools_streaming(agent)
        else:
            # Default and auto_accept: all tools
            register_task_tools_streaming(agent)

    # Maximum retries when no tools are called
    MAX_NO_TOOL_RETRIES = 2

    def _build_prompt_with_context(self, prompt: str, working_dir: str) -> str:
        """Build the prompt with working directory context."""
        return f"""Working directory: {working_dir}

When using file tools, use paths relative to or within this working directory.
For example, to read a file at {working_dir}/README.md, use file_path="{working_dir}/README.md".

User request: {prompt}"""

    def _build_retry_prompt(self, prompt: str, working_dir: str) -> str:
        """Build a stronger prompt for retry after no tool calls."""
        return f"""Working directory: {working_dir}

IMPORTANT: Your previous response was REJECTED because you did not call any tools.
You MUST call a tool (like glob_files, bash_readonly, or read_file) BEFORE responding.
DO NOT answer from memory. DO NOT fabricate information.

Call a tool NOW to gather real information, then respond based on the results.

User request: {prompt}"""

    @logged()
    async def run(
        self,
        prompt: str,
        working_dir: str | None = None,
        allowed_paths: list[str] | None = None,
        mode: PermissionMode = PermissionMode.default,
        **kwargs: Any
    ) -> str:
        """
        Run the task agent to execute a multi-step task.

        Args:
            prompt: Description of the task to execute
            working_dir: Working directory for the agent
            allowed_paths: Restrict tool access to these paths
            mode: Permission mode controlling tool access

        Returns:
            Consolidated task summary with results
        """
        effective_working_dir = working_dir or os.getcwd()

        # Get agent configured for this mode
        agent = self._get_agent_for_mode(mode)

        async with trace_span("task_agent_run"):
            retries = 0
            while retries <= self.MAX_NO_TOOL_RETRIES:
                # Create fresh context for each attempt
                ctx = TaskContext(
                    working_dir=effective_working_dir,
                    allowed_paths=allowed_paths or self._settings.effective_allowed_paths,
                    timeout_seconds=self._settings.tool_timeout_seconds,
                    mode=mode,
                )

                # Build prompt - use retry prompt if this is a retry
                if retries == 0:
                    full_prompt = self._build_prompt_with_context(prompt, effective_working_dir)
                else:
                    full_prompt = self._build_retry_prompt(prompt, effective_working_dir)
                    logger.warning(f"Retry {retries}/{self.MAX_NO_TOOL_RETRIES}: No tools called, retrying with stronger prompt")

                try:
                    result = await agent.run(full_prompt, deps=ctx)

                    # Check if tools were called
                    if ctx.tools_called == 0 and retries < self.MAX_NO_TOOL_RETRIES:
                        retries += 1
                        continue

                    if ctx.tools_called == 0:
                        logger.warning("Agent responded without calling tools after all retries")

                    return result.output
                except Exception as e:
                    logger.exception(f"Task agent error: {e}")
                    raise

            # Should not reach here, but just in case
            return result.output

    async def run_stream(
        self,
        prompt: str,
        working_dir: str | None = None,
        allowed_paths: list[str] | None = None,
        mode: PermissionMode = PermissionMode.default,
        **kwargs: Any
    ) -> AsyncIterator[StreamEvent]:
        """
        Run the task agent with structured event streaming.

        Instead of streaming raw text (which gets garbled during tool calls),
        yields structured events that clients can render appropriately.

        Args:
            prompt: Task description
            working_dir: Working directory
            allowed_paths: Restrict tool access
            mode: Permission mode controlling tool access

        Yields:
            StreamEvent objects for tool progress and final response.

        Event types:
        - tool_start: Tool execution beginning
        - tool_done: Tool execution complete with summary
        - thinking: Agent status update
        - response: Final response text
        - error: Error occurred
        - done: Stream complete
        """
        effective_working_dir = working_dir or os.getcwd()

        # Get agent configured for this mode
        agent = self._get_agent_for_mode(mode)

        async with trace_span("task_agent_stream"):
            # Emit initial thinking event
            yield StreamEvent(
                event=StreamEventType.thinking,
                message="Starting task execution..."
            )

            retries = 0
            response = ""

            while retries <= self.MAX_NO_TOOL_RETRIES:
                # Create fresh event queue and context for each attempt
                event_queue: asyncio.Queue[StreamEvent] = asyncio.Queue()

                ctx = TaskContext(
                    working_dir=effective_working_dir,
                    allowed_paths=allowed_paths or self._settings.effective_allowed_paths,
                    timeout_seconds=self._settings.tool_timeout_seconds,
                    mode=mode,
                    event_queue=event_queue,
                )

                # Build prompt - use retry prompt if this is a retry
                if retries == 0:
                    full_prompt = self._build_prompt_with_context(prompt, effective_working_dir)
                else:
                    full_prompt = self._build_retry_prompt(prompt, effective_working_dir)
                    yield StreamEvent(
                        event=StreamEventType.thinking,
                        message=f"Retrying (attempt {retries + 1})..."
                    )

                # Run agent in background task so we can yield events
                async def run_agent() -> str:
                    try:
                        result = await agent.run(full_prompt, deps=ctx)
                        return result.output
                    except Exception as e:
                        logger.exception(f"Task agent stream error: {e}")
                        raise

                agent_task = asyncio.create_task(run_agent())

                # Yield events from queue while agent runs
                try:
                    while not agent_task.done():
                        try:
                            # Check for events with timeout
                            event = await asyncio.wait_for(
                                event_queue.get(),
                                timeout=0.1
                            )
                            yield event
                        except asyncio.TimeoutError:
                            # No events, check if agent is done
                            continue

                    # Drain remaining events
                    while not event_queue.empty():
                        yield event_queue.get_nowait()

                    # Get final result
                    response = await agent_task

                    # Check if tools were called - if not, retry
                    if ctx.tools_called == 0 and retries < self.MAX_NO_TOOL_RETRIES:
                        logger.warning(f"No tools called, retrying ({retries + 1}/{self.MAX_NO_TOOL_RETRIES})")
                        retries += 1
                        continue

                    if ctx.tools_called == 0:
                        logger.warning("Agent responded without calling tools after all retries")

                    # Success - break out of retry loop
                    break

                except Exception as e:
                    logger.exception(f"Stream error: {e}")
                    yield StreamEvent(
                        event=StreamEventType.error,
                        error_message=str(e)
                    )
                    # Cancel agent if still running
                    if not agent_task.done():
                        agent_task.cancel()
                        try:
                            await agent_task
                        except asyncio.CancelledError:
                            pass
                    return

            # Yield response in chunks for streaming feel
            chunk_size = 100
            for i in range(0, len(response), chunk_size):
                chunk = response[i:i + chunk_size]
                yield StreamEvent(
                    event=StreamEventType.response,
                    text=chunk
                )
                # Small delay for streaming effect
                await asyncio.sleep(0.01)

            # Signal completion
            yield StreamEvent(
                event=StreamEventType.done,
                mode=mode.value
            )


# Create and register the singleton instance
task_agent = TaskAgentImpl()
register_agent(task_agent)


async def task(
    prompt: str,
    working_dir: str | None = None,
    **kwargs: Any
) -> str:
    """Run task execution."""
    return await task_agent.run(prompt, working_dir=working_dir, **kwargs)


async def task_stream(
    prompt: str,
    working_dir: str | None = None,
    **kwargs: Any
) -> AsyncIterator[StreamEvent]:
    """Run task execution with event streaming."""
    async for event in task_agent.run_stream(prompt, working_dir=working_dir, **kwargs):
        yield event