tatlock/tests/responses/test_streaming_delegation.py

"""
Tests for real-time think message streaming and context plumbing in
the direct delegation paths.

_stream_direct_delegation must be an async generator that yields the
"start" think message BEFORE the expert runs (so 'Allow me to consult
the archives, sir.' streams while research is in flight), and both
direct delegation paths must pass trimmed conversation history as
expert context.
"""

from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import pytest

from src.agents.delegation import DelegationResult
from src.responses.streaming import (
    ReasoningSummaryDelta,
    ReasoningSummaryDone,
    StreamingCoordinator,
)

HISTORY = [
    {"role": "user", "content": "Tell me about my homelab wiki"},
    {"role": "assistant", "content": "It documents your services, sir."},
]


def _librarian_result(output: str = "Findings.") -> DelegationResult:
    return DelegationResult(
        expert_name="librarian",
        task="task",
        success=True,
        output=output,
    )


@pytest.mark.unit
class TestStreamDirectDelegation:
    """Real-time streaming behavior of _stream_direct_delegation."""

    @pytest.mark.asyncio
    async def test_start_think_streams_before_research_runs(self):
        coordinator = StreamingCoordinator()
        tracker = AsyncMock()
        results: dict = {}

        with patch(
            "src.agents.delegation.delegate_to_librarian",
            new_callable=AsyncMock,
            return_value=_librarian_result(),
        ) as mock_delegate:
            gen = coordinator._stream_direct_delegation(
                user_message="Search for Docker info",
                recommendation=SimpleNamespace(recommended_capabilities=["librarian"]),
                tracker=tracker,
                conversation_id="conv_1",
                conversation_history=HISTORY,
                results=results,
            )

            # First event: the start think message, BEFORE any research
            first = await gen.__anext__()
            assert isinstance(first, ReasoningSummaryDelta)
            assert first.delta.strip() != ""
            assert (
                mock_delegate.await_count == 0
            ), "start think message must stream before the expert runs"

            second = await gen.__anext__()
            assert isinstance(second, ReasoningSummaryDone)
            assert mock_delegate.await_count == 0

            # Third event: completion message - research has now run
            third = await gen.__anext__()
            assert isinstance(third, ReasoningSummaryDelta)
            assert mock_delegate.await_count == 1

            remaining = [event async for event in gen]
            assert any(isinstance(e, ReasoningSummaryDone) for e in remaining)

        # Results dict is populated for Phase 2 synthesis
        assert results["expert_results"] == {"librarian": "Findings."}
        assert results["tools_called"] == ["delegate_to_librarian"]
        assert len(results["think_messages"]) == 2

    @pytest.mark.asyncio
    async def test_conversation_history_passed_as_context(self):
        coordinator = StreamingCoordinator()
        tracker = AsyncMock()

        with patch(
            "src.agents.delegation.delegate_to_librarian",
            new_callable=AsyncMock,
            return_value=_librarian_result(),
        ) as mock_delegate:
            events = [
                event
                async for event in coordinator._stream_direct_delegation(
                    user_message="And what services does it list?",
                    recommendation=SimpleNamespace(recommended_capabilities=["librarian"]),
                    tracker=tracker,
                    conversation_id="conv_1",
                    conversation_history=HISTORY,
                    results={},
                )
            ]

        assert events, "generator must yield think events"
        context = mock_delegate.await_args.kwargs["context"]
        assert "Tell me about my homelab wiki" in context
        assert "It documents your services, sir." in context

    @pytest.mark.asyncio
    async def test_failed_delegation_streams_error_think(self):
        coordinator = StreamingCoordinator()
        tracker = AsyncMock()
        results: dict = {}

        failed = DelegationResult(
            expert_name="librarian",
            task="task",
            success=False,
            output="I'm afraid the archives proved difficult to access.",
            error="The Librarian was unable to complete the task.",
        )

        with patch(
            "src.agents.delegation.delegate_to_librarian",
            new_callable=AsyncMock,
            return_value=failed,
        ):
            events = [
                event
                async for event in coordinator._stream_direct_delegation(
                    user_message="Search for Docker info",
                    recommendation=SimpleNamespace(recommended_capabilities=["librarian"]),
                    tracker=tracker,
                    conversation_id="conv_1",
                    results=results,
                )
            ]

        assert results["tools_called"] == []
        # Expert result carries the curated user-safe sentence
        assert "archives" in results["expert_results"]["librarian"]
        deltas = [e.delta for e in events if isinstance(e, ReasoningSummaryDelta)]
        assert len(deltas) == 2  # start + error think messages


@pytest.mark.unit
class TestServiceDelegationContext:
    """The non-streaming direct delegation path passes trimmed history."""

    @pytest.mark.asyncio
    async def test_direct_delegation_with_results_passes_context(self):
        from src.responses.service import _direct_delegation_with_results

        tracker = AsyncMock()

        with patch(
            "src.agents.delegation.delegate_to_librarian",
            new_callable=AsyncMock,
            return_value=_librarian_result(),
        ) as mock_delegate:
            results = await _direct_delegation_with_results(
                user_message="And what services does it list?",
                recommendation=SimpleNamespace(recommended_capabilities=["librarian"]),
                tracker=tracker,
                conversation_id="conv_1",
                conversation_history=HISTORY,
            )

        context = mock_delegate.await_args.kwargs["context"]
        assert "Tell me about my homelab wiki" in context
        assert results["expert_results"]["librarian"] == "Findings."