Skip to content

Commit 31219a1

Browse files
use socket ID to identify connections on auto-index client
1 parent 5ae38b3 commit 31219a1

File tree

3 files changed

+12
-10
lines changed

3 files changed

+12
-10
lines changed

nicegui/air.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ def _handle_handshake(data: Dict[str, Any]) -> bool:
135135
core.app.storage.copy_tab(data['old_tab_id'], data['tab_id'])
136136
client.tab_id = data['tab_id']
137137
client.on_air = True
138-
client.handle_handshake(data.get('next_message_id'))
138+
client.handle_handshake(data['sid'], data.get('next_message_id'))
139139
return True
140140

141141
@self.relay.on('client_disconnect')
@@ -144,7 +144,7 @@ def _handle_client_disconnect(data: Dict[str, Any]) -> None:
144144
client_id = data['client_id']
145145
if client_id not in Client.instances:
146146
return
147-
Client.instances[client_id].handle_disconnect()
147+
Client.instances[client_id].handle_disconnect(data['sid'])
148148

149149
@self.relay.on('connect')
150150
async def _handle_connect() -> None:

nicegui/client.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import inspect
55
import time
66
import uuid
7+
from collections import defaultdict
78
from contextlib import contextmanager
89
from pathlib import Path
910
from typing import TYPE_CHECKING, Any, Awaitable, Callable, ClassVar, Dict, Iterable, Iterator, List, Optional, Union
@@ -59,7 +60,7 @@ def __init__(self, page: page, *, request: Optional[Request]) -> None:
5960
self.environ: Optional[Dict[str, Any]] = None
6061
self.shared = request is None
6162
self.on_air = False
62-
self._num_connections = 0
63+
self._num_connections = defaultdict(int)
6364
self._delete_task: Optional[asyncio.Task] = None
6465
self._deleted = False
6566
self.tab_id: Optional[str] = None
@@ -235,10 +236,10 @@ def on_disconnect(self, handler: Union[Callable[..., Any], Awaitable]) -> None:
235236
"""Add a callback to be invoked when the client disconnects."""
236237
self.disconnect_handlers.append(handler)
237238

238-
def handle_handshake(self, next_message_id: Optional[int]) -> None:
239+
def handle_handshake(self, socket_id: str, next_message_id: Optional[int]) -> None:
239240
"""Cancel pending disconnect task and invoke connect handlers."""
240241
self._cancel_delete_task()
241-
self._num_connections += 1
242+
self._num_connections[socket_id] += 1
242243
if next_message_id is not None:
243244
self.outbox.try_rewind(next_message_id)
244245
storage.request_contextvar.set(self.request)
@@ -247,18 +248,19 @@ def handle_handshake(self, next_message_id: Optional[int]) -> None:
247248
for t in core.app._connect_handlers: # pylint: disable=protected-access
248249
self.safe_invoke(t)
249250

250-
def handle_disconnect(self) -> None:
251+
def handle_disconnect(self, socket_id: str) -> None:
251252
"""Wait for the browser to reconnect; invoke disconnect handlers if it doesn't."""
252253
self._cancel_delete_task()
253-
self._num_connections -= 1
254+
self._num_connections[socket_id] -= 1
254255
for t in self.disconnect_handlers:
255256
self.safe_invoke(t)
256257
for t in core.app._disconnect_handlers: # pylint: disable=protected-access
257258
self.safe_invoke(t)
258259
if not self.shared:
259260
async def delete_content() -> None:
260261
await asyncio.sleep(self.page.resolve_reconnect_timeout())
261-
if self._num_connections == 0:
262+
if self._num_connections[socket_id] == 0:
263+
self._num_connections.pop(socket_id)
262264
self.delete()
263265
self._delete_task = background_tasks.create(delete_content())
264266

nicegui/nicegui.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ async def _on_handshake(sid: str, data: Dict[str, Any]) -> bool:
171171
else:
172172
client.environ = sio.get_environ(sid)
173173
await sio.enter_room(sid, client.id)
174-
client.handle_handshake(data.get('next_message_id'))
174+
client.handle_handshake(sid, data.get('next_message_id'))
175175
assert client.tab_id is not None
176176
await core.app.storage._create_tab_storage(client.tab_id) # pylint: disable=protected-access
177177
return True
@@ -184,7 +184,7 @@ def _on_disconnect(sid: str) -> None:
184184
client_id = query['client_id'][0]
185185
client = Client.instances.get(client_id)
186186
if client:
187-
client.handle_disconnect()
187+
client.handle_disconnect(sid)
188188

189189

190190
@sio.on('event')

0 commit comments

Comments
 (0)