tatlock/tests/integration/test_tatlock_streaming.py

"""
Integration tests for Tatlock agent streaming through full API stack.

These tests verify the complete streaming flow from API endpoint through
StreamingCoordinator to TatlockAgent, ensuring no text duplication and
proper delta calculation.
"""

import json

import pytest
from fastapi.testclient import TestClient
from httpx import AsyncClient


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_streaming_no_duplication(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock streaming produces no text duplication.

    This test catches the bug where accumulated text from PydanticAI was
    being re-streamed multiple times by the StreamingCoordinator.
    Note: Requires running server, may xfail if server unavailable or LLM times out.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Say hello"}],
        "stream": True,
    }

    collected_deltas = []

    try:
        async with async_client.stream(
            "POST",
            "/v1/responses",
            json=request_data,
            timeout=60.0,  # Increase timeout for LLM response
        ) as response:
            if response.status_code != 200:
                pytest.xfail(f"Server returned {response.status_code}")
            assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

            async for line in response.aiter_lines():
                if not line.strip():
                    continue

                if line.startswith("event: "):
                    line[7:].strip()
                elif line.startswith("data: "):
                    data_str = line[6:].strip()
                    if data_str != "[DONE]":
                        try:
                            chunk = json.loads(data_str)

                            # Collect output text deltas
                            if chunk.get("event") == "response.output_text.delta":
                                collected_deltas.append(chunk["delta"])

                        except json.JSONDecodeError:
                            pass
    except Exception as e:
        pytest.xfail(f"Streaming request failed (server may be unavailable): {e}")

    # Reconstruct full text from deltas
    full_text = "".join(collected_deltas)

    # Verify we got some response (xfail if LLM didn't produce output)
    if len(full_text) == 0:
        pytest.xfail("No text received from streaming (LLM may have timed out)")

    # Verify no obvious duplication patterns
    # Check that common words don't appear excessively repeated
    words = full_text.lower().split()
    if len(words) > 0:
        # Check for consecutive duplicate words (sign of duplication bug)
        consecutive_dupes = sum(
            1 for i in range(len(words) - 1) if words[i] == words[i + 1] and len(words[i]) > 3
        )
        # Allow a few duplicates (natural language), but not excessive
        assert (
            consecutive_dupes < len(words) * 0.1
        ), f"Too many consecutive duplicate words: {consecutive_dupes}/{len(words)}"


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_chat_streaming_no_duplication(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock streaming through Chat Completions API.

    Tests the full stack through the chat completions wrapper to ensure
    streaming works correctly without duplication.
    """
    request_data = {
        "model": "Tatlock",
        "messages": [{"role": "user", "content": "Hello"}],
        "stream": True,
    }

    collected_content = []

    async with async_client.stream(
        "POST",
        "/v1/chat/completions",
        json=request_data,
        timeout=30.0,
    ) as response:
        assert response.status_code == 200

        async for line in response.aiter_lines():
            if not line.strip():
                continue

            if line.startswith("data: "):
                data_str = line[6:].strip()
                if data_str == "[DONE]":
                    break

                try:
                    chunk = json.loads(data_str)

                    # Collect content deltas from choices
                    if "choices" in chunk and len(chunk["choices"]) > 0:
                        delta = chunk["choices"][0].get("delta", {})
                        if "content" in delta and delta["content"]:
                            collected_content.append(delta["content"])

                except json.JSONDecodeError:
                    pass

    # Reconstruct full response
    full_response = "".join(collected_content)

    # Verify we got a response
    assert len(full_response) > 0, "Should have received response content"

    # Check for duplication patterns
    words = full_response.lower().split()
    if len(words) > 0:
        consecutive_dupes = sum(
            1 for i in range(len(words) - 1) if words[i] == words[i + 1] and len(words[i]) > 3
        )
        assert (
            consecutive_dupes < len(words) * 0.1
        ), f"Too many consecutive duplicate words in chat response: {consecutive_dupes}/{len(words)}"


@pytest.mark.integration
def test_tatlock_non_streaming_responses_api(client: TestClient):
    """
    Integration test: Verify Tatlock non-streaming through Responses API.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Say hello"}],
        "stream": False,
    }

    response = client.post("/v1/responses", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Verify response structure
    assert data["status"] == "completed"
    assert "output" in data
    assert len(data["output"]) > 0

    # Get the message content
    message_item = next((item for item in data["output"] if item["type"] == "message"), None)
    assert message_item is not None, "Should have a message output item"
    assert len(message_item["content"]) > 0

    text = message_item["content"][0]["text"]
    assert len(text) > 0, "Should have response text"


@pytest.mark.integration
def test_tatlock_non_streaming_chat_api(client: TestClient):
    """
    Integration test: Verify Tatlock non-streaming through Chat Completions API.
    """
    request_data = {
        "model": "Tatlock",
        "messages": [{"role": "user", "content": "Hello"}],
        "stream": False,
    }

    response = client.post("/v1/chat/completions", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Verify OpenAI-compatible structure
    assert "id" in data
    assert data["object"] == "chat.completion"
    assert "choices" in data
    assert len(data["choices"]) > 0

    # Verify content
    choice = data["choices"][0]
    assert choice["message"]["role"] == "assistant"
    assert len(choice["message"]["content"]) > 0


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_streaming_delta_accumulation(async_client: AsyncClient):
    """
    Integration test: Verify deltas accumulate correctly without duplication.

    This test explicitly checks that when we accumulate all deltas,
    we get a coherent response without repeated text.
    Note: Requires running server, may xfail if server unavailable or LLM times out.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Count to three"}],
        "stream": True,
    }

    collected_deltas = []
    previous_full_text = ""

    try:
        async with async_client.stream(
            "POST",
            "/v1/responses",
            json=request_data,
            timeout=60.0,
        ) as response:
            if response.status_code != 200:
                pytest.xfail(f"Server returned {response.status_code}")

            async for line in response.aiter_lines():
                if not line.strip():
                    continue

                if line.startswith("data: "):
                    data_str = line[6:].strip()
                    if data_str != "[DONE]":
                        try:
                            chunk = json.loads(data_str)

                            if chunk.get("event") == "response.output_text.delta":
                                delta = chunk["delta"]
                                collected_deltas.append(delta)

                                # Verify each delta is new content
                                current_full = "".join(collected_deltas)
                                assert current_full.startswith(
                                    previous_full_text
                                ), "Deltas should accumulate progressively"
                                previous_full_text = current_full

                        except json.JSONDecodeError:
                            pass
    except Exception as e:
        pytest.xfail(f"Streaming request failed (server may be unavailable): {e}")

    full_text = "".join(collected_deltas)
    if len(full_text) == 0:
        pytest.xfail("No text received from streaming (LLM may have timed out)")


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_with_reasoning(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock with reasoning enabled.
    Note: Requires running server, may xfail if server unavailable or LLM times out.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Hello"}],
        "reasoning": {"effort": "medium", "summary": "auto"},
        "stream": True,
    }

    has_reasoning = False
    has_output = False

    try:
        async with async_client.stream(
            "POST",
            "/v1/responses",
            json=request_data,
            timeout=60.0,
        ) as response:
            if response.status_code != 200:
                pytest.xfail(f"Server returned {response.status_code}")

            async for line in response.aiter_lines():
                if not line.strip():
                    continue

                if line.startswith("data: "):
                    data_str = line[6:].strip()
                    if data_str != "[DONE]":
                        try:
                            chunk = json.loads(data_str)

                            if chunk.get("event") == "response.reasoning_summary_text.delta":
                                has_reasoning = True
                            elif chunk.get("event") == "response.output_text.delta":
                                has_output = True

                        except json.JSONDecodeError:
                            pass
    except Exception as e:
        pytest.xfail(f"Streaming request failed (server may be unavailable): {e}")

    if not has_reasoning:
        pytest.xfail("No reasoning summary received (LLM may have timed out)")
    if not has_output:
        pytest.xfail("No output text received (LLM may have timed out)")


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_markdown_formatting_preserved(async_client: AsyncClient):
    """
    Integration test: Verify markdown formatting is preserved in responses.

    Tests that code blocks, newlines, and other markdown formatting
    are properly preserved through the streaming pipeline.
    Note: Requires running server, may xfail if server unavailable or LLM times out.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Can you give me an HTML5 boilerplate template?"}],
        "stream": True,
    }

    collected_deltas = []

    try:
        async with async_client.stream(
            "POST",
            "/v1/responses",
            json=request_data,
            timeout=90.0,  # Give extra time for code generation
        ) as response:
            if response.status_code != 200:
                pytest.xfail(f"Server returned {response.status_code}")

            async for line in response.aiter_lines():
                if not line.strip():
                    continue

                if line.startswith("data: "):
                    data_str = line[6:].strip()
                    if data_str != "[DONE]":
                        try:
                            chunk = json.loads(data_str)

                            if chunk.get("event") == "response.output_text.delta":
                                collected_deltas.append(chunk["delta"])

                        except json.JSONDecodeError:
                            pass
    except Exception as e:
        pytest.xfail(f"Streaming request failed (server may be unavailable): {e}")

    # Reconstruct full response
    full_response = "".join(collected_deltas)

    # Always print the response for debugging
    print("\n" + "=" * 80)
    print("FULL RESPONSE (repr):")
    print("=" * 80)
    print(repr(full_response))
    print("\n" + "=" * 80)
    print("FULL RESPONSE (formatted):")
    print("=" * 80)
    print(full_response)
    print("=" * 80 + "\n")

    # Verify we got a response (xfail if LLM didn't produce output)
    if len(full_response) < 100:
        pytest.xfail(f"Response too short ({len(full_response)} chars), LLM may have timed out")

    # Check for code block - xfail if not present (LLM may respond differently)
    if "```" not in full_response:
        pytest.xfail("No markdown code blocks in response (LLM response varied)")

    # Verify newlines are preserved (not all collapsed to spaces)
    newline_count = full_response.count("\n")
    if newline_count < 5:
        pytest.xfail(f"Only {newline_count} newlines, formatting may have been lost")

    # Verify code block markers are complete
    code_block_starts = full_response.count("```")
    # Should have at least opening and closing markers (even count)
    assert code_block_starts % 2 == 0, "Code blocks should have matching opening/closing markers"
    assert code_block_starts >= 2, "Should have at least one complete code block"

    # Verify HTML tags are present (indicates code block content is preserved)
    has_html = "<!DOCTYPE html>" in full_response or "<html" in full_response
    if not has_html:
        pytest.xfail("No HTML5 boilerplate in response (LLM response varied)")

    # Verify indentation is preserved (check for multiple spaces in a row)
    # This indicates that code formatting with indentation is maintained
    assert "  " in full_response, "Should preserve indentation (multiple spaces)"


@pytest.mark.integration
def test_tatlock_markdown_non_streaming(client: TestClient):
    """
    Integration test: Verify markdown in non-streaming mode.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Give me a simple Python hello world code"}],
        "stream": False,
    }

    response = client.post("/v1/responses", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Get the message content
    message_item = next((item for item in data["output"] if item["type"] == "message"), None)
    assert message_item is not None

    text = message_item["content"][0]["text"]

    # Verify markdown code block
    assert "```" in text, "Should contain code block markers"
    assert "\n" in text, "Should contain newlines"