Skip to content

Commit

Permalink
use socket ID to identify connections on auto-index client
Browse files Browse the repository at this point in the history
  • Loading branch information
falkoschindler committed Jan 31, 2025
1 parent 5ae38b3 commit 31219a1
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 10 deletions.
4 changes: 2 additions & 2 deletions nicegui/air.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ def _handle_handshake(data: Dict[str, Any]) -> bool:
core.app.storage.copy_tab(data['old_tab_id'], data['tab_id'])
client.tab_id = data['tab_id']
client.on_air = True
client.handle_handshake(data.get('next_message_id'))
client.handle_handshake(data['sid'], data.get('next_message_id'))
return True

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

@self.relay.on('connect')
async def _handle_connect() -> None:
Expand Down
14 changes: 8 additions & 6 deletions nicegui/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import inspect
import time
import uuid
from collections import defaultdict
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable, ClassVar, Dict, Iterable, Iterator, List, Optional, Union
Expand Down Expand Up @@ -59,7 +60,7 @@ def __init__(self, page: page, *, request: Optional[Request]) -> None:
self.environ: Optional[Dict[str, Any]] = None
self.shared = request is None
self.on_air = False
self._num_connections = 0
self._num_connections = defaultdict(int)
self._delete_task: Optional[asyncio.Task] = None
self._deleted = False
self.tab_id: Optional[str] = None
Expand Down Expand Up @@ -235,10 +236,10 @@ def on_disconnect(self, handler: Union[Callable[..., Any], Awaitable]) -> None:
"""Add a callback to be invoked when the client disconnects."""
self.disconnect_handlers.append(handler)

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

def handle_disconnect(self) -> None:
def handle_disconnect(self, socket_id: str) -> None:
"""Wait for the browser to reconnect; invoke disconnect handlers if it doesn't."""
self._cancel_delete_task()
self._num_connections -= 1
self._num_connections[socket_id] -= 1
for t in self.disconnect_handlers:
self.safe_invoke(t)
for t in core.app._disconnect_handlers: # pylint: disable=protected-access
self.safe_invoke(t)
if not self.shared:
async def delete_content() -> None:
await asyncio.sleep(self.page.resolve_reconnect_timeout())
if self._num_connections == 0:
if self._num_connections[socket_id] == 0:
self._num_connections.pop(socket_id)
self.delete()
self._delete_task = background_tasks.create(delete_content())

Expand Down
4 changes: 2 additions & 2 deletions nicegui/nicegui.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ async def _on_handshake(sid: str, data: Dict[str, Any]) -> bool:
else:
client.environ = sio.get_environ(sid)
await sio.enter_room(sid, client.id)
client.handle_handshake(data.get('next_message_id'))
client.handle_handshake(sid, data.get('next_message_id'))
assert client.tab_id is not None
await core.app.storage._create_tab_storage(client.tab_id) # pylint: disable=protected-access
return True
Expand All @@ -184,7 +184,7 @@ def _on_disconnect(sid: str) -> None:
client_id = query['client_id'][0]
client = Client.instances.get(client_id)
if client:
client.handle_disconnect()
client.handle_disconnect(sid)


@sio.on('event')
Expand Down

0 comments on commit 31219a1

Please sign in to comment.