Enhancement: websocket heartbeat (#13739)

---------

Co-authored-by: shamoon <4887959+shamoon@users.noreply.github.com>

src-ui/src/app/services/websocket-status.service.spec.ts

@@ -438,6 +438,21 @@ describe('ConsumerStatusService', () => {
     expect(updated).toBeTruthy()
   })
 
+  it('should ignore keep-alive heartbeat messages from the server', () => {
+    let updated = false
+    let deleted = false
+    websocketStatusService.onDocumentUpdated().subscribe(() => (updated = true))
+    websocketStatusService.onDocumentDeleted().subscribe(() => (deleted = true))
+
+    websocketStatusService.connect()
+    server.send({ type: WebsocketStatusType.HEARTBEAT })
+
+    expect(updated).toBeFalsy()
+    expect(deleted).toBeFalsy()
+    expect(websocketStatusService.getConsumerStatus()).toHaveLength(0)
+    websocketStatusService.disconnect()
+  })
+
   it('should ignore document updated events the user cannot view', () => {
     let updated = false
     websocketStatusService.onDocumentUpdated().subscribe(() => {

src-ui/src/app/services/websocket-status.service.ts

@@ -11,6 +11,7 @@ export enum WebsocketStatusType {
   STATUS_UPDATE = 'status_update',
   DOCUMENTS_DELETED = 'documents_deleted',
   DOCUMENT_UPDATED = 'document_updated',
+  HEARTBEAT = 'heartbeat',
 }
 
 // see ProgressStatusOptions in src/documents/plugins/helpers.py
@@ -207,6 +208,10 @@ export class WebsocketStatusService {
         case WebsocketStatusType.STATUS_UPDATE:
           this.handleProgressUpdate(messageData as WebsocketProgressMessage)
           break
+
+        case WebsocketStatusType.HEARTBEAT:
+          // keep-alive from the server, see paperless.consumers.StatusConsumer
+          break
       }
     }
   }

src/paperless/consumers.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+import contextlib
 import json
 from typing import TYPE_CHECKING
 
@@ -14,8 +16,13 @@ if TYPE_CHECKING:
     from documents.plugins.helpers import PermissionsData
     from documents.plugins.helpers import StatusUpdatePayload
 
+HEARTBEAT_INTERVAL = 30
+HEARTBEAT_MESSAGE = json.dumps({"type": "heartbeat"})
+
 
 class StatusConsumer(AsyncWebsocketConsumer):
+    heartbeat_task: asyncio.Task | None = None
+
     def _authenticated(self) -> bool:
         user: AbstractBaseUser | AnonymousUser | None = self.scope.get("user")
         return user is not None and user.is_authenticated
@@ -39,10 +46,27 @@ class StatusConsumer(AsyncWebsocketConsumer):
             return
         await self.channel_layer.group_add("status_updates", self.channel_name)
         await self.accept()
+        self._start_heartbeat()
 
     async def disconnect(self, code: int) -> None:
+        await self._stop_heartbeat()
         await self.channel_layer.group_discard("status_updates", self.channel_name)
 
+    def _start_heartbeat(self) -> None:
+        self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+    async def _stop_heartbeat(self) -> None:
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self.heartbeat_task
+            self.heartbeat_task = None
+
+    async def _heartbeat_loop(self) -> None:
+        while True:
+            await asyncio.sleep(HEARTBEAT_INTERVAL)
+            await self.send(HEARTBEAT_MESSAGE)
+
     async def status_update(self, event: StatusUpdatePayload) -> None:
         if not self._authenticated():
             await self.close()

src/paperless/tests/test_websockets.py

@@ -189,6 +189,23 @@ class TestWebSockets:
 
         await communicator.disconnect()
 
+    @pytest.mark.anyio
+    async def test_heartbeat(self, mocker: MockerFixture) -> None:
+        mocker.patch(
+            "paperless.consumers.StatusConsumer._authenticated",
+            return_value=True,
+        )
+        mocker.patch("paperless.consumers.HEARTBEAT_INTERVAL", 0.01)
+
+        communicator = WebsocketCommunicator(application, "/ws/status/")
+        connected, _ = await communicator.connect()
+        assert connected
+
+        assert await communicator.receive_json_from() == {"type": "heartbeat"}
+        assert await communicator.receive_json_from() == {"type": "heartbeat"}
+
+        await communicator.disconnect()
+
     def test_manager_send_progress(self, mocker: MockerFixture) -> None:
         mock_group_send = mocker.patch(
             "channels.layers.InMemoryChannelLayer.group_send",