Skip to content

Commit a893737

Browse files
dan-distributiveclaude
dan-distributive
and
claude
committed
Add a real WebSocket builtin module
pythonmonkey has never had a WebSocket implementation. dcp-client's own SocketIOTransport.buildOptions() explicitly special-cases this: every platform except pythonmonkey gets transports: ['websocket', 'polling'] and upgrades to a persistent WebSocket within the first round trip or two; pythonmonkey is hard-restricted to ['polling'] forever, since `typeof WebSocket` was previously `undefined`. That forces pythonmonkey to be the only client that ever *sustains* HTTP long-polling for an entire connection's lifetime against services like DCP's package manager - a code path essentially unexercised by any other real client, where a real, reproducible session-routing bug lives (a freshly issued session id doesn't reliably route to the same backend on the very next polling request; confirmed independent of pythonmonkey/dcp-client with plain curl, aiohttp, and Node https). Rather than work around that bug, this gives pythonmonkey a real WebSocket implementation - a JS class matching the DOM API (builtin_modules/WebSocket.js), backed by a Python aiohttp-based implementation (builtin_modules/WebSocket-internal.py), following the exact existing pattern XMLHttpRequest already uses (loaded dynamically via require() at pythonmonkey import time, no C++ rebuild needed). Two real bugs were found and fixed while building it (documented inline in WebSocket.js/WebSocket-internal.py's own docstrings): - A race between the 'open' event firing and send/close functions being available, since real clients often send their first message immediately on 'open'. - engine.io-client's real WS transport unconditionally calling `this.ws._socket.unref()` (a Node `ws`-library-shaped assumption that doesn't hold for a browser-style WebSocket), worked around with a harmless no-op `_socket` property. This alone doesn't change any existing behavior - dcp-client's platform check still needs its pythonmonkey exclusion relaxed on its own side for anything to actually use this. Verified via a real end-to-end DCP job (job.exec(), real network dispatch to the demo/dcp compute group) after also relaxing that check locally. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 943face commit a893737

3 files changed

Lines changed: 300 additions & 0 deletions

File tree

‎python/pythonmonkey/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@
1515
require("timers")
1616
require("url")
1717
require("XMLHttpRequest")
18+
require("WebSocket")

‎python/pythonmonkey/builtin_modules/WebSocket-internal.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
# @file WebSocket-internal.py
2+
# @brief internal helper functions for WebSocket, backed by aiohttp
3+
# @author (local patch, this session)
4+
# @date September 2026
5+
6+
import asyncio
7+
import aiohttp
8+
from typing import Callable, Union, List
9+
10+
11+
async def wsConnect(
12+
url: str,
13+
protocols: List[str],
14+
headers: dict,
15+
onOpen: Callable[[Callable, Callable, Callable], None],
16+
onMessage: Callable[[Union[str, bytearray], bool], None],
17+
onError: Callable[[str], None],
18+
onClose: Callable[[int, str], None],
19+
debug: Callable[[str], Callable[..., None]],
20+
/
21+
):
22+
"""
23+
NOTE: onOpen is called with (sendText, sendBinary, close) as arguments,
24+
rather than having this whole function return them once the coroutine
25+
finishes -- the JS side needs those functions available in the SAME
26+
synchronous callback that fires the 'open' event, not one microtask
27+
later. A real client (dcp-client included) commonly sends its first
28+
message immediately in reaction to 'open'; if the send/close functions
29+
only became available via this async function's eventual return value
30+
(resolved on a later microtask), that first send would silently race
31+
against them not being wired up yet and get dropped -- confirmed
32+
empirically as a silent hang (no error, no crash, connection just never
33+
progresses) the first time this was tried without this fix.
34+
"""
35+
session = aiohttp.ClientSession()
36+
37+
try:
38+
ws = await session.ws_connect(url, protocols=tuple(protocols) if protocols else (), headers=headers or {})
39+
except Exception as e:
40+
try:
41+
onError(str(e))
42+
finally:
43+
await session.close()
44+
return
45+
46+
async def sendText(data: str):
47+
if not ws.closed:
48+
await ws.send_str(data)
49+
50+
async def sendBinary(data):
51+
if not ws.closed:
52+
await ws.send_bytes(bytes(data))
53+
54+
async def closeConn(code: int, reason: str):
55+
if not ws.closed:
56+
await ws.close(code=code or 1000, message=(reason or '').encode('utf-8'))
57+
58+
onOpen(sendText, sendBinary, closeConn)
59+
60+
close_code = 1006
61+
close_reason = ''
62+
try:
63+
async for msg in ws:
64+
debug('ws:io')('received', msg.type, 'len=', len(msg.data) if hasattr(msg.data, '__len__') else None)
65+
if msg.type == aiohttp.WSMsgType.TEXT:
66+
onMessage(msg.data, False)
67+
elif msg.type == aiohttp.WSMsgType.BINARY:
68+
onMessage(bytearray(msg.data), True)
69+
elif msg.type == aiohttp.WSMsgType.ERROR:
70+
onError(str(ws.exception()))
71+
elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.CLOSED):
72+
break
73+
close_code = ws.close_code if ws.close_code is not None else 1000
74+
except Exception as e:
75+
onError(str(e))
76+
finally:
77+
try:
78+
await session.close()
79+
except Exception:
80+
pass
81+
onClose(close_code, close_reason)
82+
83+
84+
# Module exports
85+
exports['wsConnect'] = wsConnect # type: ignore

‎python/pythonmonkey/builtin_modules/WebSocket.js‎

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
/**
2+
* @file WebSocket.js
3+
* Implement the WebSocket API, backed by Python's aiohttp
4+
* WebSocket client (WebSocket-internal.py).
5+
*
6+
* LOCAL PATCH (this session): pythonmonkey has never had a
7+
* WebSocket implementation, which is why dcp-client's
8+
* SocketIOTransport.buildOptions() hard-codes
9+
* `transports: platform !== 'pythonmonkey' ? ['websocket','polling'] : ['polling']`
10+
* -- every other platform upgrades to a persistent WebSocket
11+
* almost immediately, while pythonmonkey is forced to sustain
12+
* long-polling for the entire connection lifetime. That
13+
* sustained-polling code path is essentially unexercised by any
14+
* other real client, and packages.distributed.computer's own
15+
* backend has a real, reproducible session-routing bug in it
16+
* (a freshly-issued session id 404s on the very next polling
17+
* request, confirmed with plain curl/aiohttp/Node https/a bare
18+
* engine.io-client, no dcp-client or pythonmonkey involved).
19+
* Giving pythonmonkey a real WebSocket lets it take the exact
20+
* same well-exercised upgrade path every other platform takes,
21+
* sidestepping that bug entirely instead of working around it.
22+
*
23+
* @date September 2026
24+
*/
25+
'use strict';
26+
27+
const { EventTarget, Event } = require('event-target');
28+
const { DOMException } = require('dom-exception');
29+
const { URL } = require('url');
30+
const { wsConnect } = require('WebSocket-internal');
31+
const debug = globalThis.python.eval('__import__("pythonmonkey").bootstrap.require')('debug');
32+
33+
// exposed
34+
class MessageEvent extends Event
35+
{
36+
constructor(type, eventInitDict = {})
37+
{
38+
super(type);
39+
this.data = eventInitDict.data;
40+
}
41+
}
42+
43+
// exposed
44+
class CloseEvent extends Event
45+
{
46+
constructor(type, eventInitDict = {})
47+
{
48+
super(type);
49+
this.code = eventInitDict.code ?? 1000;
50+
this.reason = eventInitDict.reason ?? '';
51+
this.wasClean = eventInitDict.wasClean ?? true;
52+
}
53+
}
54+
55+
/**
56+
* Implement the `WebSocket` API according to the spec, backed by aiohttp.
57+
* @see https://websockets.spec.whatwg.org/
58+
*/
59+
class WebSocket extends EventTarget
60+
{
61+
/** @readonly */ static CONNECTING = 0;
62+
/** @readonly */ static OPEN = 1;
63+
/** @readonly */ static CLOSING = 2;
64+
/** @readonly */ static CLOSED = 3;
65+
66+
/** @readonly */ CONNECTING = 0;
67+
/** @readonly */ OPEN = 1;
68+
/** @readonly */ CLOSING = 2;
69+
/** @readonly */ CLOSED = 3;
70+
71+
// event handlers -- EventTarget#dispatchEvent auto-invokes these
72+
onopen = null;
73+
onmessage = null;
74+
onerror = null;
75+
onclose = null;
76+
77+
#readyState = WebSocket.CONNECTING;
78+
#conn = null;
79+
#url;
80+
#protocol = '';
81+
#sendBuffer = []; // messages queued before the underlying connection is ready
82+
83+
// engine.io-client's WS transport (addEventListeners()) unconditionally
84+
// does `this.ws._socket.unref()` when its `autoUnref` option is set --
85+
// real Node `ws` library sockets expose the underlying raw net.Socket as
86+
// `._socket`, but a browser-style WebSocket has no such concept. dcp-client
87+
// sets `autoUnref: true` for pythonmonkey (it's in env.js's
88+
// referencedTimerPlatformList), so without this dummy property that call
89+
// throws "can't access property unref, this.ws._socket is undefined" the
90+
// instant the connection opens. A no-op unref() is exactly correct here:
91+
// there is nothing OS-level for pythonmonkey to unref in the first place.
92+
_socket = { unref() {}, ref() {} };
93+
94+
/**
95+
* @param {string | URL} url
96+
* @param {string | string[]} [protocols]
97+
*/
98+
constructor(url, protocols)
99+
{
100+
super();
101+
const parsedURL = new URL(url);
102+
if (!['ws:', 'wss:'].includes(parsedURL.protocol))
103+
throw new DOMException(`Invalid WebSocket URL scheme "${parsedURL.protocol}"`, 'SyntaxError');
104+
this.#url = parsedURL.href;
105+
106+
const protoArray = protocols ? (Array.isArray(protocols) ? protocols : [protocols]) : [];
107+
108+
// aiohttp's ws_connect() expects a plain http(s):// URL, not ws(s)://
109+
const httpURL = this.#url.replace(/^ws/, 'http');
110+
111+
debug('ws:connect')(`connecting to ${httpURL}`);
112+
113+
wsConnect(
114+
httpURL,
115+
protoArray,
116+
{},
117+
(sendText, sendBinary, closeFn) => // onOpen
118+
{
119+
// Wired up in the SAME synchronous callback that fires 'open' --
120+
// see WebSocket-internal.py's docstring for why this matters (a
121+
// message sent in reaction to 'open', which real clients commonly
122+
// do, must never race ahead of these being available).
123+
this.#conn = { sendText, sendBinary, close: closeFn };
124+
this.#readyState = WebSocket.OPEN;
125+
debug('ws:open')(`connected to ${this.#url}`);
126+
for (const queued of this.#sendBuffer)
127+
this.#doSend(queued);
128+
this.#sendBuffer = [];
129+
this.dispatchEvent(new Event('open'));
130+
},
131+
(data, isBinary) => // onMessage
132+
{
133+
const payload = isBinary ? new Uint8Array(data).buffer : data;
134+
this.dispatchEvent(new MessageEvent('message', { data: payload }));
135+
},
136+
(message) => // onError
137+
{
138+
debug('ws:error')(message);
139+
this.dispatchEvent(new Event('error'));
140+
},
141+
(code, reason) => // onClose
142+
{
143+
this.#readyState = WebSocket.CLOSED;
144+
debug('ws:close')(`closed, code=${code} reason=${reason}`);
145+
this.dispatchEvent(new CloseEvent('close', { code, reason, wasClean: code === 1000 }));
146+
},
147+
debug,
148+
).catch((e) =>
149+
{
150+
this.#readyState = WebSocket.CLOSED;
151+
debug('ws:error')(String(e));
152+
this.dispatchEvent(new Event('error'));
153+
this.dispatchEvent(new CloseEvent('close', { code: 1006, reason: String(e), wasClean: false }));
154+
});
155+
}
156+
157+
get readyState() { return this.#readyState; }
158+
get url() { return this.#url; }
159+
get protocol() { return this.#protocol; }
160+
get bufferedAmount() { return 0; } // not tracked
161+
162+
/**
163+
* @param {string | ArrayBuffer | ArrayBufferView} data
164+
*/
165+
send(data)
166+
{
167+
if (this.#readyState === WebSocket.CONNECTING)
168+
throw new DOMException('WebSocket is still connecting (readyState CONNECTING)', 'InvalidStateError');
169+
if (this.#readyState !== WebSocket.OPEN)
170+
return; // per spec: silently discard if not OPEN
171+
this.#doSend(data);
172+
}
173+
174+
#doSend(data)
175+
{
176+
if (!this.#conn)
177+
return;
178+
if (typeof data === 'string')
179+
this.#conn.sendText(data);
180+
else if (data instanceof ArrayBuffer)
181+
this.#conn.sendBinary(new Uint8Array(data));
182+
else
183+
this.#conn.sendBinary(data); // TypedArray/DataView
184+
}
185+
186+
/**
187+
* @param {number} [code]
188+
* @param {string} [reason]
189+
*/
190+
close(code = 1000, reason = '')
191+
{
192+
if (this.#readyState === WebSocket.CLOSING || this.#readyState === WebSocket.CLOSED)
193+
return;
194+
this.#readyState = WebSocket.CLOSING;
195+
if (this.#conn)
196+
this.#conn.close(code, reason);
197+
}
198+
}
199+
200+
/* A side-effect of loading this module is to add WebSocket and related
201+
* symbols to the global object, matching XMLHttpRequest.js's convention,
202+
* so real code (like dcp-client, once its platform check is relaxed) can
203+
* use `new WebSocket(...)` directly with no require() needed.
204+
*/
205+
if (!globalThis.WebSocket)
206+
globalThis.WebSocket = WebSocket;
207+
if (!globalThis.MessageEvent)
208+
globalThis.MessageEvent = MessageEvent;
209+
if (!globalThis.CloseEvent)
210+
globalThis.CloseEvent = CloseEvent;
211+
212+
exports.WebSocket = WebSocket;
213+
exports.MessageEvent = MessageEvent;
214+
exports.CloseEvent = CloseEvent;

0 commit comments

Comments
 (0)