From e85fb81ab00ab2232950d6d25ef7c417f0c777a4 Mon Sep 17 00:00:00 2001
From: jherrera-jump <jherrera@jumptrading.com>
Date: Wed, 12 Mar 2025 21:43:58 +0000
Subject: [PATCH] fix rpc nodes count after limiting peers msg

---
 src/api/entities.ts         |  8 ++++++++
 src/api/useSetAtomWsData.ts | 19 +++++++++++++++----
 src/atoms.ts                | 15 +++++++++++----
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/api/entities.ts b/src/api/entities.ts
index a364267..c6d200a 100644
--- a/src/api/entities.ts
+++ b/src/api/entities.ts
@@ -413,6 +413,10 @@ export const peerRemoveSchema = z.object({
   identity_pubkey: z.string(),
 });
 
+const rpcCountUpdateSchema = z.object({
+  rpc_count: z.number(),
+});
+
 const peersUpdateSchema = z.object({
   add: z.array(peerUpdateSchema).optional(),
   update: z.array(peerUpdateSchema).optional(),
@@ -424,6 +428,10 @@ export const peersSchema = z.discriminatedUnion("key", [
     key: z.literal("update"),
     value: peersUpdateSchema,
   }),
+  peersTopicSchema.extend({
+    key: z.literal("rpc"),
+    value: rpcCountUpdateSchema,
+  }),
 ]);
 
 const tsTileTimersSchema = z.object({
diff --git a/src/api/useSetAtomWsData.ts b/src/api/useSetAtomWsData.ts
index c9e64b4..c4715e6 100644
--- a/src/api/useSetAtomWsData.ts
+++ b/src/api/useSetAtomWsData.ts
@@ -39,6 +39,7 @@ import {
   deleteSlotStatusBoundsAtom,
   deleteSlotResponseBoundsAtom,
   skipRateAtom,
+  rpcCountAtom,
 } from "../atoms";
 import {
   EstimatedSlotDuration,
@@ -128,6 +129,7 @@ export function useSetAtomWsData() {
   const addPeers = useSetAtom(addPeersAtom);
   const updatePeers = useSetAtom(updatePeersAtom);
   const removePeers = useSetAtom(removePeersAtom);
+  const setRpcCount = useSetAtom(rpcCountAtom)
 
   const setBlockEngine = useSetAtom(blockEngineAtom);
 
@@ -247,10 +249,19 @@ export function useSetAtomWsData() {
             break;
         }
       } else if (topic === "peers") {
-        const { value } = peersSchema.parse(msg);
-        addPeers(value.add);
-        updatePeers(value.update);
-        removePeers(value.remove);
+        const { key, value } = peersSchema.parse(msg);
+        switch (key) {
+          case "update": {
+            addPeers(value.add);
+            updatePeers(value.update);
+            removePeers(value.remove);
+            break;
+          }
+          case "rpc": {
+            setRpcCount(value.rpc_count);
+            break;
+          }
+        }
       } else if (topic === "slot") {
         const { key, value } = slotSchema.parse(msg);
         switch (key) {
diff --git a/src/atoms.ts b/src/atoms.ts
index c7b9f4a..9eb1d13 100644
--- a/src/atoms.ts
+++ b/src/atoms.ts
@@ -366,14 +366,21 @@ export const removePeersAtom = atom(null, (get, set, peers?: PeerRemove[]) => {
   }, removePeerDelay);
 });
 
+const _rpcCountAtom = atom<number>(0);
+
+export const rpcCountAtom = atom(
+  (get) => get(_rpcCountAtom),
+  (get, set, count: number) => {
+    set(_rpcCountAtom, count);
+  }
+);
+
 export const peerStatsAtom = atom((get) => {
   const peers = get(peersAtom);
+  const rpcCount = get(rpcCountAtom);
   if (!peers) return;
 
   const activePeers = Object.values(peers).filter((p) => !p.removed);
-  const rpc = activePeers.filter(
-    (p) => p.vote.every((v) => !v.activated_stake) && !!p.gossip,
-  );
   const validators = activePeers.filter((p) =>
     p.vote.some((v) => v.activated_stake),
   );
@@ -395,7 +402,7 @@ export const peerStatsAtom = atom((get) => {
   );
 
   return {
-    rpcCount: rpc.length,
+    rpcCount: rpcCount,
     validatorCount: validators.length,
     activeStake,
     delinquentStake,