provide optional socket in page.on_disconnect

main.py

@@ -461,7 +461,7 @@ You can run a function or coroutine as a parallel task by passing it to one of t
 - `ui.on_shutdown`: Called when NiceGUI is shut down or restarted.
 - `ui.on_connect`: Called when a client connects to NiceGUI. (Optional argument: Starlette request)
 - `ui.on_page_ready`: Called when the page is ready and the websocket is connected. (Optional argument: socket)
-- `ui.on_disconnect`: Called when a client disconnects from NiceGUI.
+- `ui.on_disconnect`: Called when a client disconnects from NiceGUI. (Optional argument: socket)
 
 When NiceGUI is shut down or restarted, the startup tasks will be automatically canceled.
 '''

nicegui/elements/page.py

@@ -89,7 +89,14 @@ class Page(jp.QuasarPage):
 
     async def on_disconnect(self, websocket=None) -> None:
         for disconnect_handler in ([self.disconnect_handler] if self.disconnect_handler else []) + disconnect_handlers:
-            await disconnect_handler() if is_coroutine(disconnect_handler) else disconnect_handler()
+            arg_count = len(inspect.signature(disconnect_handler).parameters)
+            is_coro = is_coroutine(disconnect_handler)
+            if arg_count == 1:
+                await disconnect_handler(websocket) if is_coro else disconnect_handler(websocket)
+            elif arg_count == 0:
+                await disconnect_handler() if is_coro else disconnect_handler()
+            else:
+                raise ValueError(f'invalid number of arguments (0 or 1 allowed, got {arg_count})')
         await super().on_disconnect(websocket)
 
     def __enter__(self):