"""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()