dinos
8bb1d6c0e3
* refactor(stream): langgraph streaming v3 * fix: address CR comments * chore(stream): add success field to state --------- Co-authored-by: Xi Zhang <106144707+X-iZhang@users.noreply.github.com>
602 lines
21 KiB
Python
602 lines
21 KiB
Python
"""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_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()
|