fix(download): stop aria2 from leaking transfers when a download is cancelled

A cancel landing between aria2.addUri acceptance and the _transfers
registration found no tracked transfer, so DownloadManager tolerated the
"not found" and only cancelled the asyncio task — the daemon kept
downloading the file untracked while history showed the download as
cancelled.

- Register the gid in _transfers immediately after addUri returns,
  before any further await (state-store persist moved after it)
- Shield the addUri RPC so a mid-flight cancellation still learns the
  accepted gid and forceRemoves it before re-raising CancelledError
- On cancellation during the state persist, remove the daemon transfer
  unless it is paused (skip_download relies on paused gids surviving)

py/services/aria2_downloader.py

@@ -161,6 +161,11 @@ class Aria2Downloader:
         (typically an expired CivitAI signed URL): a fresh URL is resolved
         and the partial download continues.  Recovery is bounded by
         ``MAX_TRANSFER_RECOVERY_ATTEMPTS``.
+
+        Cancellation never leaks daemon transfers: the gid is tracked in
+        ``_transfers`` before any post-``addUri`` await, and a gid accepted
+        by the daemon while the caller is being cancelled is removed again
+        before the ``CancelledError`` propagates.
         """
 
         await self._ensure_process()
@@ -251,7 +256,11 @@ class Aria2Downloader:
                 await asyncio.sleep(self._poll_interval)
         finally:
             current = self._transfers.get(download_id)
-            if current is not None and current.gid == transfer.gid:
+            if (
+                transfer is not None
+                and current is not None
+                and current.gid == transfer.gid
+            ):
                 self._transfers.pop(download_id, None)
 
     async def _get_status_with_retry(
@@ -339,21 +348,43 @@ class Aria2Downloader:
             resolved_url != url,
         )
 
+        # Shield the addUri RPC from cancellation: the daemon may accept the
+        # download even when the caller is cancelled while the request is in
+        # flight.  On cancellation, wait for the RPC result so the freshly
+        # created gid can be removed instead of leaking an untracked
+        # download that keeps running in the daemon.
+        add_task = asyncio.ensure_future(
+            self._rpc_call("aria2.addUri", [[resolved_url], options])
+        )
         try:
-            gid = await self._rpc_call("aria2.addUri", [[resolved_url], options])
+            gid = await asyncio.shield(add_task)
+        except asyncio.CancelledError:
+            leaked_gid: Any = None
+            try:
+                leaked_gid = await add_task
+            except Exception:
+                leaked_gid = None
+            if isinstance(leaked_gid, str) and leaked_gid:
+                logger.info(
+                    "Removing aria2 gid %s accepted while download %s was "
+                    "being cancelled",
+                    leaked_gid,
+                    download_id,
+                )
+                try:
+                    await self._rpc_call("aria2.forceRemove", [leaked_gid])
+                except Exception as exc:
+                    logger.warning(
+                        "Failed to remove leaked aria2 gid %s for download %s: %s",
+                        leaked_gid,
+                        download_id,
+                        exc,
+                    )
+            raise
         except Exception as exc:
             raise Aria2Error(f"Failed to schedule aria2 download: {exc}") from exc
 
         logger.debug("aria2 accepted download %s with gid %s", download_id, gid)
-        await self._state_store.upsert(
-            download_id,
-            {
-                "gid": gid,
-                "save_path": save_path,
-                "status": "downloading",
-                "url": url,
-            },
-        )
         return gid
 
     async def _register_transfer(
@@ -372,7 +403,46 @@ class Aria2Downloader:
             headers=headers,
         )
         transfer = Aria2Transfer(gid=gid, save_path=os.path.abspath(save_path))
+        # Register the transfer before any further await: once the daemon
+        # holds the gid, cancel_download() must be able to find it.  An await
+        # in between would open a window where a concurrent cancel reports
+        # "Download task not found" and the daemon keeps downloading
+        # untracked.
         self._transfers[download_id] = transfer
+        try:
+            await self._state_store.upsert(
+                download_id,
+                {
+                    "gid": gid,
+                    "save_path": transfer.save_path,
+                    "status": "downloading",
+                    "url": url,
+                },
+            )
+        except asyncio.CancelledError:
+            # The task was cancelled while persisting state and the
+            # coordinator's cancel ran before the transfer was registered
+            # above.  Remove the daemon transfer unless it was deliberately
+            # paused (skip_download preserves paused transfers for resume).
+            status = None
+            try:
+                status = await self.get_status(download_id)
+            except Exception:
+                status = None
+            if status is not None and status.get("status") != "paused":
+                try:
+                    await self._rpc_call("aria2.forceRemove", [gid])
+                except Exception as exc:
+                    logger.warning(
+                        "Failed to remove aria2 gid %s for cancelled download %s: %s",
+                        gid,
+                        download_id,
+                        exc,
+                    )
+                current = self._transfers.get(download_id)
+                if current is not None and current.gid == gid:
+                    self._transfers.pop(download_id, None)
+            raise
         return transfer
 
     async def get_status(self, download_id: str) -> Optional[Dict[str, Any]]:

tests/services/test_aria2_downloader.py

@@ -889,6 +889,215 @@ async def test_cancel_download_tolerates_missing_gid(monkeypatch):
     assert await downloader._state_store.get("download-1") is None
 
 
+@pytest.mark.asyncio
+async def test_register_transfer_tracks_gid_before_state_persist_completes(
+    tmp_path, monkeypatch
+):
+    """A cancel arriving while the state store write is still in flight must
+    already find the transfer — otherwise the gid leaks and the daemon keeps
+    downloading."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    rpc_calls = []
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        rpc_calls.append((method, params))
+        if method == "aria2.addUri":
+            return "gid-1"
+        if method == "aria2.forceRemove":
+            return "OK"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+
+    persist_started = asyncio.Event()
+    persist_release = asyncio.Event()
+
+    class BlockingStore:
+        async def upsert(self, download_id, payload):
+            persist_started.set()
+            await persist_release.wait()
+
+        async def remove(self, download_id):
+            return None
+
+    monkeypatch.setattr(downloader, "_state_store", BlockingStore())
+
+    register_task = asyncio.create_task(
+        downloader._register_transfer(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        )
+    )
+    await asyncio.wait_for(persist_started.wait(), timeout=1.0)
+
+    transfer = downloader._transfers.get("download-1")
+    assert transfer is not None and transfer.gid == "gid-1"
+
+    result = await downloader.cancel_download("download-1")
+    assert result["success"] is True
+    assert ("aria2.forceRemove", ["gid-1"]) in rpc_calls
+
+    persist_release.set()
+    registered = await register_task
+    assert registered.gid == "gid-1"
+
+
+@pytest.mark.asyncio
+async def test_schedule_download_removes_gid_accepted_while_cancelled(
+    tmp_path, monkeypatch
+):
+    """Cancelling while the addUri RPC is in flight must remove the gid the
+    daemon accepted, instead of leaking an untracked download."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    add_uri_started = asyncio.Event()
+    force_removed = []
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            add_uri_started.set()
+            # The daemon processes the request while the client is cancelled.
+            await asyncio.sleep(0.05)
+            return "gid-leaked"
+        if method == "aria2.forceRemove":
+            force_removed.append(params[0])
+            return "OK"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+
+    schedule_task = asyncio.create_task(
+        downloader._schedule_download(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        )
+    )
+    await asyncio.wait_for(add_uri_started.wait(), timeout=1.0)
+    schedule_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await schedule_task
+
+    assert force_removed == ["gid-leaked"]
+    assert "download-1" not in downloader._transfers
+
+
+@pytest.mark.asyncio
+async def test_register_transfer_cancelled_during_persist_removes_active_gid(
+    tmp_path, monkeypatch
+):
+    """Cancellation landing after the gid is registered but before the state
+    store write finishes must remove the still-active daemon transfer."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    persist_started = asyncio.Event()
+    force_removed = []
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            return "gid-2"
+        if method == "aria2.tellStatus":
+            return {"gid": "gid-2", "status": "active"}
+        if method == "aria2.forceRemove":
+            force_removed.append(params[0])
+            return "OK"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+
+    class BlockingStore:
+        async def upsert(self, download_id, payload):
+            persist_started.set()
+            await asyncio.Event().wait()
+
+        async def remove(self, download_id):
+            return None
+
+    monkeypatch.setattr(downloader, "_state_store", BlockingStore())
+
+    register_task = asyncio.create_task(
+        downloader._register_transfer(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        )
+    )
+    await asyncio.wait_for(persist_started.wait(), timeout=1.0)
+    register_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await register_task
+
+    assert force_removed == ["gid-2"]
+    assert "download-1" not in downloader._transfers
+
+
+@pytest.mark.asyncio
+async def test_register_transfer_cancelled_during_persist_preserves_paused_gid(
+    tmp_path, monkeypatch
+):
+    """skip_download pauses the daemon transfer before cancelling the task;
+    the unwind cleanup must not remove a deliberately paused gid."""
+    downloader = Aria2Downloader()
+    downloader._rpc_url = "http://127.0.0.1/jsonrpc"
+    downloader._rpc_secret = "secret"
+
+    save_path = tmp_path / "downloads" / "model.safetensors"
+    persist_started = asyncio.Event()
+    force_removed = []
+
+    async def fake_rpc_call(method, params, **_kwargs):
+        if method == "aria2.addUri":
+            return "gid-3"
+        if method == "aria2.tellStatus":
+            return {"gid": "gid-3", "status": "paused"}
+        if method == "aria2.forceRemove":
+            force_removed.append(params[0])
+            return "OK"
+        raise AssertionError(f"Unexpected RPC method: {method}")
+
+    monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
+
+    class BlockingStore:
+        async def upsert(self, download_id, payload):
+            persist_started.set()
+            await asyncio.Event().wait()
+
+        async def remove(self, download_id):
+            return None
+
+    monkeypatch.setattr(downloader, "_state_store", BlockingStore())
+
+    register_task = asyncio.create_task(
+        downloader._register_transfer(
+            "https://example.com/model.safetensors",
+            str(save_path),
+            download_id="download-1",
+        )
+    )
+    await asyncio.wait_for(persist_started.wait(), timeout=1.0)
+    register_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await register_task
+
+    assert force_removed == []
+    assert downloader._transfers["download-1"].gid == "gid-3"
+
+
 @pytest.mark.asyncio
 async def test_rpc_call_suppresses_error_log_when_log_errors_false(
     monkeypatch, caplog