LLeMbas/tests/test_terminal_capture.py

"""Where one command ends and the next begins, and what is kept of it.

Three layers, tested separately because they fail separately: the byte scanner,
the bounding, and a real PTY emitting real markers.

The scanner tests need no server and no event loop -- it is a byte state
machine. That is most of the argument for parsing server-side rather than in
xterm: the thing that has to be right is testable without a browser.
"""

from __future__ import annotations

import pytest

from lembas.services.agent import capture, shell_marks

asyncssh = pytest.importorskip("asyncssh")


def _osc(body: str) -> bytes:
    return f"\033]{body}\007".encode()


def _seen(chunks) -> list[tuple[str, str]]:
    marks: list[tuple[str, str]] = []
    scanner = shell_marks.Marks(lambda kind, value: marks.append((kind, value)))
    for chunk in chunks:
        scanner.feed(chunk)
    return marks


# --- The scanner -------------------------------------------------------------
def test_a_marker_is_read_out_of_a_stream():
    marks = _seen([b"hello" + _osc("133;A") + b"$ "])

    assert marks == [("A", "")]


def test_a_marker_split_across_writes_is_still_seen():
    """The test that justifies parsing server-side being safe at all.

    The pump is handed 64KB at a time on no particular boundary, so the ESC and
    the `]` land in different frames often enough to matter. Fed one byte at a
    time, which is the worst case and the cheapest way to prove it.
    """
    stream = b"before" + _osc("633;E;make -j8") + b"after"

    marks = _seen([stream[i : i + 1] for i in range(len(stream))])

    assert marks == [("E", "make -j8")]


def test_a_marker_terminated_with_st_rather_than_bel_is_seen():
    """Both terminators are legal and shells in the wild use both."""
    marks = _seen([b"\033]133;D;2\033\\"])

    assert marks == [("D", "2")]


def test_a_lone_escape_bracket_in_binary_output_does_not_swallow_the_session():
    """`cat` of a binary file produces stray ESC ] regularly. Without the bound
    one of them would buffer the rest of the session and never emit."""
    junk = b"\033]" + b"x" * (shell_marks.MAX_MARKER_BYTES + 10)

    marks = _seen([junk, _osc("133;A")])

    assert marks == [("A", "")]


def test_an_unterminated_marker_does_not_eat_the_next_one():
    marks = _seen([b"\033]133;A" + _osc("133;C")])

    assert ("C", "") in marks


def test_an_escape_that_is_not_an_osc_is_ignored():
    """A colour change is `ESC [`, not `ESC ]`, and there is a lot of it."""
    assert _seen([b"\033[31mred\033[0m"]) == []


def test_a_semicolon_in_a_command_survives_the_escaping():
    """The payload may contain no raw `;` or the fields split, so the shell
    escapes it and this puts it back."""
    marks = _seen([_osc("633;E;cd /tmp \\x3b ls")])

    assert marks == [("E", "cd /tmp \\x3b ls")]
    assert shell_marks.unescape(marks[0][1]) == "cd /tmp ; ls"


def test_an_unrelated_osc_is_not_ours():
    """Setting the window title is OSC 0 and happens constantly."""
    assert _seen([_osc("0;some title")]) == []


# --- The bounding ------------------------------------------------------------
def test_a_flood_keeps_the_head_and_the_tail_and_says_what_it_dropped():
    """Either half alone is the wrong half: a build that fails ten megabytes in
    has the invocation at the top and the error at the bottom."""
    found = capture.Capture(command="make")
    found.absorb(b"START\n")
    found.absorb(b"x" * (capture.CAPTURE_HEAD_BYTES * 2))
    found.absorb(b"END\n")

    output = found.output()

    assert output.startswith("START")
    assert output.rstrip().endswith("END")
    assert "dropped" in output


def test_short_output_is_kept_whole():
    found = capture.Capture(command="ls")
    found.absorb(b"one\ntwo\n")

    assert found.output() == "one\ntwo"


def test_a_progress_bar_collapses_to_its_last_state():
    """The highest-value transform here. Only the last state of a line was ever
    on screen, and keeping every one turns two megabytes of `pip install` into
    two megabytes of spinner in somebody's prompt."""
    found = capture.Capture(command="pip install")
    found.absorb(b"".join(f"\r{n}%".encode() for n in range(500)) + b"\ndone\n")

    output = found.output()

    assert output == "499%\ndone"


def test_escape_sequences_are_stripped():
    found = capture.Capture(command="ls --color")
    found.absorb(b"\033[31mred\033[0m\n")

    assert found.output() == "red"


def test_a_split_multibyte_character_degrades_rather_than_raising():
    """Head/tail slicing splits UTF-8 at will, so the decode has to replace."""
    found = capture.Capture(command="cat")
    found.absorb("héllo".encode()[:3])

    assert isinstance(found.output(), str)


def test_a_fence_cannot_be_ended_early_by_the_output():
    """A real injection route: output containing three backticks would close
    the block, and everything after it would read to the model as prose rather
    than as what a machine printed."""
    block = capture.fenced("here are ``` three")

    assert block.startswith("````")
    assert block.rstrip().endswith("````")


def test_the_attribution_sits_outside_the_fence():
    """So nothing the far side printed can forge it. The `$ ` line is
    synthesised here too -- what the shell echoed carries readline's editing
    escapes and is not the command."""
    found = capture.Capture(command="pytest -q", cwd="/srv/work", exit_status=1, ended=1.0)
    found.absorb(b"1 failed\n")

    text = found.as_text(label="Container")

    assert text.startswith("Ran in the terminal on Container, in /srv/work — exit 1")
    assert "```console\n$ pytest -q" in text


def test_output_still_running_is_capturable():
    """"Copy the last command" while `make` is going should give what has been
    printed so far, marked as running -- not "nothing yet"."""
    found = capture.Capture(command="make")
    found.absorb(b"compiling\n")

    assert found.running
    assert "still running" in found.as_text(label="Box")
    assert "compiling" in found.as_text(label="Box")


def test_a_very_long_line_is_cut():
    found = capture.Capture(command="cat bundle.min.js")
    found.absorb(b"z" * (capture.MAX_LINE_CHARS * 3) + b"\n")

    assert len(found.output()) <= capture.MAX_LINE_CHARS + 1


# --- The command string ------------------------------------------------------
def test_the_command_branches_on_the_shell_name():
    command = shell_marks.command_for("/srv/work")

    assert "case ${SHELL##*/} in" in command
    assert "bash --rcfile" in command
    assert "ZDOTDIR=" in command


def test_a_single_quote_in_the_directory_cannot_end_the_quoting():
    command = shell_marks.command_for("/tmp/it's here; rm -rf /")

    assert command.startswith("cd '/tmp/it'\\''s here; rm -rf /'")


def test_switching_integration_off_gives_back_exactly_what_was_there_before():
    assert shell_marks.command_for("", integrate=False) is None
    assert (
        shell_marks.command_for("/srv/work", integrate=False)
        == "cd '/srv/work' 2>/dev/null; exec ${SHELL:-/bin/sh} -l"
    )


def test_the_integration_writes_nothing_to_the_terminals_input_side():
    """Which is the whole reason this mechanism was chosen over feeding
    `source …` in as keystrokes. Nothing is echoed because nothing is typed, so
    there is no setup to hide from the scrollback and no fan-out gate."""
    command = shell_marks.command_for("/srv/work")

    assert "printf %s" in command  # written to a file
    assert "$__L/rc" in command


# --- Against a real PTY ------------------------------------------------------
# A fake shell that speaks the markers rather than a real bash: what is under
# test here is the session's bookkeeping, and a real bash would make every
# assertion depend on whoever's dotfiles the machine happens to carry. The rc
# files themselves are verified by hand, once, per shell.
import asyncio  # noqa: E402

from lembas.services.agent import terminal as terminal_service  # noqa: E402


class _Server(asyncssh.SSHServer):
    def begin_auth(self, username: str) -> bool:
        return False


@pytest.fixture
async def marking_host():
    seen: dict = {}

    def osc(body: str) -> str:
        return f"\033]{body}\007"

    async def handler(process):
        seen["command"] = process.command
        process.stdout.write(osc("633;LEMBAS;bash;1"))
        process.stdout.write(osc("633;P;Cwd=/srv/work") + osc("133;A") + "$ ")
        while True:
            try:
                line = (await process.stdin.readline()).rstrip("\n")
            except asyncssh.TerminalSizeChanged:
                continue
            except Exception:  # noqa: BLE001
                break
            if not line or line == "exit":
                break

            process.stdout.write(osc("633;E;" + line.replace(";", "\\x3b")))
            process.stdout.write(osc("133;C"))
            status = 0
            if line.startswith("fail"):
                process.stdout.write("boom\n")
                status = 3
            elif line.startswith("spin"):
                for n in range(200):
                    process.stdout.write(f"\r{n}%")
                process.stdout.write("\ndone\n")
            else:
                process.stdout.write(f"out:{line}\n")
            process.stdout.write(osc(f"133;D;{status}"))
            process.stdout.write(osc("633;P;Cwd=/srv/work") + osc("133;A") + "$ ")
        process.exit(0)

    server = await asyncssh.create_server(
        _Server,
        "127.0.0.1",
        0,
        server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
        process_factory=handler,
    )
    port = next(iter(server.sockets)).getsockname()[1]
    from lembas.services.agent import ssh as ssh_service

    line, _fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
    try:
        yield {"port": port, "host_key": line, "seen": seen}
    finally:
        await terminal_service.shutdown()
        server.close()
        await server.wait_closed()


def _spec(host) -> dict:
    return {
        "host": "127.0.0.1",
        "port": host["port"],
        "username": "tester",
        "auth": "password",
        "password": "",
        "private_key": "",
        "key_passphrase": "",
        "host_key": host["host_key"],
        "connect_timeout": 10,
    }


async def _session(host, **kwargs):
    return await terminal_service.open_session(
        "chat-1",
        owner_id="user-1",
        profile_id="profile-1",
        label="Container",
        spec=_spec(host),
        **kwargs,
    )


async def _settle(session, want, *, timeout=5.0):
    """Wait until `want(session)` holds, or give up."""
    async with asyncio.timeout(timeout):
        while not want(session):
            await asyncio.sleep(0.02)


async def test_the_shell_says_it_loaded_the_integration(marking_host):
    session = await _session(marking_host)
    session.attach()

    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)
    assert session.shell == "bash"
    await session.close()


async def test_a_command_and_its_output_are_captured_between_the_markers(marking_host):
    session = await _session(marking_host)
    session.attach()
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)

    await session.send(b"ls -la\n")
    await _settle(session, lambda s: s.last is not None)

    assert session.last.command == "ls -la"
    assert "out:ls -la" in session.last.output()
    assert session.last.cwd == "/srv/work"
    await session.close()


async def test_the_exit_code_comes_back_with_the_command(marking_host):
    session = await _session(marking_host)
    session.attach()
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)

    await session.send(b"fail now\n")
    await _settle(session, lambda s: s.last is not None)

    assert session.last.exit_status == 3
    assert not session.last.running
    await session.close()


async def test_a_semicolon_in_the_command_survives_the_round_trip(marking_host):
    session = await _session(marking_host)
    session.attach()
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)

    await session.send(b"cd /tmp ; ls\n")
    await _settle(session, lambda s: s.last is not None)

    assert session.last.command == "cd /tmp ; ls"
    await session.close()


async def test_a_progress_bar_is_collapsed_before_it_reaches_a_prompt(marking_host):
    session = await _session(marking_host)
    session.attach()
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)

    await session.send(b"spin\n")
    await _settle(session, lambda s: s.last is not None)

    output = session.last.output()
    assert output.count("%") == 1
    assert output.rstrip().endswith("done")
    await session.close()


async def test_the_markers_still_reach_the_browser(marking_host):
    """Fanned out unchanged. xterm consumes an OSC it has no handler for and
    never draws it, and rewriting frames on the hot path would break the
    "nothing decodes, so nothing splits" property the pump depends on."""
    session = await _session(marking_host)
    viewer = session.attach()
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_LIVE)

    seen = b""
    while not viewer.queue.empty():
        chunk = viewer.queue.get_nowait()
        if isinstance(chunk, bytes):
            seen += chunk
    assert b"\033]133;A" in seen
    await session.close()


async def test_a_shell_that_never_marks_ends_up_reported_as_none(marking_host, monkeypatch):
    """The fallback path's own test, and the one that matters most: without it
    the buttons would sit greyed out forever with no explanation. This is what
    catches a `.bashrc` that ends in `exec tmux`."""
    monkeypatch.setattr(terminal_service, "INTEGRATION_GRACE", 0.0)
    session = await _session(marking_host)
    session.integration = terminal_service.INTEGRATION_LOADING
    session._marks = shell_marks.Marks(lambda kind, value: None)  # deaf on purpose
    session.attach()

    await session.send(b"anything\n")
    await _settle(session, lambda s: s.integration == terminal_service.INTEGRATION_NONE)
    await session.close()