Handling downloads cancelled on the browser

Author: darrenburns

src/textual_serve/app_service.py

@@ -187,6 +187,10 @@ async def start(self, width: int, height: int) -> None:
     async def stop(self) -> None:
         """Stop the process and wait for it to complete."""
         if self._task is not None:
+            await self._download_manager.cancel_app_downloads(
+                app_service_id=self.app_service_id
+            )
+
             await self.send_meta({"type": "quit"})
             await self._task
             self._task = None

src/textual_serve/download_manager.py

@@ -116,7 +116,8 @@ async def download(self, delivery_key: str) -> AsyncGenerator[bytes, None]:
                 chunk = None
 
             if not chunk:
-                # Empty chunk - the app process has finished sending the file.
+                # Empty chunk - the app process has finished sending the file
+                # or the download has been cancelled.
                 incoming_chunks.task_done()
                 del self._active_downloads[delivery_key]
                 break
@@ -131,7 +132,14 @@ async def chunk_received(self, delivery_key: str, chunk: bytes | str) -> None:
             delivery_key: The delivery key that the chunk was received for.
             chunk: The chunk that was received.
         """
-        download = self._active_downloads[delivery_key]
+
+        download = self._active_downloads.get(delivery_key)
+        if not download:
+            # The download may have been cancelled - e.g. the websocket
+            # was closed before the download could complete.
+            log.debug("Chunk received for cancelled download %r", delivery_key)
+            return
+
         if isinstance(chunk, str):
             chunk = chunk.encode(download.encoding or "utf-8")
         await download.incoming_chunks.put(chunk)
@@ -155,3 +163,13 @@ async def get_download_metadata(self, delivery_key: str) -> Download:
             delivery_key: The delivery key to get the metadata for.
         """
         return self._active_downloads[delivery_key]
+
+    async def cancel_app_downloads(self, app_service_id: str) -> None:
+        """Cancel all downloads for the given app service.
+
+        Args:
+            app_service_id: The app service ID to cancel downloads for.
+        """
+        for download in self._active_downloads.values():
+            if download.app_service.app_service_id == app_service_id:
+                await download.incoming_chunks.put(None)