From 7cf80289c3e7603bdb7bce7f736a4ed7b6958a29 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 4 Jun 2024 09:52:27 -0700
Subject: [PATCH] fix: active users

---
 backend/apps/socket/main.py  | 20 ++++++++++++++++++++
 src/routes/auth/+page.svelte |  4 +++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/backend/apps/socket/main.py b/backend/apps/socket/main.py
index a8fc50efa..58f48025b 100644
--- a/backend/apps/socket/main.py
+++ b/backend/apps/socket/main.py
@@ -31,6 +31,26 @@ async def connect(sid, environ, auth):
             await sio.emit("user-count", {"count": len(set(USER_POOL))})
 
 
+@sio.on("user-join")
+async def user_join(sid, data):
+    print("user-join", sid, data)
+
+    auth = data["auth"] if "auth" in data else None
+
+    if auth and "token" in auth:
+        data = decode_token(auth["token"])
+
+        if data is not None and "id" in data:
+            user = Users.get_user_by_id(data["id"])
+
+        if user:
+            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)
diff --git a/src/routes/auth/+page.svelte b/src/routes/auth/+page.svelte
index f836a5cc9..321ded45f 100644
--- a/src/routes/auth/+page.svelte
+++ b/src/routes/auth/+page.svelte
@@ -3,7 +3,7 @@
 	import { userSignIn, userSignUp } from '$lib/apis/auths';
 	import Spinner from '$lib/components/common/Spinner.svelte';
 	import { WEBUI_API_BASE_URL, WEBUI_BASE_URL } from '$lib/constants';
-	import { WEBUI_NAME, config, user } from '$lib/stores';
+	import { WEBUI_NAME, config, user, socket } from '$lib/stores';
 	import { onMount, getContext } from 'svelte';
 	import { toast } from 'svelte-sonner';
 	import { generateInitialsImage, canvasPixelTest } from '$lib/utils';
@@ -22,6 +22,8 @@
 			console.log(sessionUser);
 			toast.success($i18n.t(`You're now logged in.`));
 			localStorage.token = sessionUser.token;
+
+			$socket.emit('user-join', { auth: { token: sessionUser.token } });
 			await user.set(sessionUser);
 			goto('/');
 		}