@nevuamarkets/poly-websockets 0.0.1

package/dist/modules/GroupRegistry.js

@@ -0,0 +1,261 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.GroupRegistry = void 0;
+const async_mutex_1 = require("async-mutex");
+const lodash_1 = __importDefault(require("lodash"));
+const uuid_1 = require("uuid");
+const WebSocketSubscriptions_1 = require("../types/WebSocketSubscriptions");
+const logger_1 = require("../logger");
+/*
+ * Global group store and mutex, intentionally **not** exported anymore to prevent
+ * accidental external mutation.  All access should go through the helper methods
+ * on GroupRegistry instead.
+ */
+const wsGroups = [];
+const wsGroupsMutex = new async_mutex_1.Mutex();
+class GroupRegistry {
+    /**
+     * Atomic mutate helper.
+     *
+     * @param fn - The function to run atomically.
+     * @returns The result of the function.
+     */
+    async mutate(fn) {
+        const release = await wsGroupsMutex.acquire();
+        try {
+            return await fn(wsGroups);
+        }
+        finally {
+            release();
+        }
+    }
+    /**
+     * Read-only copy of the registry.
+     *
+     * Only to be used in test suite.
+     */
+    snapshot() {
+        return wsGroups.map(group => ({
+            ...group,
+            assetIds: new Set(group.assetIds),
+        }));
+    }
+    /**
+     * Find the first group with capacity to hold new assets.
+     *
+     * Returns the groupId if found, otherwise null.
+     */
+    findGroupWithCapacity(newAssetLen, maxPerWS) {
+        for (const group of wsGroups) {
+            if (group.assetIds.size === 0)
+                continue;
+            if (group.assetIds.size + newAssetLen <= maxPerWS)
+                return group.groupId;
+        }
+        return null;
+    }
+    /**
+     * Get the indices of all groups that contain the asset.
+     *
+     * Returns an array of indices.
+     */
+    getGroupIndicesForAsset(assetId) {
+        var _a;
+        const indices = [];
+        for (let i = 0; i < wsGroups.length; i++) {
+            if ((_a = wsGroups[i]) === null || _a === void 0 ? void 0 : _a.assetIds.has(assetId))
+                indices.push(i);
+        }
+        return indices;
+    }
+    /**
+     * Check if any group contains the asset.
+     */
+    hasAsset(assetId) {
+        return wsGroups.some(group => group.assetIds.has(assetId));
+    }
+    /**
+     * Find the group by groupId.
+     *
+     * Returns the group if found, otherwise undefined.
+     */
+    findGroupById(groupId) {
+        return wsGroups.find(g => g.groupId === groupId);
+    }
+    /**
+     * Atomically remove **all** groups from the registry and return them so the
+     * caller can perform any asynchronous cleanup (closing sockets, etc.)
+     * outside the lock.
+     *
+     * Returns the removed groups.
+     */
+    async clearAllGroups() {
+        let removed = [];
+        await this.mutate(groups => {
+            removed = [...groups];
+            groups.length = 0;
+        });
+        return removed;
+    }
+    /**
+     * Add new asset subscriptions.
+     *
+     * – Ignores assets that are already subscribed.
+     * – Either reuses an existing group with capacity or creates new groups (size ≤ maxPerWS).
+     * – If appending to a group:
+     *  - A new group is created with the updated assetIds.
+     *  - The old group is marked for cleanup.
+     *  - The group is added to the list of groups to connect.
+     *
+     * @param assetIds - The assetIds to add.
+     * @param maxPerWS - The maximum number of assets per WebSocket group.
+     * @returns An array of *new* groupIds that need websocket connections.
+     */
+    async addAssets(assetIds, maxPerWS) {
+        const groupIdsToConnect = [];
+        let newAssetIds = [];
+        await this.mutate(groups => {
+            newAssetIds = assetIds.filter(id => !groups.some(g => g.assetIds.has(id)));
+            if (newAssetIds.length === 0)
+                return;
+            const existingGroupId = this.findGroupWithCapacity(newAssetIds.length, maxPerWS);
+            /*
+                If no existing group with capacity is found, create new groups.
+            */
+            if (existingGroupId === null) {
+                const chunks = lodash_1.default.chunk(newAssetIds, maxPerWS);
+                for (const chunk of chunks) {
+                    const groupId = (0, uuid_1.v4)();
+                    groups.push({
+                        groupId,
+                        assetIds: new Set(chunk),
+                        wsClient: null,
+                        status: WebSocketSubscriptions_1.WebSocketStatus.PENDING
+                    });
+                    groupIdsToConnect.push(groupId);
+                }
+                /*
+                    If an existing group with capacity is found, update the group.
+                */
+            }
+            else {
+                const existingGroup = groups.find(g => g.groupId === existingGroupId);
+                if (!existingGroup) {
+                    // Should never happen
+                    throw new Error(`Group with capacity not found for ${newAssetIds.join(', ')}`);
+                }
+                const updatedAssetIds = new Set([...existingGroup.assetIds, ...newAssetIds]);
+                // Mark old group ready for cleanup
+                existingGroup.assetIds = new Set();
+                existingGroup.status = WebSocketSubscriptions_1.WebSocketStatus.CLEANUP;
+                const groupId = (0, uuid_1.v4)();
+                groups.push({
+                    groupId,
+                    assetIds: updatedAssetIds,
+                    wsClient: null,
+                    status: WebSocketSubscriptions_1.WebSocketStatus.PENDING
+                });
+                groupIdsToConnect.push(groupId);
+            }
+        });
+        if (newAssetIds.length > 0) {
+            logger_1.logger.info({
+                message: `Added ${newAssetIds.length} new asset(s)`
+            });
+        }
+        return groupIdsToConnect;
+    }
+    /**
+     * Remove asset subscriptions from every group that contains the asset.
+     *
+     * It should be only one group that contains the asset, we search all of them
+     * regardless.
+     *
+     * Returns the list of assetIds that were removed.
+     */
+    async removeAssets(assetIds, bookCache) {
+        const removedAssetIds = [];
+        await this.mutate(groups => {
+            groups.forEach(group => {
+                if (group.assetIds.size === 0)
+                    return;
+                assetIds.forEach(id => {
+                    if (group.assetIds.delete(id)) {
+                        bookCache.clear(id);
+                        removedAssetIds.push(id);
+                    }
+                });
+            });
+        });
+        if (removedAssetIds.length > 0) {
+            logger_1.logger.info({
+                message: `Removed ${removedAssetIds.length} asset(s)`
+            });
+        }
+        return removedAssetIds;
+    }
+    /**
+     * Disconnect a group.
+     */
+    disconnectGroup(group) {
+        var _a;
+        (_a = group.wsClient) === null || _a === void 0 ? void 0 : _a.close();
+        group.wsClient = null;
+        logger_1.logger.info({
+            message: 'Disconnected group',
+            groupId: group.groupId,
+            assetIds: Array.from(group.assetIds),
+        });
+    }
+    ;
+    /**
+     * Check status of groups and reconnect or cleanup as needed.
+     *
+     * – Empty groups are removed from the global array and returned.
+     * – Dead (but non-empty) groups are reset so that caller can reconnect them.
+     * – Pending groups are returned so that caller can connect them.
+     *
+     * Returns an array of group IDs that need to be reconnected, after cleaning up empty and cleanup-marked groups.
+     */
+    async getGroupsToReconnectAndCleanup() {
+        const reconnectIds = [];
+        await this.mutate(groups => {
+            const groupsToRemove = new Set();
+            for (const group of groups) {
+                if (group.assetIds.size === 0) {
+                    groupsToRemove.add(group.groupId);
+                    continue;
+                }
+                if (group.status === WebSocketSubscriptions_1.WebSocketStatus.ALIVE) {
+                    continue;
+                }
+                if (group.status === WebSocketSubscriptions_1.WebSocketStatus.DEAD) {
+                    this.disconnectGroup(group);
+                    reconnectIds.push(group.groupId);
+                }
+                if (group.status === WebSocketSubscriptions_1.WebSocketStatus.CLEANUP) {
+                    groupsToRemove.add(group.groupId);
+                    group.assetIds = new Set();
+                    continue;
+                }
+                if (group.status === WebSocketSubscriptions_1.WebSocketStatus.PENDING) {
+                    reconnectIds.push(group.groupId);
+                }
+            }
+            if (groupsToRemove.size > 0) {
+                groups.forEach(group => {
+                    if (groupsToRemove.has(group.groupId)) {
+                        this.disconnectGroup(group);
+                    }
+                });
+                const remaining = groups.filter(group => !groupsToRemove.has(group.groupId));
+                groups.splice(0, groups.length, ...remaining);
+            }
+        });
+        return reconnectIds;
+    }
+}
+exports.GroupRegistry = GroupRegistry;

package/dist/modules/GroupSocket.d.ts

@@ -0,0 +1,22 @@
+import Bottleneck from 'bottleneck';
+import { WebSocketGroup } from '../types/WebSocketSubscriptions';
+import { OrderBookCache } from './OrderBookCache';
+import { WebSocketHandlers } from '../types/PolymarketWebSocket';
+export declare class GroupSocket {
+    private group;
+    private limiter;
+    private bookCache;
+    private handlers;
+    private pingInterval?;
+    constructor(group: WebSocketGroup, limiter: Bottleneck, bookCache: OrderBookCache, handlers: WebSocketHandlers);
+    /**
+     * Establish the websocket connection using the provided Bottleneck limiter.
+     *
+     */
+    connect(): Promise<void>;
+    private setupEventHandlers;
+    private handleBookEvents;
+    private handleTickEvents;
+    private handlePriceChangeEvents;
+    private handleLastTradeEvents;
+}

package/dist/modules/GroupSocket.js

@@ -0,0 +1,311 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.GroupSocket = void 0;
+const ws_1 = __importDefault(require("ws"));
+const logger_1 = require("../logger");
+const WebSocketSubscriptions_1 = require("../types/WebSocketSubscriptions");
+const PolymarketWebSocket_1 = require("../types/PolymarketWebSocket");
+const lodash_1 = __importDefault(require("lodash"));
+const ms_1 = __importDefault(require("ms"));
+const crypto_1 = require("crypto");
+const CLOB_WSS_URL = 'wss://ws-subscriptions-clob.polymarket.com/ws/market';
+class GroupSocket {
+    constructor(group, limiter, bookCache, handlers) {
+        this.group = group;
+        this.limiter = limiter;
+        this.bookCache = bookCache;
+        this.handlers = handlers;
+    }
+    /**
+     * Establish the websocket connection using the provided Bottleneck limiter.
+     *
+     */
+    async connect() {
+        if (this.group.assetIds.size === 0) {
+            this.group.status = WebSocketSubscriptions_1.WebSocketStatus.CLEANUP;
+            return;
+        }
+        try {
+            logger_1.logger.info({
+                message: 'Connecting to CLOB WebSocket',
+                groupId: this.group.groupId,
+                assetIdsLength: this.group.assetIds.size,
+            });
+            this.group.wsClient = await this.limiter.schedule({ priority: 0 }, async () => {
+                const ws = new ws_1.default(CLOB_WSS_URL);
+                /*
+                    This handler will be replaced by the handlers in setupEventHandlers
+                */
+                ws.on('error', (err) => {
+                    logger_1.logger.warn({
+                        message: 'Error connecting to CLOB WebSocket',
+                        error: err,
+                        groupId: this.group.groupId,
+                        assetIdsLength: this.group.assetIds.size,
+                    });
+                });
+                return ws;
+            });
+        }
+        catch (err) {
+            this.group.status = WebSocketSubscriptions_1.WebSocketStatus.DEAD;
+            throw err; // caller responsible for error handler
+        }
+        this.setupEventHandlers();
+    }
+    setupEventHandlers() {
+        const group = this.group;
+        const handlers = this.handlers;
+        /*
+            Define handlers within this scope to capture 'this' context
+        */
+        const handleOpen = async () => {
+            var _a;
+            if (group.assetIds.size === 0) {
+                group.status = WebSocketSubscriptions_1.WebSocketStatus.CLEANUP;
+                return;
+            }
+            group.status = WebSocketSubscriptions_1.WebSocketStatus.ALIVE;
+            group.wsClient.send(JSON.stringify({ assets_ids: Array.from(group.assetIds), type: 'market' }));
+            await ((_a = handlers.onWSOpen) === null || _a === void 0 ? void 0 : _a.call(handlers, group.groupId, Array.from(group.assetIds)));
+            this.pingInterval = setInterval(() => {
+                if (group.assetIds.size === 0) {
+                    clearInterval(this.pingInterval);
+                    group.status = WebSocketSubscriptions_1.WebSocketStatus.CLEANUP;
+                    return;
+                }
+                if (!group.wsClient) {
+                    clearInterval(this.pingInterval);
+                    group.status = WebSocketSubscriptions_1.WebSocketStatus.DEAD;
+                    return;
+                }
+                group.wsClient.ping();
+            }, (0, crypto_1.randomInt)((0, ms_1.default)('15s'), (0, ms_1.default)('25s')));
+        };
+        const handleMessage = async (data) => {
+            var _a, _b;
+            let events = [];
+            try {
+                const parsedData = JSON.parse(data.toString());
+                events = Array.isArray(parsedData) ? parsedData : [parsedData];
+            }
+            catch (err) {
+                await ((_a = handlers.onError) === null || _a === void 0 ? void 0 : _a.call(handlers, new Error(`Not JSON: ${data.toString()}`)));
+                return;
+            }
+            events = lodash_1.default.filter(events, (event) => lodash_1.default.size(event.asset_id) > 0);
+            const bookEvents = [];
+            const lastTradeEvents = [];
+            const tickEvents = [];
+            const priceChangeEvents = [];
+            for (const event of events) {
+                /*
+                    Skip events for asset ids that are not in the group to ensure that
+                    we don't get stale events for assets that were removed.
+                */
+                if (!group.assetIds.has(event.asset_id)) {
+                    continue;
+                }
+                if ((0, PolymarketWebSocket_1.isBookEvent)(event)) {
+                    bookEvents.push(event);
+                }
+                else if ((0, PolymarketWebSocket_1.isLastTradePriceEvent)(event)) {
+                    lastTradeEvents.push(event);
+                }
+                else if ((0, PolymarketWebSocket_1.isTickSizeChangeEvent)(event)) {
+                    tickEvents.push(event);
+                }
+                else if ((0, PolymarketWebSocket_1.isPriceChangeEvent)(event)) {
+                    priceChangeEvents.push(event);
+                }
+                else {
+                    await ((_b = handlers.onError) === null || _b === void 0 ? void 0 : _b.call(handlers, new Error(`Unknown event: ${JSON.stringify(event)}`)));
+                }
+            }
+            await this.handleBookEvents(bookEvents);
+            await this.handleTickEvents(tickEvents);
+            await this.handlePriceChangeEvents(priceChangeEvents);
+            await this.handleLastTradeEvents(lastTradeEvents);
+        };
+        const handlePong = () => {
+            group.groupId;
+        };
+        const handleError = (err) => {
+            var _a;
+            group.status = WebSocketSubscriptions_1.WebSocketStatus.DEAD;
+            clearInterval(this.pingInterval);
+            (_a = handlers.onError) === null || _a === void 0 ? void 0 : _a.call(handlers, new Error(`WebSocket error for group ${group.groupId}: ${err.message}`));
+        };
+        const handleClose = async (code, reason) => {
+            var _a;
+            group.status = WebSocketSubscriptions_1.WebSocketStatus.DEAD;
+            clearInterval(this.pingInterval);
+            await ((_a = handlers.onWSClose) === null || _a === void 0 ? void 0 : _a.call(handlers, group.groupId, code, reason.toString()));
+        };
+        if (group.wsClient) {
+            // Remove any existing handlers
+            group.wsClient.removeAllListeners();
+            // Add the handlers
+            group.wsClient.on('open', handleOpen);
+            group.wsClient.on('message', handleMessage);
+            group.wsClient.on('pong', handlePong);
+            group.wsClient.on('error', handleError);
+            group.wsClient.on('close', handleClose);
+        }
+        if (group.assetIds.size === 0) {
+            group.status = WebSocketSubscriptions_1.WebSocketStatus.CLEANUP;
+            return;
+        }
+        if (!group.wsClient) {
+            group.status = WebSocketSubscriptions_1.WebSocketStatus.DEAD;
+            return;
+        }
+    }
+    async handleBookEvents(bookEvents) {
+        var _a, _b;
+        if (bookEvents.length) {
+            for (const event of bookEvents) {
+                this.bookCache.replaceBook(event);
+            }
+            await ((_b = (_a = this.handlers).onBook) === null || _b === void 0 ? void 0 : _b.call(_a, bookEvents));
+        }
+    }
+    async handleTickEvents(tickEvents) {
+        var _a, _b;
+        if (tickEvents.length) {
+            await ((_b = (_a = this.handlers).onTickSizeChange) === null || _b === void 0 ? void 0 : _b.call(_a, tickEvents));
+        }
+    }
+    async handlePriceChangeEvents(priceChangeEvents) {
+        var _a, _b, _c, _d, _e, _f, _g, _h;
+        if (priceChangeEvents.length) {
+            await ((_b = (_a = this.handlers).onPriceChange) === null || _b === void 0 ? void 0 : _b.call(_a, priceChangeEvents));
+            for (const event of priceChangeEvents) {
+                try {
+                    this.bookCache.upsertPriceChange(event);
+                }
+                catch (err) {
+                    logger_1.logger.warn({
+                        message: `Skipping price_change: book not found for asset`,
+                        asset_id: event.asset_id,
+                        event: event,
+                        error: err
+                    });
+                    continue;
+                }
+                let spreadOver10Cents;
+                try {
+                    spreadOver10Cents = this.bookCache.spreadOver(event.asset_id, 0.1);
+                }
+                catch (err) {
+                    logger_1.logger.warn({
+                        message: 'Skipping price_change: error calculating spread',
+                        asset_id: event.asset_id,
+                        event: event,
+                        error: err
+                    });
+                    await ((_d = (_c = this.handlers).onError) === null || _d === void 0 ? void 0 : _d.call(_c, err));
+                    continue;
+                }
+                if (!spreadOver10Cents) {
+                    let newPrice;
+                    try {
+                        newPrice = this.bookCache.midpoint(event.asset_id);
+                    }
+                    catch (err) {
+                        logger_1.logger.warn({
+                            message: 'Skipping price_change: error calculating midpoint',
+                            asset_id: event.asset_id,
+                            event: event,
+                            error: err
+                        });
+                        await ((_f = (_e = this.handlers).onError) === null || _f === void 0 ? void 0 : _f.call(_e, err));
+                        continue;
+                    }
+                    const bookEntry = this.bookCache.getBookEntry(event.asset_id);
+                    if (!bookEntry) {
+                        logger_1.logger.warn({
+                            message: 'Skipping price_change: book not found for asset',
+                            asset_id: event.asset_id,
+                            event: event,
+                        });
+                        continue;
+                    }
+                    if (newPrice !== bookEntry.price) {
+                        bookEntry.price = newPrice;
+                        const priceUpdateEvent = {
+                            asset_id: event.asset_id,
+                            event_type: 'price_update',
+                            triggeringEvent: event,
+                            timestamp: event.timestamp,
+                            book: { bids: bookEntry.bids, asks: bookEntry.asks },
+                            price: newPrice,
+                            midpoint: bookEntry.midpoint || '',
+                            spread: bookEntry.spread || '',
+                        };
+                        await ((_h = (_g = this.handlers).onPolymarketPriceUpdate) === null || _h === void 0 ? void 0 : _h.call(_g, [priceUpdateEvent]));
+                    }
+                }
+            }
+        }
+    }
+    async handleLastTradeEvents(lastTradeEvents) {
+        var _a, _b, _c, _d, _e, _f;
+        if (lastTradeEvents.length) {
+            /*
+                Note: There is no need to edit the book here. According to the docs, a separate
+                book event is sent when a trade affects the book.
+
+                See: https://docs.polymarket.com/developers/CLOB/websocket/market-channel#book-message
+            */
+            await ((_b = (_a = this.handlers).onLastTradePrice) === null || _b === void 0 ? void 0 : _b.call(_a, lastTradeEvents));
+            for (const event of lastTradeEvents) {
+                let spreadOver10Cents;
+                try {
+                    spreadOver10Cents = this.bookCache.spreadOver(event.asset_id, 0.1);
+                }
+                catch (err) {
+                    logger_1.logger.warn({
+                        message: 'Skipping last_trade_price: error calculating spread',
+                        asset_id: event.asset_id,
+                        event: event,
+                        error: err
+                    });
+                    await ((_d = (_c = this.handlers).onError) === null || _d === void 0 ? void 0 : _d.call(_c, err));
+                    continue;
+                }
+                if (spreadOver10Cents) {
+                    // Ensure no trailing zeros
+                    const newPrice = parseFloat(event.price).toString();
+                    const bookEntry = this.bookCache.getBookEntry(event.asset_id);
+                    if (!bookEntry) {
+                        logger_1.logger.warn({
+                            message: 'Skipping last_trade_price: book not found for asset',
+                            asset_id: event.asset_id,
+                            event: event,
+                        });
+                        continue;
+                    }
+                    if (newPrice !== bookEntry.price) {
+                        bookEntry.price = newPrice;
+                        const priceUpdateEvent = {
+                            asset_id: event.asset_id,
+                            event_type: 'price_update',
+                            triggeringEvent: event,
+                            timestamp: event.timestamp,
+                            book: { bids: bookEntry.bids, asks: bookEntry.asks },
+                            price: newPrice,
+                            midpoint: bookEntry.midpoint || '',
+                            spread: bookEntry.spread || '',
+                        };
+                        await ((_f = (_e = this.handlers).onPolymarketPriceUpdate) === null || _f === void 0 ? void 0 : _f.call(_e, [priceUpdateEvent]));
+                    }
+                }
+            }
+        }
+    }
+}
+exports.GroupSocket = GroupSocket;

package/dist/modules/OrderBookCache.d.ts

@@ -0,0 +1,49 @@
+import { BookEvent, PriceChangeEvent, PriceLevel } from '../types/PolymarketWebSocket';
+export interface BookEntry {
+    bids: PriceLevel[];
+    asks: PriceLevel[];
+    price: string | null;
+    midpoint: string | null;
+    spread: string | null;
+}
+export declare function sortDescendingInPlace(bookSide: PriceLevel[]): void;
+export declare class OrderBookCache {
+    private bookCache;
+    constructor();
+    /**
+     * Replace full book (after a `book` event)
+     */
+    replaceBook(event: BookEvent): void;
+    /**
+     * Update a cached book from a `price_change` event.
+     *
+     * Returns true if the book was updated.
+     * Throws if the book is not found.
+     */
+    upsertPriceChange(event: PriceChangeEvent): void;
+    /**
+     * Return `true` if best-bid/best-ask spread exceeds `cents`.
+     *
+     * Side effect: updates the book's spread
+     *
+     * Throws if either side of the book is empty.
+     */
+    spreadOver(assetId: string, cents?: number): boolean;
+    /**
+     * Calculate the midpoint of the book, rounded to 3dp, no trailing zeros
+     *
+     * Side effect: updates the book's midpoint
+     *
+     * Throws if
+     * - the book is not found or missing either bid or ask
+     * - the midpoint is NaN.
+    */
+    midpoint(assetId: string): string;
+    clear(assetId?: string): void;
+    /**
+     * Get a book entry by asset id.
+     *
+     * Return null if the book is not found.
+     */
+    getBookEntry(assetId: string): BookEntry | null;
+}