From fb610c489832dcde52d18346c6ddf7d569106642 Mon Sep 17 00:00:00 2001
From: MingxuanGame <MingxuanGame@outlook.com>
Date: Sun, 27 Jul 2025 16:23:47 +0000
Subject: [PATCH] feat(signalr): support on_client_connect hook

---
 app/signalr/hub/hub.py | 8 +++++++-
 app/signalr/router.py  | 4 ++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/app/signalr/hub/hub.py b/app/signalr/hub/hub.py
index 02c5537..40e75da 100644
--- a/app/signalr/hub/hub.py
+++ b/app/signalr/hub/hub.py
@@ -48,7 +48,7 @@ class Client:
         message = await self.connection.receive()
         d = message.get("bytes") or message.get("text", "").encode()
         if not d:
-            raise WebSocketDisconnect(code=1008, reason="Empty message received.")
+            return PingPacket()  # FIXME: Graceful empty message handling
         return self.procotol.decode(d)
 
     async def _ping(self):
@@ -104,6 +104,10 @@ class Hub:
         client._ping_task = task
         return client
 
+    async def on_connect(self, client: Client) -> None:
+        if method := getattr(self, "on_client_connect", None):
+            await method(client)
+
     async def remove_client(self, connection_id: str) -> None:
         if client := self.clients.get(connection_id):
             del self.clients[connection_id]
@@ -138,6 +142,8 @@ class Hub:
                 task = asyncio.create_task(self._handle_packet(client, packet))
                 self.tasks.add(task)
                 task.add_done_callback(self.tasks.discard)
+            except StopIteration:
+                pass
             except WebSocketDisconnect as e:
                 print(
                     f"Client {client.connection_id} disconnected: {e.code}, {e.reason}"
diff --git a/app/signalr/router.py b/app/signalr/router.py
index 3d70931..5c2f08f 100644
--- a/app/signalr/router.py
+++ b/app/signalr/router.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import time
 from typing import Literal
@@ -92,4 +93,7 @@ async def connect(
     if error or not client:
         await websocket.close(code=1008)
         return
+    task = asyncio.create_task(hub_.on_connect(client))
+    hub_.tasks.add(task)
+    task.add_done_callback(hub_.tasks.discard)
     await hub_._listen_client(client)