v0.3.0 #15

Merged
pptx704 merged 64 commits from dev into main 2026-06-25 22:11:26 +00:00
17 changed files with 2680 additions and 3717 deletions
Showing only changes of commit f3fb5a59c4 - Show all commits

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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