fix(download): self-heal aria2 transfers lost on daemon restart

py/services/aria2_downloader.py

@@ -24,6 +24,10 @@ from .settings_manager import get_settings_manager
 
 logger = logging.getLogger(__name__)
 
+# Maximum times the download poll loop will re-schedule a transfer after it
+# is lost (daemon restart / RPC outage) before failing the download.
+MAX_TRANSFER_RECOVERY_ATTEMPTS = 2
+
 def _try_certifi_ca_path() -> str | None:
     """Return the certifi CA bundle path if available, else None."""
     try:
@@ -85,6 +89,7 @@ class Aria2Downloader:
         self._rpc_session: Optional[aiohttp.ClientSession] = None
         self._rpc_session_lock = asyncio.Lock()
         self._process_lock = asyncio.Lock()
+        self._register_lock = asyncio.Lock()
         self._transfers: Dict[str, Aria2Transfer] = {}
         self._poll_interval = 0.5
         self._state_store = Aria2TransferStateStore()
@@ -103,26 +108,58 @@ class Aria2Downloader:
         progress_callback=None,
         headers: Optional[Dict[str, str]] = None,
     ) -> Tuple[bool, str]:
-        """Download a file using aria2 RPC and wait for completion."""
+        """Download a file using aria2 RPC and wait for completion.
+
+        The poll loop is self-healing: when the in-memory transfer entry
+        disappears (e.g. another download restarted the daemon and
+        ``close()`` cleared ``_transfers``) or the RPC becomes unreachable,
+        the transfer is re-scheduled with ``continue=true`` so the download
+        resumes from the on-disk ``.aria2`` control file. Recovery is bounded
+        by ``MAX_TRANSFER_RECOVERY_ATTEMPTS``.
+        """
 
         await self._ensure_process()
         save_path = os.path.abspath(save_path)
-        transfer = self._transfers.get(download_id)
-        if transfer is None or os.path.abspath(transfer.save_path) != save_path:
-            gid = await self._schedule_download(
-                url,
-                save_path,
-                download_id=download_id,
-                headers=headers,
-            )
-            transfer = Aria2Transfer(gid=gid, save_path=save_path)
-            self._transfers[download_id] = transfer
 
+        async with self._register_lock:
+            transfer = self._transfers.get(download_id)
+            if transfer is None or os.path.abspath(transfer.save_path) != save_path:
+                transfer = await self._register_transfer(
+                    url,
+                    save_path,
+                    download_id=download_id,
+                    headers=headers,
+                )
+
+        recovery_attempts = 0
         try:
             while True:
-                status = await self._get_status_with_retry(download_id)
+                try:
+                    status = await self._get_status_with_retry(download_id)
+                except Aria2Error:
+                    status = None
+
                 if status is None:
-                    return False, "aria2 download not found"
+                    if recovery_attempts >= MAX_TRANSFER_RECOVERY_ATTEMPTS:
+                        return False, "aria2 download not found"
+                    recovery_attempts += 1
+                    logger.warning(
+                        "aria2 transfer %s lost; re-scheduling with resume "
+                        "(attempt %d/%d)",
+                        download_id,
+                        recovery_attempts,
+                        MAX_TRANSFER_RECOVERY_ATTEMPTS,
+                    )
+                    await asyncio.sleep(1.0)
+                    await self._ensure_process()
+                    async with self._register_lock:
+                        transfer = await self._register_transfer(
+                            url,
+                            save_path,
+                            download_id=download_id,
+                            headers=headers,
+                        )
+                    continue
 
                 snapshot = self._build_progress_snapshot(status)
                 if progress_callback is not None:
@@ -139,7 +176,9 @@ class Aria2Downloader:
 
                 await asyncio.sleep(self._poll_interval)
         finally:
-            self._transfers.pop(download_id, None)
+            current = self._transfers.get(download_id)
+            if current is not None and current.gid == transfer.gid:
+                self._transfers.pop(download_id, None)
 
     async def _get_status_with_retry(
         self, download_id: str, *, max_retries: int = 4, retry_delay: float = 3.0
@@ -242,6 +281,25 @@ class Aria2Downloader:
         )
         return gid
 
+    async def _register_transfer(
+        self,
+        url: str,
+        save_path: str,
+        *,
+        download_id: str,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Aria2Transfer:
+        """Schedule a download and track it in the in-memory transfer registry."""
+        gid = await self._schedule_download(
+            url,
+            save_path,
+            download_id=download_id,
+            headers=headers,
+        )
+        transfer = Aria2Transfer(gid=gid, save_path=os.path.abspath(save_path))
+        self._transfers[download_id] = transfer
+        return transfer
+
     async def get_status(self, download_id: str) -> Optional[Dict[str, Any]]:
         """Return the raw aria2 status payload for a known download."""
 

tests/services/test_aria2_downloader.py

@@ -6,7 +6,12 @@ from unittest.mock import AsyncMock
 
 import pytest
 
-from py.services.aria2_downloader import Aria2Downloader, Aria2Error, Aria2Transfer
+from py.services.aria2_downloader import (
+    Aria2Downloader,
+    Aria2Error,
+    Aria2Transfer,
+    MAX_TRANSFER_RECOVERY_ATTEMPTS,
+)
 from py.services.aria2_transfer_state import Aria2TransferStateStore
 from py.services import aria2_transfer_state
 
@@ -246,6 +251,253 @@ async def test_download_file_reuses_existing_transfer_without_add_uri(
     assert [call[0] for call in rpc_calls] == ["aria2.tellStatus", "aria2.tellStatus"]
 
 
+@pytest.mark.asyncio
+async def test_download_file_recovers_when_transfer_lost_mid_poll(
+    tmp_path, monkeypatch
+):
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+    poll_count = {"n": 0}
+
+    async def fake_rpc_call(method, params):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return "gid-1" if add_uri_count["n"] == 1 else "gid-2"
+        if method == "aria2.tellStatus":
+            poll_count["n"] += 1
+            if poll_count["n"] == 1:
+                # Simulate a concurrent close() wiping the transfer mid-poll.
+                downloader._transfers.pop("download-1", None)
+                return {
+                    "gid": "gid-1",
+                    "status": "active",
+                    "completedLength": "5",
+                    "totalLength": "10",
+                    "downloadSpeed": "25",
+                }
+            return {
+                "gid": "gid-2",
+                "status": "complete",
+                "completedLength": "10",
+                "totalLength": "10",
+                "downloadSpeed": "0",
+                "files": [{"path": str(save_path)}],
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is True
+    assert result == str(save_path)
+    assert add_uri_count["n"] == 2
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_recovers_when_rpc_fails_mid_poll(tmp_path, monkeypatch):
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+    poll_count = {"n": 0}
+
+    async def fake_rpc_call(method, params):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return "gid-1" if add_uri_count["n"] == 1 else "gid-2"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    async def fake_get_status_with_retry(download_id):
+        poll_count["n"] += 1
+        if poll_count["n"] == 1:
+            raise Aria2Error(
+                "Failed to query aria2 download status after 4 attempts: boom"
+            )
+        return {
+            "gid": "gid-2",
+            "status": "complete",
+            "completedLength": "10",
+            "totalLength": "10",
+            "downloadSpeed": "0",
+            "files": [{"path": str(save_path)}],
+        }
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr(downloader, "_get_status_with_retry", fake_get_status_with_retry)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is True
+    assert result == str(save_path)
+    assert add_uri_count["n"] == 2
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_fails_after_recovery_attempts_exhausted(
+    tmp_path, monkeypatch
+):
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+
+    async def fake_rpc_call(method, params):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return f"gid-{add_uri_count['n']}"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    async def fake_get_status(download_id):
+        return None  # transfer never tracked / always lost
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr(downloader, "get_status", fake_get_status)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is False
+    assert result == "aria2 download not found"
+    assert add_uri_count["n"] == 1 + MAX_TRANSFER_RECOVERY_ATTEMPTS
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_concurrent_same_id_schedules_once(tmp_path, monkeypatch):
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_count = {"n": 0}
+    poll_count = {"n": 0}
+
+    async def fake_rpc_call(method, params):
+        if method == "aria2.addUri":
+            add_uri_count["n"] += 1
+            return "gid-1"
+        if method == "aria2.tellStatus":
+            poll_count["n"] += 1
+            if poll_count["n"] < 4:
+                return {
+                    "gid": "gid-1",
+                    "status": "active",
+                    "completedLength": "5",
+                    "totalLength": "10",
+                    "downloadSpeed": "25",
+                }
+            return {
+                "gid": "gid-1",
+                "status": "complete",
+                "completedLength": "10",
+                "totalLength": "10",
+                "downloadSpeed": "0",
+                "files": [{"path": str(save_path)}],
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    results = await asyncio.gather(
+        downloader.download_file(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        ),
+        downloader.download_file(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        ),
+    )
+
+    assert all(success for success, _ in results)
+    assert all(result == str(save_path) for _, result in results)
+    assert add_uri_count["n"] <= 2
+    assert downloader._transfers == {}
+
+
+@pytest.mark.asyncio
+async def test_download_file_cleanup_preserves_newer_registration(tmp_path, monkeypatch):
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    poll_count = {"n": 0}
+
+    async def fake_rpc_call(method, params):
+        if method == "aria2.addUri":
+            return "gid-1"
+        if method == "aria2.tellStatus":
+            poll_count["n"] += 1
+            if poll_count["n"] == 1:
+                # Simulate another invocation registering its own transfer.
+                downloader._transfers["download-1"] = Aria2Transfer(
+                    gid="gid-new", save_path=str(save_path)
+                )
+                return {
+                    "gid": "gid-1",
+                    "status": "active",
+                    "completedLength": "5",
+                    "totalLength": "10",
+                    "downloadSpeed": "25",
+                }
+            return {
+                "gid": "gid-new",
+                "status": "complete",
+                "completedLength": "10",
+                "totalLength": "10",
+                "downloadSpeed": "0",
+                "files": [{"path": str(save_path)}],
+            }
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+    monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
+
+    success, result = await downloader.download_file(
+        "https://example.com/model.safetensors",
+        str(save_path),
+        download_id="download-1",
+    )
+
+    assert success is True
+    assert result == str(save_path)
+    assert downloader._transfers["download-1"].gid == "gid-new"
+
+
 def test_build_progress_snapshot_normalizes_numeric_fields():
     downloader = Aria2Downloader()