nanobot/tests/test_agent_loop_tool_result.py

"""Tests for agent loop handling of ToolResult and CLIResult objects."""

import pytest
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.anthropic.base import ToolResult, CLIResult
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse, ToolCallRequest


@pytest.fixture
def mock_provider():
    """Create mock LLM provider."""
    provider = MagicMock()
    provider.chat = AsyncMock()
    provider.thinking_budget = 0
    return provider


@pytest.fixture
def mock_session_manager():
    """Create mock session manager."""
    session_mgr = MagicMock()
    session_mgr.load = AsyncMock(return_value={
        "messages": [],
        "metadata": {},
    })
    session_mgr.save = MagicMock()  # Synchronous in production, not async
    return session_mgr


@pytest.fixture
def mock_bus():
    """Create mock message bus."""
    bus = MagicMock(spec=MessageBus)
    bus.publish = AsyncMock()
    return bus


@pytest.fixture
def agent_loop(mock_provider, mock_session_manager, mock_bus, tmp_path):
    """Create agent loop for testing."""
    return AgentLoop(
        provider=mock_provider,
        session_manager=mock_session_manager,
        bus=mock_bus,
        workspace=tmp_path,
        max_iterations=5,
    )


@pytest.mark.asyncio
async def test_tool_result_with_output(agent_loop, mock_provider):
    """Test handling ToolResult with output field."""
    # Mock LLM responses
    mock_provider.chat.side_effect = [
        # First call: request tool
        LLMResponse(
            content="Using tool",
            tool_calls=[ToolCallRequest(id="call_1", name="test_tool", arguments={})],
        ),
        # Second call: final response
        LLMResponse(content="Done"),
    ]

    # Mock tool that returns ToolResult
    tool_result = ToolResult(output="Tool executed successfully")
    agent_loop.tools.execute = AsyncMock(return_value=tool_result)

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    # Verify tool result was added to messages
    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]

    # Find the tool result message
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")
    assert tool_msg["content"] == "Tool executed successfully"


@pytest.mark.asyncio
async def test_tool_result_with_error(agent_loop, mock_provider):
    """Test handling ToolResult with error field."""
    mock_provider.chat.side_effect = [
        LLMResponse(
            content="Using tool",
            tool_calls=[ToolCallRequest(id="call_1", name="test_tool", arguments={})],
        ),
        LLMResponse(content="Error handled"),
    ]

    tool_result = ToolResult(error="Command failed: exit code 1")
    agent_loop.tools.execute = AsyncMock(return_value=tool_result)

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")
    assert "Error:" in tool_msg["content"]
    assert "Command failed: exit code 1" in tool_msg["content"]


@pytest.mark.asyncio
async def test_tool_result_with_base64_image(agent_loop, mock_provider):
    """Test handling ToolResult with base64_image field."""
    mock_provider.chat.side_effect = [
        LLMResponse(
            content="Taking screenshot",
            tool_calls=[ToolCallRequest(id="call_1", name="screenshot", arguments={})],
        ),
        LLMResponse(content="Screenshot analyzed"),
    ]

    tool_result = ToolResult(
        output="Screenshot taken",
        base64_image="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==",
    )
    agent_loop.tools.execute = AsyncMock(return_value=tool_result)

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")

    # Should contain both text and image
    assert isinstance(tool_msg["content"], list)
    assert len(tool_msg["content"]) == 2

    # Text content
    text_part = next(p for p in tool_msg["content"] if p["type"] == "text")
    assert text_part["text"] == "Screenshot taken"

    # Image content
    image_part = next(p for p in tool_msg["content"] if p["type"] == "image")
    assert image_part["source"]["type"] == "base64"
    assert image_part["source"]["media_type"] == "image/png"
    assert "iVBORw0KGgoAAAANS" in image_part["source"]["data"]


@pytest.mark.asyncio
async def test_cli_result_handling(agent_loop, mock_provider):
    """Test handling CLIResult from text editor tools."""
    mock_provider.chat.side_effect = [
        LLMResponse(
            content="Editing file",
            tool_calls=[ToolCallRequest(id="call_1", name="edit", arguments={})],
        ),
        LLMResponse(content="File edited"),
    ]

    cli_result = CLIResult(
        exit_code=0,
        output="File updated successfully",
        error="",
    )
    agent_loop.tools.execute = AsyncMock(return_value=cli_result)

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")
    assert tool_msg["content"] == "File updated successfully"


@pytest.mark.asyncio
async def test_legacy_string_result(agent_loop, mock_provider):
    """Test backward compatibility with string results from function tools."""
    mock_provider.chat.side_effect = [
        LLMResponse(
            content="Using tool",
            tool_calls=[ToolCallRequest(id="call_1", name="legacy_tool", arguments={})],
        ),
        LLMResponse(content="Done"),
    ]

    # Legacy tool returns plain string
    agent_loop.tools.execute = AsyncMock(return_value="Plain text result")

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")
    assert tool_msg["content"] == "Plain text result"


@pytest.mark.asyncio
async def test_tool_result_output_and_error(agent_loop, mock_provider):
    """Test handling ToolResult with both output and error."""
    mock_provider.chat.side_effect = [
        LLMResponse(
            content="Running command",
            tool_calls=[ToolCallRequest(id="call_1", name="bash", arguments={})],
        ),
        LLMResponse(content="Handled"),
    ]

    tool_result = ToolResult(
        output="Partial output before error",
        error="Unexpected termination",
    )
    agent_loop.tools.execute = AsyncMock(return_value=tool_result)

    message = InboundMessage(
        channel="test",
        chat_id="123",
        sender_id="user1",
        content="Test message",
    )

    response = await agent_loop._process_message(message)

    calls = mock_provider.chat.call_args_list
    second_call_messages = calls[1][1]["messages"]
    tool_msg = next(m for m in second_call_messages if m.get("role") == "tool")

    # Should contain both output and error
    content = tool_msg["content"]
    assert "Partial output before error" in content
    assert "Error:" in content
    assert "Unexpected termination" in content