Add a real WebSocket builtin module

python/pythonmonkey/__init__.py

@@ -15,3 +15,4 @@
 require("timers")
 require("url")
 require("XMLHttpRequest")
+require("WebSocket")

python/pythonmonkey/builtin_modules/WebSocket-internal.py

@@ -0,0 +1,85 @@
+# @file     WebSocket-internal.py
+# @brief    internal helper functions for WebSocket, backed by aiohttp
+# @author   (local patch, this session)
+# @date     September 2026
+
+import asyncio
+import aiohttp
+from typing import Callable, Union, List
+
+
+async def wsConnect(
+    url: str,
+    protocols: List[str],
+    headers: dict,
+    onOpen: Callable[[Callable, Callable, Callable], None],
+    onMessage: Callable[[Union[str, bytearray], bool], None],
+    onError: Callable[[str], None],
+    onClose: Callable[[int, str], None],
+    debug: Callable[[str], Callable[..., None]],
+    /
+):
+  """
+  NOTE: onOpen is called with (sendText, sendBinary, close) as arguments,
+  rather than having this whole function return them once the coroutine
+  finishes -- the JS side needs those functions available in the SAME
+  synchronous callback that fires the 'open' event, not one microtask
+  later. A real client (dcp-client included) commonly sends its first
+  message immediately in reaction to 'open'; if the send/close functions
+  only became available via this async function's eventual return value
+  (resolved on a later microtask), that first send would silently race
+  against them not being wired up yet and get dropped -- confirmed
+  empirically as a silent hang (no error, no crash, connection just never
+  progresses) the first time this was tried without this fix.
+  """
+  session = aiohttp.ClientSession()
+
+  try:
+    ws = await session.ws_connect(url, protocols=tuple(protocols) if protocols else (), headers=headers or {})
+  except Exception as e:
+    try:
+      onError(str(e))
+    finally:
+      await session.close()
+    return
+
+  async def sendText(data: str):
+    if not ws.closed:
+      await ws.send_str(data)
+
+  async def sendBinary(data):
+    if not ws.closed:
+      await ws.send_bytes(bytes(data))
+
+  async def closeConn(code: int, reason: str):
+    if not ws.closed:
+      await ws.close(code=code or 1000, message=(reason or '').encode('utf-8'))
+
+  onOpen(sendText, sendBinary, closeConn)
+
+  close_code = 1006
+  close_reason = ''
+  try:
+    async for msg in ws:
+      debug('ws:io')('received', msg.type, 'len=', len(msg.data) if hasattr(msg.data, '__len__') else None)
+      if msg.type == aiohttp.WSMsgType.TEXT:
+        onMessage(msg.data, False)
+      elif msg.type == aiohttp.WSMsgType.BINARY:
+        onMessage(bytearray(msg.data), True)
+      elif msg.type == aiohttp.WSMsgType.ERROR:
+        onError(str(ws.exception()))
+      elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.CLOSED):
+        break
+    close_code = ws.close_code if ws.close_code is not None else 1000
+  except Exception as e:
+    onError(str(e))
+  finally:
+    try:
+      await session.close()
+    except Exception:
+      pass
+    onClose(close_code, close_reason)
+
+
+# Module exports
+exports['wsConnect'] = wsConnect  # type: ignore

python/pythonmonkey/builtin_modules/WebSocket.js

@@ -0,0 +1,214 @@
+/**
+ * @file     WebSocket.js
+ *           Implement the WebSocket API, backed by Python's aiohttp
+ *           WebSocket client (WebSocket-internal.py).
+ *
+ *           LOCAL PATCH (this session): pythonmonkey has never had a
+ *           WebSocket implementation, which is why dcp-client's
+ *           SocketIOTransport.buildOptions() hard-codes
+ *           `transports: platform !== 'pythonmonkey' ? ['websocket','polling'] : ['polling']`
+ *           -- every other platform upgrades to a persistent WebSocket
+ *           almost immediately, while pythonmonkey is forced to sustain
+ *           long-polling for the entire connection lifetime. That
+ *           sustained-polling code path is essentially unexercised by any
+ *           other real client, and packages.distributed.computer's own
+ *           backend has a real, reproducible session-routing bug in it
+ *           (a freshly-issued session id 404s on the very next polling
+ *           request, confirmed with plain curl/aiohttp/Node https/a bare
+ *           engine.io-client, no dcp-client or pythonmonkey involved).
+ *           Giving pythonmonkey a real WebSocket lets it take the exact
+ *           same well-exercised upgrade path every other platform takes,
+ *           sidestepping that bug entirely instead of working around it.
+ *
+ * @date     September 2026
+ */
+'use strict';
+
+const { EventTarget, Event } = require('event-target');
+const { DOMException } = require('dom-exception');
+const { URL } = require('url');
+const { wsConnect } = require('WebSocket-internal');
+const debug = globalThis.python.eval('__import__("pythonmonkey").bootstrap.require')('debug');
+
+// exposed
+class MessageEvent extends Event
+{
+  constructor(type, eventInitDict = {})
+  {
+    super(type);
+    this.data = eventInitDict.data;
+  }
+}
+
+// exposed
+class CloseEvent extends Event
+{
+  constructor(type, eventInitDict = {})
+  {
+    super(type);
+    this.code = eventInitDict.code ?? 1000;
+    this.reason = eventInitDict.reason ?? '';
+    this.wasClean = eventInitDict.wasClean ?? true;
+  }
+}
+
+/**
+ * Implement the `WebSocket` API according to the spec, backed by aiohttp.
+ * @see https://websockets.spec.whatwg.org/
+ */
+class WebSocket extends EventTarget
+{
+  /** @readonly */ static CONNECTING = 0;
+  /** @readonly */ static OPEN       = 1;
+  /** @readonly */ static CLOSING    = 2;
+  /** @readonly */ static CLOSED     = 3;
+
+  /** @readonly */ CONNECTING = 0;
+  /** @readonly */ OPEN       = 1;
+  /** @readonly */ CLOSING    = 2;
+  /** @readonly */ CLOSED     = 3;
+
+  // event handlers -- EventTarget#dispatchEvent auto-invokes these
+  onopen = null;
+  onmessage = null;
+  onerror = null;
+  onclose = null;
+
+  #readyState = WebSocket.CONNECTING;
+  #conn = null;
+  #url;
+  #protocol = '';
+  #sendBuffer = []; // messages queued before the underlying connection is ready
+
+  // engine.io-client's WS transport (addEventListeners()) unconditionally
+  // does `this.ws._socket.unref()` when its `autoUnref` option is set --
+  // real Node `ws` library sockets expose the underlying raw net.Socket as
+  // `._socket`, but a browser-style WebSocket has no such concept. dcp-client
+  // sets `autoUnref: true` for pythonmonkey (it's in env.js's
+  // referencedTimerPlatformList), so without this dummy property that call
+  // throws "can't access property unref, this.ws._socket is undefined" the
+  // instant the connection opens. A no-op unref() is exactly correct here:
+  // there is nothing OS-level for pythonmonkey to unref in the first place.
+  _socket = { unref() {}, ref() {} };
+
+  /**
+   * @param {string | URL} url
+   * @param {string | string[]} [protocols]
+   */
+  constructor(url, protocols)
+  {
+    super();
+    const parsedURL = new URL(url);
+    if (!['ws:', 'wss:'].includes(parsedURL.protocol))
+      throw new DOMException(`Invalid WebSocket URL scheme "${parsedURL.protocol}"`, 'SyntaxError');
+    this.#url = parsedURL.href;
+
+    const protoArray = protocols ? (Array.isArray(protocols) ? protocols : [protocols]) : [];
+
+    // aiohttp's ws_connect() expects a plain http(s):// URL, not ws(s)://
+    const httpURL = this.#url.replace(/^ws/, 'http');
+
+    debug('ws:connect')(`connecting to ${httpURL}`);
+
+    wsConnect(
+      httpURL,
+      protoArray,
+      {},
+      (sendText, sendBinary, closeFn) => // onOpen
+      {
+        // Wired up in the SAME synchronous callback that fires 'open' --
+        // see WebSocket-internal.py's docstring for why this matters (a
+        // message sent in reaction to 'open', which real clients commonly
+        // do, must never race ahead of these being available).
+        this.#conn = { sendText, sendBinary, close: closeFn };
+        this.#readyState = WebSocket.OPEN;
+        debug('ws:open')(`connected to ${this.#url}`);
+        for (const queued of this.#sendBuffer)
+          this.#doSend(queued);
+        this.#sendBuffer = [];
+        this.dispatchEvent(new Event('open'));
+      },
+      (data, isBinary) => // onMessage
+      {
+        const payload = isBinary ? new Uint8Array(data).buffer : data;
+        this.dispatchEvent(new MessageEvent('message', { data: payload }));
+      },
+      (message) => // onError
+      {
+        debug('ws:error')(message);
+        this.dispatchEvent(new Event('error'));
+      },
+      (code, reason) => // onClose
+      {
+        this.#readyState = WebSocket.CLOSED;
+        debug('ws:close')(`closed, code=${code} reason=${reason}`);
+        this.dispatchEvent(new CloseEvent('close', { code, reason, wasClean: code === 1000 }));
+      },
+      debug,
+    ).catch((e) =>
+    {
+      this.#readyState = WebSocket.CLOSED;
+      debug('ws:error')(String(e));
+      this.dispatchEvent(new Event('error'));
+      this.dispatchEvent(new CloseEvent('close', { code: 1006, reason: String(e), wasClean: false }));
+    });
+  }
+
+  get readyState() { return this.#readyState; }
+  get url() { return this.#url; }
+  get protocol() { return this.#protocol; }
+  get bufferedAmount() { return 0; } // not tracked
+
+  /**
+   * @param {string | ArrayBuffer | ArrayBufferView} data
+   */
+  send(data)
+  {
+    if (this.#readyState === WebSocket.CONNECTING)
+      throw new DOMException('WebSocket is still connecting (readyState CONNECTING)', 'InvalidStateError');
+    if (this.#readyState !== WebSocket.OPEN)
+      return; // per spec: silently discard if not OPEN
+    this.#doSend(data);
+  }
+
+  #doSend(data)
+  {
+    if (!this.#conn)
+      return;
+    if (typeof data === 'string')
+      this.#conn.sendText(data);
+    else if (data instanceof ArrayBuffer)
+      this.#conn.sendBinary(new Uint8Array(data));
+    else
+      this.#conn.sendBinary(data); // TypedArray/DataView
+  }
+
+  /**
+   * @param {number} [code]
+   * @param {string} [reason]
+   */
+  close(code = 1000, reason = '')
+  {
+    if (this.#readyState === WebSocket.CLOSING || this.#readyState === WebSocket.CLOSED)
+      return;
+    this.#readyState = WebSocket.CLOSING;
+    if (this.#conn)
+      this.#conn.close(code, reason);
+  }
+}
+
+/* A side-effect of loading this module is to add WebSocket and related
+ * symbols to the global object, matching XMLHttpRequest.js's convention,
+ * so real code (like dcp-client, once its platform check is relaxed) can
+ * use `new WebSocket(...)` directly with no require() needed.
+ */
+if (!globalThis.WebSocket)
+  globalThis.WebSocket = WebSocket;
+if (!globalThis.MessageEvent)
+  globalThis.MessageEvent = MessageEvent;
+if (!globalThis.CloseEvent)
+  globalThis.CloseEvent = CloseEvent;
+
+exports.WebSocket = WebSocket;
+exports.MessageEvent = MessageEvent;
+exports.CloseEvent = CloseEvent;