Collab: Add RoomCount in Websocket

Author: wr1159

collab-service/app/server.js

@@ -21,6 +21,7 @@ const io = new Server(server, {
   },
 });
 
+// Partially Referenced from: https://github.com/yjs/y-monaco/blob/master/demo/monaco-demo.js
 const yjsWs = new WebSocketServer({ noServer: true });
 yjsWs.on("connection", (conn, req) => {
   setupWSConnection(conn, req, {
@@ -36,18 +37,23 @@ setInterval(() => {
   const stats = {
     conns,
     docs: docs.size,
-    websocket: `ws://localhost:${PORT}`,
-    http: `http://localhost:${PORT}`,
   };
   console.log(`${new Date().toISOString()} Stats: ${JSON.stringify(stats)}`);
 }, 10000);
 
+const roomUserCount = new Map();
 io.on("connection", (socket) => {
   console.log("User connected to Socket.IO");
 
   // Join a room based on roomId
   socket.on("joinRoom", (roomId) => {
     socket.join(roomId);
+    if (roomUserCount.has(roomId)) {
+      roomUserCount.set(roomId, roomUserCount.get(roomId) + 1);
+    } else {
+      roomUserCount.set(roomId, 1);
+    }
+    io.to(roomId).emit("roomCount", roomUserCount.get(roomId));
     console.log(`User joined room: ${roomId}`);
   });
 
@@ -60,6 +66,18 @@ io.on("connection", (socket) => {
     io.to(roomId).emit("chatMessage", newMessage);
   });
 
+  socket.on("disconnecting", () => {
+    socket.rooms.forEach((roomId) => {
+      if (roomId !== socket.id) {
+        if (roomUserCount.has(roomId)) {
+          const currentCount = roomUserCount.get(roomId);
+          roomUserCount.set(roomId, currentCount - 1);
+          io.to(roomId).emit("roomCount", roomUserCount.get(roomId));
+        }
+      }
+    });
+  });
+
   socket.on("disconnect", () => {
     console.log("User disconnected from Socket.IO");
   });

frontend/components/collab/chat.tsx

@@ -55,6 +55,7 @@ export default function Chat({
   const [aiMessages, setAiMessages] = useState<Message[]>([]);
   const [isConnected, setIsConnected] = useState(false);
   const [isLoading, setIsLoading] = useState(true);
+  const [roomCount, setRoomCount] = useState(0);
   const lastMessageRef = useRef<HTMLDivElement | null>(null);
 
   useEffect(() => {
@@ -145,6 +146,14 @@ export default function Chat({
       setIsConnected(false);
     });
 
+    socketInstance.on("roomCount", (roomCount: number) => {
+      toast({
+        title: "Room Count",
+        description: "Current room count: " + roomCount,
+      });
+      setRoomCount(roomCount);
+    });
+
     socketInstance.on("chatMessage", (message: Message) => {
       setPartnerMessages((prev) => {
         const exists = prev.some(
@@ -171,7 +180,7 @@ export default function Chat({
     return () => {
       socketInstance.disconnect();
     };
-  }, [roomId, own_user_id, auth?.user?.id]);
+  }, [roomId, own_user_id, auth?.user?.id, toast]);
 
   useEffect(() => {
     const scrollWithDelay = () => {
@@ -282,7 +291,7 @@ export default function Chat({
     <Card className="flex flex-col h-full">
       <CardHeader>
         <CardTitle className="flex justify-between items-center">
-          Chat
+          Chat (Users: {roomCount})
           <span
             className={`h-2 w-2 rounded-full ${isConnected ? "bg-green-500" : "bg-red-500"}`}
           />