SneakyCode/app/agent/loop.py

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

import json
from typing import Any

from app.agent.context import SessionContext
from app.models.config import 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 (
    print_error,
    print_iteration_header,
    print_tool_call,
    print_tool_result,
    print_token_usage,
    print_warning,
)
from app.utils.logging import get_logger

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,
    ) -> None:
        self._config = config
        self._ctx = ctx
        self._client = client
        self._handler = handler
        self._registry = registry
        self._permissions = permissions
        self._tools_schema = registry.get_openai_tools_schema()
        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]
        return (
            "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)."
        )

    def _get_messages_with_system_prompt(self) -> list[Message]:
        """Prepend the system prompt to conversation history."""
        system_msg = Message(role="system", content=self._system_prompt)
        return [system_msg] + self._ctx.get_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
        for iteration in range(1, max_iter + 1):
            if self._cancelled:
                print_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:
                    print_warning(f"Token budget pressure: dropped {dropped} oldest messages.")
                else:
                    print_warning("Token budget exceeded, cannot truncate further. Stopping.")
                    break

            if iteration > 1:
                print_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:
                total = self._ctx.token_counter.cumulative_usage.total_tokens
                if total == 0:
                    total = self._ctx.estimated_tokens
                print_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()

                if reasoning_only_streak >= _MAX_REASONING_RETRIES:
                    # Nudge the model by injecting a user hint
                    print_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. Do not just think — provide your actual response.",
                    )
                    reasoning_only_streak = 0
                else:
                    print_warning("Model produced reasoning but no response. Retrying...")
                continue

            # Successful response — reset streak
            reasoning_only_streak = 0

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

            # Execute tool calls
            results = 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,
                )

            # Check if finish tool was called
            if any(r.tool_name == "finish" for r in results):
                break
        else:
            print_warning(f"Agent reached maximum iterations ({max_iter}). Stopping.")

    async def _llm_step(self) -> 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.

        Returns:
            The assistant Message, or None if an error occurred.
        """
        messages = self._get_messages_with_system_prompt()
        try:
            chunk_iter = self._client.stream_chat_with_retry(messages, tools=self._tools_schema)
            return await self._handler.process_stream(chunk_iter)
        except KeyboardInterrupt:
            print_warning("Response interrupted.")
            self._handler.reset()
            return None
        except (LLMConnectionError, LLMStreamError) as e:
            partial = self._handler.get_partial_message()
            if partial is not None:
                print_warning(f"Stream interrupted ({e}), returning partial response.")
                return partial
            print_error(f"Connection error: {e}")
            return None
        except LLMError as e:
            print_error(f"LLM error: {e}")
            return None

    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:
                print_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:
                    print_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:
                    print_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 self._permissions.check(name, description=desc):
                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:
                    print_tool_result(name, result.error or "", is_error=True)
                continue

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

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

        return results