Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635"""One reader owns the socket. It never awaits a consumer or user callback."""
from __future__ import annotations
from typing import SupportsIndex
import asyncioimport base64import copyfrom dataclasses import dataclassimport hashlibimport jsonimport mathimport osimport structimport urllib.parsefrom collections.abc import Callablefrom typing import Any
from .errors import ( BrowserError, CommandTimeout, ConnectionLost, EventOverflow, ProtocolError, WaitTimeout,)
JSON = dict[str, Any]
def positive(value: float, name: str) -> float: if isinstance(value, bool) or not math.isfinite(value) or value <= 0: raise ValueError(f"{name} must be finite and positive") return value
@dataclassclass CloseFrame: code: int reason: str
class ConnectionClosed(Exception): def __init__(self, code: int = 1006, reason: str = "") -> None: super().__init__(f"WebSocket closed: {code} {reason}") self.code = code self.reason = reason self.rcvd = CloseFrame(code, reason) self.sent: CloseFrame | None = None
class SimpleWebSocketClient: def __init__( self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, max_message_bytes: int = 16 * 1024 * 1024, ) -> None: self.reader = reader self.writer = writer self.max_message_bytes = max_message_bytes self._closed = False self._close_sent = False self.close_code: int | None = None self.close_reason: str = ""
@classmethod async def connect( cls, url: str, *, headers: dict[str, str] | None = None, timeout: float = 15, max_message_bytes: int = 16 * 1024 * 1024, ) -> SimpleWebSocketClient: parsed = urllib.parse.urlsplit(url) if parsed.scheme not in ("ws", "wss"): raise ValueError(f"Expected ws:// or wss:// URL, got: {url!r}") hostname = parsed.hostname or "127.0.0.1" port = parsed.port or (443 if parsed.scheme == "wss" else 80) ssl_ctx = True if parsed.scheme == "wss" else None
async with asyncio.timeout(timeout): reader, writer = await asyncio.open_connection( hostname, port, ssl=ssl_ctx, limit=max_message_bytes ) key = base64.b64encode(os.urandom(16)).decode("ascii") path = parsed.path or "/" if parsed.query: path = f"{path}?{parsed.query}"
host_header = f"{hostname}:{port}" if parsed.port else hostname lines = [ f"GET {path} HTTP/1.1", f"Host: {host_header}", "Upgrade: websocket", "Connection: Upgrade", f"Sec-WebSocket-Key: {key}", "Sec-WebSocket-Version: 13", ] if headers: for k, v in headers.items(): lines.append(f"{k}: {v}") raw_req = "\r\n".join(lines) + "\r\n\r\n" writer.write(raw_req.encode("ascii")) await writer.drain()
status_line = await reader.readline() if not status_line: writer.close() raise ConnectionClosed(1006, "EOF during WebSocket handshake") parts = status_line.decode("latin1").split(" ", 2) if len(parts) < 2 or parts[1] != "101": body = await reader.read(4096) writer.close() raise ConnectionClosed( 1002, f"Handshake failed: {status_line.decode('latin1').strip()} {body.decode('latin1', errors='replace')}", )
resp_headers: dict[str, str] = {} while True: line = await reader.readline() if line in (b"\r\n", b"\n", b""): break hparts = line.decode("latin1").split(":", 1) if len(hparts) == 2: resp_headers[hparts[0].strip().lower()] = hparts[1].strip()
expected_accept = base64.b64encode( hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii") ).digest() ).decode("ascii") accept = resp_headers.get("sec-websocket-accept") if accept != expected_accept: writer.close() raise ConnectionClosed( 1002, f"Sec-WebSocket-Accept mismatch: expected {expected_accept}, got {accept}", )
return cls(reader, writer, max_message_bytes=max_message_bytes)
async def send(self, data: str) -> None: if self._closed: raise ConnectionClosed(self.close_code or 1000, self.close_reason) payload = data.encode("utf-8") mask_key = os.urandom(4) length = len(payload)
header = bytearray([0x80 | 0x1]) # FIN=1, Opcode=TEXT if length <= 125: header.append(0x80 | length) elif length <= 65535: header.append(0x80 | 126) header.extend(struct.pack("!H", length)) else: header.append(0x80 | 127) header.extend(struct.pack("!Q", length)) header.extend(mask_key)
masked = bytearray(payload) for i in range(length): masked[i] ^= mask_key[i % 4]
self.writer.write(header + masked) await self.writer.drain()
async def recv_message(self) -> str: fragments: list[bytes] = [] total_len = 0
while True: if self._closed: raise ConnectionClosed(self.close_code or 1000, self.close_reason)
try: head = await self.reader.readexactly(2) except ( asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, ) as exc: self._closed = True raise ConnectionClosed(1006, "Connection closed abruptly") from exc
b1, b2 = head[0], head[1] fin = bool(b1 & 0x80) opcode = b1 & 0x0F is_masked = bool(b2 & 0x80) payload_len = b2 & 0x7F
if payload_len == 126: ext = await self.reader.readexactly(2) payload_len = struct.unpack("!H", ext)[0] elif payload_len == 127: ext = await self.reader.readexactly(8) payload_len = struct.unpack("!Q", ext)[0]
total_len += payload_len if total_len > self.max_message_bytes: await self.close(1009, "Message exceeds size limit") exc = ConnectionClosed( 1009, f"Message exceeded max_message_bytes {self.max_message_bytes}" ) exc.sent = CloseFrame(1009, "Message exceeds size limit") raise exc
if is_masked: mask_key = await self.reader.readexactly(4)
payload = await self.reader.readexactly(payload_len) if is_masked: unmasked = bytearray(payload) for i in range(len(unmasked)): unmasked[i] ^= mask_key[i % 4] payload = bytes(unmasked)
if opcode == 0x8: # CLOSE code = 1000 reason = "" if len(payload) >= 2: code = struct.unpack("!H", payload[:2])[0] reason = payload[2:].decode("utf-8", errors="replace") self.close_code = code self.close_reason = reason self._closed = True if not self._close_sent: try: self._close_sent = True close_resp = bytearray([0x80 | 0x8, 0x80]) mask = os.urandom(4) close_resp.extend(mask) self.writer.write(close_resp) await self.writer.drain() self.writer.close() except Exception: pass raise ConnectionClosed(code, reason) elif opcode == 0x9: # PING pong = bytearray([0x80 | 0xA, 0x80 | (len(payload) & 0x7F)]) mask = os.urandom(4) pong.extend(mask) masked_pong = bytearray(payload) for i in range(len(masked_pong)): masked_pong[i] ^= mask[i % 4] self.writer.write(pong + masked_pong) await self.writer.drain() continue elif opcode == 0xA: # PONG continue elif opcode in (0x1, 0x2): # TEXT or BINARY fragments.append(payload) if fin: break elif opcode == 0x0: # CONTINUATION fragments.append(payload) if fin: break else: raise ConnectionClosed(1002, f"Unknown WebSocket opcode: {opcode}")
raw_bytes = b"".join(fragments) return raw_bytes.decode("utf-8", errors="replace")
def __aiter__(self) -> SimpleWebSocketClient: return self
async def __anext__(self) -> str: try: return await self.recv_message() except ConnectionClosed: raise StopAsyncIteration
async def close(self, code: int = 1000, reason: str = "") -> None: if self._closed and self._close_sent: return self._closed = True if not self._close_sent: self._close_sent = True try: reason_bytes = reason.encode("utf-8")[:123] payload = struct.pack("!H", code) + reason_bytes mask = os.urandom(4) length = len(payload) header = bytearray([0x80 | 0x8, 0x80 | length]) header.extend(mask) masked = bytearray(payload) for i in range(length): masked[i] ^= mask[i % 4] self.writer.write(header + masked) await self.writer.drain() self.writer.close() await self.writer.wait_closed() except Exception: pass
class Subscription: def __init__( self, connection: Connection, method: str, session_id: str | None, capacity: int, max_bytes: int, ) -> None: if capacity < 1 or max_bytes < 1: raise ValueError("event capacity and max_bytes must be positive") self.connection = connection self.method = method self.session_id = session_id self.capacity = capacity self.max_bytes = max_bytes self._queue: asyncio.Queue[tuple[JSON, int] | None] = asyncio.Queue(capacity) self._bytes = 0 self._error: BrowserError | None = None self._closed = False
def _push(self, params: JSON, size: int) -> None: if self._closed: return if self._queue.full() or self._bytes + size > self.max_bytes: self._fail( EventOverflow( "Event buffer overflowed; this subscription is no longer complete.", method=self.method, capacity=self.capacity, max_bytes=self.max_bytes, recovery="Create a new subscription; do not assume no event occurred.", ) ) return self._queue.put_nowait((copy.deepcopy(params), size)) self._bytes += size
def _fail(self, error: BrowserError) -> None: self._error = error self.close()
def close(self) -> None: self._closed = True self.connection._subscriptions.discard(self) while not self._queue.empty(): self._queue.get_nowait() self._bytes = 0 self._queue.put_nowait(None)
async def next( self, *, where: Callable[[JSON], bool] | None = None, timeout: float = 10, ) -> JSON: """Return one payload. A predicate consumes unmatched events on this stream.""" positive(timeout, "timeout") try: async with asyncio.timeout(timeout): while True: if self._error is not None: raise self._error if self._closed: raise ConnectionLost("Event subscription is closed.") item = await self._queue.get() if item is None: if self._error is not None: raise self._error raise ConnectionLost("Event subscription is closed.") params, size = item self._bytes -= size if where is None or where(params): return params except TimeoutError as exc: raise WaitTimeout( "No matching event before the deadline.", method=self.method ) from exc
def drain(self) -> list[JSON]: """Take the currently buffered payloads without waiting.""" if self._error is not None: raise self._error result = [] while not self._queue.empty(): item = self._queue.get_nowait() if item is not None: params, size = item self._bytes -= size result.append(params) return result
def __reduce_ex__(self, protocol: SupportsIndex, /) -> Any: raise TypeError( "Live event subscriptions cannot be saved; drain them to ordinary data." )
class Connection: def __init__( self, socket: SimpleWebSocketClient, *, timeout: float, max_message_bytes: int = 16 * 1024 * 1024, ) -> None: self.socket = socket self.max_message_bytes = max_message_bytes self._detached_sessions: set[str] = set() self.timeout = positive(timeout, "timeout") self._next_id = 0 self._pending: dict[int, asyncio.Future[JSON]] = {} self._subscriptions: set[Subscription] = set() self._pending_sessions: dict[int, str | None] = {} self._closed: ConnectionLost | None = None self.on_event: Callable[[str, JSON, str | None], None] | None = None self._reader = asyncio.create_task(self._read(), name="browser-cdp-reader")
@classmethod async def open( cls, url: str, *, timeout: float = 15, headers: dict[str, str] | None = None, max_message_bytes: int = 16 * 1024 * 1024, ) -> Connection: positive(timeout, "timeout") if max_message_bytes < 1024: raise ValueError("max_message_bytes must be at least 1024") try: socket = await SimpleWebSocketClient.connect( url, headers=headers, timeout=timeout, max_message_bytes=max_message_bytes, ) except Exception as exc: raise ConnectionLost( "Could not open the CDP WebSocket.", cause=type(exc).__name__ ) from exc return cls(socket, timeout=timeout, max_message_bytes=max_message_bytes)
@property def closed(self) -> bool: return self._closed is not None
def subscribe( self, method: str, session_id: str | None, *, capacity: int = 256, max_bytes: int = 2 * 1024 * 1024, ) -> Subscription: if self._closed is not None: raise self._closed if session_id in self._detached_sessions: raise ConnectionLost( "Target session detached; use await page.reattach().", session_id=session_id, ) if len(self._subscriptions) >= 64: raise ValueError("At most 64 live event subscriptions per connection") sub = Subscription(self, method, session_id, capacity, max_bytes) self._subscriptions.add(sub) return sub
async def send( self, method: str, params: JSON | None = None, *, session_id: str | None = None, timeout: float | None = None, ) -> JSON: duration = positive(self.timeout if timeout is None else timeout, "timeout") if self._closed is not None: raise self._closed if session_id in self._detached_sessions: raise ConnectionLost( "Target session detached; use await page.reattach().", session_id=session_id, ) if not isinstance(method, str) or "." not in method: raise ValueError("method must be a CDP Domain.command name") if len(self._pending) >= 256: raise ValueError("At most 256 in-flight CDP commands per connection") self._next_id += 1 request_id = self._next_id request = {"id": request_id, "method": method, "params": params or {}} if session_id is not None: request["sessionId"] = session_id encoded = json.dumps(request, allow_nan=False) future: asyncio.Future[JSON] = asyncio.get_running_loop().create_future() self._pending[request_id] = future self._pending_sessions[request_id] = session_id sent = False try: async with asyncio.timeout(duration): sent = True await self.socket.send(encoded) reply = await future if "error" in reply: error = reply["error"] raise ProtocolError( error.get("message", "CDP command failed"), method=method, code=error.get("code"), data=error.get("data"), ) return reply.get("result", {}) except TimeoutError as exc: raise CommandTimeout( "CDP command exceeded its deadline; it was not replayed.", method=method, may_have_executed=sent, timeout=duration, ) from exc except ConnectionClosed as exc: raise ConnectionLost( "CDP connection closed.", method=method, may_have_executed=sent ) from exc finally: self._pending.pop(request_id, None) self._pending_sessions.pop(request_id, None) if not future.done(): future.cancel() elif not future.cancelled(): future.exception()
def detach_session(self, session_id: str | None) -> None: if session_id is None: return self._detached_sessions.add(session_id) error = ConnectionLost( "Target session detached; use await page.reattach().", session_id=session_id ) for request_id, sid in tuple(self._pending_sessions.items()): if sid == session_id: future = self._pending.get(request_id) if future is not None and not future.done(): future.set_exception(error) for sub in tuple(self._subscriptions): if sub.session_id == session_id: sub._fail(error)
def _fail(self, error: ConnectionLost) -> None: if self._closed is not None: return self._closed = error for future in self._pending.values(): if not future.done(): future.set_exception(error) for sub in tuple(self._subscriptions): sub._fail(error)
async def _read(self) -> None: try: async for raw in self.socket: message = json.loads(raw) if not isinstance(message, dict): raise ValueError("CDP message is not an object") if "id" in message: future = self._pending.get(message["id"]) if future is not None and not future.done(): future.set_result(message) continue method = message.get("method", "") params = message.get("params", {}) session_id = message.get("sessionId") if self.on_event is not None: self.on_event(method, params, session_id) for sub in tuple(self._subscriptions): if sub.method == method and sub.session_id == session_id: sub._push(params, len(raw)) except asyncio.CancelledError: raise except Exception as exc: details: JSON = { "cause": type(exc).__name__, "max_message_bytes": self.max_message_bytes, } if isinstance(exc, ConnectionClosed): close = ( exc.sent if exc.sent and exc.sent.code == 1009 else (exc.rcvd or exc.sent) ) if close is not None: details.update(close_code=close.code, close_reason=close.reason) if details.get("close_code") == 1009: message = ( "CDP message exceeded max_message_bytes; event history has a gap. " "Reattach explicitly with a larger finite max_message_bytes budget." ) else: message = "CDP reader stopped; event history has a gap. Use await page.reattach()." self._fail(ConnectionLost(message, **details)) finally: self._fail(ConnectionLost("CDP disconnected; event history has a gap.")) await self.socket.close()
async def close(self) -> None: self._fail(ConnectionLost("CDP connection explicitly disconnected.")) try: await self.socket.close() finally: self._reader.cancel() await asyncio.gather(self._reader, return_exceptions=True)
def __reduce_ex__(self, protocol: SupportsIndex, /) -> Any: raise TypeError( "A live CDP connection cannot be saved; save a reconnect descriptor." )
class CDPSession: """A session-bound escape hatch. Raw commands bypass high-level safety checks."""
def __init__(self, connection: Connection, session_id: str | None = None) -> None: self.connection = connection self.session_id = session_id
async def send( self, method: str, params: JSON | None = None, *, timeout: float | None = None, ) -> JSON: return await self.connection.send( method, params, session_id=self.session_id, timeout=timeout )