A terminal panel beside an agent chat
A real shell on the chat's own connection, opened and closed like the inspector and never beside it. The modes govern the model; what a person types is theirs, since they hold the credential and could open the same shell with an ssh client. The model cannot see the panel -- a button copies the output you choose into the composer. The session outlives the socket: closing the panel leaves a build running, and coming back reattaches with the scrollback. Two tabs share one shell and the smaller window decides the size. It ends on an idle timeout, on deleting the chat, on disabling, moving or deleting the connection, and on a restart -- which says why rather than quietly opening a fresh shell that has lost the working directory. The nginx template's `Connection ""` is right for SSE and fails every WebSocket handshake, so `location /` now uses a `map $http_upgrade`; update.sh grows a drift check for it, because the only symptom on a stale vhost is a panel that cannot connect. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
@@ -80,6 +80,28 @@ def fresh_generation_registry() -> Iterator[None]:
|
||||
generation_service._TASKS.clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fresh_terminal_registry() -> Iterator[None]:
|
||||
"""Empty the open-shell registry between tests, for the same reason.
|
||||
|
||||
A leaked entry holds an asyncssh connection belonging to an event loop that
|
||||
has since closed, and the reaper task is module-level too -- one left
|
||||
running would wake up inside the next test's loop.
|
||||
"""
|
||||
from lembas.services.agent import terminal as terminal_service
|
||||
|
||||
def _clear() -> None:
|
||||
reaper = terminal_service._REAPER
|
||||
if reaper is not None:
|
||||
reaper.cancel()
|
||||
terminal_service._REAPER = None
|
||||
terminal_service._SESSIONS.clear()
|
||||
|
||||
_clear()
|
||||
yield
|
||||
_clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db() -> Iterator[Session]:
|
||||
session = get_session_factory()()
|
||||
|
||||
@@ -117,7 +117,7 @@ def test_the_dangerous_defaults_are_all_passed_explicitly():
|
||||
"""Every LLeMbas user shares one unix account, so "whatever the account has
|
||||
lying around" is never the right answer. This is the most important test in
|
||||
the file and it needs no server at all."""
|
||||
kwargs = ssh._connect_kwargs(_spec(22, "127.0.0.1 ssh-ed25519 AAAA\n"))
|
||||
kwargs = ssh.connect_kwargs(_spec(22, "127.0.0.1 ssh-ed25519 AAAA\n"))
|
||||
|
||||
# Bytes, never None: None turns host key checking off entirely.
|
||||
assert isinstance(kwargs["known_hosts"], bytes)
|
||||
@@ -132,11 +132,11 @@ def test_the_dangerous_defaults_are_all_passed_explicitly():
|
||||
|
||||
def test_a_profile_with_no_confirmed_host_key_refuses_to_connect():
|
||||
with pytest.raises(ExecError, match="host key has not been confirmed"):
|
||||
ssh._connect_kwargs(_spec(22, ""))
|
||||
ssh.connect_kwargs(_spec(22, ""))
|
||||
|
||||
|
||||
def test_a_password_profile_sends_a_password_and_no_keys():
|
||||
kwargs = ssh._connect_kwargs(
|
||||
kwargs = ssh.connect_kwargs(
|
||||
_spec(22, "h ssh-ed25519 AAAA\n", auth="password", password="hunter2")
|
||||
)
|
||||
assert kwargs["password"] == "hunter2"
|
||||
@@ -144,7 +144,7 @@ def test_a_password_profile_sends_a_password_and_no_keys():
|
||||
|
||||
|
||||
def test_a_key_profile_sends_no_password():
|
||||
kwargs = ssh._connect_kwargs(_spec(22, "h ssh-ed25519 AAAA\n", password="stale"))
|
||||
kwargs = ssh.connect_kwargs(_spec(22, "h ssh-ed25519 AAAA\n", password="stale"))
|
||||
assert kwargs["password"] is None
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,437 @@
|
||||
"""The shell behind the terminal panel.
|
||||
|
||||
Everything here runs against a real SSH server with a real PTY, because the
|
||||
whole point of this module is what a pseudo-terminal does and a stub would agree
|
||||
with any design at all.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from lembas.services.agent import terminal as terminal_service
|
||||
from lembas.services.agent.base import ExecError
|
||||
from lembas.services.agent.terminal import Session
|
||||
|
||||
asyncssh = pytest.importorskip("asyncssh")
|
||||
|
||||
|
||||
# --- A machine with a shell on it ----------------------------------------------
|
||||
class _Server(asyncssh.SSHServer):
|
||||
def begin_auth(self, username: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def shell_host():
|
||||
"""A server whose sessions behave enough like a shell to be tested against.
|
||||
|
||||
It announces itself, echoes what is typed at it, and remembers what the PTY
|
||||
was asked for -- which is the part the client side has to get right.
|
||||
"""
|
||||
seen: dict = {}
|
||||
|
||||
async def handler(process):
|
||||
seen["command"] = process.command
|
||||
seen["term_type"] = process.get_terminal_type()
|
||||
seen["term_size"] = process.get_terminal_size()
|
||||
seen["process"] = process
|
||||
process.stdout.write("READY\n")
|
||||
while True:
|
||||
try:
|
||||
line = await process.stdin.readline()
|
||||
except asyncssh.TerminalSizeChanged:
|
||||
# A window change interrupts the read rather than arriving as
|
||||
# data. A real shell redraws and carries on.
|
||||
continue
|
||||
except Exception: # noqa: BLE001 - the session ending is not a failure
|
||||
break
|
||||
if not line or line.rstrip("\n") == "exit":
|
||||
break
|
||||
process.stdout.write(f"echo:{line.rstrip()}\n")
|
||||
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:
|
||||
# Before the server, always: a client connection still open makes
|
||||
# `wait_closed` wait for it, and a test that failed mid-way has one.
|
||||
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 _open(host, *, chat_id="chat-1", owner="user-1", project_dir="", **kwargs) -> Session:
|
||||
return await terminal_service.open_session(
|
||||
chat_id,
|
||||
owner_id=owner,
|
||||
profile_id="profile-1",
|
||||
label="Test box",
|
||||
spec=_spec(host),
|
||||
project_dir=project_dir,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def _read(viewer, *, timeout: float = 5.0) -> bytes:
|
||||
return await asyncio.wait_for(viewer.queue.get(), timeout=timeout)
|
||||
|
||||
|
||||
async def _read_until(viewer, needle: bytes, *, timeout: float = 5.0) -> bytes:
|
||||
"""Frames are whatever the far side wrote, so a line can arrive in pieces.
|
||||
|
||||
Starts from the snapshot, because a viewer that attached after the shell
|
||||
had already said something finds it there rather than on the queue -- which
|
||||
is the whole design, and would otherwise make this hang.
|
||||
"""
|
||||
seen = viewer.snapshot
|
||||
while needle not in seen:
|
||||
chunk = await _read(viewer, timeout=timeout)
|
||||
assert chunk is not None, f"the session closed before {needle!r} arrived"
|
||||
seen += chunk
|
||||
return seen
|
||||
|
||||
|
||||
# --- The PTY --------------------------------------------------------------------
|
||||
async def test_the_pty_is_asked_for_with_a_term_type_and_a_size(shell_host):
|
||||
session = await _open(shell_host, cols=100, rows=30)
|
||||
viewer = session.attach(cols=100, rows=30)
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
assert shell_host["seen"]["term_type"] == "xterm-256color"
|
||||
assert shell_host["seen"]["term_size"][:2] == (100, 30)
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_a_project_directory_becomes_a_cd_before_the_shell(shell_host):
|
||||
session = await _open(shell_host, project_dir="/srv/work")
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
command = shell_host["seen"]["command"]
|
||||
assert command is not None
|
||||
assert "cd '/srv/work'" in command
|
||||
assert "exec" in command
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_no_project_directory_means_the_plain_login_shell(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
assert shell_host["seen"]["command"] is None
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_a_single_quote_in_the_directory_cannot_end_the_quoting(shell_host):
|
||||
session = await _open(shell_host, project_dir="/tmp/it's here; rm -rf /")
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
command = shell_host["seen"]["command"]
|
||||
assert command.startswith("cd '/tmp/it'\\''s here; rm -rf /'")
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_what_is_typed_reaches_the_shell_and_comes_back(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
await session.send(b"hello\n")
|
||||
assert b"echo:hello" in await _read_until(viewer, b"echo:hello")
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_a_resize_reaches_the_far_side(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach(cols=80, rows=24)
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
session.resize(viewer, 120, 40)
|
||||
process = shell_host["seen"]["process"]
|
||||
for _ in range(50):
|
||||
if process.get_terminal_size()[:2] == (120, 40):
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
assert process.get_terminal_size()[:2] == (120, 40)
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_a_silly_size_is_clamped_rather_than_forwarded(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach(cols=100_000, rows=100_000)
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
assert viewer.cols == terminal_service.MAX_COLS
|
||||
assert viewer.rows == terminal_service.MAX_ROWS
|
||||
await session.close()
|
||||
|
||||
|
||||
# --- Two tabs, one shell ---------------------------------------------------------
|
||||
async def test_two_attachments_share_one_shell(shell_host):
|
||||
session = await _open(shell_host)
|
||||
first = session.attach()
|
||||
await _read_until(first, b"READY")
|
||||
second = session.attach()
|
||||
|
||||
await session.send(b"both\n")
|
||||
assert b"echo:both" in await _read_until(first, b"echo:both")
|
||||
assert b"echo:both" in await _read_until(second, b"echo:both")
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_opening_twice_returns_the_same_session(shell_host):
|
||||
first = await _open(shell_host)
|
||||
second = await _open(shell_host)
|
||||
assert first is second
|
||||
await first.close()
|
||||
|
||||
|
||||
async def test_the_smaller_viewer_decides_the_size(shell_host):
|
||||
session = await _open(shell_host)
|
||||
wide = session.attach(cols=200, rows=60)
|
||||
await _read_until(wide, b"READY")
|
||||
session.attach(cols=90, rows=25)
|
||||
|
||||
assert (session.cols, session.rows) == (90, 25)
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_the_size_lifts_again_when_the_smaller_window_goes(shell_host):
|
||||
session = await _open(shell_host)
|
||||
wide = session.attach(cols=200, rows=60)
|
||||
await _read_until(wide, b"READY")
|
||||
narrow = session.attach(cols=90, rows=25)
|
||||
|
||||
session.detach(narrow)
|
||||
assert (session.cols, session.rows) == (200, 60)
|
||||
await session.close()
|
||||
|
||||
|
||||
# --- Lifetime --------------------------------------------------------------------
|
||||
async def test_detaching_leaves_the_session_running(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
session.detach(viewer)
|
||||
|
||||
assert terminal_service.get("chat-1") is session
|
||||
assert not session.closed
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_reattaching_replays_the_scrollback(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
await session.send(b"before\n")
|
||||
await _read_until(viewer, b"echo:before")
|
||||
session.detach(viewer)
|
||||
|
||||
returning = session.attach()
|
||||
assert b"READY" in returning.snapshot
|
||||
assert b"echo:before" in returning.snapshot
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_scrollback_is_bounded(shell_host, monkeypatch):
|
||||
monkeypatch.setattr(terminal_service, "SCROLLBACK_BYTES", 1024)
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
for index in range(200):
|
||||
session._remember(f"line {index} ".encode() + b"y" * 100)
|
||||
assert session._scrollback_bytes <= 1024 + 110
|
||||
assert len(session.attach().snapshot) <= 1024 + 110
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_a_slow_reader_is_dropped_rather_than_buffered(monkeypatch):
|
||||
"""No server needed: fanning out is a plain function over the viewers.
|
||||
|
||||
The alternative design -- a pump that waits for a full queue -- would stall
|
||||
every other viewer behind the one that stopped reading, and buffer without
|
||||
bound while it did.
|
||||
"""
|
||||
monkeypatch.setattr(terminal_service, "VIEWER_QUEUE", 4)
|
||||
session = Session("chat-1", owner_id="user-1", profile_id="p", label="box")
|
||||
keeping_up = session.attach()
|
||||
stalled = session.attach()
|
||||
|
||||
for index in range(20):
|
||||
session._fan_out(f"chunk {index}\n".encode())
|
||||
while not keeping_up.queue.empty():
|
||||
keeping_up.queue.get_nowait()
|
||||
|
||||
assert stalled.dropped
|
||||
assert stalled.id not in session.viewers
|
||||
assert keeping_up.id in session.viewers
|
||||
assert not session.closed
|
||||
# Woken rather than left waiting: the sentinel is how the socket learns to
|
||||
# reconnect, and the scrollback is what makes that free.
|
||||
assert stalled.queue.get_nowait() is None
|
||||
|
||||
|
||||
async def test_the_shell_exiting_closes_the_session(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
await session.send(b"exit\n")
|
||||
for _ in range(200):
|
||||
if session.closed:
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
assert session.closed
|
||||
assert session.closed_reason == terminal_service.CLOSED_EXITED
|
||||
# The viewer is woken with the sentinel rather than left waiting.
|
||||
frame = await _read(viewer)
|
||||
while frame is not None:
|
||||
frame = await _read(viewer)
|
||||
|
||||
|
||||
async def test_a_closed_session_lingers_with_its_reason(shell_host):
|
||||
session = await _open(shell_host)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
await session.close(terminal_service.CLOSED_IDLE)
|
||||
|
||||
assert terminal_service.get("chat-1") is None
|
||||
lingering = terminal_service.peek("chat-1")
|
||||
assert lingering is not None
|
||||
assert lingering.closed_reason == terminal_service.CLOSED_IDLE
|
||||
|
||||
|
||||
async def test_an_idle_session_is_reaped(shell_host, monkeypatch):
|
||||
monkeypatch.setattr(terminal_service, "REAP_INTERVAL", 0.05)
|
||||
session = await _open(shell_host, idle_timeout=0.1)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
session.detach(viewer)
|
||||
|
||||
for _ in range(200):
|
||||
if session.closed:
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
assert session.closed
|
||||
assert session.closed_reason == terminal_service.CLOSED_IDLE
|
||||
|
||||
|
||||
async def test_a_watched_session_is_never_idle(shell_host, monkeypatch):
|
||||
monkeypatch.setattr(terminal_service, "REAP_INTERVAL", 0.05)
|
||||
session = await _open(shell_host, idle_timeout=0.1)
|
||||
viewer = session.attach()
|
||||
await _read_until(viewer, b"READY")
|
||||
|
||||
await asyncio.sleep(0.4)
|
||||
assert not session.closed
|
||||
assert session.idle_for == 0.0
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_closing_a_profile_ends_its_shells(shell_host):
|
||||
first = await _open(shell_host, chat_id="chat-1")
|
||||
second = await _open(shell_host, chat_id="chat-2")
|
||||
for session in (first, second):
|
||||
await _read_until(session.attach(), b"READY")
|
||||
|
||||
assert await terminal_service.close_for_profile("profile-1") == 2
|
||||
assert first.closed and second.closed
|
||||
assert first.closed_reason == terminal_service.CLOSED_REVOKED
|
||||
|
||||
|
||||
async def test_closing_a_chat_ends_only_its_shell(shell_host):
|
||||
first = await _open(shell_host, chat_id="chat-1")
|
||||
second = await _open(shell_host, chat_id="chat-2")
|
||||
|
||||
assert await terminal_service.close_chat("chat-1")
|
||||
assert first.closed
|
||||
assert not second.closed
|
||||
await second.close()
|
||||
|
||||
|
||||
async def test_shutdown_closes_every_shell(shell_host):
|
||||
first = await _open(shell_host, chat_id="chat-1")
|
||||
second = await _open(shell_host, chat_id="chat-2")
|
||||
|
||||
await terminal_service.shutdown()
|
||||
assert first.closed and second.closed
|
||||
assert first.closed_reason == terminal_service.CLOSED_SHUTDOWN
|
||||
assert terminal_service.count() == 0
|
||||
|
||||
|
||||
# --- Caps -------------------------------------------------------------------------
|
||||
async def test_the_instance_cap_refuses_a_new_shell(shell_host):
|
||||
await _open(shell_host, chat_id="chat-1", max_sessions=1)
|
||||
with pytest.raises(ExecError, match="as many terminals open as it allows"):
|
||||
await _open(shell_host, chat_id="chat-2", max_sessions=1)
|
||||
|
||||
|
||||
async def test_the_per_person_cap_counts_only_that_person(shell_host):
|
||||
await _open(shell_host, chat_id="chat-1", owner="user-1", max_per_user=1)
|
||||
with pytest.raises(ExecError, match="already have 1 terminal"):
|
||||
await _open(shell_host, chat_id="chat-2", owner="user-1", max_per_user=1)
|
||||
|
||||
# Somebody else is unaffected.
|
||||
other = await _open(shell_host, chat_id="chat-3", owner="user-2", max_per_user=1)
|
||||
assert other is terminal_service.get("chat-3")
|
||||
await terminal_service.shutdown()
|
||||
|
||||
|
||||
async def test_a_closed_session_does_not_count_against_the_cap(shell_host):
|
||||
session = await _open(shell_host, chat_id="chat-1", max_per_user=1)
|
||||
await terminal_service.close_chat("chat-1")
|
||||
assert session.closed
|
||||
|
||||
replacement = await _open(shell_host, chat_id="chat-2", max_per_user=1)
|
||||
assert not replacement.closed
|
||||
await replacement.close()
|
||||
|
||||
|
||||
# --- Refusals ---------------------------------------------------------------------
|
||||
async def test_an_unconfirmed_host_key_is_refused_before_anything_is_sent(shell_host):
|
||||
spec = _spec(shell_host)
|
||||
spec["host_key"] = ""
|
||||
session = Session("chat-1", owner_id="user-1", profile_id="p", label="Test box")
|
||||
|
||||
with pytest.raises(ExecError, match="host key has not been confirmed"):
|
||||
await session.start(spec)
|
||||
|
||||
|
||||
async def test_a_different_host_key_is_refused(shell_host):
|
||||
spec = _spec(shell_host)
|
||||
spec["host_key"] = f"[127.0.0.1]:{shell_host['port']} ssh-ed25519 {'A' * 68}\n"
|
||||
session = Session("chat-1", owner_id="user-1", profile_id="p", label="Test box")
|
||||
|
||||
with pytest.raises(ExecError):
|
||||
await session.start(spec)
|
||||
@@ -0,0 +1,520 @@
|
||||
"""The WebSocket behind the terminal panel: who gets one, and what it carries.
|
||||
|
||||
The refusals are most of this file on purpose. It is the one endpoint where
|
||||
getting past the door means a shell on somebody's machine, so every "no" is
|
||||
pinned -- including *where* it happens, since a browser can read a frame and
|
||||
cannot read a rejected handshake.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from lembas.db.models import KIND_AGENT, KIND_CHAT, Chat, Connection, Model, SshProfile, User
|
||||
from lembas.services import settings_store
|
||||
from lembas.services.agent import policy
|
||||
from lembas.services.agent import terminal as terminal_service
|
||||
|
||||
asyncssh = pytest.importorskip("asyncssh")
|
||||
|
||||
|
||||
class _Server(asyncssh.SSHServer):
|
||||
def begin_auth(self, username: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shell_host():
|
||||
"""A real SSH server with a shell-ish session, on its own thread and loop.
|
||||
|
||||
Its own loop matters: `websocket_connect` is synchronous and blocks the
|
||||
test's loop, so a server sharing it could never accept the connection the
|
||||
endpoint is trying to make.
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
||||
thread.start()
|
||||
seen: dict = {}
|
||||
|
||||
async def handler(process):
|
||||
seen["command"] = process.command
|
||||
seen["term_size"] = process.get_terminal_size()
|
||||
process.stdout.write("READY\n")
|
||||
while True:
|
||||
try:
|
||||
line = await process.stdin.readline()
|
||||
except asyncssh.TerminalSizeChanged:
|
||||
# A window change interrupts the read rather than arriving as
|
||||
# data. A real shell redraws and carries on, and so must this
|
||||
# one, or every resize would look like the shell exiting.
|
||||
seen["term_size"] = process.get_terminal_size()
|
||||
continue
|
||||
except Exception: # noqa: BLE001 - the session ending is not a failure
|
||||
break
|
||||
if not line:
|
||||
break
|
||||
text = line.rstrip("\n")
|
||||
if text == "exit":
|
||||
break
|
||||
seen.setdefault("typed", []).append(text)
|
||||
process.stdout.write(f"echo:{text}\n")
|
||||
process.exit(0)
|
||||
|
||||
async def start():
|
||||
from lembas.services.agent import ssh as ssh_service
|
||||
|
||||
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]
|
||||
line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
|
||||
return server, port, line, fingerprint
|
||||
|
||||
server, port, host_key, fingerprint = asyncio.run_coroutine_threadsafe(start(), loop).result(10)
|
||||
try:
|
||||
yield {"port": port, "host_key": host_key, "fingerprint": fingerprint, "seen": seen}
|
||||
finally:
|
||||
|
||||
async def stop():
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
asyncio.run_coroutine_threadsafe(stop(), loop).result(10)
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def close_terminals(client, shell_host):
|
||||
"""End every shell before the server goes.
|
||||
|
||||
A session outlives the socket on purpose, so a test that opens one leaves it
|
||||
open -- and `wait_closed` then waits on a client connection that nothing is
|
||||
going to close. It runs on the TestClient's portal because that is the loop
|
||||
the asyncssh connection belongs to.
|
||||
"""
|
||||
yield
|
||||
client.portal.call(terminal_service.shutdown)
|
||||
|
||||
|
||||
def _agent_chat(db, user_id, shell_host, *, kind=KIND_AGENT, project_dir="") -> Chat:
|
||||
"""A chat pointed at the test server, with everything switched on."""
|
||||
settings_store.update(
|
||||
db, {"enabled": True, "terminal_enabled": True}, key=settings_store.AGENTS
|
||||
)
|
||||
|
||||
profile = SshProfile(
|
||||
owner_id=user_id,
|
||||
name="Test box",
|
||||
host="127.0.0.1",
|
||||
port=shell_host["port"],
|
||||
username="tester",
|
||||
host_key=shell_host["host_key"],
|
||||
host_fingerprint=shell_host["fingerprint"],
|
||||
)
|
||||
db.add(profile)
|
||||
db.commit()
|
||||
|
||||
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
|
||||
db.add(connection)
|
||||
db.commit()
|
||||
db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
|
||||
db.commit()
|
||||
|
||||
chat = Chat(
|
||||
user_id=user_id,
|
||||
model_id="m",
|
||||
connection_id=connection.id,
|
||||
kind=kind,
|
||||
ssh_profile_id=profile.id,
|
||||
project_dir=project_dir,
|
||||
agent_mode=policy.MODE_MANUAL,
|
||||
)
|
||||
db.add(chat)
|
||||
db.commit()
|
||||
return chat
|
||||
|
||||
|
||||
def _url(chat: Chat, **params) -> str:
|
||||
query = "".join(f"&{k}={v}" for k, v in params.items())
|
||||
return f"/api/chats/{chat.id}/terminal/ws?{query.lstrip('&')}"
|
||||
|
||||
|
||||
def _headers(client: TestClient) -> dict:
|
||||
"""What a browser on this origin sends. The endpoint requires both."""
|
||||
return {"origin": str(client.base_url).rstrip("/"), "host": client.base_url.host}
|
||||
|
||||
|
||||
def _first_json(socket) -> dict:
|
||||
return json.loads(socket.receive_text())
|
||||
|
||||
|
||||
def _read_until(socket, needle: bytes, *, frames: int = 40) -> bytes:
|
||||
seen = b""
|
||||
for _ in range(frames):
|
||||
message = socket.receive()
|
||||
if message.get("bytes") is not None:
|
||||
seen += message["bytes"]
|
||||
if needle in seen:
|
||||
return seen
|
||||
elif message.get("text"):
|
||||
payload = json.loads(message["text"])
|
||||
raise AssertionError(f"the socket said {payload} before {needle!r} arrived")
|
||||
raise AssertionError(f"{needle!r} never arrived; saw {seen!r}")
|
||||
|
||||
|
||||
# --- Getting in -------------------------------------------------------------------
|
||||
def test_a_shell_opens_and_says_hello(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
hello = _first_json(socket)
|
||||
assert hello["t"] == "ready"
|
||||
assert hello["label"] == "Test box"
|
||||
assert b"READY" in _read_until(socket, b"READY")
|
||||
|
||||
|
||||
def test_what_is_typed_reaches_the_shell(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
assert _first_json(socket)["t"] == "ready"
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_bytes(b"hello\n")
|
||||
assert b"echo:hello" in _read_until(socket, b"echo:hello")
|
||||
|
||||
|
||||
def test_the_requested_size_reaches_the_pty(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(
|
||||
_url(chat, cols=120, rows=40), headers=_headers(client)
|
||||
) as socket:
|
||||
hello = _first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
|
||||
assert hello["cols"] == 120
|
||||
assert shell_host["seen"]["term_size"][:2] == (120, 40)
|
||||
|
||||
|
||||
def test_the_project_directory_is_where_the_shell_starts(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
chat = _agent_chat(db, user_id, shell_host, project_dir="/srv/work")
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
assert _first_json(socket)["dir"] == "/srv/work"
|
||||
_read_until(socket, b"READY")
|
||||
|
||||
assert "cd '/srv/work'" in shell_host["seen"]["command"]
|
||||
|
||||
|
||||
def test_the_mode_does_not_apply_to_what_a_person_types(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
"""Plan mode stops the *model* running anything. It is not a keyboard lock.
|
||||
|
||||
Pinned because it looks like a bug to anybody reading `policy.py` next to
|
||||
this, and "fixing" it would make the panel useless in the mode people spend
|
||||
the most time in.
|
||||
"""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
chat.agent_mode = policy.MODE_PLAN
|
||||
db.commit()
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
assert _first_json(socket)["t"] == "ready"
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_bytes(b"rm -rf /tmp/nothing\n")
|
||||
_read_until(socket, b"echo:rm -rf /tmp/nothing")
|
||||
|
||||
assert "rm -rf /tmp/nothing" in shell_host["seen"]["typed"]
|
||||
|
||||
|
||||
# --- Staying out ------------------------------------------------------------------
|
||||
def test_a_stranger_is_refused_before_the_socket_is_accepted(client, db, user_id, shell_host):
|
||||
"""No cookie, no handshake. Accepting first would mean an unauthenticated
|
||||
socket existed at all, however briefly."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
client.cookies.clear()
|
||||
# noqa: B017 - starlette raises on a rejected handshake, and the class differs
|
||||
with pytest.raises(Exception), client.websocket_connect( # noqa: B017
|
||||
_url(chat), headers=_headers(client)
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_a_cross_site_origin_is_refused_before_the_socket_is_accepted(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
headers = _headers(client) | {"origin": "https://evil.example"}
|
||||
with pytest.raises(Exception), client.websocket_connect( # noqa: B017
|
||||
_url(chat), headers=headers
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_an_absent_origin_is_refused(client, db, registered, user_id, shell_host):
|
||||
"""Required rather than checked-when-present: nothing without an Origin is
|
||||
a browser, and a non-browser client has no business here."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with pytest.raises(Exception), client.websocket_connect( # noqa: B017
|
||||
_url(chat), headers={"host": client.base_url.host}
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_someone_elses_chat_is_refused_with_a_reason(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
stranger = User(name="Sam", email="sam@shire.test", password_hash="x")
|
||||
db.add(stranger)
|
||||
db.commit()
|
||||
chat.user_id = stranger.id
|
||||
db.commit()
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "no longer exists" in payload["message"]
|
||||
|
||||
|
||||
def test_an_ordinary_chat_has_no_machine_to_open(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host, kind=KIND_CHAT)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "ordinary chat" in payload["message"]
|
||||
|
||||
|
||||
def test_the_instance_switch_refuses_with_a_reason(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
settings_store.update(db, {"terminal_enabled": False}, key=settings_store.AGENTS)
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "switched off" in payload["message"]
|
||||
|
||||
|
||||
def test_a_disabled_connection_refuses_with_a_reason(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
profile = db.get(SshProfile, chat.ssh_profile_id)
|
||||
profile.enabled = False
|
||||
db.commit()
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "not usable" in payload["message"]
|
||||
|
||||
|
||||
def test_an_unconfirmed_host_key_says_so_rather_than_failing_silently(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
"""The one refusal that is genuinely actionable: press Check and accept."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
profile = db.get(SshProfile, chat.ssh_profile_id)
|
||||
profile.host_key = ""
|
||||
db.commit()
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "host key has not been confirmed" in payload["message"]
|
||||
|
||||
|
||||
def test_without_the_permission_there_is_no_terminal(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
user = db.get(User, user_id)
|
||||
user.role = "user" # an admin bypasses every permission, deliberately
|
||||
db.commit()
|
||||
settings_store.update(
|
||||
db,
|
||||
{"default_permissions": {"agent.terminal": False, "tools.agent": True, "agent.ssh": True}},
|
||||
)
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
payload = _first_json(socket)
|
||||
assert payload["t"] == "error"
|
||||
assert "permission" in payload["message"]
|
||||
|
||||
|
||||
# --- Lifetime ---------------------------------------------------------------------
|
||||
def test_closing_the_socket_leaves_the_shell_running(client, db, registered, user_id, shell_host):
|
||||
"""The whole reason the session is not owned by the socket."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_bytes(b"remember\n")
|
||||
_read_until(socket, b"echo:remember")
|
||||
|
||||
session = terminal_service.get(chat.id)
|
||||
assert session is not None
|
||||
assert not session.closed
|
||||
|
||||
|
||||
def test_reconnecting_replays_what_was_missed(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_bytes(b"before\n")
|
||||
_read_until(socket, b"echo:before")
|
||||
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
assert _first_json(socket)["t"] == "ready"
|
||||
# Arrives as the scrollback, in one frame, before anything new.
|
||||
assert b"echo:before" in _read_until(socket, b"echo:before")
|
||||
|
||||
|
||||
def test_the_shell_exiting_is_announced(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_bytes(b"exit\n")
|
||||
|
||||
payload = None
|
||||
for _ in range(40):
|
||||
message = socket.receive()
|
||||
if message.get("text"):
|
||||
payload = json.loads(message["text"])
|
||||
break
|
||||
assert payload is not None
|
||||
assert payload["t"] == "closed"
|
||||
assert payload["reason"] == terminal_service.CLOSED_EXITED
|
||||
|
||||
|
||||
def test_a_resize_frame_reaches_the_far_side(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
socket.send_text(json.dumps({"t": "resize", "cols": 132, "rows": 43}))
|
||||
# Round-tripped through the shell so the resize has certainly been read.
|
||||
socket.send_bytes(b"after\n")
|
||||
_read_until(socket, b"echo:after")
|
||||
|
||||
assert shell_host["seen"]["term_size"][:2] == (132, 43)
|
||||
|
||||
|
||||
# --- The panel on the page --------------------------------------------------------
|
||||
def test_the_panel_and_its_vendored_terminal_are_on_an_agent_chat(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
body = client.get(f"/chat/{chat.id}").text
|
||||
|
||||
assert 'id="terminal"' in body
|
||||
assert "vendor/xterm.js" in body
|
||||
assert 'data-toggle-group="side"' in body
|
||||
|
||||
|
||||
def test_an_ordinary_chat_loads_none_of_it(client, db, registered, user_id, shell_host):
|
||||
"""280KB of terminal on a page that could never use it."""
|
||||
chat = _agent_chat(db, user_id, shell_host, kind=KIND_CHAT)
|
||||
body = client.get(f"/chat/{chat.id}").text
|
||||
|
||||
assert 'id="terminal"' not in body
|
||||
assert "xterm" not in body
|
||||
|
||||
|
||||
def test_the_panel_is_absent_when_the_instance_switch_is_off(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
settings_store.update(db, {"terminal_enabled": False}, key=settings_store.AGENTS)
|
||||
|
||||
assert 'id="terminal"' not in client.get(f"/chat/{chat.id}").text
|
||||
|
||||
|
||||
def test_the_panel_is_absent_without_ssh_installed(
|
||||
client, db, registered, user_id, shell_host, monkeypatch
|
||||
):
|
||||
"""A button whose only outcome is an error frame is worse than no button."""
|
||||
from lembas.services.agent import ssh as ssh_service
|
||||
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
monkeypatch.setattr(ssh_service, "available", lambda: ssh_service.INSTALL_HINT)
|
||||
|
||||
assert 'id="terminal"' not in client.get(f"/chat/{chat.id}").text
|
||||
|
||||
|
||||
# --- Things that must reach a shell already open ----------------------------------
|
||||
def _open_and_leave(client, chat) -> None:
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
assert terminal_service.get(chat.id) is not None
|
||||
|
||||
|
||||
def test_deleting_the_chat_closes_its_terminal(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
_open_and_leave(client, chat)
|
||||
|
||||
assert client.delete(f"/api/chats/{chat.id}").status_code == 204
|
||||
assert terminal_service.get(chat.id) is None
|
||||
|
||||
|
||||
def test_disabling_the_connection_closes_its_terminal(
|
||||
client, db, registered, user_id, shell_host
|
||||
):
|
||||
"""The model stops on the next reply anyway. A shell already open would not."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
_open_and_leave(client, chat)
|
||||
|
||||
profile = db.get(SshProfile, chat.ssh_profile_id)
|
||||
client.post(
|
||||
f"/api/agents/{profile.id}",
|
||||
data={
|
||||
"name": profile.name,
|
||||
"host": profile.host,
|
||||
"port": str(profile.port),
|
||||
"username": profile.username,
|
||||
"auth": profile.auth,
|
||||
"connect_timeout": "15",
|
||||
},
|
||||
follow_redirects=False,
|
||||
)
|
||||
db.expire_all()
|
||||
assert db.get(SshProfile, profile.id).enabled is False
|
||||
assert terminal_service.get(chat.id) is None
|
||||
|
||||
|
||||
def test_deleting_the_connection_closes_its_terminal(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
_open_and_leave(client, chat)
|
||||
|
||||
profile_id = chat.ssh_profile_id
|
||||
client.post(f"/api/agents/{profile_id}/delete", follow_redirects=False)
|
||||
assert terminal_service.get(chat.id) is None
|
||||
|
||||
|
||||
def test_forgetting_the_host_key_closes_its_terminal(client, db, registered, user_id, shell_host):
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
_open_and_leave(client, chat)
|
||||
|
||||
client.post(f"/api/agents/{chat.ssh_profile_id}/forget")
|
||||
assert terminal_service.get(chat.id) is None
|
||||
|
||||
|
||||
def test_nonsense_control_frames_are_ignored(client, db, registered, user_id, shell_host):
|
||||
"""A frame is text somebody could forge; none of it may crash the socket."""
|
||||
chat = _agent_chat(db, user_id, shell_host)
|
||||
with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
|
||||
_first_json(socket)
|
||||
_read_until(socket, b"READY")
|
||||
for frame in ("not json", "[]", '{"t":"unknown"}', '{"t":"resize"}'):
|
||||
socket.send_text(frame)
|
||||
socket.send_bytes(b"alive\n")
|
||||
assert b"echo:alive" in _read_until(socket, b"echo:alive")
|
||||
Reference in New Issue
Block a user