From bbfa54a6b9d097342ea9458ef10db90697c71dfd Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 4 Jun 2024 01:10:31 -0700
Subject: [PATCH] feat: active user count

---
 backend/apps/socket/main.py                   | 27 +++++++++--------
 .../components/layout/Sidebar/UserMenu.svelte | 30 +++++++++++++++++--
 src/lib/stores/index.ts                       |  1 +
 src/routes/+layout.svelte                     |  9 ++++--
 4 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/backend/apps/socket/main.py b/backend/apps/socket/main.py
index 3fbfb8ae8..a8fc50efa 100644
--- a/backend/apps/socket/main.py
+++ b/backend/apps/socket/main.py
@@ -7,6 +7,8 @@ sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi")
 app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
 
 # Dictionary to maintain the user pool
+
+
 USER_POOL = {}
 
 
@@ -22,24 +24,25 @@ async def connect(sid, environ, auth):
             user = Users.get_user_by_id(data["id"])
 
         if user:
-            USER_POOL[sid] = {
-                "id": user.id,
-                "name": user.name,
-                "email": user.email,
-                "role": user.role,
-            }
+            USER_POOL[sid] = user.id
             print(f"user {user.name}({user.id}) connected with session ID {sid}")
 
+            print(len(set(USER_POOL)))
+            await sio.emit("user-count", {"count": len(set(USER_POOL))})
+
+
+@sio.on("user-count")
+async def user_count(sid):
+    print("user-count", sid)
+    await sio.emit("user-count", {"count": len(set(USER_POOL))})
+
 
 @sio.event
-def disconnect(sid):
+async def disconnect(sid):
     if sid in USER_POOL:
         disconnected_user = USER_POOL.pop(sid)
         print(f"user {disconnected_user} disconnected with session ID {sid}")
+
+        await sio.emit("user-count", {"count": len(USER_POOL)})
     else:
         print(f"Unknown session ID {sid} disconnected")
-
-
-@sio.event
-def disconnect(sid):
-    print("disconnect", sid)
diff --git a/src/lib/components/layout/Sidebar/UserMenu.svelte b/src/lib/components/layout/Sidebar/UserMenu.svelte
index 7d571555d..1525727c0 100644
--- a/src/lib/components/layout/Sidebar/UserMenu.svelte
+++ b/src/lib/components/layout/Sidebar/UserMenu.svelte
@@ -1,11 +1,11 @@
 <script lang="ts">
 	import { DropdownMenu } from 'bits-ui';
-	import { createEventDispatcher, getContext } from 'svelte';
+	import { createEventDispatcher, getContext, onMount } from 'svelte';
 
 	import { flyAndScale } from '$lib/utils/transitions';
 	import { goto } from '$app/navigation';
 	import ArchiveBox from '$lib/components/icons/ArchiveBox.svelte';
-	import { showSettings } from '$lib/stores';
+	import { showSettings, activeUserCount } from '$lib/stores';
 	import { fade, slide } from 'svelte/transition';
 
 	const i18n = getContext('i18n');
@@ -107,7 +107,7 @@
 				</button>
 			{/if}
 
-			<hr class=" dark:border-gray-800 my-2 p-0" />
+			<hr class=" dark:border-gray-800 my-1.5 p-0" />
 
 			<button
 				class="flex rounded-md py-2 px-3 w-full hover:bg-gray-50 dark:hover:bg-gray-800 transition"
@@ -139,6 +139,30 @@
 				<div class=" self-center font-medium">{$i18n.t('Sign Out')}</div>
 			</button>
 
+			{#if $activeUserCount}
+				<hr class=" dark:border-gray-800 my-1.5 p-0" />
+
+				<div class="flex rounded-md py-1.5 px-3 text-xs gap-2.5 items-center">
+					<div class=" flex items-center">
+						<span class="relative flex size-2">
+							<span
+								class="animate-ping absolute inline-flex h-full w-full rounded-full bg-green-400 opacity-75"
+							/>
+							<span class="relative inline-flex rounded-full size-2 bg-green-500" />
+						</span>
+					</div>
+
+					<div class=" translate-y-[0.25px]">
+						<span class=" font-medium">
+							{$i18n.t('Active Users')}:
+						</span>
+						<span class=" font-semibold">
+							{$activeUserCount}
+						</span>
+					</div>
+				</div>
+			{/if}
+
 			<!-- <DropdownMenu.Item class="flex items-center px-3 py-2 text-sm  font-medium">
 				<div class="flex items-center">Profile</div>
 			</DropdownMenu.Item> -->
diff --git a/src/lib/stores/index.ts b/src/lib/stores/index.ts
index 619b175d9..90bf1358c 100644
--- a/src/lib/stores/index.ts
+++ b/src/lib/stores/index.ts
@@ -15,6 +15,7 @@ export const MODEL_DOWNLOAD_POOL = writable({});
 export const mobile = writable(false);
 
 export const socket: Writable<null | Socket> = writable(null);
+export const activeUserCount: Writable<null | number> = writable(null);
 
 export const theme = writable('system');
 export const chatId = writable('');
diff --git a/src/routes/+layout.svelte b/src/routes/+layout.svelte
index ccccccec3..3057752d4 100644
--- a/src/routes/+layout.svelte
+++ b/src/routes/+layout.svelte
@@ -2,7 +2,7 @@
 	import { io } from 'socket.io-client';
 
 	import { onMount, tick, setContext } from 'svelte';
-	import { config, user, theme, WEBUI_NAME, mobile, socket } from '$lib/stores';
+	import { config, user, theme, WEBUI_NAME, mobile, socket, activeUserCount } from '$lib/stores';
 	import { goto } from '$app/navigation';
 	import { Toaster, toast } from 'svelte-sonner';
 
@@ -69,7 +69,12 @@
 					console.log('connected');
 				});
 
-				socket.set(_socket);
+				await socket.set(_socket);
+
+				_socket.on('user-count', (data) => {
+					console.log('user-count', data);
+					activeUserCount.set(data.count);
+				});
 
 				if (localStorage.token) {
 					// Get Session User Info