fix: fully drain terminal paste writes to PTY
- Make the host PTY master non-blocking and wait for write readiness. - Retry partial writes so a large bracketed paste always delivers its closing marker instead of leaving pi in paste mode. - Add regression coverage for partial PTY writes and resolve diagnostics. Quality gates: ruff, mypy, focused pytest (3 passed)
This commit is contained in:
Developer
@@ -19,6 +19,8 @@ from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_session_counters_by_instance: dict[str, int] = {}
|
||||
|
||||
|
||||
class TerminalSession:
|
||||
"""Manages a single terminal session with event-driven PTY I/O.
|
||||
@@ -46,9 +48,6 @@ class TerminalSession:
|
||||
# Max WebSocket frame size
|
||||
MAX_FRAME_SIZE = 64 * 1024
|
||||
|
||||
# Session number counter per instance_id for auto-naming
|
||||
_instance_counters: dict[str, int] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
@@ -99,11 +98,11 @@ class TerminalSession:
|
||||
# Ack timeout fallback
|
||||
self._ack_timeout_handle: asyncio.TimerHandle | None = None
|
||||
|
||||
@classmethod
|
||||
def _generate_name(cls, instance_id: str) -> str:
|
||||
@staticmethod
|
||||
def _generate_name(instance_id: str) -> str:
|
||||
"""Generate an auto-incremented session name for the instance."""
|
||||
count = cls._instance_counters.get(instance_id, 0) + 1
|
||||
cls._instance_counters[instance_id] = count
|
||||
count = _session_counters_by_instance.get(instance_id, 0) + 1
|
||||
_session_counters_by_instance[instance_id] = count
|
||||
return f"Session {count}"
|
||||
|
||||
async def start(self, startup_command: str | None = None) -> None:
|
||||
@@ -121,8 +120,12 @@ class TerminalSession:
|
||||
if self.container_user:
|
||||
exec_cmd.extend(["--user", self.container_user])
|
||||
|
||||
# Create a pseudo-terminal on the host
|
||||
# Create a pseudo-terminal on the host. The master must be
|
||||
# non-blocking: a browser paste can be larger than the PTY input
|
||||
# buffer, and write_input() drains it asynchronously without dropping
|
||||
# the closing bracketed-paste marker.
|
||||
self._master_fd, slave_fd = pty.openpty()
|
||||
os.set_blocking(self._master_fd, False)
|
||||
|
||||
# Put the host PTY into raw mode so it behaves as a pass-through
|
||||
# pipe. openpty() leaves the slave in canonical mode by default,
|
||||
@@ -204,6 +207,9 @@ class TerminalSession:
|
||||
|
||||
try:
|
||||
data = os.read(self._master_fd, 4096)
|
||||
except BlockingIOError:
|
||||
# The readiness notification raced with another callback.
|
||||
return
|
||||
except (OSError, IOError) as exc:
|
||||
logger.debug("PTY read error for session %s: %s", self.session_id, exc)
|
||||
self._handle_eof()
|
||||
@@ -343,12 +349,42 @@ class TerminalSession:
|
||||
pass
|
||||
logger.info("Session %s EOF handled, websockets closed", self.session_id)
|
||||
|
||||
async def _wait_for_write_ready(self, fd: int) -> None:
|
||||
"""Wait until a non-blocking PTY master can accept more input."""
|
||||
loop = asyncio.get_running_loop()
|
||||
writable = loop.create_future()
|
||||
|
||||
def mark_writable() -> None:
|
||||
if not writable.done():
|
||||
writable.set_result(None)
|
||||
|
||||
loop.add_writer(fd, mark_writable)
|
||||
try:
|
||||
await writable
|
||||
finally:
|
||||
loop.remove_writer(fd)
|
||||
|
||||
async def write_input(self, data: bytes) -> None:
|
||||
"""Write input to the PTY master."""
|
||||
"""Write all terminal input bytes to the PTY master in order."""
|
||||
if self._master_fd is None or self._closed:
|
||||
return
|
||||
|
||||
fd = self._master_fd
|
||||
remaining = memoryview(data)
|
||||
try:
|
||||
os.write(self._master_fd, data)
|
||||
while remaining and not self._closed and self._master_fd == fd:
|
||||
try:
|
||||
written = os.write(fd, remaining)
|
||||
except BlockingIOError:
|
||||
await self._wait_for_write_ready(fd)
|
||||
continue
|
||||
|
||||
if written == 0:
|
||||
await self._wait_for_write_ready(fd)
|
||||
continue
|
||||
|
||||
remaining = remaining[written:]
|
||||
|
||||
self.last_activity = time.time()
|
||||
except (OSError, IOError) as exc:
|
||||
logger.debug("PTY write error for session %s: %s", self.session_id, exc)
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
"""Unit tests for TerminalSession docker exec invocation."""
|
||||
"""Unit tests for TerminalSession PTY handling."""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
import pytest # pyright: ignore[reportMissingImports]
|
||||
|
||||
from src.services.terminal.terminal_session import TerminalSession
|
||||
|
||||
@@ -19,17 +19,20 @@ async def test_start_passes_container_user_to_docker_exec() -> None:
|
||||
container_user="dev",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.services.terminal.terminal_session.pty.openpty",
|
||||
return_value=(1, 2),
|
||||
with (
|
||||
patch(
|
||||
"src.services.terminal.terminal_session.pty.openpty",
|
||||
return_value=(1, 2),
|
||||
),
|
||||
patch("src.services.terminal.terminal_session.os.set_blocking") as set_blocking,
|
||||
patch("src.services.terminal.terminal_session.tty.setraw"),
|
||||
patch(
|
||||
"src.services.terminal.terminal_session.asyncio.create_subprocess_exec",
|
||||
new=AsyncMock(),
|
||||
) as mock_exec,
|
||||
patch("src.services.terminal.terminal_session.os.close"),
|
||||
):
|
||||
with patch("src.services.terminal.terminal_session.tty.setraw"):
|
||||
with patch(
|
||||
"src.services.terminal.terminal_session.asyncio.create_subprocess_exec",
|
||||
new=AsyncMock(),
|
||||
) as mock_exec:
|
||||
with patch("src.services.terminal.terminal_session.os.close"):
|
||||
await session.start()
|
||||
await session.start()
|
||||
|
||||
args, _kwargs = mock_exec.call_args
|
||||
assert "docker" in args
|
||||
@@ -38,6 +41,7 @@ async def test_start_passes_container_user_to_docker_exec() -> None:
|
||||
user_index = args.index("--user")
|
||||
assert args[user_index + 1] == "dev"
|
||||
assert "container-123" in args
|
||||
set_blocking.assert_called_once_with(1, False)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@@ -50,17 +54,44 @@ async def test_start_omits_user_when_not_configured() -> None:
|
||||
container_id="container-123",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.services.terminal.terminal_session.pty.openpty",
|
||||
return_value=(1, 2),
|
||||
with (
|
||||
patch(
|
||||
"src.services.terminal.terminal_session.pty.openpty",
|
||||
return_value=(1, 2),
|
||||
),
|
||||
patch("src.services.terminal.terminal_session.os.set_blocking"),
|
||||
patch("src.services.terminal.terminal_session.tty.setraw"),
|
||||
patch(
|
||||
"src.services.terminal.terminal_session.asyncio.create_subprocess_exec",
|
||||
new=AsyncMock(),
|
||||
) as mock_exec,
|
||||
patch("src.services.terminal.terminal_session.os.close"),
|
||||
):
|
||||
with patch("src.services.terminal.terminal_session.tty.setraw"):
|
||||
with patch(
|
||||
"src.services.terminal.terminal_session.asyncio.create_subprocess_exec",
|
||||
new=AsyncMock(),
|
||||
) as mock_exec:
|
||||
with patch("src.services.terminal.terminal_session.os.close"):
|
||||
await session.start()
|
||||
await session.start()
|
||||
|
||||
args, _kwargs = mock_exec.call_args
|
||||
assert "--user" not in args
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_input_retries_partial_pty_writes() -> None:
|
||||
"""A large paste is fully written even when the PTY accepts it in chunks."""
|
||||
session = TerminalSession(
|
||||
session_id=str(uuid.uuid4()),
|
||||
instance_id=uuid.uuid4(),
|
||||
container_id="container-123",
|
||||
)
|
||||
session._master_fd = 42
|
||||
|
||||
with patch(
|
||||
"src.services.terminal.terminal_session.os.write",
|
||||
side_effect=[2, 2, 1],
|
||||
) as write:
|
||||
await session.write_input(b"hello")
|
||||
|
||||
assert [bytes(call.args[1]) for call in write.call_args_list] == [
|
||||
b"hello",
|
||||
b"llo",
|
||||
b"o",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user