299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""Unit tests for the AgentLoop ReAct-style tool-call loop."""
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from app.agent.context import SessionContext
|
|
from app.agent.loop import AgentLoop
|
|
from app.models.config import (
|
|
AgentConfig,
|
|
AppConfig,
|
|
DisplayConfig,
|
|
LLMConfig,
|
|
PermissionsConfig,
|
|
ToolsConfig,
|
|
)
|
|
from app.models.message import Message
|
|
from app.models.tool_call import ToolCall, ToolCallFunction, ToolResult, ToolResultStatus
|
|
from app.services.llm import LLMClient
|
|
from app.services.permissions import PermissionsService
|
|
from app.services.streaming import StreamHandler
|
|
from app.tools.registry import ToolRegistry, create_default_registry
|
|
from app.utils.display import DisplayAdapter
|
|
|
|
|
|
@pytest.fixture
|
|
def config() -> AppConfig:
|
|
return AppConfig(
|
|
llm=LLMConfig(
|
|
model="test-model",
|
|
endpoint="http://localhost:11434",
|
|
),
|
|
agent=AgentConfig(
|
|
max_iterations=5,
|
|
max_conversation_tokens=32000,
|
|
workspace_root=Path("/tmp/test-workspace"),
|
|
),
|
|
permissions=PermissionsConfig(
|
|
auto_approve=["read_file", "list_dir", "grep_files", "find_files", "finish"],
|
|
),
|
|
display=DisplayConfig(
|
|
show_tool_calls=True,
|
|
show_token_usage=False,
|
|
stream_output=False,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def ctx(config: AppConfig) -> SessionContext:
|
|
return SessionContext(config)
|
|
|
|
|
|
@pytest.fixture
|
|
def client() -> MagicMock:
|
|
return MagicMock(spec=LLMClient)
|
|
|
|
|
|
@pytest.fixture
|
|
def handler() -> MagicMock:
|
|
mock = MagicMock(spec=StreamHandler)
|
|
mock.usage = None
|
|
mock.had_reasoning_only = False
|
|
mock.reset = MagicMock()
|
|
mock.get_partial_message = MagicMock(return_value=None)
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def registry(config: AppConfig) -> ToolRegistry:
|
|
return create_default_registry(config.agent.workspace_root, config)
|
|
|
|
|
|
@pytest.fixture
|
|
def permissions(config: AppConfig) -> PermissionsService:
|
|
return PermissionsService(config.permissions)
|
|
|
|
|
|
@pytest.fixture
|
|
def display() -> MagicMock:
|
|
return MagicMock(spec=DisplayAdapter)
|
|
|
|
|
|
@pytest.fixture
|
|
def agent(
|
|
config: AppConfig,
|
|
ctx: SessionContext,
|
|
client: MagicMock,
|
|
handler: MagicMock,
|
|
registry: ToolRegistry,
|
|
permissions: PermissionsService,
|
|
display: MagicMock,
|
|
) -> AgentLoop:
|
|
return AgentLoop(config, ctx, client, handler, registry, permissions, display)
|
|
|
|
|
|
def _make_text_message(content: str) -> Message:
|
|
"""Helper: create an assistant message with text only (no tool calls)."""
|
|
return Message(role="assistant", content=content, tool_calls=None)
|
|
|
|
|
|
def _make_tool_call_message(
|
|
tool_name: str,
|
|
arguments: str,
|
|
tc_id: str = "call_001",
|
|
content: str | None = None,
|
|
) -> Message:
|
|
"""Helper: create an assistant message with a single tool call."""
|
|
return Message(
|
|
role="assistant",
|
|
content=content,
|
|
tool_calls=[
|
|
ToolCall(
|
|
id=tc_id,
|
|
type="function",
|
|
function=ToolCallFunction(name=tool_name, arguments=arguments),
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
class TestAgentLoop:
|
|
@pytest.mark.asyncio
|
|
async def test_plain_text_response(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""LLM returns text with no tool calls — loop completes in 1 iteration."""
|
|
handler.process_stream = AsyncMock(return_value=_make_text_message("Hello!"))
|
|
|
|
await agent.run_turn("Hi there")
|
|
|
|
assert handler.process_stream.call_count == 1
|
|
# History: user + assistant
|
|
history = ctx.get_history()
|
|
assert len(history) == 2
|
|
assert history[0].role == "user"
|
|
assert history[1].role == "assistant"
|
|
assert history[1].content == "Hello!"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_tool_call(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""LLM calls a tool, then responds with text — 2 LLM calls."""
|
|
handler.process_stream = AsyncMock(
|
|
side_effect=[
|
|
_make_tool_call_message("list_dir", '{"directory_path": "."}'),
|
|
_make_text_message("Here are the files."),
|
|
]
|
|
)
|
|
|
|
await agent.run_turn("List files")
|
|
|
|
assert handler.process_stream.call_count == 2
|
|
history = ctx.get_history()
|
|
# user, assistant (tool_call), tool (result), assistant (text)
|
|
assert len(history) == 4
|
|
assert history[0].role == "user"
|
|
assert history[1].role == "assistant"
|
|
assert history[1].tool_calls is not None
|
|
assert history[2].role == "tool"
|
|
assert history[3].role == "assistant"
|
|
assert history[3].content == "Here are the files."
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finish_tool_breaks_loop(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""Calling the finish tool terminates the loop immediately."""
|
|
handler.process_stream = AsyncMock(
|
|
return_value=_make_tool_call_message("finish", '{"message": "All done!"}'),
|
|
)
|
|
|
|
await agent.run_turn("Do something")
|
|
|
|
assert handler.process_stream.call_count == 1
|
|
history = ctx.get_history()
|
|
# user, assistant (finish call), tool (finish result)
|
|
assert len(history) == 3
|
|
assert history[2].role == "tool"
|
|
assert "All done!" in (history[2].content or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_iterations(self, agent: AgentLoop, handler: MagicMock, config: AppConfig) -> None:
|
|
"""Loop stops at max_iterations when LLM keeps calling tools."""
|
|
handler.process_stream = AsyncMock(
|
|
return_value=_make_tool_call_message("list_dir", '{"directory_path": "."}'),
|
|
)
|
|
|
|
await agent.run_turn("Keep going")
|
|
|
|
# Should call LLM max_iterations times
|
|
assert handler.process_stream.call_count == config.agent.max_iterations
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bad_json_arguments(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""Invalid JSON in tool arguments produces an error result, no exception."""
|
|
handler.process_stream = AsyncMock(
|
|
side_effect=[
|
|
_make_tool_call_message("list_dir", "not valid json{{{"),
|
|
_make_text_message("Sorry about that."),
|
|
]
|
|
)
|
|
|
|
await agent.run_turn("Bad args")
|
|
|
|
history = ctx.get_history()
|
|
# user, assistant (bad call), tool (error), assistant (apology)
|
|
assert len(history) == 4
|
|
tool_msg = history[2]
|
|
assert tool_msg.role == "tool"
|
|
assert "Invalid JSON" in (tool_msg.content or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unknown_tool(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""Unknown tool name produces an error result listing available tools."""
|
|
handler.process_stream = AsyncMock(
|
|
side_effect=[
|
|
_make_tool_call_message("nonexistent_tool", "{}"),
|
|
_make_text_message("I'll try something else."),
|
|
]
|
|
)
|
|
|
|
await agent.run_turn("Use fake tool")
|
|
|
|
history = ctx.get_history()
|
|
tool_msg = history[2]
|
|
assert tool_msg.role == "tool"
|
|
assert "Unknown tool" in (tool_msg.content or "")
|
|
assert "nonexistent_tool" in (tool_msg.content or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_permission_denied(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext, config: AppConfig) -> None:
|
|
"""Denied tool produces an error result."""
|
|
# Add list_dir to deny list
|
|
config.permissions.deny.append("list_dir")
|
|
# Recreate permissions service with updated config
|
|
agent._permissions = PermissionsService(config.permissions)
|
|
|
|
handler.process_stream = AsyncMock(
|
|
side_effect=[
|
|
_make_tool_call_message("list_dir", '{"directory_path": "."}'),
|
|
_make_text_message("Permission was denied."),
|
|
]
|
|
)
|
|
|
|
await agent.run_turn("List files")
|
|
|
|
history = ctx.get_history()
|
|
tool_msg = history[2]
|
|
assert tool_msg.role == "tool"
|
|
assert "Permission denied" in (tool_msg.content or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_connection_error_stops_loop(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""LLM connection error terminates the loop gracefully."""
|
|
from app.services.llm import LLMConnectionError
|
|
|
|
handler.process_stream = AsyncMock(side_effect=LLMConnectionError("Connection refused"))
|
|
|
|
await agent.run_turn("Hello")
|
|
|
|
# Only the user message should be in history (no assistant message added)
|
|
history = ctx.get_history()
|
|
assert len(history) == 1
|
|
assert history[0].role == "user"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_tool_calls_in_single_response(self, agent: AgentLoop, handler: MagicMock, ctx: SessionContext) -> None:
|
|
"""Multiple tool calls in one response are all executed."""
|
|
multi_tc_msg = Message(
|
|
role="assistant",
|
|
content=None,
|
|
tool_calls=[
|
|
ToolCall(
|
|
id="call_001",
|
|
type="function",
|
|
function=ToolCallFunction(name="list_dir", arguments='{"directory_path": "."}'),
|
|
),
|
|
ToolCall(
|
|
id="call_002",
|
|
type="function",
|
|
function=ToolCallFunction(name="find_files", arguments='{"pattern": "*.py"}'),
|
|
),
|
|
],
|
|
)
|
|
handler.process_stream = AsyncMock(
|
|
side_effect=[
|
|
multi_tc_msg,
|
|
_make_text_message("Found everything."),
|
|
]
|
|
)
|
|
|
|
await agent.run_turn("List and find files")
|
|
|
|
history = ctx.get_history()
|
|
# user, assistant (2 tool calls), tool (result 1), tool (result 2), assistant (text)
|
|
assert len(history) == 5
|
|
assert history[2].role == "tool"
|
|
assert history[2].tool_call_id == "call_001"
|
|
assert history[3].role == "tool"
|
|
assert history[3].tool_call_id == "call_002"
|
|
assert history[4].content == "Found everything."
|