fix(api): log socket disconnects at the same level as connects

invokeai/app/api/sockets.py

@@ -235,6 +235,7 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b
                     "user_id": token_data.user_id,
                     "is_admin": is_admin,
                     "token_epoch": token_epoch,
+                    "authenticated": True,
                 }
                 logger.info(f"Socket {sid} connected with user_id: {token_data.user_id}, is_admin: {is_admin}")
                 await self._sio.enter_room(sid, f"user:{token_data.user_id}")
@@ -255,6 +256,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")
@@ -277,21 +279,48 @@ 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:
-            user_id = self._socket_users[sid].get("user_id")
-            del self._socket_users[sid]
-            # Forget the user's revalidation failures once their last socket is gone. The
-            # count stands in for how long *these* sockets have gone unchecked, so a user
-            # who reconnects must start over. The sweep drops orphaned counters too, but it
-            # cannot be relied on here: `_revalidation_loop` still skips the sweep entirely
-            # when nothing is live, which a departing last socket can leave it. Without
-            # this, a user whose only socket closed part-way through a database outage
-            # would come back one failed read away from losing the admin room.
-            if user_id is not None and not any(info.get("user_id") == user_id for info in self._socket_users.values()):
-                self._revalidation_failures.pop(user_id, None)
-            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
+
+        user_id = user_info.get("user_id")
+        # Forget the user's revalidation failures once their last socket is gone. The
+        # count stands in for how long these sockets have gone unchecked, so a user who
+        # reconnects must start over.
+        if user_id is not None and not any(info.get("user_id") == user_id for info in self._socket_users.values()):
+            self._revalidation_failures.pop(user_id, None)
+
+        # `.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)
 
     def start(self) -> None:
         """Start background work. Called from the app's startup hook.
@@ -713,6 +742,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"]

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
@@ -185,3 +185,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._sio._trigger_event("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]