feat: add terminal support for guest execution agent and enhance lifecycle management

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
   "fastapi>=0.115,<1",
   "pydantic>=2,<3",
   "uvicorn>=0.30,<1",
+  "websockets>=13,<17",
 ]
 
 [project.optional-dependencies]

uvm/guest_agent.py

@@ -6,6 +6,9 @@ The host also imports its wire-format helpers; importing it never starts a serve
 
 from __future__ import annotations
 
+import base64
+import errno
+import fcntl
 import json
 import math
 import os
@@ -18,6 +21,7 @@ import stat
 import struct
 import subprocess
 import sys
+import termios
 import threading
 import time
 import uuid
@@ -35,6 +39,8 @@ MAX_RESPONSE_BYTES = 6 * MAX_OUTPUT_BYTES + 4096
 CONNECT_TIMEOUT_SECONDS = 5.0
 CLEANUP_SECONDS = 5.0
 MAX_CONCURRENT_COMMANDS = 8
+MAX_TERMINAL_BYTES = 8192
+MAX_TERMINAL_FRAME = 16384
 CONFIG_PATH = Path("/etc/uvm-agent.json")
 
 
@@ -48,6 +54,23 @@ class ExecutionFailure(Exception):
         self.outcome_unknown = outcome_unknown
 
 
+def validate_terminal_size(rows: object, cols: object) -> None:
+    if type(rows) is not int or type(cols) is not int or not (2 <= rows <= 200 and 2 <= cols <= 400):
+        raise ProtocolError("terminal dimensions must be integers: 2–200 rows and 2–400 columns")
+
+
+def terminal_data(value: object) -> bytes:
+    if not isinstance(value, str) or len(value) > 4 * ((MAX_TERMINAL_BYTES + 2) // 3):
+        raise ProtocolError("terminal data exceeds its size limit")
+    try:
+        data = base64.b64decode(value, validate=True)
+    except (ValueError, UnicodeError) as error:
+        raise ProtocolError("invalid terminal data") from error
+    if len(data) > MAX_TERMINAL_BYTES:
+        raise ProtocolError("terminal data exceeds its size limit")
+    return data
+
+
 def validate_command(command: object, sudo: object, timeout_seconds: object) -> None:
     if not isinstance(command, str) or not command.strip() or "\0" in command:
         raise ProtocolError("command must be non-empty text without NUL characters")
@@ -167,7 +190,9 @@ def _cleanup_job(process: subprocess.Popen | None, scope: JobScope) -> None:
             failed = True
         finally:
             for name in ("stdout", "stderr"):
-                getattr(process, name).close()
+                stream = getattr(process, name)
+                if stream is not None:
+                    stream.close()
     try:
         scope.close()
     except (OSError, ExecutionFailure):
@@ -239,9 +264,12 @@ def run_command(command: str, username: str, timeout_seconds: float, *, scope_fa
     }
 
 
-def child_main(scope: Path, username: str, command: str) -> None:
+def child_main(scope: Path, username: str, command: str | None) -> None:
     """Enter the job cgroup and account in a fresh interpreter, then replace it."""
     (scope / "cgroup.procs").write_text(str(os.getpid()), encoding="ascii")
+    if command is None:
+        # Popen's start_new_session already performed setsid in this fresh child.
+        fcntl.ioctl(0, termios.TIOCSCTTY, 0)
     account = pwd.getpwnam(username)
     if os.geteuid() == 0:
         os.initgroups(account.pw_name, account.pw_gid)
@@ -256,15 +284,111 @@ def child_main(scope: Path, username: str, command: str) -> None:
         "SHELL": "/bin/sh", "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
         "LANG": "C.UTF-8",
     }
-    os.execve("/bin/sh", ["sh", "-c", command], environment)
+    if command is None:
+        shell = account.pw_shell or "/bin/sh"
+        environment.update(SHELL=shell, TERM="xterm-256color", COLORTERM="truecolor")
+        os.execve(shell, [shell, "-i"], environment)
+    else:
+        os.execve("/bin/sh", ["sh", "-c", command], environment)
 
 
-def serve_connection(connection: socket.socket, config: dict, *, runner=run_command) -> None:
+def run_terminal(connection: socket.socket, username: str, rows: int, cols: int, *, scope_factory=JobScope) -> None:
+    """Bridge one real guest PTY until shell exit, explicit close, or host loss."""
+    validate_terminal_size(rows, cols)
+    scope = scope_factory()
+    process = None
+    master = slave = None
+    try:
+        master, slave = os.openpty()
+        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
+        process = subprocess.Popen(
+            [sys.executable, "-I", str(Path(__file__).resolve()), "--terminal-child", str(scope.path), username],
+            stdin=slave, stdout=slave, stderr=slave, start_new_session=True,
+            env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8"},
+        )
+        os.close(slave)
+        slave = None
+        os.set_blocking(master, False)
+        send_frame(connection, {"type": "ready"}, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+        pending = bytearray()
+        with selectors.DefaultSelector() as selector:
+            selector.register(connection, selectors.EVENT_READ, "host")
+            selector.register(master, selectors.EVENT_READ, "pty")
+            finished = False
+            while not finished:
+                for key, events in selector.select(0.1):
+                    if key.data == "host":
+                        message = receive_frame(connection, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+                        if message == {"type": "close"}:
+                            finished = True
+                            break
+                        if set(message) == {"type", "data"} and message["type"] == "input":
+                            pending.extend(terminal_data(message["data"]))
+                            if len(pending) > 65536:
+                                raise ProtocolError("terminal input backlog exceeded")
+                        elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize":
+                            validate_terminal_size(message["rows"], message["cols"])
+                            fcntl.ioctl(master, termios.TIOCSWINSZ, struct.pack("HHHH", message["rows"], message["cols"], 0, 0))
+                        else:
+                            raise ProtocolError("invalid terminal control message")
+                        selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty")
+                    else:
+                        if events & selectors.EVENT_READ:
+                            try:
+                                data = os.read(master, MAX_TERMINAL_BYTES)
+                            except OSError as error:
+                                if error.errno != errno.EIO:
+                                    raise
+                                data = b""
+                            if not data:
+                                finished = True
+                                break
+                            send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")},
+                                       MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+                        if events & selectors.EVENT_WRITE and pending:
+                            try:
+                                del pending[:os.write(master, pending)]
+                            except BlockingIOError:
+                                pass
+                            selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty")
+                if _exited(process):
+                    # Drain the leader's final output before terminating any background jobs.
+                    while True:
+                        try:
+                            data = os.read(master, MAX_TERMINAL_BYTES)
+                        except (BlockingIOError, OSError):
+                            break
+                        if not data:
+                            break
+                        send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")},
+                                   MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+                    break
+    finally:
+        try:
+            _cleanup_job(process, scope)
+        finally:
+            for descriptor in (master, slave):
+                if descriptor is not None:
+                    os.close(descriptor)
+    send_frame(connection, {"type": "exit", "exit_code": process.returncode}, MAX_TERMINAL_FRAME,
+               time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+
+
+def serve_connection(connection: socket.socket, config: dict, *, runner=run_command, terminal_runner=run_terminal) -> None:
     with connection:
         try:
             deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
             send_frame(connection, {"version": PROTOCOL_VERSION, **config}, 4096, deadline)
             request = receive_frame(connection, MAX_REQUEST_BYTES, deadline)
+            if set(request) == {"version", "terminal"}:
+                if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION:
+                    raise ProtocolError("unsupported guest protocol version")
+                size = request["terminal"]
+                if not isinstance(size, dict) or set(size) != {"rows", "cols"}:
+                    raise ProtocolError("invalid terminal request")
+                validate_terminal_size(size["rows"], size["cols"])
+                terminal_runner(connection, config["username"], size["rows"], size["cols"])
+                return
             if set(request) != {"version", "command", "sudo", "timeout_seconds"}:
                 raise ProtocolError("invalid execution request fields")
             if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION:
@@ -305,6 +429,13 @@ def load_config(path: Path = CONFIG_PATH) -> dict:
 
 
 def main() -> None:
+    if len(sys.argv) == 4 and sys.argv[1] == "--terminal-child":
+        try:
+            child_main(Path(sys.argv[2]), sys.argv[3], None)
+        except (OSError, KeyError) as error:
+            print(f"uvm-agent: could not launch guest terminal: {error}", file=sys.stderr)
+            sys.exit(126)
+        return
     if len(sys.argv) == 5 and sys.argv[1] == "--child":
         try:
             child_main(Path(sys.argv[2]), sys.argv[3], sys.argv[4])

uvm/guest_terminal.py

@@ -0,0 +1,89 @@
+"""Interactive, process-pinned guest terminal transport over the existing vsock agent."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+import struct
+import time
+
+from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
+from .guest_agent import (
+    CONNECT_TIMEOUT_SECONDS, MAX_TERMINAL_FRAME, PROTOCOL_VERSION, ProtocolError,
+    receive_frame, send_frame, terminal_data, validate_terminal_size,
+)
+from .guest_exec import GuestSession
+
+
+def validate_size(rows: object, cols: object) -> None:
+    try:
+        validate_terminal_size(rows, cols)
+    except ProtocolError as error:
+        raise ValidationError(str(error)) from error
+
+
+class GuestTerminal(GuestSession):
+    def start(self, rows: int, cols: int) -> None:
+        validate_size(rows, cols)
+        try:
+            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
+            send_frame(self._connection, {"version": PROTOCOL_VERSION, "terminal": {"rows": rows, "cols": cols}},
+                       MAX_TERMINAL_FRAME, deadline)
+            response = receive_frame(self._connection, MAX_TERMINAL_FRAME, deadline)
+            if response != {"type": "ready"}:
+                raise ProtocolError("terminal support is unavailable")
+            self._connection.setblocking(False)
+            self._write_lock = asyncio.Lock()
+        except (OSError, EOFError, ProtocolError) as error:
+            self.close()
+            raise GuestAgentUnavailableError(
+                "Interactive terminal support is unavailable. Stop this VM, use Install / update terminal support "
+                "in its Sessions view, then start it and launch a new session."
+            ) from error
+
+    async def _receive_exact(self, size: int) -> bytes:
+        data = bytearray()
+        while len(data) < size:
+            chunk = await asyncio.get_running_loop().sock_recv(self._connection, size - len(data))
+            if not chunk:
+                raise EOFError("guest terminal connection closed")
+            data.extend(chunk)
+        return bytes(data)
+
+    async def read(self) -> dict:
+        try:
+            # An idle shell has no output deadline; partial messages are bounded.
+            first = await self._receive_exact(1)
+            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
+                size, = struct.unpack("!I", first + await self._receive_exact(3))
+                if not 0 < size <= MAX_TERMINAL_FRAME:
+                    raise ProtocolError("invalid terminal frame length")
+                message = json.loads(await self._receive_exact(size))
+            if isinstance(message, dict):
+                if set(message) == {"type", "data"} and message["type"] == "output":
+                    return {"type": "output", "data": terminal_data(message["data"])}
+                if (set(message) == {"type", "exit_code"} and message["type"] == "exit"
+                        and type(message["exit_code"]) is int and -64 <= message["exit_code"] <= 255):
+                    return message
+            raise ProtocolError("invalid guest terminal response")
+        except (OSError, EOFError, ValueError, UnicodeError, RecursionError) as error:
+            raise GuestExecutionError("The guest terminal connection ended unexpectedly.", outcome_unknown=True) from error
+
+    async def _send(self, message: dict) -> None:
+        payload = json.dumps(message, separators=(",", ":")).encode("ascii")
+        if len(payload) > MAX_TERMINAL_FRAME:
+            raise ValidationError("terminal input exceeds its size limit")
+        async with self._write_lock:
+            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
+                await asyncio.get_running_loop().sock_sendall(self._connection, struct.pack("!I", len(payload)) + payload)
+
+    async def write(self, data: bytes) -> None:
+        await self._send({"type": "input", "data": base64.b64encode(data).decode("ascii")})
+
+    async def resize(self, rows: int, cols: int) -> None:
+        validate_size(rows, cols)
+        await self._send({"type": "resize", "rows": rows, "cols": cols})
+
+    async def end(self) -> None:
+        await self._send({"type": "close"})

uvm/images.py

@@ -317,7 +317,9 @@ class ImageStore:
                 if not install:
                     marker = self._read_guest_file(disk, "/etc/uvm-image", files / "marker",
                                                    required=False, regular_only=True)
-                    if marker is None or marker.path.read_text(encoding="utf-8") != "uvm-ubuntu-jammy-v4\n":
+                    if marker is None or marker.path.read_text(encoding="utf-8") not in {
+                        "uvm-ubuntu-jammy-v4\n", "uvm-ubuntu-jammy-v5\n",
+                    }:
                         return False
                 descriptor, candidate = tempfile.mkstemp(prefix=".rootfs-agent-", suffix=".ext4", dir=disk.parent)
                 os.close(descriptor)

uvm/install.py

@@ -21,7 +21,7 @@ from .system import CommandRunner, check_kvm, ensure_data_directories, require_r
 
 
 _UBUNTU_RELEASE = "jammy"
-_ROOTFS_RECIPE = "uvm-ubuntu-jammy-v4\n"
+_ROOTFS_RECIPE = "uvm-ubuntu-jammy-v5\n"
 _GUEST_PACKAGES = (
     "ca-certificates", "openssh-server", "sudo", "systemd-sysv", "udev",
     "dbus", "libpam-systemd", "iproute2", "iputils-ping", "curl", "python3",

uvm/lifecycle.py

@@ -19,6 +19,7 @@ from .firecracker.config import build_config, write_config
 from .firecracker.process import FirecrackerProcessManager, ProcessInfo
 from .guest_agent import DEFAULT_TIMEOUT_SECONDS
 from .guest_exec import ExecResult, GuestSession, validate_exec
+from .guest_terminal import GuestTerminal, validate_size
 from .images import ImageStore
 from .integrity import load_manifest, verify_file
 from .known_hosts import KnownHostsManager
@@ -61,6 +62,7 @@ class LifecycleService:
         known_hosts: KnownHostsManager,
         client_factory: Callable[[Path, float], FirecrackerClient] = FirecrackerClient,
         guest_connector: Callable[[Path, str, str, int], GuestSession] = GuestSession.connect,
+        terminal_connector: Callable[[Path, str, str, int], GuestTerminal] = GuestTerminal.connect,
         clock: Callable[[], float] = time.time,
     ) -> None:
         self._settings = settings
@@ -71,6 +73,7 @@ class LifecycleService:
         self._known_hosts = known_hosts
         self._client_factory = client_factory
         self._guest_connector = guest_connector
+        self._terminal_connector = terminal_connector
         self._clock = clock
 
     def create(self, spec: VmSpec) -> VmRecord:
@@ -520,6 +523,24 @@ class LifecycleService:
         finally:
             session.close()
 
+    def open_terminal(self, vm_id: str, *, rows: int = 24, cols: int = 80) -> tuple[GuestTerminal, str]:
+        """Connect to the exact running VMM under the lock, then open its guest PTY."""
+        require_root()
+        validate_size(rows, cols)
+        with self._state_store.operation_lock():
+            vm = self._load_for_operation(vm_id, exact_id=True)
+            self._not_deleting(vm)
+            if (vm.status != "running" or vm.operation is not None or vm.pending_resize is not None
+                    or self._process.observe(vm) != "alive"):
+                raise ConflictError("terminal sessions require a running VM with no pending lifecycle operation")
+            terminal = self._terminal_connector(self._settings.guest_socket_path(vm.id), vm.id, vm.username, vm.pid)
+        try:
+            terminal.start(rows, cols)
+            return terminal, vm.username
+        except BaseException:
+            terminal.close()
+            raise
+
     def install_guest_agent(self, vm_id: str) -> VmRecord:
         """Explicitly install/upgrade the agent on a cleanly stopped private disk."""
         require_root()