diff --git a/invokeai/app/services/download/download_default.py b/invokeai/app/services/download/download_default.py index 56197dcfb35..bc7cfb7b0df 100644 --- a/invokeai/app/services/download/download_default.py +++ b/invokeai/app/services/download/download_default.py @@ -433,11 +433,14 @@ def _do_download(self, job: DownloadJob) -> None: ) if resp.status_code == 416 and resume_from > 0: # Range not satisfiable - local partial is already complete - expected = job.expected_total_bytes or job.total_bytes or resume_from - if resume_from == expected: + match = re.fullmatch(r"bytes \*/(\d+)", resp.headers.get("Content-Range", ""), flags=re.IGNORECASE) + expected = int(match.group(1)) if match else None + if expected is not None and resume_from == expected: job.total_bytes = expected + job.expected_total_bytes = expected job.bytes = resume_from job.download_path = job.download_path or job.dest + self._in_progress_path(job.download_path).rename(job.download_path) self._signal_job_started(job) self._signal_job_complete(job) return diff --git a/tests/app/services/download/test_download_queue.py b/tests/app/services/download/test_download_queue.py index bfd2e00fdbb..90e2ba1e3df 100644 --- a/tests/app/services/download/test_download_queue.py +++ b/tests/app/services/download/test_download_queue.py @@ -9,7 +9,7 @@ import pytest from pydantic.networks import AnyHttpUrl from requests.sessions import Session -from requests_testadapter import TestAdapter +from requests_testadapter import TestAdapter, TestSession from invokeai.app.services.config import get_config from invokeai.app.services.config.config_default import URLRegexTokenPair @@ -81,6 +81,64 @@ def test_errors(tmp_path: Path, mm2_session: Session) -> None: queue.stop() +@pytest.mark.timeout(timeout=10, method="thread") +def test_completed_resume_with_416_promotes_in_progress_file(tmp_path: Path) -> None: + source = AnyHttpUrl("https://test.com/complete.safetensors") + content = b"complete" + destination = tmp_path / "complete.safetensors" + in_progress_path = destination.with_name(destination.name + ".downloading") + in_progress_path.write_bytes(content) + + session = TestSession() + session.mount( + str(source), + TestAdapter(b"", status=416, headers={"Content-Range": f"bytes */{len(content)}"}), + ) + completed_files: list[bool] = [] + queue = DownloadQueueService(requests_session=session) + queue.start() + try: + job = queue.download( + source=source, + dest=destination, + on_complete=lambda completed_job: completed_files.append(completed_job.download_path.exists()), + ) + queue.join() + finally: + queue.stop() + + assert job.status == DownloadJobStatus.COMPLETED + assert destination.read_bytes() == content + assert not in_progress_path.exists() + assert completed_files == [True] + + +@pytest.mark.timeout(timeout=10, method="thread") +def test_mismatched_416_resume_keeps_in_progress_file(tmp_path: Path) -> None: + source = AnyHttpUrl("https://test.com/stale.safetensors") + destination = tmp_path / "stale.safetensors" + in_progress_path = destination.with_name(destination.name + ".downloading") + in_progress_path.write_bytes(b"stale data") + + session = TestSession() + session.mount( + str(source), + TestAdapter(b"", status=416, headers={"Content-Range": "bytes */8"}), + ) + queue = DownloadQueueService(requests_session=session) + queue.start() + try: + job = queue.download(source=source, dest=destination) + queue.join() + finally: + queue.stop() + + assert job.status == DownloadJobStatus.PAUSED + assert job.resume_required + assert not destination.exists() + assert in_progress_path.exists() + + @pytest.mark.timeout(timeout=10, method="thread") def test_event_bus(tmp_path: Path, mm2_session: Session) -> None: event_bus = TestEventService()