Heartbeat mechanism to check if receiver is connected

Author: codeSamuraii

lib/store.py

@@ -117,6 +117,14 @@ async def get_progress(self) -> int:
         progress = await self.redis.get(self._k_progress)
         return int(progress) if progress else 0
 
+    async def set_receiver_active(self) -> None:
+        """Mark receiver as actively downloading with TTL."""
+        await self.redis.set(self.key('receiver_active'), '1', ex=5)
+
+    async def is_receiver_active(self) -> bool:
+        """Check if receiver is actively downloading."""
+        return bool(await self.redis.exists(self.key('receiver_active')))
+
     async def cleanup(self) -> None:
         """Delete all transfer data."""
         pattern = self.key('*')

lib/transfer.py

@@ -60,6 +60,11 @@ def _format_uid(uid: str):
     def get_file_info(self):
         return self.file.name, self.file.size, self.file.type
 
+    @property
+    async def receiver_connected(self) -> bool:
+        """Check if a receiver is actively downloading."""
+        return await self.store.is_receiver_active()
+
     async def notify_receiver_connected(self):
         """Notify sender that receiver connected."""
         await self.store.set_event('receiver_connected')
@@ -109,6 +114,8 @@ async def produce_download(self, on_error: Callable[[Exception | str], Awaitable
             self.info(f"▼ Resuming from byte {self.bytes_downloaded}")
 
         try:
+            await self.store.set_receiver_active()
+
             async for chunk in self.store.stream_chunks():
                 if chunk == self.DEAD_FLAG:
                     raise TransferError("Sender disconnected")
@@ -120,6 +127,7 @@ async def produce_download(self, on_error: Callable[[Exception | str], Awaitable
 
                 self.bytes_downloaded += len(chunk)
                 await self.store.save_progress(self.bytes_downloaded)
+                await self.store.set_receiver_active()
                 yield chunk
 
         except TransferError as e:

views/http.py

@@ -108,9 +108,9 @@ async def http_download(request: Request, uid: str):
 
     if not is_curl and not request.query_params.get('download'):
         log.info(f"▼ Browser request detected, serving download page")
-        progress = await transfer.store.get_progress()
+        receiver_connected = await transfer.receiver_connected
         return templates.TemplateResponse(request, "download.html",
-            transfer.file.to_readable_dict() | {'receiver_connected': progress > 0})
+            transfer.file.to_readable_dict() | {'receiver_connected': receiver_connected})
 
     await transfer.notify_receiver_connected()