diff --git a/src/plyngent/cli/display.py b/src/plyngent/cli/display.py index 65eabe4..b89a83f 100644 --- a/src/plyngent/cli/display.py +++ b/src/plyngent/cli/display.py @@ -5,6 +5,7 @@ import json import os import re import shlex +import shutil import sys from contextvars import ContextVar from dataclasses import dataclass @@ -495,25 +496,58 @@ def _echo_stream(text: str) -> None: _ = sys.stdout.flush() -def _clear_streamed_lines(line_count: int) -> None: - """Move cursor up and clear the streamed plain-text region (TTY only).""" +def _clear_streamed_lines(line_count: int, *, resume_above: bool = False) -> None: + """Clear exactly *line_count* terminal lines ending at the cursor line. + + The cursor ends at column 0 of the topmost cleared line, so a re-render + overwrites the erased block in place instead of appending below it. With + ``resume_above`` the cursor steps one line further up: the next + ``"\\n"``-prefixed write then reuses the last cleared line — needed by the + parallel-batch fallback, which must not touch the line above the prefix. + + ``line_count`` counts *physical* rows; callers fold wrapped lines via + :func:`_line_count_for_clear`. + """ if line_count <= 0: return - # Clear current line, then each previous line of the streamed block. - for _ in range(line_count): + # Clear the cursor line, then each line above it while moving up. + for _ in range(line_count - 1): _ = sys.stdout.write("\r\033[2K\033[1A") _ = sys.stdout.write("\r\033[2K") + if resume_above: + _ = sys.stdout.write("\033[1A") with contextlib.suppress(OSError): _ = sys.stdout.flush() +def _terminal_columns() -> int: + """Terminal width for wrap-aware row counts (fallback: 80).""" + try: + columns = shutil.get_terminal_size().columns + except OSError, ValueError: + return 80 + return columns if columns > 0 else 80 + + +def _physical_line_count(text: str) -> int: + """Rows streamed *text* occupies on the terminal, counting wrapping. + + ``_echo_stream`` writes without a trailing newline, so the cursor sits at + the end of the last row and owns one row even when the text ends in ``"\\n"``. + A row filled to the exact width is counted once: the wrap of the final + column is pending, so erasing must not claim the row below it. + """ + columns = _terminal_columns() + return sum(max(1, -(-len(line) // columns)) for line in text.split("\n")) + + def _line_count_for_clear(label: str, body: str) -> int: - """Approximate terminal lines used by ``label\\n + body`` for cursor erase.""" + """Physical terminal rows used by ``label\\n + body`` for cursor erase.""" if not body and not label: return 0 - # Label is on its own line; body may contain newlines. + # Label is on its own line; body may contain newlines and wrap. text = f"{label}\n{body}" if label else body - return text.count("\n") + 1 + return _physical_line_count(text) class _PrettyToolStream: @@ -542,7 +576,9 @@ class _PrettyToolStream: return if self._open: # Parallel batch: the open prefix could never receive its detail. - _clear_streamed_lines(1) + # Erase only the prefix line — the line above may be streamed + # reasoning — and let the fallback's leading newline reuse it. + _clear_streamed_lines(1, resume_above=True) self._open = False self._whole_lines = True return @@ -591,7 +627,8 @@ def _flush_assistant_markdown(body: str, *, pretty: bool) -> None: click.echo() return if pretty: - lines = _line_count_for_clear("assistant:", body) + # +1: the blank separator ``begin_assistant`` printed above the label. + lines = _line_count_for_clear("assistant:", body) + 1 _clear_streamed_lines(lines) print_markdown(body, label="assistant:") click.echo() diff --git a/tests/test_cli/test_display.py b/tests/test_cli/test_display.py index 7803c4f..e97f3bf 100644 --- a/tests/test_cli/test_display.py +++ b/tests/test_cli/test_display.py @@ -4,6 +4,8 @@ from typing import TYPE_CHECKING from plyngent.agent import ErrorEvent, ReasoningDeltaEvent, TextDeltaEvent, ToolCallEvent, ToolResultEvent from plyngent.cli.display import ( + _clear_streamed_lines, + _line_count_for_clear, get_markdown_enabled, markdown_render_available, print_markdown, @@ -637,6 +639,93 @@ async def test_pretty_parallel_batch_erases_open_prefix( assert "* Read 'b.txt' L1-9 (done)\n" in out +async def test_pretty_parallel_batch_keeps_the_line_above_the_prefix( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """Regression: erasing the open prefix must not eat the line above it. + + Reasoning streams without a trailing newline, so the prefix line sits right + below the reasoning text; the erase used to clear one row too many and wiped + that reasoning line out of the screen. + """ + monkeypatch.setattr("plyngent.cli.display.interactive_terminal", lambda: True) + await render_events( + _aiter( + [ + ReasoningDeltaEvent(content="think"), + _pretty_call("read_file", '{"path": "a.txt"}'), + _pretty_call("read_file", '{"path": "b.txt"}'), + _result("L1-4\none\n"), + _result("L1-9\ntwo\n"), + ] + ) + ) + out = capsys.readouterr().out + assert out == ( + "\nreasoning:\nthink" + "\n* Read 'a.txt' " + "\r\x1b[2K\x1b[1A" + "\n* Read 'a.txt' L1-4 (done)\n" + "\n* Read 'b.txt' L1-9 (done)\n" + "\n\n" + ) + + +def test_clear_streamed_lines_clears_exactly_the_given_rows(capsys: pytest.CaptureFixture[str]) -> None: + """Two rows means two clears; the cursor lands on the topmost cleared row.""" + _clear_streamed_lines(2) + assert capsys.readouterr().out == "\r\x1b[2K\x1b[1A\r\x1b[2K" + + +def test_clear_streamed_lines_resume_above_steps_back_one_row(capsys: pytest.CaptureFixture[str]) -> None: + """``resume_above`` clears the cursor row only, then steps above it.""" + _clear_streamed_lines(1, resume_above=True) + assert capsys.readouterr().out == "\r\x1b[2K\x1b[1A" + + +def test_clear_streamed_lines_ignores_non_positive_counts(capsys: pytest.CaptureFixture[str]) -> None: + _clear_streamed_lines(0) + _clear_streamed_lines(-1) + assert capsys.readouterr().out == "" + + +def test_line_count_for_clear_counts_wrapped_rows(monkeypatch: pytest.MonkeyPatch) -> None: + """A streamed line wider than the terminal owns more than one row.""" + monkeypatch.setenv("COLUMNS", "20") + assert _line_count_for_clear("assistant:", "x" * 20) == 2 # label + full row + assert _line_count_for_clear("assistant:", "x" * 21) == 3 # label + two rows + assert _line_count_for_clear("assistant:", "a\nb") == 3 + assert _line_count_for_clear("", "") == 0 + + +async def test_flush_markdown_erases_wrapped_assistant_rows( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """A wrapped assistant segment is fully erased before the markdown re-render.""" + monkeypatch.setattr("plyngent.cli.display.interactive_terminal", lambda: True) + monkeypatch.setattr("plyngent.cli.display.markdown_render_available", lambda: True) + monkeypatch.setenv("COLUMNS", "20") + set_markdown_enabled(True) + try: + await render_events( + _aiter( + [ + TextDeltaEvent(content="x" * 21), + _pretty_call("read_file", '{"path": "a.txt"}'), + _result("L1-4\none\n"), + ] + ) + ) + finally: + set_markdown_enabled(True) + out = capsys.readouterr().out + # label + two wrapped body rows + the blank separator above the label. + assert out.count("\x1b[2K") == 4 + assert out.count("\x1b[1A") == 3 + + async def test_pretty_prefix_closed_before_error_line( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str],