EvoScientist-Multi/tests/test_stream_display.py

"""Tests for Rich streaming display helpers."""

from typing import Any, cast

from rich.console import Console
from rich.markdown import Markdown

from EvoScientist.stream.display import (
    _fix_markdown_heading_spacing,
    create_streaming_display,
    resolve_final_status_footer,
)
from EvoScientist.stream.state import SubAgentState


def _render_text(renderable) -> str:
    console = Console(record=True, width=100, color_system=None)
    console.print(renderable)
    return console.export_text()


def test_resolve_final_status_footer_hides_footer_for_interactive_cli():
    """Interactive CLI hides the final status footer (the prompt redraws it)."""
    assert resolve_final_status_footer(True, lambda: "footer") is None


def test_resolve_final_status_footer_keeps_footer_for_noninteractive():
    """Non-interactive output keeps the footer so callers see the trailing status."""
    assert resolve_final_status_footer(False, lambda: "footer") == "footer"


def test_final_display_respects_disabled_thinking():
    """A final-frame preference must not override the user visibility flag."""
    renderable = create_streaming_display(
        thinking_text="private reasoning",
        show_thinking=False,
        is_final=True,
        final_show_thinking=True,
    )

    rendered = _render_text(renderable)

    assert "private reasoning" not in rendered
    assert "Thinking" not in rendered


def test_streaming_display_keeps_narration_visible_with_pending_memory_tool():
    """Profile-memory reads still block, while lead-in text remains visible."""
    narration = "Here is the answer."
    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "read_file",
                "args": {"path": "/memories/profile/USER_PROFILE.md"},
            }
        ],
        tool_results=[],
    )

    rendered = _render_text(renderable)

    assert "Here is the answer." in rendered
    assert "Reading memory" in rendered
    assert "Running" in rendered
    assert rendered.index("Here is the answer.") < rendered.index("Reading memory")


def test_streaming_display_keeps_narration_visible_while_processing_tool_result():
    """Completed tools still block while their result is being processed."""
    narration = "Here is the answer."
    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "read_file",
                "args": {"path": "/memories/profile/USER_PROFILE.md"},
            }
        ],
        tool_results=[
            {
                "id": "tc1",
                "name": "read_file",
                "content": "# User profile\n\n- Likes concise updates.",
            }
        ],
        is_processing=True,
    )

    rendered = _render_text(renderable)

    assert "Here is the answer." in rendered
    assert "Reading memory" in rendered
    assert "Analyzing results" in rendered
    assert rendered.index("Here is the answer.") < rendered.index("Reading memory")


def test_streaming_display_pairs_root_tool_results_by_id():
    """Out-of-order same-name root tool results should not pair by list index."""
    renderable = create_streaming_display(
        tool_calls=[
            {
                "id": "tc-a",
                "name": "execute",
                "args": {"command": "python a.py"},
            },
            {
                "id": "tc-b",
                "name": "execute",
                "args": {"command": "python b.py"},
            },
        ],
        tool_results=[
            {
                "id": "tc-b",
                "name": "execute",
                "content": "Error: result from b",
            }
        ],
    )

    rendered = _render_text(renderable)

    assert "execute(python a.py)" in rendered
    assert "execute(python b.py)" in rendered
    assert rendered.index("execute(python a.py)") < rendered.index("Running")
    assert rendered.index("execute(python b.py)") < rendered.index("result from b")


def test_streaming_display_keeps_narration_visible_with_pending_normal_tool():
    """Ordinary tools use the same pending-tool behavior as memory reads."""
    narration = "I will inspect the files first."
    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "execute",
                "args": {"command": "rg -n TODO ."},
            }
        ],
        tool_results=[],
    )

    rendered = _render_text(renderable)

    assert "execute(rg -n TODO .)" in rendered
    assert "Running" in rendered
    assert "I will inspect the files first." in rendered
    assert rendered.index("I will inspect the files first.") < rendered.index(
        "execute(rg -n TODO .)"
    )


def test_streaming_display_keeps_narration_separate_when_answer_streams():
    """Post-tool answers should not re-render pre-tool narration as Markdown."""
    narration = "I will inspect the files first.\n"
    answer = "The delayed check completed."
    renderable = create_streaming_display(
        response_text=f"{narration}{answer}",
        latest_text=answer,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "execute",
                "args": {"command": "python check.py"},
            }
        ],
        tool_results=[
            {
                "id": "tc1",
                "name": "execute",
                "content": "check complete",
            }
        ],
        response_markdown=Markdown("SHOULD NOT RENDER"),
    )

    rendered = _render_text(renderable)

    assert "I will inspect the files first." in rendered
    assert "The delayed check completed." in rendered
    assert "SHOULD NOT RENDER" not in rendered
    assert rendered.index("I will inspect the files first.") < rendered.index(
        "execute(python check.py)"
    )
    assert rendered.index("execute(python check.py)") < rendered.index(
        "The delayed check completed."
    )


def test_streaming_display_keeps_narration_separate_in_final_frame():
    """The final Rich frame should preserve narration without one concat block."""
    narration = "I will inspect the files first.\n"
    answer = "The delayed check completed."
    renderable = create_streaming_display(
        response_text=f"{narration}{answer}",
        latest_text=answer,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "execute",
                "args": {"command": "python check.py"},
            }
        ],
        tool_results=[
            {
                "id": "tc1",
                "name": "execute",
                "content": "check complete",
            }
        ],
        is_final=True,
        response_markdown=Markdown("SHOULD NOT RENDER"),
    )

    rendered = _render_text(renderable)

    assert "I will inspect the files first." in rendered
    assert "The delayed check completed." in rendered
    assert "SHOULD NOT RENDER" not in rendered
    assert rendered.index("I will inspect the files first.") < rendered.index(
        "execute(python check.py)"
    )
    assert rendered.index("execute(python check.py)") < rendered.index(
        "The delayed check completed."
    )


def test_streaming_display_preserves_narration_for_pending_final_tool():
    """Stopped/error final frames keep narration attached to pending tools."""
    narration = "I will inspect the files first.\n"
    answer = "[Stopped.]"
    renderable = create_streaming_display(
        response_text=f"{narration}{answer}",
        latest_text=answer,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "tc1",
                "name": "execute",
                "args": {"command": "sleep 30"},
            }
        ],
        tool_results=[],
        is_final=True,
    )

    rendered = _render_text(renderable)

    assert "I will inspect the files first." in rendered
    assert "execute(sleep 30)" in rendered
    assert "[Stopped.]" in rendered
    assert rendered.index("I will inspect the files first.") < rendered.index(
        "execute(sleep 30)"
    )
    assert rendered.index("execute(sleep 30)") < rendered.index("[Stopped.]")


def test_streaming_display_interleaves_multiple_narration_segments():
    """Multiple narrated segments should stay attached to their following tools."""
    first = "I will inspect the files first.\n"
    second = "I found one file, now I will run it.\n"
    answer = "The delayed check completed."
    renderable = create_streaming_display(
        response_text=f"{first}{second}{answer}",
        latest_text=answer,
        narrated_response_end=len(first) + len(second),
        narration_segments=[
            (0, first),
            (1, second),
        ],
        tool_calls=[
            {
                "id": "tc1",
                "name": "execute",
                "args": {"command": "rg -n delayed ."},
            },
            {
                "id": "tc2",
                "name": "execute",
                "args": {"command": "python check.py"},
            },
        ],
        tool_results=[
            {
                "id": "tc1",
                "name": "execute",
                "content": "check.py",
            },
            {
                "id": "tc2",
                "name": "execute",
                "content": "check complete",
            },
        ],
        is_final=True,
    )

    rendered = _render_text(renderable)

    assert rendered.index("I will inspect the files first.") < rendered.index(
        "execute(rg -n delayed .)"
    )
    assert rendered.index("execute(rg -n delayed .)") < rendered.index(
        "I found one file, now I will run it."
    )
    assert rendered.index("I found one file, now I will run it.") < rendered.index(
        "execute(python check.py)"
    )
    assert rendered.index("execute(python check.py)") < rendered.index(
        "The delayed check completed."
    )


def test_streaming_display_preserves_narration_for_collapsed_completed_tool():
    """Live collapsed completed summaries keep narration from hidden tools."""
    narration = "I will check the early result.\n"
    tool_calls = [
        {
            "id": f"tc{i}",
            "name": "execute",
            "args": {"command": f"python step_{i}.py"},
        }
        for i in range(5)
    ]
    tool_results = [
        {
            "id": f"tc{i}",
            "name": "execute",
            "content": f"step {i} complete",
        }
        for i in range(5)
    ]

    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=tool_calls,
        tool_results=tool_results,
    )

    rendered = _render_text(renderable)

    assert "I will check the early result." in rendered
    assert "1 completed" in rendered
    assert "execute(python step_0.py)" not in rendered
    assert rendered.index("I will check the early result.") < rendered.index(
        "1 completed"
    )


def test_streaming_display_preserves_narration_for_collapsed_running_tool():
    """Live collapsed running summaries keep narration from hidden tools."""
    narration = "I will start the long-running check.\n"
    tool_calls = [
        {
            "id": f"tc{i}",
            "name": "execute",
            "args": {"command": f"sleep {i + 1}"},
        }
        for i in range(4)
    ]

    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=tool_calls,
        tool_results=[],
    )

    rendered = _render_text(renderable)

    assert "I will start the long-running check." in rendered
    assert "1 more running" in rendered
    assert "execute(sleep 1)" not in rendered
    assert rendered.index("I will start the long-running check.") < rendered.index(
        "1 more running"
    )


def test_streaming_display_preserves_task_narration_while_subagent_runs():
    """Narration before a task call should stay attached to the task section."""
    narration = "I'll ask a specialist to inspect this.\n"
    subagent = SubAgentState("code-agent", "inspect this", "task:code", "task1")
    subagent.is_active = True
    subagent.add_tool_call("execute", {"command": "rg -n TODO ."}, "sa1")

    renderable = create_streaming_display(
        response_text=narration,
        narrated_response_end=len(narration),
        narration_segments=[(0, narration)],
        tool_calls=[
            {
                "id": "task1",
                "name": "task",
                "args": {
                    "subagent_type": "code-agent",
                    "description": "inspect this",
                },
            }
        ],
        tool_results=[],
        subagents=[subagent],
    )

    rendered = _render_text(renderable)

    assert "I'll ask a specialist to inspect this." in rendered
    assert "Cooking with code-agent" in rendered
    assert "execute(rg -n TODO .)" in rendered
    assert rendered.index("I'll ask a specialist to inspect this.") < rendered.index(
        "Cooking with code-agent"
    )


def test_streaming_display_orders_final_task_narration_by_tool_index():
    """Task narration should not move after regular-tool narration in final frames."""
    first = "I'll ask a specialist to inspect this.\n"
    second = "Now I will run the result locally.\n"
    answer = "The local run passed."
    subagent = SubAgentState("code-agent", "inspect this", "task:code", "task1")
    subagent.is_active = False
    subagent.add_tool_call("execute", {"command": "rg -n TODO ."}, "sa1")
    subagent.add_tool_result("execute", "todo.py", True, "sa1")

    renderable = create_streaming_display(
        response_text=f"{first}{second}{answer}",
        latest_text=answer,
        narrated_response_end=len(first) + len(second),
        narration_segments=[
            (0, first),
            (1, second),
        ],
        tool_calls=[
            {
                "id": "task1",
                "name": "task",
                "args": {
                    "subagent_type": "code-agent",
                    "description": "inspect this",
                },
            },
            {
                "id": "tc2",
                "name": "execute",
                "args": {"command": "python todo.py"},
            },
        ],
        tool_results=[
            {
                "id": "task1",
                "name": "task",
                "content": "todo.py",
            },
            {
                "id": "tc2",
                "name": "execute",
                "content": "passed",
            },
        ],
        subagents=[subagent],
        is_final=True,
    )

    rendered = _render_text(renderable)

    assert rendered.index("I'll ask a specialist to inspect this.") < rendered.index(
        "Cooking with code-agent"
    )
    assert rendered.index("Cooking with code-agent") < rendered.index(
        "Now I will run the result locally."
    )
    assert rendered.index("Now I will run the result locally.") < rendered.index(
        "execute(python todo.py)"
    )
    assert rendered.index("execute(python todo.py)") < rendered.index(
        "The local run passed."
    )


class TestFixMarkdownHeadingSpacing:
    """Pure-helper tests: heading levels, idempotence, EOS / CRLF / fenced
    code. The display-copy-only contract at call sites is covered by
    ``TestAssistantMessageBufferContract``.
    """

    def test_inserts_missing_space(self):
        """Inserts a space after `#`-marker for all 6 ATX heading levels."""
        assert _fix_markdown_heading_spacing("#Bar") == "# Bar"
        assert _fix_markdown_heading_spacing("##Bar") == "## Bar"
        assert _fix_markdown_heading_spacing("###Bar") == "### Bar"
        assert _fix_markdown_heading_spacing("####Bar") == "#### Bar"
        assert _fix_markdown_heading_spacing("#####Bar") == "##### Bar"
        assert _fix_markdown_heading_spacing("######Bar") == "###### Bar"

    def test_idempotent_on_valid_headings(self):
        """Already-spaced markers are unchanged; ``f(f(x)) == f(x)``."""
        assert _fix_markdown_heading_spacing("### Foo") == "### Foo"
        assert _fix_markdown_heading_spacing("# Bar\n## Baz") == "# Bar\n## Baz"
        # Running twice gives the same result as running once.
        once = _fix_markdown_heading_spacing("###Foo\n##Bar")
        twice = _fix_markdown_heading_spacing(once)
        assert once == twice == "### Foo\n## Bar"

    def test_multiline_mixed(self):
        """Each line of a multiline string is normalised independently."""
        src = "###A\n## B\n#C\n#### D"
        assert _fix_markdown_heading_spacing(src) == "### A\n## B\n# C\n#### D"

    def test_indented_and_blockquote_unchanged(self):
        """Lines whose `#` is not at column 0 are left alone (`^` requires col 0)."""
        # Indented lines (treated as code by CommonMark) — `^` only matches
        # column 0, so the helper leaves them alone.
        assert _fix_markdown_heading_spacing("   ###Indented") == "   ###Indented"
        # Blockquote-prefixed lines — the `>` shifts the heading away from
        # column 0; helper does not touch them. Documents accepted edge case.
        assert _fix_markdown_heading_spacing("> ###Quoted") == "> ###Quoted"
        # Empty string and whitespace-only inputs are unchanged.
        assert _fix_markdown_heading_spacing("") == ""
        assert _fix_markdown_heading_spacing("\n\n") == "\n\n"

    def test_bare_hash_at_end_of_string_unchanged(self):
        """A trailing `#` (e.g. mid-stream chunk) must not gain a spurious
        space. Positive lookahead requires a real non-excluded char to
        follow, so EOS naturally fails the match.
        """
        assert _fix_markdown_heading_spacing("#") == "#"
        assert _fix_markdown_heading_spacing("##") == "##"
        assert _fix_markdown_heading_spacing("######") == "######"
        # Trailing hash on a non-trailing line is also untouched (the line
        # has no follow-up content yet).
        assert _fix_markdown_heading_spacing("Foo\n###") == "Foo\n###"

    def test_crlf_line_endings(self):
        """CRLF (`\\r\\n`) line endings: `\\r` is in the exclusion set so
        empty CRLF heading lines are unchanged, and a real CRLF heading
        gets a space inserted in front of the carriage-return-free part.
        """
        # Empty CRLF heading — must not become `# \r\n`.
        assert _fix_markdown_heading_spacing("#\r\n") == "#\r\n"
        assert _fix_markdown_heading_spacing("###\r\n") == "###\r\n"
        # Multi-line mixed CRLF — both lines fixed.
        assert _fix_markdown_heading_spacing("###A\r\n##B") == "### A\r\n## B"
        # Trailing CRLF after content — fix applies, line ending preserved.
        assert _fix_markdown_heading_spacing("###Foo\r\n") == "### Foo\r\n"

    def test_fenced_code_block_known_limitation(self):
        """The regex is context-free, so `###define` at column 0 inside a
        backtick fence WILL get a space in the display copy. This test
        documents (and locks in) the accepted trade-off — flip these
        assertions if a future fix gates on fence parsing.
        """
        src = "```c\n###define X 1\n```"
        # Currently DOES alter the line inside the fence.
        assert _fix_markdown_heading_spacing(src) == "```c\n### define X 1\n```"


class TestAssistantMessageBufferContract:
    """Regression guard: ``AssistantMessage`` flush/mount must apply the
    heading fix to a display copy and leave ``self._content`` untouched.
    """

    def _make_widget(self, initial: str = ""):
        from unittest.mock import MagicMock

        from EvoScientist.cli.widgets.assistant_message import AssistantMessage

        msg = AssistantMessage(initial_content=initial)
        fake_md = MagicMock()
        cast(Any, msg).query_one = MagicMock(return_value=fake_md)
        return msg, fake_md

    def test_flush_markdown_does_not_mutate_buffer(self):
        msg, fake_md = self._make_widget()
        msg._content = "###Foo\n##Bar"

        msg._flush_markdown()

        # Raw streaming buffer is preserved verbatim.
        assert msg._content == "###Foo\n##Bar"
        # The Textual Markdown widget receives the fixed display copy.
        fake_md.update.assert_called_once_with("### Foo\n## Bar")
        # Flush latch is cleared.
        assert msg._flush_pending is False

    def test_on_mount_does_not_mutate_initial_content(self):
        msg, fake_md = self._make_widget(initial="###Hello")

        msg.on_mount()

        assert msg._content == "###Hello"
        fake_md.update.assert_called_once_with("### Hello")

    def test_on_mount_no_op_when_initial_empty(self):
        msg, fake_md = self._make_widget(initial="")

        msg.on_mount()

        assert msg._content == ""
        fake_md.update.assert_not_called()