From bed9d9220d272bf9838fd4b07025f621332e7363 Mon Sep 17 00:00:00 2001
From: ciaranbor <ciaranborourke-dev@proton.me>
Date: Wed, 1 Apr 2026 19:52:54 +0100
Subject: [PATCH] Update download progress on pause

---
 src/exo/download/coordinator.py                |  7 +++++++
 src/exo/download/tests/test_cancel_download.py | 11 ++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
index a7a77064..39741d94 100644
--- a/src/exo/download/coordinator.py
+++ b/src/exo/download/coordinator.py
@@ -158,10 +158,17 @@ class DownloadCoordinator:
             logger.info(f"Cancelling download for {model_id}")
             self.active_downloads[model_id].cancel()
             current_status = self.download_status[model_id]
+            downloaded = Memory()
+            total = Memory()
+            if isinstance(current_status, DownloadOngoing):
+                downloaded = current_status.download_progress.downloaded
+                total = current_status.download_progress.total
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
                 model_directory=self._default_model_dir(model_id),
+                downloaded=downloaded,
+                total=total,
             )
             self.download_status[model_id] = pending
             await self.event_sender.send(
diff --git a/src/exo/download/tests/test_cancel_download.py b/src/exo/download/tests/test_cancel_download.py
index 740e28bc..3d02d65f 100644
--- a/src/exo/download/tests/test_cancel_download.py
+++ b/src/exo/download/tests/test_cancel_download.py
@@ -189,6 +189,14 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
         # Wait for the download to actually start (blocking in ensure_shard)
         await asyncio.wait_for(slow_downloader.download_started.wait(), timeout=2.0)
 
+        # Drain any events emitted before the cancel (initial DownloadPending, DownloadOngoing)
+        while True:
+            try:
+                async with asyncio.timeout(0.1):
+                    await event_recv.receive()
+            except TimeoutError:
+                break
+
         # Cancel the download
         await cmd_send.send(
             ForwarderDownloadCommand(
@@ -197,10 +205,11 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
             )
         )
 
-        # Should receive a DownloadPending event
+        # Should receive a DownloadPending event with preserved progress
         pending = await _wait_for_pending(event_recv, MODEL_ID)
         assert pending is not None, "Cancel should emit DownloadPending"
         assert pending.shard_metadata.model_card.model_id == MODEL_ID
+        assert pending.total == Memory.from_mb(100), "Should preserve total bytes"
 
         # Give coordinator time to clean up
         await asyncio.sleep(0.05)