v0.3.0 #15
@ -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\": \"<base64-encoded bytes>\"}\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\": \"<base64-encoded PTY\
|
||||
\ bytes>\"}\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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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,13 +157,20 @@ class PtySession:
|
||||
return self
|
||||
|
||||
def __next__(self) -> PtyEvent:
|
||||
while True:
|
||||
if self._done:
|
||||
raise StopIteration
|
||||
try:
|
||||
raw = self._ws.receive_text()
|
||||
msg = self._ws.receive()
|
||||
except _WS_CLOSED:
|
||||
raise StopIteration
|
||||
event = _parse_pty_event(json.loads(raw))
|
||||
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
|
||||
@ -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,13 +302,20 @@ class AsyncPtySession:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> PtyEvent:
|
||||
while True:
|
||||
if self._done:
|
||||
raise StopAsyncIteration
|
||||
try:
|
||||
raw = await self._ws.receive_text()
|
||||
msg = await self._ws.receive()
|
||||
except _WS_CLOSED:
|
||||
raise StopAsyncIteration
|
||||
event = _parse_pty_event(json.loads(raw))
|
||||
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
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user