tatlock/tests/integration/test_steward_streaming.py

"""
Integration tests for Steward + Tatlock streaming.

Tests the complete streaming flow with Steward preprocessing.
"""

from unittest.mock import patch

import pytest

from src.core.startup import initialize_application
from src.responses.schemas import ResponseRequest
from src.responses.streaming import StreamEventType, StreamingCoordinator


@pytest.fixture(scope="module", autouse=True)
def setup_household_registry():
    """Initialize household registry before running tests."""
    initialize_application()


class TestStewardStreaming:
    """Test Steward + Tatlock streaming integration."""

    @pytest.mark.asyncio
    async def test_stream_with_steward_basic(self):
        """Test basic streaming with Steward preprocessing."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "What's 2 + 2?"}],
            stream=True,
        )

        # Mock the Steward analysis
        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            # Mock the streaming method (async generator)
            with patch(
                "src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream"
            ) as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                # Mock Steward recommendation
                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="Math calculation requires tatlock_core",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                )

                # Mock Tatlock streaming response as async generator
                async def mock_stream(*args, **kwargs):
                    yield "Certainly, sir. "
                    yield "2 + 2 equals 4."

                mock_tatlock_stream.return_value = mock_stream()

                # Execute streaming
                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Verify event sequence
                event_types = [e.event for e in events]

                # Should have reasoning summary deltas
                assert StreamEventType.REASONING_SUMMARY_DELTA in event_types
                assert StreamEventType.REASONING_SUMMARY_DONE in event_types

                # Should have output text deltas
                assert StreamEventType.OUTPUT_TEXT_DELTA in event_types
                assert StreamEventType.OUTPUT_TEXT_DONE in event_types

                # Should end with response.done
                assert events[-1].event == StreamEventType.RESPONSE_DONE

                # Verify Steward and Tatlock were called
                assert mock_steward.called
                assert mock_tatlock_stream.called

    @pytest.mark.asyncio
    async def test_stream_with_conversation_history(self):
        """Test streaming with conversation history."""
        request = ResponseRequest(
            model="tatlock",
            input=[
                {"role": "user", "content": "What's 5 times 3?"},
                {"role": "assistant", "content": "That equals 15, sir."},
                {"role": "user", "content": "And divided by 3?"},
            ],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch(
                "src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream"
            ) as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="Follow-up calculation based on previous result of 15",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(
                        has_previous_context=True,
                        relevant_turns=[0],
                        context_summary="Previous calculation in turn 0",
                    ),
                )

                async def mock_stream(*args, **kwargs):
                    yield "15 divided by 3 equals 5, sir."

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Verify conversation history was passed to Steward
                call_kwargs = mock_steward.call_args[1]
                assert "conversation_history" in call_kwargs
                assert len(call_kwargs["conversation_history"]) == 2  # First Q&A pair

                # Verify final response includes both reasoning and message
                final_event = events[-1]
                assert final_event.event == StreamEventType.RESPONSE_DONE
                assert len(final_event.response.output) == 2  # Reasoning + Message

    @pytest.mark.asyncio
    async def test_stream_reasoning_contains_steward_analysis(self):
        """Test that reasoning summary contains Steward's analysis."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "Test request"}],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch(
                "src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream"
            ) as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="This is a test analysis with specific markers",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                )

                async def mock_stream(*args, **kwargs):
                    yield "Test response"

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                reasoning_deltas = []

                async for event in coordinator.stream_response_with_steward(request):
                    if event.event == StreamEventType.REASONING_SUMMARY_DELTA:
                        reasoning_deltas.append(event.delta)

                # Combine all reasoning deltas
                full_reasoning = "".join(reasoning_deltas)

                # Should contain Steward's analysis
                assert "test analysis" in full_reasoning.lower()
                assert len(reasoning_deltas) > 0, "Should have streamed reasoning deltas"

    @pytest.mark.asyncio
    async def test_stream_with_missing_capabilities(self):
        """Test streaming when Steward detects missing capabilities."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "Generate an image of a sunset"}],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch(
                "src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream"
            ) as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=[],
                    reasoning="Image generation not available in current toolset",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                    missing_capabilities="Image generation capability would be needed",
                )

                async def mock_stream(*args, **kwargs):
                    yield "I'm afraid I don't have image generation capabilities, sir."

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Should complete successfully even with missing capabilities
                assert events[-1].event == StreamEventType.RESPONSE_DONE

                # Verify empty scoped tools were passed to stream method
                tatlock_kwargs = mock_tatlock_stream.call_args[1]
                assert "scoped_tools" in tatlock_kwargs
                assert tatlock_kwargs["scoped_tools"] == []