diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index c1cc753c88a..380e17fe792 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -186,10 +186,13 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b logger.warning(f"Rejecting socket {sid}: unable to verify user record") return False - # Store user_id and is_admin in socket users dict + # Store user_id and is_admin in socket users dict. `authenticated` records + # which half of this handler admitted the socket, so the disconnect message + # can be logged at the same level as the connect message below. self._socket_users[sid] = { "user_id": token_data.user_id, "is_admin": token_data.is_admin, + "authenticated": True, } logger.info( f"Socket {sid} connected with user_id: {token_data.user_id}, is_admin: {token_data.is_admin}" @@ -212,6 +215,7 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b self._socket_users[sid] = { "user_id": "system", "is_admin": True, + "authenticated": False, } logger.debug(f"Socket {sid} connected as system admin (single-user mode)") await self._sio.enter_room(sid, "user:system") @@ -234,11 +238,41 @@ def _is_multiuser_enabled() -> bool: # so we never accidentally admit an anonymous socket. return True - async def _handle_disconnect(self, sid: str) -> None: - """Handle socket disconnection and cleanup user info.""" - if sid in self._socket_users: - del self._socket_users[sid] - logger.debug(f"Socket {sid} disconnected and cleaned up") + async def _handle_disconnect(self, sid: str, reason: str | None = None) -> None: + """Handle socket disconnection and cleanup user info. + + Logged at the same level as the matching connect message. When only the connect + half is visible at INFO, a client that reconnects in a loop — a backgrounded tab + whose timers the browser has throttled, a flaky network — is indistinguishable + from an unbounded pile of accumulating sockets. + + `reason` is python-socketio's disconnect reason — one of `ping timeout`, + `transport close`, `transport error`, `client disconnect`, `server disconnect` — + which is the first thing worth knowing when sockets are churning. Versions from + 5.12 always pass one; `python-socketio` is unpinned, so an older install calls this + with `sid` alone and `reason` falls back to its default. + + Two constraints on this body, both imposed by how python-socketio invokes it: + + - It must not raise. `AsyncServer._handle_disconnect` does not guard the + `_trigger_event` call, so an exception here skips `manager.disconnect()` and + leaves the sid in its rooms and in `server.environ` for the life of the process. + - The `pop` must stay first. `_trigger_event` retries `disconnect` handlers on + `TypeError` with one fewer argument, and that retry wraps the *await of the + handler*, not just the argument binding — so a `TypeError` raised anywhere below + would silently re-enter this method. + """ + user_info = self._socket_users.pop(sid, None) + if user_info is None: + return + + # `.get` throughout: entries are populated by convention, not by a schema, and a + # KeyError here would cost the caller its cleanup (see above). + message = f"Socket {sid} disconnected (user_id: {user_info.get('user_id')}, reason: {reason or 'unknown'})" + if user_info.get("authenticated"): + logger.info(message) + else: + logger.debug(message) async def _handle_sub_queue(self, sid: str, data: Any) -> None: """Handle queue subscription and add socket to both queue and user-specific rooms.""" @@ -258,6 +292,7 @@ async def _handle_sub_queue(self, sid: str, data: Any) -> None: self._socket_users[sid] = { "user_id": "system", "is_admin": True, + "authenticated": False, } user_id = self._socket_users[sid]["user_id"] diff --git a/tests/app/test_workflow_socketio.py b/tests/app/test_workflow_socketio.py index ba317e35886..c092145a918 100644 --- a/tests/app/test_workflow_socketio.py +++ b/tests/app/test_workflow_socketio.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import ANY, AsyncMock +from unittest.mock import ANY, AsyncMock, Mock import pytest from fastapi import FastAPI @@ -178,3 +178,77 @@ async def test_shared_to_private_transition_emits_access_revoked_to_shared_room( data={"workflow_id": "wf-1", "user_id": "owner-1", "timestamp": ANY}, room="workflows:shared", ) + + +@pytest.mark.anyio +async def test_authenticated_socket_logs_disconnect_at_same_level_as_connect(monkeypatch: pytest.MonkeyPatch) -> None: + """A connect logged at INFO must be matched by a disconnect logged at INFO. + + Otherwise a client that reconnects in a loop looks exactly like sockets piling up. + """ + socketio = SocketIO(FastAPI()) + socketio._sio.enter_room = AsyncMock() + _patch_multiuser_context(monkeypatch, user_id="user-1", is_admin=False) + await socketio._handle_connect("sid-1", {}, {"token": "valid-token"}) + + log = SimpleNamespace(info=Mock(), debug=Mock(), warning=Mock(), error=Mock()) + monkeypatch.setattr("invokeai.app.api.sockets.logger", log) + + await socketio._handle_disconnect("sid-1", "ping timeout") + + log.info.assert_called_once() + message = log.info.call_args.args[0] + assert "sid-1" in message + assert "user-1" in message + assert "ping timeout" in message + log.debug.assert_not_called() + assert "sid-1" not in socketio._socket_users + + +@pytest.mark.anyio +async def test_single_user_socket_logs_disconnect_at_debug(monkeypatch: pytest.MonkeyPatch) -> None: + """The single-user connect is logged at DEBUG, so its disconnect must be too.""" + socketio = SocketIO(FastAPI()) + socketio._sio.enter_room = AsyncMock() + _patch_single_user_context(monkeypatch) + await socketio._handle_connect("sid-1", {}, None) + + log = SimpleNamespace(info=Mock(), debug=Mock(), warning=Mock(), error=Mock()) + monkeypatch.setattr("invokeai.app.api.sockets.logger", log) + + await socketio._handle_disconnect("sid-1", "transport close") + + log.debug.assert_called_once() + log.info.assert_not_called() + assert "sid-1" not in socketio._socket_users + + +@pytest.mark.anyio +async def test_disconnect_of_unknown_socket_is_silent(monkeypatch: pytest.MonkeyPatch) -> None: + """An unknown sid must not raise: an exception here would cost python-socketio its own + cleanup (`AsyncServer._handle_disconnect` does not guard the handler call, so the sid + would stay in its rooms and in `server.environ` for the life of the process).""" + socketio = SocketIO(FastAPI()) + log = SimpleNamespace(info=Mock(), debug=Mock(), warning=Mock(), error=Mock()) + monkeypatch.setattr("invokeai.app.api.sockets.logger", log) + + await socketio._handle_disconnect("never-connected") + + log.info.assert_not_called() + log.debug.assert_not_called() + + +@pytest.mark.anyio +async def test_disconnect_without_reason_is_accepted(monkeypatch: pytest.MonkeyPatch) -> None: + """`python-socketio` is unpinned; versions before 5.12 call the handler with `sid` alone.""" + socketio = SocketIO(FastAPI()) + socketio._sio.enter_room = AsyncMock() + _patch_multiuser_context(monkeypatch, user_id="user-1", is_admin=False) + await socketio._handle_connect("sid-1", {}, {"token": "valid-token"}) + + log = SimpleNamespace(info=Mock(), debug=Mock(), warning=Mock(), error=Mock()) + monkeypatch.setattr("invokeai.app.api.sockets.logger", log) + + await socketio._handle_disconnect("sid-1") + + assert "unknown" in log.info.call_args.args[0]