SneakyCode/app/agent/loop.py

"""AgentLoop — ReAct-style tool-call loop for autonomous task execution."""

from __future__ import annotations

import json
import time
from typing import TYPE_CHECKING, Any

from app.agent.context import SessionContext
from app.models.config import AgentMode, AppConfig
from app.models.message import Message
from app.models.tool_call import ToolCall, ToolResult, ToolResultStatus
from app.services.llm import LLMClient, LLMConnectionError, LLMError, LLMStreamError
from app.services.permissions import PermissionsService
from app.services.streaming import StreamHandler
from app.tools.registry import ToolRegistry
from app.utils.display import DisplayAdapter
from app.utils.logging import get_logger

if TYPE_CHECKING:
    from app.services.debug_log import DebugLogger
    from app.services.skill_runner import SkillRunner
    from app.services.skills import SkillsManager

logger = get_logger(__name__)

_MAX_REASONING_RETRIES = 2


class AgentLoop:
    """ReAct-style agent loop that streams LLM responses and executes tool calls.

    The loop sends conversation history to the LLM, parses tool calls from the
    response, executes them with permission checks, feeds results back, and
    repeats until the LLM produces a plain-text response or calls ``finish``.
    """

    def __init__(
        self,
        config: AppConfig,
        ctx: SessionContext,
        client: LLMClient,
        handler: StreamHandler,
        registry: ToolRegistry,
        permissions: PermissionsService,
        display: DisplayAdapter | None = None,
        debug_logger: DebugLogger | None = None,
        skills_manager: SkillsManager | None = None,
        skill_runner: SkillRunner | None = None,
    ) -> None:
        self._config = config
        self._ctx = ctx
        self._client = client
        self._handler = handler
        self._registry = registry
        self._permissions = permissions
        self._display = display
        self._debug = debug_logger
        self._skills = skills_manager
        self._skill_runner = skill_runner
        self._tools_schema = registry.get_openai_tools_schema()
        if self._permissions.mode == AgentMode.PLAN:
            read_only = PermissionsService.READ_ONLY_TOOLS
            self._tools_schema = [
                t for t in self._tools_schema
                if t["function"]["name"] in read_only
            ]
        self._system_prompt = self._build_system_prompt()
        self._cancelled = False

    def cancel(self) -> None:
        """Request cancellation of the current agent turn."""
        self._cancelled = True

    def _build_system_prompt(self) -> str:
        """Build the system prompt including tool schemas and agent instructions."""
        tool_names = [t["function"]["name"] for t in self._tools_schema]
        prompt = (
            "You are SneakyCode, a local AI coding agent. "
            "You help users with software engineering tasks by reading files, "
            "searching code, and answering questions about their project.\n\n"
            f"Workspace root: {self._config.agent.workspace_root}\n"
            "IMPORTANT: All tool path arguments must be RELATIVE to the workspace root. "
            'Use "." for the root, "app/main.py" for files, "app/" for subdirectories. '
            "Never pass absolute paths or the workspace root path itself.\n\n"
            "Available tools: " + ", ".join(tool_names) + "\n\n"
            "When you have fully completed the user's request, call the `finish` tool "
            "with a brief summary. If you can answer directly without tools, just respond "
            "with text (no tool call needed)."
        )
        if self._skills:
            prompt += self._skills.get_system_prompt_snippet()
        if self._skill_runner and self._skill_runner.is_active:
            prompt += (
                f"\n\nCurrently active skill: {self._skill_runner.active_skill_name}. "
                "When the skill's objective is complete, call the `finish_skill` tool."
            )
        if self._permissions.mode == AgentMode.PLAN:
            prompt += (
                "\n\nYou are in PLAN mode. You may only use read-only tools: "
                "read_file, list_dir, grep_files, find_files, finish. "
                "Do NOT attempt to write files, edit code, or run commands. "
                "Instead, describe what changes you would make, which files "
                "you would modify, and provide the reasoning for each change."
            )
        return prompt

    # Models whose chat templates understand /no_think directives.
    _THINKING_MODEL_PREFIXES = ("qwen", "qwq")

    def _model_supports_no_think(self) -> bool:
        """Check if the current model uses a thinking chat template."""
        model_lower = self._config.llm.model.lower()
        return any(model_lower.startswith(p) for p in self._THINKING_MODEL_PREFIXES)

    def _get_messages_with_system_prompt(self) -> list[Message]:
        """Prepend the system prompt to conversation history.

        When thinking is disabled on a model that supports it, appends a
        system-level /no_think directive after the last user message so
        Qwen 3.x (and similar) chat templates see it.
        """
        system_msg = Message(role="system", content=self._system_prompt)
        history = self._ctx.get_history()

        if not self._config.llm.thinking and self._model_supports_no_think() and history:
            history = list(history)
            # Find last user message and insert a system hint after it
            for i in range(len(history) - 1, -1, -1):
                if history[i].role == "user":
                    no_think_msg = Message(
                        role="system",
                        content="/no_think",
                    )
                    history.insert(i + 1, no_think_msg)
                    break

        return [system_msg] + history

    async def run_turn(self, user_input: str) -> None:
        """Execute one full agent turn: add user message, loop until done.

        Args:
            user_input: The user's message text.
        """
        self._ctx.add_message("user", user_input)
        self._cancelled = False

        max_iter = self._config.agent.max_iterations
        reasoning_only_streak = 0
        empty_streak = 0
        for iteration in range(1, max_iter + 1):
            if self._cancelled:
                if self._display:
                    self._display.write_warning("Agent loop cancelled.")
                break

            # Check token budget — try truncation before giving up
            if self._ctx.token_counter.is_over_budget():
                system_tokens = self._ctx.token_counter.estimate_tokens(self._system_prompt)
                dropped = self._ctx.truncate_history(system_tokens)
                if dropped > 0:
                    if self._display:
                        self._display.write_warning(f"Token budget pressure: dropped {dropped} oldest messages.")
                else:
                    if self._display:
                        self._display.write_warning("Token budget exceeded, cannot truncate further. Stopping.")
                    break

            if iteration > 1:
                if self._display:
                    self._display.write_iteration_header(iteration, max_iter)

            # Stream LLM response
            assistant_msg = await self._llm_step()
            if assistant_msg is None:
                break

            # Record assistant message
            self._ctx.add_message(
                "assistant",
                assistant_msg.content,
                tool_calls=assistant_msg.tool_calls,
            )

            # Detect reasoning-only response (model thought but produced nothing)
            reasoning_only = self._handler.had_reasoning_only

            # Record token usage
            if self._handler.usage:
                self._ctx.token_counter.count_usage(self._handler.usage)

            if self._config.display.show_token_usage and self._display:
                total = self._ctx.token_counter.cumulative_usage.total_tokens
                if total == 0:
                    total = self._ctx.estimated_tokens
                self._display.write_token_usage(total, self._ctx.token_counter.budget)

            self._handler.reset()

            # Reasoning-only: model produced thinking tokens but no content or tool calls.
            if reasoning_only:
                reasoning_only_streak += 1
                self._ctx.pop_last_message()

                # When thinking is disabled, reasoning-only is expected model noise.
                # Nudge immediately and silently to avoid wasting iterations.
                thinking_disabled = not self._config.llm.thinking

                # If the last context messages are tool errors, nudge immediately
                # rather than wasting retries — the model is likely confused by the error.
                has_recent_tool_error = any(
                    m.role == "tool" and m.content and m.content.startswith("Unknown ")
                    for m in self._ctx.get_history()[-3:]
                )

                should_nudge = (
                    thinking_disabled
                    or has_recent_tool_error
                    or reasoning_only_streak >= _MAX_REASONING_RETRIES
                )

                if should_nudge:
                    if not thinking_disabled and self._display:
                        self._display.write_warning(
                            f"Model produced reasoning but no response {reasoning_only_streak} times. "
                            "Nudging model to respond..."
                        )
                    self._ctx.add_message(
                        "user",
                        "Please respond with your answer. If a tool call failed, briefly explain what happened and continue.",
                    )
                    reasoning_only_streak = 0
                else:
                    if self._display:
                        self._display.write_warning("Model produced reasoning but no response. Retrying...")
                continue

            # Successful response — reset streak
            reasoning_only_streak = 0

            # Detect completely empty response (no content, no tool calls)
            if not assistant_msg.content and not assistant_msg.tool_calls:
                empty_streak += 1
                self._ctx.pop_last_message()  # Don't keep empty messages
                if empty_streak >= 2:
                    if self._display:
                        self._display.write_warning(
                            "Model returned repeated empty responses — "
                            "try a different model or check Ollama logs."
                        )
                    break
                if self._display:
                    self._display.write_warning("Model returned empty response. Retrying without tools...")
                # Retry without tool schemas — some models return empty when
                # tools are in the payload but the model can't handle them.
                assistant_msg = await self._llm_step(skip_tools=True)
                if assistant_msg is None:
                    break
                if assistant_msg.content:
                    self._ctx.add_message("assistant", assistant_msg.content)
                    if self._display:
                        self._display.write_assistant_message(assistant_msg.content)
                    self._handler.reset()
                    break
                # Still empty even without tools
                self._handler.reset()
                continue

            empty_streak = 0  # reset on successful non-empty response

            # Display any assistant text content (even if tool calls follow)
            if self._display and assistant_msg.content:
                self._display.write_assistant_message(assistant_msg.content)

            # No tool calls → task complete (plain text response)
            if not assistant_msg.tool_calls:
                break

            # Execute tool calls
            results = await self._execute_tool_calls(assistant_msg.tool_calls)

            # Add tool results to context
            for result in results:
                content = result.output if result.status == ToolResultStatus.SUCCESS else (result.error or "Unknown error")
                self._ctx.add_message(
                    "tool",
                    content,
                    tool_call_id=result.tool_call_id,
                    name=result.tool_name,
                )

            # Rebuild tools schema and system prompt if skill state may have changed
            if any(r.tool_name in ("load_skill", "finish_skill") for r in results):
                self._tools_schema = self._registry.get_openai_tools_schema()
                self._system_prompt = self._build_system_prompt()

            # Check if finish tool was called (finish_skill does NOT break the loop)
            if any(r.tool_name == "finish" for r in results):
                break
        else:
            if self._display:
                self._display.write_warning(f"Agent reached maximum iterations ({max_iter}). Stopping.")

    async def _llm_step(self, *, skip_tools: bool = False) -> Message | None:
        """Stream one LLM response and return the accumulated Message.

        Uses retry-enabled streaming. On mid-stream errors, attempts to recover
        partial content if available.

        Args:
            skip_tools: If True, send the request without tool schemas (fallback mode).

        Returns:
            The assistant Message, or None if an error occurred.
        """
        messages = self._get_messages_with_system_prompt()
        if self._debug:
            self._debug.log_request(messages, self._config.llm.model)
        tools = None if skip_tools else self._tools_schema
        t0 = time.monotonic()
        try:
            chunk_iter = self._client.stream_chat_with_retry(messages, tools=tools)
            result = await self._handler.process_stream(chunk_iter)
            if result and self._debug:
                elapsed = (time.monotonic() - t0) * 1000
                self._debug.log_response(result, self._handler.usage, elapsed)
            return result
        except KeyboardInterrupt:
            if self._display:
                self._display.write_warning("Response interrupted.")
            self._handler.reset()
            return None
        except (LLMConnectionError, LLMStreamError) as e:
            partial = self._handler.get_partial_message()
            if partial is not None:
                if self._display:
                    self._display.write_warning(f"Stream interrupted ({e}), returning partial response.")
                return partial
            if self._display:
                self._display.write_error(f"Connection error: {e}")
            return None
        except LLMError as e:
            if self._display:
                self._display.write_error(f"LLM error: {e}")
            return None

    async def _execute_tool_calls(self, tool_calls: list[ToolCall]) -> list[ToolResult]:
        """Execute a list of tool calls with permission checks.

        Args:
            tool_calls: Tool calls from the LLM response.

        Returns:
            List of ToolResult objects (one per tool call).
        """
        results: list[ToolResult] = []
        available_names = list(self._registry.get_all().keys())

        for tc in tool_calls:
            name = tc.function.name
            tc_id = tc.id

            # Display the tool call
            if self._config.display.show_tool_calls and self._display:
                self._display.write_tool_call(name, tc.function.arguments)

            # Parse arguments
            try:
                parsed_args: dict[str, Any] = json.loads(tc.function.arguments) if tc.function.arguments else {}
            except json.JSONDecodeError as e:
                result = ToolResult(
                    tool_call_id=tc_id,
                    tool_name=name,
                    status=ToolResultStatus.ERROR,
                    error=f"Invalid JSON in arguments: {e}",
                )
                results.append(result)
                if self._config.display.show_tool_calls and self._display:
                    self._display.write_tool_result(name, result.error or "", is_error=True)
                continue

            # Look up tool
            tool = self._registry.get(name)
            if tool is None:
                result = ToolResult(
                    tool_call_id=tc_id,
                    tool_name=name,
                    status=ToolResultStatus.ERROR,
                    error=f"Unknown tool '{name}'. Available: {available_names}",
                )
                results.append(result)
                if self._config.display.show_tool_calls and self._display:
                    self._display.write_tool_result(name, result.error or "", is_error=True)
                continue

            # Check permissions (truncate args for display in prompt)
            desc = tc.function.arguments[:120] + "..." if len(tc.function.arguments) > 120 else tc.function.arguments
            if not await self._permissions.check(name, description=desc, arguments=tc.function.arguments):
                result = ToolResult(
                    tool_call_id=tc_id,
                    tool_name=name,
                    status=ToolResultStatus.ERROR,
                    error=f"Permission denied for tool '{name}'",
                )
                results.append(result)
                if self._config.display.show_tool_calls and self._display:
                    self._display.write_tool_result(name, result.error or "", is_error=True)
                continue

            # Execute tool (BaseTool.run never raises)
            tool_t0 = time.monotonic()
            result = tool.run(tc_id, parsed_args)
            results.append(result)

            if self._debug:
                tool_elapsed = (time.monotonic() - tool_t0) * 1000
                self._debug.log_tool_execution(name, result.status.value, tool_elapsed)

            if self._config.display.show_tool_calls and self._display:
                is_error = result.status == ToolResultStatus.ERROR
                output = result.error if is_error else result.output
                self._display.write_tool_result(name, output or "", is_error=is_error)

        return results