wrenn/python-sdk
2
0
Fork 1
Code Issues Pull Requests Projects Wiki Activity

fix: send and receive PTY data over binary frames
All checks were successful
ci/woodpecker/push/unit Pipeline was successful
Details

Browse Source 
PTY input and output now travel as raw WebSocket binary frames to
match the updated server protocol. Control messages stay as JSON.
This commit is contained in:
Rafeed M. Bhuiyan
2026-06-23 04:34:43 +06:00
parent 875af78e7a
commit f3fb5a59c4
4 changed files with 129 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
api/openapi.yaml
View File

@ -987,13 +987,17 @@ paths:
\ \"rows\": 24,\n \"envs\": {\"TERM\": \"xterm-256color\"},\n \"cwd\": \"/home/user\",\n \"user\": \"user\"\n}\n\ \ \"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\ ```\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\ \ 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\": \"<base64-encoded bytes>\"}\n\ \ sends** (control messages, after session is established — JSON text frames):\n```json\n{\"type\": \"resize\", \"\
{\"type\": \"resize\", \"cols\": 120, \"rows\": 40}\n{\"type\": \"kill\"}\n```\n\n**Server sends**:\n```json\n{\"\ cols\": 120, \"rows\": 40}\n{\"type\": \"kill\"}\n```\n\nKeystroke input is sent as **raw binary WebSocket frames**\
type\": \"started\", \"tag\": \"pty-abc123de\", \"pid\": 42}\n{\"type\": \"output\", \"data\": \"<base64-encoded PTY\ \ containing the\nterminal bytes verbatim — not JSON, not base64. Every binary frame from\nthe client is treated as\
\ bytes>\"}\n{\"type\": \"exit\", \"exit_code\": 0}\n{\"type\": \"error\", \"data\": \"description\", \"fatal\": true}\n\ \ PTY input.\n\n**Server sends** (control messages — JSON text frames):\n```json\n{\"type\": \"started\", \"tag\"\
{\"type\": \"ping\"}\n```\n\nPTY data (input and output) is base64-encoded because it contains raw\nterminal bytes\ : \"pty-abc123de\", \"pid\": 42}\n{\"type\": \"exit\", \"exit_code\": 0}\n{\"type\": \"error\", \"data\": \"description\"\
\ (escape sequences, control codes) that are not valid UTF-8.\n\nSessions persist across WebSocket disconnections\ , \"fatal\": true}\n{\"type\": \"ping\"}\n```\n\nPTY output is sent as **raw binary WebSocket frames** containing\
\ — the process keeps\nrunning in the capsule. Use the `tag` from the \"started\" response to\nreconnect later.\n" \ 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: responses:
'101': '101':
description: WebSocket upgrade description: WebSocket upgrade

2
src/wrenn/models/_generated.py
View File

@ -1,6 +1,6 @@
# generated by datamodel-codegen: # generated by datamodel-codegen:
# filename: openapi.yaml # filename: openapi.yaml
# timestamp: 2026-06-22T19:04:03+00:00 # timestamp: 2026-06-22T22:24:45+00:00
from __future__ import annotations from __future__ import annotations
from typing import Annotated from typing import Annotated

115
src/wrenn/pty.py
View File

@ -1,6 +1,5 @@
from __future__ import annotations from __future__ import annotations
import base64
import json import json
from collections.abc import AsyncIterator, Iterator from collections.abc import AsyncIterator, Iterator
from enum import StrEnum from enum import StrEnum
@ -8,6 +7,7 @@ from typing import Any
import httpx_ws import httpx_ws
from pydantic import BaseModel from pydantic import BaseModel
from wsproto.events import BytesMessage, TextMessage
# A clean (``WebSocketDisconnect``) or abrupt (``WebSocketNetworkError``) close # A clean (``WebSocketDisconnect``) or abrupt (``WebSocketNetworkError``) close
# both mean the PTY stream has ended; iteration must stop on either. # both mean the PTY stream has ended; iteration must stop on either.
@ -31,7 +31,8 @@ class PtyEvent(BaseModel):
fatal: bool | None = None 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", "") msg_type = raw.get("type", "")
if msg_type == "started": if msg_type == "started":
return PtyEvent( return PtyEvent(
@ -39,10 +40,6 @@ def _parse_pty_event(raw: dict[str, Any]) -> PtyEvent:
pid=raw.get("pid"), pid=raw.get("pid"),
tag=raw.get("tag"), 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": if msg_type == "exit":
return PtyEvent(type=PtyEventType.exit, exit_code=raw.get("exit_code", -1)) return PtyEvent(type=PtyEventType.exit, exit_code=raw.get("exit_code", -1))
if msg_type == "error": if msg_type == "error":
@ -133,10 +130,10 @@ class PtySession:
"""Send raw bytes to the PTY stdin. """Send raw bytes to the PTY stdin.
Args: 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_bytes(data)
self._ws.send_text(json.dumps({"type": "input", "data": encoded}))
def resize(self, cols: int, rows: int) -> None: def resize(self, cols: int, rows: int) -> None:
"""Resize the PTY terminal. """Resize the PTY terminal.
@ -160,27 +157,34 @@ class PtySession:
return self return self
def __next__(self) -> PtyEvent: def __next__(self) -> PtyEvent:
if self._done: while True:
raise StopIteration if self._done:
try: raise StopIteration
raw = self._ws.receive_text() try:
except _WS_CLOSED: msg = self._ws.receive()
raise StopIteration except _WS_CLOSED:
event = _parse_pty_event(json.loads(raw)) raise StopIteration
if event.type == PtyEventType.started: if isinstance(msg, BytesMessage):
if event.tag is not None: return PtyEvent(type=PtyEventType.output, data=msg.data)
self._tag = event.tag if isinstance(msg, TextMessage):
if event.pid is not None: event = _parse_control_event(json.loads(msg.data))
self._pid = event.pid else:
if event.type == PtyEventType.ping: # Ignore other wsproto events (Ping/Pong/etc handled by httpx_ws).
self._send_pong() continue
if event.type == PtyEventType.exit: if event.type == PtyEventType.started:
self._done = True 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 return event
if event.type == PtyEventType.error and event.fatal:
self._done = True
return event
return event
def __enter__(self) -> PtySession: def __enter__(self) -> PtySession:
return self return self
@ -269,10 +273,10 @@ class AsyncPtySession:
"""Send raw bytes to the PTY stdin. """Send raw bytes to the PTY stdin.
Args: 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_bytes(data)
await self._ws.send_text(json.dumps({"type": "input", "data": encoded}))
async def resize(self, cols: int, rows: int) -> None: async def resize(self, cols: int, rows: int) -> None:
"""Resize the PTY terminal. """Resize the PTY terminal.
@ -298,27 +302,34 @@ class AsyncPtySession:
return self return self
async def __anext__(self) -> PtyEvent: async def __anext__(self) -> PtyEvent:
if self._done: while True:
raise StopAsyncIteration if self._done:
try: raise StopAsyncIteration
raw = await self._ws.receive_text() try:
except _WS_CLOSED: msg = await self._ws.receive()
raise StopAsyncIteration except _WS_CLOSED:
event = _parse_pty_event(json.loads(raw)) raise StopAsyncIteration
if event.type == PtyEventType.started: if isinstance(msg, BytesMessage):
if event.tag is not None: return PtyEvent(type=PtyEventType.output, data=msg.data)
self._tag = event.tag if isinstance(msg, TextMessage):
if event.pid is not None: event = _parse_control_event(json.loads(msg.data))
self._pid = event.pid else:
if event.type == PtyEventType.ping: # Ignore other wsproto events (Ping/Pong/etc handled by httpx_ws).
await self._send_pong() continue
if event.type == PtyEventType.exit: if event.type == PtyEventType.started:
self._done = True 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 return event
if event.type == PtyEventType.error and event.fatal:
self._done = True
return event
return event
async def __aenter__(self) -> AsyncPtySession: async def __aenter__(self) -> AsyncPtySession:
return self return self

103
tests/test_filesystem_pty.py
View File

@ -1,11 +1,11 @@
from __future__ import annotations from __future__ import annotations
import base64
import json import json
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
import respx import respx
from wsproto.events import BytesMessage, TextMessage
from wrenn.capsule import Capsule from wrenn.capsule import Capsule
from wrenn.models import FileEntry from wrenn.models import FileEntry
@ -13,9 +13,18 @@ from wrenn.pty import (
AsyncPtySession, AsyncPtySession,
PtyEventType, PtyEventType,
PtySession, 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" BASE = "https://app.wrenn.dev/api"
@ -226,50 +235,37 @@ class TestFilesExists:
class TestPtyEventParsing: class TestPtyEventParsing:
def test_started_event(self): def test_started_event(self):
raw = {"type": "started", "tag": "pty-a1b2c3d4", "pid": 42} 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.type == PtyEventType.started
assert event.pid == 42 assert event.pid == 42
assert event.tag == "pty-a1b2c3d4" 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): def test_exit_event(self):
raw = {"type": "exit", "exit_code": 0} raw = {"type": "exit", "exit_code": 0}
event = _parse_pty_event(raw) event = _parse_control_event(raw)
assert event.type == PtyEventType.exit assert event.type == PtyEventType.exit
assert event.exit_code == 0 assert event.exit_code == 0
def test_error_event(self): def test_error_event(self):
raw = {"type": "error", "data": "process not found", "fatal": True} 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.type == PtyEventType.error
assert event.data == "process not found" assert event.data == "process not found"
assert event.fatal is True assert event.fatal is True
def test_ping_event(self): def test_ping_event(self):
raw = {"type": "ping"} raw = {"type": "ping"}
event = _parse_pty_event(raw) event = _parse_control_event(raw)
assert event.type == PtyEventType.ping assert event.type == PtyEventType.ping
class TestPtySessionWrite: class TestPtySessionWrite:
def test_write_sends_base64_input(self): def test_write_sends_binary_frame(self):
ws = MagicMock() ws = MagicMock()
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
session.write(b"ls -la\n") session.write(b"ls -la\n")
sent = json.loads(ws.send_text.call_args[0][0]) ws.send_bytes.assert_called_once_with(b"ls -la\n")
assert sent["type"] == "input" ws.send_text.assert_not_called()
assert base64.b64decode(sent["data"]) == b"ls -la\n"
class TestPtySessionResize: class TestPtySessionResize:
@ -303,12 +299,11 @@ class TestPtySessionKill:
class TestPtySessionIteration: class TestPtySessionIteration:
def test_iter_yields_events_until_exit(self): def test_iter_yields_events_until_exit(self):
ws = MagicMock() ws = MagicMock()
messages = [ ws.receive.side_effect = [
json.dumps({"type": "started", "tag": "pty-abc12345", "pid": 1}), _text({"type": "started", "tag": "pty-abc12345", "pid": 1}),
json.dumps({"type": "output", "data": base64.b64encode(b"hello").decode()}), _bytes(b"hello"),
json.dumps({"type": "exit", "exit_code": 0}), _text({"type": "exit", "exit_code": 0}),
] ]
ws.receive_text.side_effect = messages
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
events = list(session) events = list(session)
assert len(events) == 3 assert len(events) == 3
@ -320,12 +315,22 @@ class TestPtySessionIteration:
assert events[2].type == PtyEventType.exit assert events[2].type == PtyEventType.exit
assert events[2].exit_code == 0 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): def test_iter_stops_on_fatal_error(self):
ws = MagicMock() ws = MagicMock()
messages = [ ws.receive.side_effect = [
json.dumps({"type": "error", "data": "fatal", "fatal": True}), _text({"type": "error", "data": "fatal", "fatal": True}),
] ]
ws.receive_text.side_effect = messages
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
events = list(session) events = list(session)
assert len(events) == 1 assert len(events) == 1
@ -335,7 +340,7 @@ class TestPtySessionIteration:
import httpx_ws import httpx_ws
ws = MagicMock() ws = MagicMock()
ws.receive_text.side_effect = httpx_ws.WebSocketDisconnect() ws.receive.side_effect = httpx_ws.WebSocketDisconnect()
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
events = list(session) events = list(session)
assert events == [] assert events == []
@ -344,9 +349,9 @@ class TestPtySessionIteration:
class TestPtySessionPong: class TestPtySessionPong:
def test_ping_triggers_pong(self): def test_ping_triggers_pong(self):
ws = MagicMock() ws = MagicMock()
ws.receive_text.side_effect = [ ws.receive.side_effect = [
json.dumps({"type": "ping"}), _text({"type": "ping"}),
json.dumps({"type": "exit", "exit_code": 0}), _text({"type": "exit", "exit_code": 0}),
] ]
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
events = list(session) events = list(session)
@ -356,9 +361,9 @@ class TestPtySessionPong:
def test_no_pong_without_ping(self): def test_no_pong_without_ping(self):
ws = MagicMock() ws = MagicMock()
ws.receive_text.side_effect = [ ws.receive.side_effect = [
json.dumps({"type": "output", "data": ""}), _bytes(b""),
json.dumps({"type": "exit", "exit_code": 0}), _text({"type": "exit", "exit_code": 0}),
] ]
session = PtySession(ws, "cl-abc") session = PtySession(ws, "cl-abc")
list(session) list(session)
@ -431,13 +436,12 @@ class TestPtySessionSendConnect:
class TestAsyncPtySession: class TestAsyncPtySession:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_write_sends_base64(self): async def test_async_write_sends_binary_frame(self):
ws = AsyncMock() ws = AsyncMock()
session = AsyncPtySession(ws, "cl-abc") session = AsyncPtySession(ws, "cl-abc")
await session.write(b"hello") await session.write(b"hello")
sent = json.loads(ws.send_text.call_args[0][0]) ws.send_bytes.assert_awaited_once_with(b"hello")
assert sent["type"] == "input" ws.send_text.assert_not_called()
assert base64.b64decode(sent["data"]) == b"hello"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_resize(self): async def test_async_resize(self):
@ -486,9 +490,9 @@ class TestAsyncPtySession:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_ping_triggers_pong(self): async def test_async_ping_triggers_pong(self):
ws = AsyncMock() ws = AsyncMock()
ws.receive_text.side_effect = [ ws.receive.side_effect = [
json.dumps({"type": "ping"}), _text({"type": "ping"}),
json.dumps({"type": "exit", "exit_code": 0}), _text({"type": "exit", "exit_code": 0}),
] ]
session = AsyncPtySession(ws, "cl-abc") session = AsyncPtySession(ws, "cl-abc")
events = [e async for e in session] events = [e async for e in session]
@ -508,12 +512,11 @@ class TestAsyncPtySession:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_iteration(self): async def test_async_iteration(self):
ws = AsyncMock() ws = AsyncMock()
messages = [ ws.receive.side_effect = [
json.dumps({"type": "started", "tag": "pty-xyz", "pid": 5}), _text({"type": "started", "tag": "pty-xyz", "pid": 5}),
json.dumps({"type": "output", "data": base64.b64encode(b"hi").decode()}), _bytes(b"hi"),
json.dumps({"type": "exit", "exit_code": 0}), _text({"type": "exit", "exit_code": 0}),
] ]
ws.receive_text.side_effect = messages
session = AsyncPtySession(ws, "cl-abc") session = AsyncPtySession(ws, "cl-abc")
events = [] events = []
async for event in session: async for event in session:
@ -522,6 +525,8 @@ class TestAsyncPtySession:
assert events[0].type == PtyEventType.started assert events[0].type == PtyEventType.started
assert session.tag == "pty-xyz" assert session.tag == "pty-xyz"
assert session.pid == 5 assert session.pid == 5
assert events[1].type == PtyEventType.output
assert events[1].data == b"hi"
assert events[2].type == PtyEventType.exit assert events[2].type == PtyEventType.exit