refactor: will this work better

scripts/export-dids.ts

@@ -185,9 +185,17 @@ let firehoseCursor: number | undefined = state?.firehose.cursor;
 	console.log(`  connecting to ${JETSTREAM_URL}`);
 	console.log(`  starting ${cursor || `<root>`}`);
 
-	const ws = createWebSocketStream<JetstreamEvent>(
-		() => JETSTREAM_URL + `?cursor=${cursor}` + `&wantedCollections=invalid.nsid.record`,
-	);
+	const ws = createWebSocketStream<JetstreamEvent>({
+		url() {
+			return JETSTREAM_URL + `?cursor=${cursor}` + `&wantedCollections=invalid.nsid.record`;
+		},
+		onOpen() {
+			console.log(`  connected`);
+		},
+		onClose() {
+			console.log(`  closed`);
+		},
+	});
 
 	for await (const data of ws) {
 		if (data.time_us > cursor) {

src/utils/websocket.ts

@@ -1,17 +1,39 @@
 import { WebSocket } from 'partysocket';
 
-export const createWebSocketStream = <T = any>(url: () => string) => {
+export const createWebSocketStream = <T = any>({
+	url,
+	onOpen,
+	onClose,
+}: {
+	url: () => string;
+	onOpen?: () => void;
+	onClose?: () => void;
+}) => {
 	let ws: WebSocket | undefined;
+	let closed = false;
 
 	return new ReadableStream<T>({
 		start(controller) {
 			ws = new WebSocket(url, null, { maxRetries: Infinity });
+			closed = false;
 
+			ws.addEventListener('open', () => {
+				onOpen?.();
+			});
+			ws.addEventListener('close', (ev) => {
+				onClose?.();
+
+				if (!closed && ev.wasClean) {
+					closed = true;
+					controller.close();
+				}
+			});
 			ws.addEventListener('message', (ev) => {
 				controller.enqueue(JSON.parse(ev.data));
 			});
 		},
 		cancel() {
+			closed = true;
 			ws?.close();
 		},
 	});