diff --git a/api/openapi.yaml b/api/openapi.yaml index 7ed1e94..f2d1c9d 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -987,13 +987,17 @@ paths: \ \"rows\": 24,\n \"envs\": {\"TERM\": \"xterm-256color\"},\n \"cwd\": \"/home/user\",\n \"user\": \"user\"\n}\n\ ```\nAll fields except `type` are optional. Defaults: cmd=\"/bin/bash\", cols=80, rows=24.\n\n**Client sends** (first\ \ message — reconnect to existing PTY):\n```json\n{\"type\": \"connect\", \"tag\": \"pty-abc123de\"}\n```\n\n**Client\ - \ sends** (after session is established):\n```json\n{\"type\": \"input\", \"data\": \"\"}\n\ - {\"type\": \"resize\", \"cols\": 120, \"rows\": 40}\n{\"type\": \"kill\"}\n```\n\n**Server sends**:\n```json\n{\"\ - type\": \"started\", \"tag\": \"pty-abc123de\", \"pid\": 42}\n{\"type\": \"output\", \"data\": \"\"}\n{\"type\": \"exit\", \"exit_code\": 0}\n{\"type\": \"error\", \"data\": \"description\", \"fatal\": true}\n\ - {\"type\": \"ping\"}\n```\n\nPTY data (input and output) is base64-encoded because it contains raw\nterminal bytes\ - \ (escape sequences, control codes) that are not valid UTF-8.\n\nSessions persist across WebSocket disconnections\ - \ — the process keeps\nrunning in the capsule. Use the `tag` from the \"started\" response to\nreconnect later.\n" + \ sends** (control messages, after session is established — JSON text frames):\n```json\n{\"type\": \"resize\", \"\ + cols\": 120, \"rows\": 40}\n{\"type\": \"kill\"}\n```\n\nKeystroke input is sent as **raw binary WebSocket frames**\ + \ containing the\nterminal bytes verbatim — not JSON, not base64. Every binary frame from\nthe client is treated as\ + \ PTY input.\n\n**Server sends** (control messages — JSON text frames):\n```json\n{\"type\": \"started\", \"tag\"\ + : \"pty-abc123de\", \"pid\": 42}\n{\"type\": \"exit\", \"exit_code\": 0}\n{\"type\": \"error\", \"data\": \"description\"\ + , \"fatal\": true}\n{\"type\": \"ping\"}\n```\n\nPTY output is sent as **raw binary WebSocket frames** containing\ + \ the\nterminal bytes verbatim — not JSON, not base64. Sending input and output\nas binary avoids the ~33% inflation\ + \ of base64 over the JSON path.\n\nThe connection negotiates permessage-deflate (RFC 7692) compression\nautomatically;\ + \ rapid output (e.g. full-screen TUI repaints) is coalesced\ninto fewer, larger frames before compression.\n\nSessions\ + \ persist across WebSocket disconnections — the process keeps\nrunning in the capsule. Use the `tag` from the \"started\"\ + \ response to\nreconnect later.\n" responses: '101': description: WebSocket upgrade diff --git a/src/wrenn/models/_generated.py b/src/wrenn/models/_generated.py index 2eeb5a3..002dd13 100644 --- a/src/wrenn/models/_generated.py +++ b/src/wrenn/models/_generated.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: openapi.yaml -# timestamp: 2026-06-22T19:04:03+00:00 +# timestamp: 2026-06-22T22:24:45+00:00 from __future__ import annotations from typing import Annotated diff --git a/src/wrenn/pty.py b/src/wrenn/pty.py index 0b7ff77..bd33e4f 100644 --- a/src/wrenn/pty.py +++ b/src/wrenn/pty.py @@ -1,6 +1,5 @@ from __future__ import annotations -import base64 import json from collections.abc import AsyncIterator, Iterator from enum import StrEnum @@ -8,6 +7,7 @@ from typing import Any import httpx_ws from pydantic import BaseModel +from wsproto.events import BytesMessage, TextMessage # A clean (``WebSocketDisconnect``) or abrupt (``WebSocketNetworkError``) close # both mean the PTY stream has ended; iteration must stop on either. @@ -31,7 +31,8 @@ class PtyEvent(BaseModel): fatal: bool | None = None -def _parse_pty_event(raw: dict[str, Any]) -> PtyEvent: +def _parse_control_event(raw: dict[str, Any]) -> PtyEvent: + """Parse a JSON control frame from the server.""" msg_type = raw.get("type", "") if msg_type == "started": return PtyEvent( @@ -39,10 +40,6 @@ def _parse_pty_event(raw: dict[str, Any]) -> PtyEvent: pid=raw.get("pid"), tag=raw.get("tag"), ) - if msg_type == "output": - raw_data = raw.get("data", "") - decoded = base64.b64decode(raw_data) if raw_data else b"" - return PtyEvent(type=PtyEventType.output, data=decoded) if msg_type == "exit": return PtyEvent(type=PtyEventType.exit, exit_code=raw.get("exit_code", -1)) if msg_type == "error": @@ -133,10 +130,10 @@ class PtySession: """Send raw bytes to the PTY stdin. Args: - data: Raw bytes to send. Base64-encoded internally. + data: Raw bytes to send. Delivered as a binary WebSocket frame + verbatim — no JSON wrapping, no base64. """ - encoded = base64.b64encode(data).decode("ascii") - self._ws.send_text(json.dumps({"type": "input", "data": encoded})) + self._ws.send_bytes(data) def resize(self, cols: int, rows: int) -> None: """Resize the PTY terminal. @@ -160,27 +157,34 @@ class PtySession: return self def __next__(self) -> PtyEvent: - if self._done: - raise StopIteration - try: - raw = self._ws.receive_text() - except _WS_CLOSED: - raise StopIteration - event = _parse_pty_event(json.loads(raw)) - if event.type == PtyEventType.started: - if event.tag is not None: - self._tag = event.tag - if event.pid is not None: - self._pid = event.pid - if event.type == PtyEventType.ping: - self._send_pong() - if event.type == PtyEventType.exit: - self._done = True + while True: + if self._done: + raise StopIteration + try: + msg = self._ws.receive() + except _WS_CLOSED: + raise StopIteration + if isinstance(msg, BytesMessage): + return PtyEvent(type=PtyEventType.output, data=msg.data) + if isinstance(msg, TextMessage): + event = _parse_control_event(json.loads(msg.data)) + else: + # Ignore other wsproto events (Ping/Pong/etc handled by httpx_ws). + continue + if event.type == PtyEventType.started: + if event.tag is not None: + self._tag = event.tag + if event.pid is not None: + self._pid = event.pid + if event.type == PtyEventType.ping: + self._send_pong() + if event.type == PtyEventType.exit: + self._done = True + return event + if event.type == PtyEventType.error and event.fatal: + self._done = True + return event return event - if event.type == PtyEventType.error and event.fatal: - self._done = True - return event - return event def __enter__(self) -> PtySession: return self @@ -269,10 +273,10 @@ class AsyncPtySession: """Send raw bytes to the PTY stdin. Args: - data: Raw bytes to send. Base64-encoded internally. + data: Raw bytes to send. Delivered as a binary WebSocket frame + verbatim — no JSON wrapping, no base64. """ - encoded = base64.b64encode(data).decode("ascii") - await self._ws.send_text(json.dumps({"type": "input", "data": encoded})) + await self._ws.send_bytes(data) async def resize(self, cols: int, rows: int) -> None: """Resize the PTY terminal. @@ -298,27 +302,34 @@ class AsyncPtySession: return self async def __anext__(self) -> PtyEvent: - if self._done: - raise StopAsyncIteration - try: - raw = await self._ws.receive_text() - except _WS_CLOSED: - raise StopAsyncIteration - event = _parse_pty_event(json.loads(raw)) - if event.type == PtyEventType.started: - if event.tag is not None: - self._tag = event.tag - if event.pid is not None: - self._pid = event.pid - if event.type == PtyEventType.ping: - await self._send_pong() - if event.type == PtyEventType.exit: - self._done = True + while True: + if self._done: + raise StopAsyncIteration + try: + msg = await self._ws.receive() + except _WS_CLOSED: + raise StopAsyncIteration + if isinstance(msg, BytesMessage): + return PtyEvent(type=PtyEventType.output, data=msg.data) + if isinstance(msg, TextMessage): + event = _parse_control_event(json.loads(msg.data)) + else: + # Ignore other wsproto events (Ping/Pong/etc handled by httpx_ws). + continue + if event.type == PtyEventType.started: + if event.tag is not None: + self._tag = event.tag + if event.pid is not None: + self._pid = event.pid + if event.type == PtyEventType.ping: + await self._send_pong() + if event.type == PtyEventType.exit: + self._done = True + return event + if event.type == PtyEventType.error and event.fatal: + self._done = True + return event return event - if event.type == PtyEventType.error and event.fatal: - self._done = True - return event - return event async def __aenter__(self) -> AsyncPtySession: return self diff --git a/tests/test_filesystem_pty.py b/tests/test_filesystem_pty.py index 1c963b9..855bc49 100644 --- a/tests/test_filesystem_pty.py +++ b/tests/test_filesystem_pty.py @@ -1,11 +1,11 @@ from __future__ import annotations -import base64 import json from unittest.mock import AsyncMock, MagicMock import pytest import respx +from wsproto.events import BytesMessage, TextMessage from wrenn.capsule import Capsule from wrenn.models import FileEntry @@ -13,9 +13,18 @@ from wrenn.pty import ( AsyncPtySession, PtyEventType, PtySession, - _parse_pty_event, + _parse_control_event, ) + +def _text(payload: dict) -> TextMessage: + return TextMessage(data=json.dumps(payload)) + + +def _bytes(data: bytes) -> BytesMessage: + return BytesMessage(data=data) + + BASE = "https://app.wrenn.dev/api" @@ -226,50 +235,37 @@ class TestFilesExists: class TestPtyEventParsing: def test_started_event(self): raw = {"type": "started", "tag": "pty-a1b2c3d4", "pid": 42} - event = _parse_pty_event(raw) + event = _parse_control_event(raw) assert event.type == PtyEventType.started assert event.pid == 42 assert event.tag == "pty-a1b2c3d4" - def test_output_event_base64(self): - encoded = base64.b64encode(b"ls -la\n").decode() - raw = {"type": "output", "data": encoded} - event = _parse_pty_event(raw) - assert event.type == PtyEventType.output - assert event.data == b"ls -la\n" - - def test_output_event_empty(self): - raw = {"type": "output", "data": ""} - event = _parse_pty_event(raw) - assert event.data == b"" - def test_exit_event(self): raw = {"type": "exit", "exit_code": 0} - event = _parse_pty_event(raw) + event = _parse_control_event(raw) assert event.type == PtyEventType.exit assert event.exit_code == 0 def test_error_event(self): raw = {"type": "error", "data": "process not found", "fatal": True} - event = _parse_pty_event(raw) + event = _parse_control_event(raw) assert event.type == PtyEventType.error assert event.data == "process not found" assert event.fatal is True def test_ping_event(self): raw = {"type": "ping"} - event = _parse_pty_event(raw) + event = _parse_control_event(raw) assert event.type == PtyEventType.ping class TestPtySessionWrite: - def test_write_sends_base64_input(self): + def test_write_sends_binary_frame(self): ws = MagicMock() session = PtySession(ws, "cl-abc") session.write(b"ls -la\n") - sent = json.loads(ws.send_text.call_args[0][0]) - assert sent["type"] == "input" - assert base64.b64decode(sent["data"]) == b"ls -la\n" + ws.send_bytes.assert_called_once_with(b"ls -la\n") + ws.send_text.assert_not_called() class TestPtySessionResize: @@ -303,12 +299,11 @@ class TestPtySessionKill: class TestPtySessionIteration: def test_iter_yields_events_until_exit(self): ws = MagicMock() - messages = [ - json.dumps({"type": "started", "tag": "pty-abc12345", "pid": 1}), - json.dumps({"type": "output", "data": base64.b64encode(b"hello").decode()}), - json.dumps({"type": "exit", "exit_code": 0}), + ws.receive.side_effect = [ + _text({"type": "started", "tag": "pty-abc12345", "pid": 1}), + _bytes(b"hello"), + _text({"type": "exit", "exit_code": 0}), ] - ws.receive_text.side_effect = messages session = PtySession(ws, "cl-abc") events = list(session) assert len(events) == 3 @@ -320,12 +315,22 @@ class TestPtySessionIteration: assert events[2].type == PtyEventType.exit assert events[2].exit_code == 0 + def test_iter_yields_empty_binary_frame(self): + ws = MagicMock() + ws.receive.side_effect = [ + _bytes(b""), + _text({"type": "exit", "exit_code": 0}), + ] + session = PtySession(ws, "cl-abc") + events = list(session) + assert events[0].type == PtyEventType.output + assert events[0].data == b"" + def test_iter_stops_on_fatal_error(self): ws = MagicMock() - messages = [ - json.dumps({"type": "error", "data": "fatal", "fatal": True}), + ws.receive.side_effect = [ + _text({"type": "error", "data": "fatal", "fatal": True}), ] - ws.receive_text.side_effect = messages session = PtySession(ws, "cl-abc") events = list(session) assert len(events) == 1 @@ -335,7 +340,7 @@ class TestPtySessionIteration: import httpx_ws ws = MagicMock() - ws.receive_text.side_effect = httpx_ws.WebSocketDisconnect() + ws.receive.side_effect = httpx_ws.WebSocketDisconnect() session = PtySession(ws, "cl-abc") events = list(session) assert events == [] @@ -344,9 +349,9 @@ class TestPtySessionIteration: class TestPtySessionPong: def test_ping_triggers_pong(self): ws = MagicMock() - ws.receive_text.side_effect = [ - json.dumps({"type": "ping"}), - json.dumps({"type": "exit", "exit_code": 0}), + ws.receive.side_effect = [ + _text({"type": "ping"}), + _text({"type": "exit", "exit_code": 0}), ] session = PtySession(ws, "cl-abc") events = list(session) @@ -356,9 +361,9 @@ class TestPtySessionPong: def test_no_pong_without_ping(self): ws = MagicMock() - ws.receive_text.side_effect = [ - json.dumps({"type": "output", "data": ""}), - json.dumps({"type": "exit", "exit_code": 0}), + ws.receive.side_effect = [ + _bytes(b""), + _text({"type": "exit", "exit_code": 0}), ] session = PtySession(ws, "cl-abc") list(session) @@ -431,13 +436,12 @@ class TestPtySessionSendConnect: class TestAsyncPtySession: @pytest.mark.asyncio - async def test_async_write_sends_base64(self): + async def test_async_write_sends_binary_frame(self): ws = AsyncMock() session = AsyncPtySession(ws, "cl-abc") await session.write(b"hello") - sent = json.loads(ws.send_text.call_args[0][0]) - assert sent["type"] == "input" - assert base64.b64decode(sent["data"]) == b"hello" + ws.send_bytes.assert_awaited_once_with(b"hello") + ws.send_text.assert_not_called() @pytest.mark.asyncio async def test_async_resize(self): @@ -486,9 +490,9 @@ class TestAsyncPtySession: @pytest.mark.asyncio async def test_async_ping_triggers_pong(self): ws = AsyncMock() - ws.receive_text.side_effect = [ - json.dumps({"type": "ping"}), - json.dumps({"type": "exit", "exit_code": 0}), + ws.receive.side_effect = [ + _text({"type": "ping"}), + _text({"type": "exit", "exit_code": 0}), ] session = AsyncPtySession(ws, "cl-abc") events = [e async for e in session] @@ -508,12 +512,11 @@ class TestAsyncPtySession: @pytest.mark.asyncio async def test_async_iteration(self): ws = AsyncMock() - messages = [ - json.dumps({"type": "started", "tag": "pty-xyz", "pid": 5}), - json.dumps({"type": "output", "data": base64.b64encode(b"hi").decode()}), - json.dumps({"type": "exit", "exit_code": 0}), + ws.receive.side_effect = [ + _text({"type": "started", "tag": "pty-xyz", "pid": 5}), + _bytes(b"hi"), + _text({"type": "exit", "exit_code": 0}), ] - ws.receive_text.side_effect = messages session = AsyncPtySession(ws, "cl-abc") events = [] async for event in session: @@ -522,6 +525,8 @@ class TestAsyncPtySession: assert events[0].type == PtyEventType.started assert session.tag == "pty-xyz" assert session.pid == 5 + assert events[1].type == PtyEventType.output + assert events[1].data == b"hi" assert events[2].type == PtyEventType.exit