SneakyCode/tests/integration/test_agent_workflows.py

"""Integration tests for end-to-end agent workflows with mocked LLM."""

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

import pytest

from app.agent.context import SessionContext
from app.agent.loop import AgentLoop
from app.models.config import AgentConfig, AppConfig, LLMConfig
from app.services.llm import LLMConnectionError
from app.services.permissions import PermissionsService
from app.services.session import SessionManager
from app.services.streaming import StreamHandler
from app.tools.registry import create_default_registry

from .conftest import MockLLMClient, make_text_chunks, make_tool_call_chunks


@pytest.fixture
def agent_factory(test_config: AppConfig, tmp_workspace: Path):
    """Factory that creates an AgentLoop wired to a MockLLMClient."""

    def create(responses):
        ctx = SessionContext(test_config)
        mock_client = MockLLMClient(responses)
        handler = StreamHandler(test_config.display)
        registry = create_default_registry(test_config.agent.workspace_root, test_config)
        permissions = PermissionsService(test_config.permissions)
        agent = AgentLoop(test_config, ctx, mock_client, handler, registry, permissions)
        return agent, ctx, mock_client

    return create


class TestAgentWorkflows:
    @pytest.mark.asyncio
    async def test_multi_turn_read_workflow(self, agent_factory, tmp_workspace: Path) -> None:
        """Agent reads a file then responds with text — full 2-turn workflow."""
        responses = [
            make_tool_call_chunks("read_file", {"file_path": "hello.txt"}),
            make_text_chunks("The file contains: Hello, world!"),
        ]
        agent, ctx, mock_client = agent_factory(responses)

        await agent.run_turn("What's in hello.txt?")

        assert mock_client.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 "Hello, world!" in (history[2].content or "")
        assert history[3].role == "assistant"
        assert "Hello, world!" in (history[3].content or "")

    @pytest.mark.asyncio
    async def test_token_budget_truncation(self, test_config: AppConfig, tmp_workspace: Path) -> None:
        """When token budget is exceeded, truncation drops messages instead of stopping."""
        # Use a tiny budget to trigger truncation
        test_config.agent.max_conversation_tokens = 100
        test_config.agent.truncation_keep_recent = 2
        test_config.agent.truncation_threshold = 0.5

        ctx = SessionContext(test_config)

        # Fill history with enough to exceed budget
        for i in range(10):
            ctx.add_message("user", f"Message {i} " * 20)
            ctx.add_message("assistant", f"Response {i} " * 20)

        # Force token counter over budget
        from app.utils.token_counter import TokenUsage
        ctx.token_counter.count_usage(TokenUsage(total_tokens=100))

        original_count = len(ctx.get_history())
        dropped = ctx.truncate_history()

        assert dropped > 0
        assert len(ctx.get_history()) < original_count
        # Recent messages should still be present
        assert len(ctx.get_history()) >= 2

    @pytest.mark.asyncio
    async def test_session_save_and_restore(self, test_config: AppConfig, tmp_workspace: Path) -> None:
        """Session can be saved after a turn and restored into a fresh context."""
        responses = [make_text_chunks("Hello from the agent!")]
        ctx = SessionContext(test_config)
        mock_client = MockLLMClient(responses)
        handler = StreamHandler(test_config.display)
        registry = create_default_registry(test_config.agent.workspace_root, test_config)
        permissions = PermissionsService(test_config.permissions)
        agent = AgentLoop(test_config, ctx, mock_client, handler, registry, permissions)

        await agent.run_turn("Hi")

        # Save session
        session_mgr = SessionManager(
            test_config.session, test_config.agent.workspace_root, test_config.llm.model
        )
        path = session_mgr.save(ctx)
        assert path.exists()

        # Restore into fresh context
        fresh_ctx = SessionContext(test_config)
        loaded = session_mgr.load_latest()
        assert loaded is not None
        session_mgr.restore(loaded, fresh_ctx)

        assert fresh_ctx.message_count == ctx.message_count
        original_history = ctx.get_history()
        restored_history = fresh_ctx.get_history()
        for orig, restored in zip(original_history, restored_history):
            assert orig.role == restored.role
            assert orig.content == restored.content

    @pytest.mark.asyncio
    async def test_retry_on_transient_error(self, test_config: AppConfig, tmp_workspace: Path) -> None:
        """Agent recovers from transient LLM errors via retry."""
        from app.services.llm import LLMClient

        ctx = SessionContext(test_config)

        # Create a real client but mock stream_chat to fail then succeed
        client = LLMClient(test_config.llm)
        call_count = 0

        async def flaky_stream(*args, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise LLMConnectionError("Temporary failure")
            for chunk in make_text_chunks("Recovered!"):
                yield chunk

        client.stream_chat = flaky_stream  # type: ignore[assignment]

        handler = StreamHandler(test_config.display)
        registry = create_default_registry(test_config.agent.workspace_root, test_config)
        permissions = PermissionsService(test_config.permissions)
        agent = AgentLoop(test_config, ctx, client, handler, registry, permissions)

        with patch("app.services.llm.asyncio.sleep", new_callable=AsyncMock):
            await agent.run_turn("Test retry")

        history = ctx.get_history()
        # Should have succeeded: user + assistant
        assert len(history) == 2
        assert history[1].content == "Recovered!"
        assert call_count == 2

        await client.close()