ppai/polymarket/libs/bestBidAskSizeWatcher.js

import { EventEmitter } from "node:events";

import { MarketWsClient } from "./polymarketClient.js";

const normalizeTokenIds = (tokenIds) => {
  if (!Array.isArray(tokenIds)) {
    throw new Error("tokenIds must be an array", { cause: 400 });
  }
  return [...new Set(tokenIds.map(String).filter(Boolean))];
};

const normalizeSize = (size) => {
  const value = Number(size);
  return Number.isFinite(value) ? value : 0;
};

const applyLevels = (levels = [], target) => {
  target.clear();
  levels.forEach(({ price, size }) => {
    if (!price) {
      return;
    }
    const normalizedSize = normalizeSize(size);
    if (normalizedSize > 0) {
      target.set(String(price), String(size));
    }
  });
};

const updateLevel = (target, price, size) => {
  if (!price) {
    return;
  }
  if (normalizeSize(size) <= 0) {
    target.delete(String(price));
    return;
  }
  target.set(String(price), String(size));
};

const getBestLevel = (levels, compare) => {
  let bestPrice;
  let bestSize;

  levels.forEach((size, price) => {
    if (bestPrice === undefined || compare(Number(price), Number(bestPrice))) {
      bestPrice = price;
      bestSize = size;
    }
  });

  if (bestPrice === undefined) {
    return { price: null, size: null };
  }
  return { price: bestPrice, size: bestSize };
};

const parseSide = (side) => {
  const normalizedSide = String(side || "").toLowerCase();
  if (["ask", "asks", "sell"].includes(normalizedSide)) {
    return "asks";
  }
  if (["bid", "bids", "buy"].includes(normalizedSide)) {
    return "bids";
  }
  return "";
};

class BestBidAskSizeWatcher extends EventEmitter {
  #tokenIds = [];
  #client;
  #books = new Map();
  #started = false;

  constructor(tokenIds = []) {
    super();
    this.#client = new MarketWsClient();
    this.subscribe(tokenIds);
  }

  get tokenIds() {
    return [...this.#tokenIds];
  }

  start() {
    if (this.#started) {
      return this;
    }
    this.#started = true;
    this.#client.on("open", () => {
      if (this.#tokenIds.length > 0) {
        this.#client.subscribeToTokensIds(this.#tokenIds);
      }
      this.emit("open");
    });
    this.#client.on("message", data => this.#handleMessage(data));
    this.#client.on("error", error => this.emit("error", error));
    this.#client.on("close", event => this.emit("close", event));
    this.#client.connect();
    return this;
  }

  stop() {
    this.#client.close();
    this.#started = false;
    return this;
  }

  subscribe(tokenIds) {
    const ids = normalizeTokenIds(tokenIds);
    const newIds = ids.filter(tokenId => !this.#books.has(tokenId));
    if (newIds.length === 0) {
      return this;
    }

    this.#tokenIds = [...this.#tokenIds, ...newIds];
    newIds.forEach(tokenId => {
      this.#books.set(tokenId, { asks: new Map(), bids: new Map() });
    });

    if (this.#started) {
      this.#client.subscribeToTokensIds(newIds);
    }
    return this;
  }

  unsubscribe(tokenIds) {
    const ids = normalizeTokenIds(tokenIds);
    const removeIds = ids.filter(tokenId => this.#books.has(tokenId));
    if (removeIds.length === 0) {
      return this;
    }

    const removeIdSet = new Set(removeIds);
    this.#tokenIds = this.#tokenIds.filter(tokenId => !removeIdSet.has(tokenId));
    removeIds.forEach(tokenId => {
      this.#books.delete(tokenId);
    });

    if (this.#started) {
      this.#client.unsubscribeToTokensIds(removeIds);
    }
    return this;
  }

  snapshot() {
    return Object.fromEntries(
      this.#tokenIds.map(tokenId => [tokenId, this.#getBestBidAskSize(tokenId)])
    );
  }

  #getBook(tokenId) {
    const id = String(tokenId);
    if (!this.#books.has(id)) {
      this.#books.set(id, { asks: new Map(), bids: new Map() });
    }
    return this.#books.get(id);
  }

  #getBestBidAskSize(tokenId) {
    const book = this.#getBook(tokenId);
    const ask = getBestLevel(book.asks, (price, bestPrice) => price < bestPrice);
    const bid = getBestLevel(book.bids, (price, bestPrice) => price > bestPrice);
    return {
      tokenId: String(tokenId),
      best_ask: ask.price,
      best_ask_size: ask.size,
      best_bid: bid.price,
      best_bid_size: bid.size,
    };
  }

  #emitUpdate(tokenId) {
    const data = this.#getBestBidAskSize(tokenId);
    this.emit("update", data);
    this.emit(String(tokenId), data);
  }

  #handleBook({ asset_id, asks = [], bids = [] }) {
    const tokenId = String(asset_id);
    if (!this.#books.has(tokenId)) {
      return;
    }
    const book = this.#getBook(tokenId);
    applyLevels(asks, book.asks);
    applyLevels(bids, book.bids);
    this.#emitUpdate(tokenId);
  }

  #handlePriceChanges(priceChanges = []) {
    const changedTokenIds = new Set();
    priceChanges.forEach(({ asset_id, price, size, side }) => {
      const tokenId = String(asset_id);
      if (!this.#books.has(tokenId)) {
        return;
      }
      const parsedSide = parseSide(side);
      if (!parsedSide) {
        return;
      }
      const book = this.#getBook(tokenId);
      updateLevel(book[parsedSide], price, size);
      changedTokenIds.add(tokenId);
    });
    changedTokenIds.forEach(tokenId => this.#emitUpdate(tokenId));
  }

  #handleMessage(data) {
    switch (data?.event_type) {
      case "book":
        this.#handleBook(data);
        break;
      case "price_change":
        this.#handlePriceChanges(data.price_changes);
        break;
    }
  }
}

export const watchBestBidAskSizes = (tokenIds) => {
  return new BestBidAskSizeWatcher(tokenIds).start();
};

export default BestBidAskSizeWatcher;