diff --git a/comfy_cli/command/models/models.py b/comfy_cli/command/models/models.py index 789fc3264..664ea3821 100644 --- a/comfy_cli/command/models/models.py +++ b/comfy_cli/command/models/models.py @@ -1048,6 +1048,11 @@ def on_progress(completed: int, total: int | None) -> None: _worker_headers(state), downloader=state.downloader, progress_callback=on_progress, + # Tag the `.part` temp with this download's id so a cancel of a + # sibling download targeting the same destination can't unlink the + # temp this worker is streaming into. `state.id` is 12 lowercase + # hex (validated filename-safe before any path was built). + part_tag=state.id, ) except DownloadCancelled: state.status = "cancelled" @@ -1750,7 +1755,9 @@ def download_cancel( # find the temp it is streaming into deleted underneath it. reclaimed = 0 if state.status != "completed" and not download_state.worker_alive(state): - reclaimed = cleanup_partials(pathlib.Path(state.dest)) + # Scope the sweep to this download's tagged temps (plus the legacy + # untagged shape) so it can't reclaim a sibling download's live temp. + reclaimed = cleanup_partials(pathlib.Path(state.dest), tag=state.id) if reclaimed: state.completed_bytes = 0 with contextlib.suppress(OSError, ValueError): @@ -1845,8 +1852,10 @@ def download_cancel( # onto the destination once the transfer completes, so a worker killed # mid-flight leaves its gigabytes *there*, not at `dest`. Without this # sweep the cancel would report success and reclaim nothing — the exact - # hand-cleanup this command exists to spare the user. - if cleanup_partials(partial): + # hand-cleanup this command exists to spare the user. Scope it to this + # download's id so a sibling download to the same destination keeps its + # own live temp. + if cleanup_partials(partial, tag=state.id): removed = True state.status = "cancelled" state.error = None if stopped else "worker may still be running; partial file left in place" diff --git a/comfy_cli/download_state.py b/comfy_cli/download_state.py index 968bf098c..faed0d128 100644 --- a/comfy_cli/download_state.py +++ b/comfy_cli/download_state.py @@ -591,16 +591,20 @@ def list_all(workspace: Path) -> list[DownloadState]: return states -def _has_partials(dest: str) -> bool: +def _has_partials(dest: str, tag: str | None = None) -> bool: """True while a ``.part`` sibling of ``dest`` still holds bytes on disk. This is the same set of files ``download-cancel`` reclaims, so it is also - the only handle a user has left on those bytes once a download has failed. + the only handle a user has left on those bytes once a download has failed — + which is why the ``tag`` must match what the cancel path sweeps. A background + worker streams into a temp tagged with its download id, so pruning has to look + for that same tagged shape (plus the legacy untagged one) or it would evict a + failed record whose gigabytes are still on disk, orphaning them. """ from comfy_cli import file_utils try: - return bool(file_utils.partial_paths_for(Path(dest))) + return bool(file_utils.partial_paths_for(Path(dest), tag)) except (OSError, ValueError): # An unreadable parent or a nonsense dest tells us nothing about the # partial; assume there is one rather than deleting the record that @@ -731,7 +735,7 @@ def prune(workspace: Path) -> int: updated = _parse_iso(state.updated_at) if updated is None or updated >= cutoff: continue - if state.status in ("failed", "cancelled") and _has_partials(state.dest): + if state.status in ("failed", "cancelled") and _has_partials(state.dest, state.id): continue if _remove_record(path): removed += 1 diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index fe6a57b8e..f2ce91011 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -516,10 +516,19 @@ def _cleanup_partial(filepath: pathlib.Path) -> None: # The httpx downloader streams into a sibling of the destination named -# ``..part`` and renames it onto the destination only -# once the last byte has landed. The suffix is public in the sense that a killed -# transfer leaves one on disk, so `download-cancel` has to be able to find it — -# hence `partial_paths_for` below rather than an ad-hoc glob at the call site. +# ``..part`` (or ``...part`` +# when a ``part_tag`` is passed) and renames it onto the destination only once the +# last byte has landed. The suffix is public in the sense that a killed transfer +# leaves one on disk, so `download-cancel` has to be able to find it — hence +# `partial_paths_for` below rather than an ad-hoc glob at the call site. +# +# The optional ```` segment scopes a temp to one download so a cancel of A +# cannot reach into a *sibling* download B streaming to the same destination: two +# background workers can legitimately target one path (the submit-time guard is +# only ``local_filepath.exists()``), and without the tag `download-cancel ` +# unlinks B's live temp, whose closing ``os.replace`` then dies with a +# non-retriable ``FileNotFoundError``. Callers pass the download id as the tag; +# foreground/untagged callers keep the shorter shape. _PART_SUFFIX = ".part" # ``tempfile.mkstemp`` fills the middle with exactly 8 characters from this # alphabet (``tempfile._RandomNameSequence``). Matching its shape — not just the @@ -540,33 +549,90 @@ def _cleanup_partial(filepath: pathlib.Path) -> None: _PART_STEM_MAX = _NAME_MAX - len(".") - _MKSTEMP_TOKEN_LEN - len(_PART_SUFFIX) -def _part_prefix(name: str) -> str: +def _part_prefix(name: str, tag: str | None = None) -> str: """The mkstemp ``prefix`` used for ``name``'s ``.part`` siblings. - Normally just ``name + "."``. A destination name too long to also carry the - token and suffix is truncated to fit — on a *byte* basis, since NAME_MAX + Untagged (``tag is None``) this is just ``name + "."``. With a ``tag`` it is + ``..``, which scopes the temp to one download so a cancel of a + sibling transfer to the same destination cannot claim it. A destination name + too long to also carry the token and suffix (and, when tagged, the tag and + its own separator) is truncated to fit — on a *byte* basis, since NAME_MAX counts bytes, but on a character boundary so the result stays valid UTF-8. The trailing ``"."`` is appended after the cut, so the prefix always ends in the separator :func:`partial_paths_for` slices on. - Two destination names agreeing for that many bytes then share a temp - namespace, which only affects which temps :func:`cleanup_partials` claims — - a far smaller problem than being unable to name a temp at all. + The tagged shape reserves the tag, an extra separator, the token and the + suffix (228 bytes of stem for a 12-char id vs. the untagged 241), so for a + long name the tagged and untagged prefixes truncate the stem to *different* + lengths and neither is a prefix-extension of the other. + + The "shared temp namespace" caveat only ever bites the untagged/foreground + shape now, plus the residual case of two names colliding *within* the legacy + (untagged) truncated stem: two destination names agreeing for that many bytes + share the untagged namespace, which only affects which temps + :func:`cleanup_partials` claims — a far smaller problem than being unable to + name a temp at all. A tag makes even that collision id-scoped. """ + if tag is None: + stem_max = _PART_STEM_MAX + else: + # ``..<8 token>.part`` — the tag, its trailing separator, the + # token and the suffix all come out of NAME_MAX before the stem does. + stem_max = ( + _NAME_MAX + - len(".") + - len(tag.encode("utf-8", "surrogatepass")) + - len(".") + - _MKSTEMP_TOKEN_LEN + - len(_PART_SUFFIX) + ) encoded = name.encode("utf-8", "surrogatepass") - if len(encoded) > _PART_STEM_MAX: - name = encoded[:_PART_STEM_MAX].decode("utf-8", "ignore") - return name + "." + if len(encoded) > stem_max: + name = encoded[:stem_max].decode("utf-8", "ignore") + if tag is None: + return name + "." + return f"{name}.{tag}." -def partial_paths_for(local_filepath: pathlib.Path) -> list[pathlib.Path]: +def partial_paths_for(local_filepath: pathlib.Path, tag: str | None = None) -> list[pathlib.Path]: """Every ``.part`` sibling this module would have created for ``local_filepath``. A transfer killed uncleanly (SIGKILL, OOM, power loss) never gets to run its own cleanup, so its ``.part`` file outlives it. This is how the cancel path finds those bytes; nothing else on disk is ever matched. + + ``tag=None`` matches only the untagged shape ``.<8 token>.part`` — the + behaviour before download ids scoped temps, and what foreground transfers + still produce. + + With a ``tag`` the result is the **union** of two shapes: + + * the tag's own temps, ``..<8 token>.part`` — the ones a worker + running this code writes; matching only these is what makes a tag-scoped + cancel skip a *sibling* download's live temp; and + * the legacy untagged shape ``.<8 token>.part`` — included so a ``.part`` + left by a foreground transfer or a pre-change binary is still reclaimed. + + The two prefixes are computed independently (they truncate the stem to + different lengths for a long name, so one is not a prefix-extension of the + other) and the shapes never cross-match: under the untagged prefix a tagged + temp's "token" slice is ``.<8>`` — too long and it contains a ``.``, so + the length/charset check below rejects it; under the tagged prefix an untagged + temp is too short to even start with ``..``. + + **Known residual:** the legacy-inclusion arm means a tag-scoped sweep can + still unlink a *foreground* (untagged) transfer's live temp, or one written by + a pre-change binary, that shares this destination. That exposure is strictly + narrower than the un-tagged behaviour it replaces and disappears once no + untagged producers remain; the companion destination-reservation change + removes the same-destination precondition entirely. """ - prefix = _part_prefix(local_filepath.name) + prefixes = [_part_prefix(local_filepath.name)] + if tag is not None: + # Tagged prefix first so the tag's own temps are the primary match; the + # untagged prefix is the legacy-inclusion fallback. + prefixes.insert(0, _part_prefix(local_filepath.name, tag)) + try: entries = list(local_filepath.parent.iterdir()) except OSError: @@ -575,23 +641,28 @@ def partial_paths_for(local_filepath: pathlib.Path) -> list[pathlib.Path]: matches = [] for entry in entries: name = entry.name - if not name.startswith(prefix) or not name.endswith(_PART_SUFFIX): + if not name.endswith(_PART_SUFFIX): continue - token = name[len(prefix) : -len(_PART_SUFFIX)] - if len(token) != _MKSTEMP_TOKEN_LEN or not set(token) <= _MKSTEMP_TOKEN_CHARS: - continue - matches.append(entry) + for prefix in prefixes: + if not name.startswith(prefix): + continue + token = name[len(prefix) : -len(_PART_SUFFIX)] + if len(token) != _MKSTEMP_TOKEN_LEN or not set(token) <= _MKSTEMP_TOKEN_CHARS: + continue + matches.append(entry) + break return sorted(matches) -def cleanup_partials(local_filepath: pathlib.Path) -> int: +def cleanup_partials(local_filepath: pathlib.Path, tag: str | None = None) -> int: """Best-effort removal of every ``.part`` sibling; returns how many went away. Used by `download-cancel`, which promises to reclaim the disk a dead worker - was using. The destination itself is never touched here. + was using. The destination itself is never touched here. See + :func:`partial_paths_for` for what ``tag`` matches (and its known residual). """ removed = 0 - for partial in partial_paths_for(local_filepath): + for partial in partial_paths_for(local_filepath, tag): try: partial.unlink() except OSError: @@ -730,6 +801,7 @@ def _download_file_httpx( *, state: dict | None = None, progress_callback: ProgressCallback | None = None, + part_tag: str | None = None, ) -> None: """Download a file using httpx streaming. Raises on HTTP or network errors. @@ -791,7 +863,7 @@ def _download_file_httpx( # pre-planted symlink can't redirect the write (CWE-377). fd, tmp_name = tempfile.mkstemp( dir=str(local_filepath.parent), - prefix=_part_prefix(local_filepath.name), + prefix=_part_prefix(local_filepath.name, part_tag), suffix=_PART_SUFFIX, ) try: @@ -844,9 +916,17 @@ def download_file( headers: dict | None = None, downloader: str = "httpx", progress_callback: ProgressCallback | None = None, + part_tag: str | None = None, ): """Helper function to download a file. + ``part_tag`` (optional) scopes the httpx ``.part`` temp to one download — + ``...part`` instead of ``..part`` — so a + concurrent ``download-cancel`` of a *sibling* transfer to the same + destination cannot sweep away this transfer's live temp (see + :func:`partial_paths_for`). Callers pass their download id. The aria2 branch + ignores it: aria2 writes to the destination directly, with no ``.part``. + ``progress_callback`` (optional) is invoked with ``(completed_bytes, total_bytes)`` as the transfer advances; ``total_bytes`` is None until the size is known. When a retry discards a partial file the callback is reset to @@ -883,7 +963,14 @@ def download_file( state["file_opened"] = False state["part_path"] = None try: - _download_file_httpx(url, local_filepath, headers, state=state, progress_callback=progress_callback) + _download_file_httpx( + url, + local_filepath, + headers, + state=state, + progress_callback=progress_callback, + part_tag=part_tag, + ) return except _retriable_exceptions() as exc: last_exc = exc diff --git a/tests/comfy_cli/command/test_model_download_background.py b/tests/comfy_cli/command/test_model_download_background.py index da2b5ca31..c495eebe6 100644 --- a/tests/comfy_cli/command/test_model_download_background.py +++ b/tests/comfy_cli/command/test_model_download_background.py @@ -202,6 +202,27 @@ def test_a_partial_pins_a_failed_or_cancelled_record(self, workspace, status): assert download_state.prune(workspace) == 1 assert download_state.read(workspace, state.id) is None + @pytest.mark.parametrize("status", ["failed", "cancelled"]) + def test_a_tagged_partial_pins_a_failed_or_cancelled_record(self, workspace, status): + """A background worker streams into a temp tagged with its download id, so + prune's pin check has to find *that* shape — not just the legacy untagged + one — or it would evict a record whose gigabytes are still on disk.""" + dest = workspace / "models" / "m.safetensors" + state = _record(workspace, status=status, age_s=_OLD_S, dest=dest) + dest.parent.mkdir(parents=True, exist_ok=True) + tagged = dest.parent / f"{dest.name}.{state.id}.ab3d9f01.part" + tagged.write_bytes(b"partial bytes") + # The fixture is only meaningful if the untagged matcher would have missed it. + assert file_utils.partial_paths_for(dest) == [] + assert file_utils.partial_paths_for(dest, tag=state.id) == [tagged] + + assert download_state.prune(workspace) == 0 + assert download_state.read(workspace, state.id) is not None + + tagged.unlink() + assert download_state.prune(workspace) == 1 + assert download_state.read(workspace, state.id) is None + def test_a_partial_does_not_pin_a_completed_record(self, workspace): """The carve-out is about unreclaimed bytes; a completed download's are at `dest`, and any leftover `.part` is unrelated debris.""" @@ -560,7 +581,7 @@ def stream(*args, **kwargs): original = _download_file_httpx - def flaky(url, path, headers=None, *, state=None, progress_callback=None): + def flaky(url, path, headers=None, *, state=None, progress_callback=None, part_tag=None): if attempts["n"] == 0: attempts["n"] = 1 if state is not None: @@ -569,7 +590,7 @@ def flaky(url, path, headers=None, *, state=None, progress_callback=None): progress_callback(0, 8) progress_callback(4, 8) raise httpx.ReadTimeout("boom") - return original(url, path, headers, state=state, progress_callback=progress_callback) + return original(url, path, headers, state=state, progress_callback=progress_callback, part_tag=part_tag) with ( patch("comfy_cli.file_utils._download_file_httpx", side_effect=flaky), @@ -653,7 +674,7 @@ def test_progress_writes_are_throttled_but_terminal_always_lands(self, workspace clock = {"t": 1000.0} monkeypatch.setattr(models.time, "monotonic", lambda: clock["t"]) - def fake_download_file(url, filepath, headers, downloader, progress_callback): + def fake_download_file(url, filepath, headers, downloader, progress_callback, part_tag=None): for completed, tick in [(3, 0.1), (6, 0.2), (9, 5.0)]: clock["t"] += tick progress_callback(completed, 9) @@ -692,7 +713,7 @@ def test_worker_writes_its_own_pid_on_startup(self, workspace, monkeypatch, tmp_ recorded = {} - def capture(url, filepath, headers, downloader, progress_callback): + def capture(url, filepath, headers, downloader, progress_callback, part_tag=None): recorded.update(download_state.read(workspace, state.id).to_dict()) filepath.write_bytes(b"ok") @@ -719,7 +740,7 @@ def test_worker_rederives_auth_headers_from_config_not_state(self, workspace, mo monkeypatch.setattr(models, "_civitai_headers", lambda: {"Authorization": "Bearer from-config"}) seen = {} - def capture(url, filepath, headers, downloader, progress_callback): + def capture(url, filepath, headers, downloader, progress_callback, part_tag=None): seen["headers"] = headers filepath.write_bytes(b"ok") @@ -1666,7 +1687,7 @@ def test_a_completed_transfer_releases_it(self, workspace, monkeypatch, tmp_path monkeypatch.setattr( models, "download_file", - lambda url, filepath, headers, downloader, progress_callback: filepath.write_bytes(b"ok"), + lambda url, filepath, headers, downloader, progress_callback, part_tag=None: filepath.write_bytes(b"ok"), ) models._download_worker(state_file=str(path)) @@ -1704,7 +1725,7 @@ def test_a_mid_transfer_cancel_releases_it(self, workspace, monkeypatch, tmp_pat state, path, claim = self._prepare(workspace, tmp_path) marker = download_state.cancel_marker_for(path) - def transfer(url, filepath, headers, downloader, progress_callback): + def transfer(url, filepath, headers, downloader, progress_callback, part_tag=None): marker.touch() monkeypatch.setattr(models.time, "monotonic", lambda: 1e9) progress_callback(1, 2) @@ -1722,7 +1743,7 @@ def test_a_claim_that_is_no_longer_ours_is_left_alone(self, workspace, monkeypat that window. An unconditional unlink would delete a live claim.""" state, path, claim = self._prepare(workspace, tmp_path) - def transfer(url, filepath, headers, downloader, progress_callback): + def transfer(url, filepath, headers, downloader, progress_callback, part_tag=None): filepath.write_bytes(b"ok") claim.unlink() assert download_state.acquire_claim(claim, download_id="ffffffffffff", dest=state.dest) @@ -1743,7 +1764,7 @@ def test_a_finished_download_does_not_wedge_its_destination(self, workspace, mon monkeypatch.setattr( models, "download_file", - lambda url, filepath, headers, downloader, progress_callback: filepath.write_bytes(b"ok"), + lambda url, filepath, headers, downloader, progress_callback, part_tag=None: filepath.write_bytes(b"ok"), ) def submit(): @@ -2975,7 +2996,7 @@ def test_worker_aborts_mid_transfer_and_clears_the_partial(self, workspace, monk state = _state(dest=str(dest), status="starting", downloader="aria2", pid=None) path = download_state.write(workspace, state) - def transfer(url, filepath, headers, downloader, progress_callback): + def transfer(url, filepath, headers, downloader, progress_callback, part_tag=None): filepath.write_bytes(b"partial") # The cancel lands after the transfer is already under way. download_state.request_cancel(download_state.cancel_path(workspace, state.id)) @@ -3002,7 +3023,7 @@ def test_worker_cancel_on_httpx_leaves_the_destination_alone(self, workspace, mo state = _state(dest=str(dest), status="starting", pid=None) path = download_state.write(workspace, state) - def transfer(url, filepath, headers, downloader, progress_callback): + def transfer(url, filepath, headers, downloader, progress_callback, part_tag=None): download_state.request_cancel(download_state.cancel_path(workspace, state.id)) progress_callback(7, 4096) @@ -3023,7 +3044,7 @@ def test_a_transfer_that_beat_the_cancel_keeps_its_file(self, workspace, json_re state = _state(dest=str(dest), status="starting", pid=None) path = download_state.write(workspace, state) - def transfer(url, filepath, headers, downloader, progress_callback): + def transfer(url, filepath, headers, downloader, progress_callback, part_tag=None): filepath.write_bytes(b"done") download_state.request_cancel(download_state.cancel_path(workspace, state.id)) @@ -3136,6 +3157,57 @@ def stop(_state, **_kwargs): assert json_renderer()["data"]["status"] == "completed" +class TestCancelIsolatesSiblingDownloads: + """Two background downloads can legitimately target one destination (the only + submit-time guard is ``local_filepath.exists()``). Each streams into its own + ``...part``, so cancelling A must reclaim A's temp (and any + legacy untagged debris) while leaving B's live temp untouched — otherwise B + keeps writing into a deleted inode and its closing rename dies with a + non-retriable ``FileNotFoundError``. + """ + + def _siblings(self, workspace, tmp_path, *, a_status): + dest = tmp_path / "m.safetensors" + a = _state(dest=str(dest), status=a_status, pid=5150, total_bytes=13000, completed_bytes=3600) + b = _state(dest=str(dest), status="downloading", pid=6000, total_bytes=13000, completed_bytes=1000) + download_state.write(workspace, a) + download_state.write(workspace, b) + + a_part = tmp_path / f"m.safetensors.{a.id}.a1b2c3d4.part" + a_part.write_bytes(b"A" * 3600) + b_part = tmp_path / f"m.safetensors.{b.id}.b2c3d4e5.part" + b_part.write_bytes(b"B" * 1000) + legacy = tmp_path / "m.safetensors.c3d4e5f6.part" + legacy.write_bytes(b"legacy") + return dest, a, b, a_part, b_part, legacy + + def test_terminal_status_cancel_of_A_spares_Bs_partial(self, workspace, json_renderer, tmp_path): + """The ``models.py`` terminal-status sweep: A is already ``failed`` (its + first status poll persisted reconcile's verdict) and its worker is gone.""" + dest, a, b, a_part, b_part, legacy = self._siblings(workspace, tmp_path, a_status="failed") + + with patch.object(download_state, "is_worker_process", return_value=False): + models.download_cancel(None, download_id=a.id) + + assert not a_part.exists(), "A's own tagged temp must be reclaimed" + assert not legacy.exists(), "an untagged legacy temp is still reclaimed" + assert b_part.read_bytes() == b"B" * 1000, "B's live temp must survive A's cancel" + assert json_renderer()["changed"] is True + + def test_live_cancel_of_A_spares_Bs_partial(self, workspace, json_renderer, monkeypatch, tmp_path): + """The ``models.py`` live-cancel sweep: A is still ``downloading`` and is + stopped before its `.part` is swept.""" + dest, a, b, a_part, b_part, legacy = self._siblings(workspace, tmp_path, a_status="downloading") + + monkeypatch.setattr(download_state, "stop_worker", lambda *_a, **_k: True) + models.download_cancel(None, download_id=a.id) + + assert not a_part.exists(), "A's own tagged temp must be reclaimed" + assert not legacy.exists(), "an untagged legacy temp is still reclaimed" + assert b_part.read_bytes() == b"B" * 1000, "B's live temp must survive A's cancel" + assert json_renderer()["data"]["status"] == "cancelled" + + class TestStateFilePermissions: """A resolved url can carry a presigned/SAS query token, so these files are secrets on a shared host — and a writable one is an attack surface.""" diff --git a/tests/test_file_utils_network.py b/tests/test_file_utils_network.py index 884cd44b5..2a6c00770 100644 --- a/tests/test_file_utils_network.py +++ b/tests/test_file_utils_network.py @@ -988,6 +988,162 @@ def killed_iter(): assert cleanup_partials(dest) == 1 +class TestTaggedPartials: + """A ``part_tag`` scopes a `.part` temp to one download, so a cancel of one + download can't reach a *sibling* download streaming to the same destination. + + The matcher then answers a tag-scoped query with the union of that tag's own + temps and the legacy untagged shape (so a `.part` from a foreground transfer + or a pre-change binary is still reclaimable). + """ + + # A 12-char lowercase-hex id, the exact shape `download_state.new_id` mints. + TAG_X = "abc123def456" + TAG_Y = "0f1e2d3c4b5a" + + def _tagged(self, dest, tag, token="a1b2c3d4", data=b"bytes"): + p = dest.parent / f"{dest.name}.{tag}.{token}.part" + p.write_bytes(data) + return p + + def _untagged(self, dest, token="c3d4e5f6", data=b"legacy"): + p = dest.parent / f"{dest.name}.{token}.part" + p.write_bytes(data) + return p + + @patch("httpx.stream") + def test_a_failed_tagged_download_leaves_a_tag_shaped_temp(self, mock_stream, tmp_path): + """Case 1: a mid-stream failure leaves `..<8>.part`, and the + tag-scoped matcher/cleaner find exactly it.""" + dest = tmp_path / "model.safetensors" + + def killed_iter(): + yield b"partial" + raise KeyboardInterrupt() + + resp = Mock() + resp.status_code = 200 + resp.headers = {} + resp.iter_bytes = Mock(side_effect=killed_iter) + resp.__enter__ = Mock(return_value=resp) + resp.__exit__ = Mock(return_value=None) + mock_stream.return_value = resp + + with ( + patch("comfy_cli.file_utils.ui.prompt_confirm_action", return_value=False), + pytest.raises(KeyboardInterrupt), + ): + download_file("http://example.com/model.safetensors", dest, part_tag=self.TAG_X) + + parts = partial_paths_for(dest, tag=self.TAG_X) + assert [p.read_bytes() for p in parts] == [b"partial"] + name = parts[0].name + prefix = f"{dest.name}.{self.TAG_X}." + assert name.startswith(prefix) and name.endswith(".part") + token = name[len(prefix) : -len(".part")] + assert len(token) == 8 and set(token) <= file_utils._MKSTEMP_TOKEN_CHARS + assert cleanup_partials(dest, tag=self.TAG_X) == 1 + + @patch("httpx.stream") + def test_a_successful_tagged_download_lands_via_rename(self, mock_stream, tmp_path): + """Case 1 (success half): the tagged temp is consumed by the rename onto + the destination, leaving nothing behind.""" + mock_stream.return_value = _make_ok_response(content=b"full model", content_length=10) + dest = tmp_path / "model.safetensors" + + renames = [] + real_replace = os.replace + + def spy(src, dst, *args, **kwargs): + renames.append((str(src), str(dst))) + return real_replace(src, dst, *args, **kwargs) + + with patch("comfy_cli.file_utils.os.replace", side_effect=spy): + download_file("http://example.com/model.safetensors", dest, part_tag=self.TAG_X) + + assert dest.read_bytes() == b"full model" + src, dst = renames[0] + assert dst == str(dest) + assert src.startswith(f"{dest}.{self.TAG_X}.") and src.endswith(".part") + assert partial_paths_for(dest, tag=self.TAG_X) == [] + assert sorted(p.name for p in tmp_path.iterdir()) == ["model.safetensors"] + + def test_a_tag_scoped_cleanup_removes_only_that_tags_temp(self, tmp_path): + """Case 2: two downloads to one destination, and cancelling X must not + touch Y's in-flight temp.""" + dest = tmp_path / "model.safetensors" + x_part = self._tagged(dest, self.TAG_X, token="a1b2c3d4", data=b"X") + y_part = self._tagged(dest, self.TAG_Y, token="b2c3d4e5", data=b"Y") + + assert partial_paths_for(dest, tag=self.TAG_X) == [x_part] + assert cleanup_partials(dest, tag=self.TAG_X) == 1 + assert not x_part.exists() + assert y_part.read_bytes() == b"Y" + + def test_a_tag_scoped_cleanup_still_reclaims_a_legacy_untagged_temp(self, tmp_path): + """Case 3: a `.part` left by a foreground transfer or a pre-change binary + (untagged shape) is still swept by a tag-scoped cancel.""" + dest = tmp_path / "model.safetensors" + legacy = self._untagged(dest, token="c3d4e5f6") + + assert partial_paths_for(dest, tag=self.TAG_X) == [legacy] + assert cleanup_partials(dest, tag=self.TAG_X) == 1 + assert not legacy.exists() + + def test_the_shapes_never_cross_match(self, tmp_path): + """Case 4: an untagged sweep ignores a tagged temp (its 21-char "token" + contains a `.`), and a tag-scoped sweep ignores a temp whose token slice + is 21 chars — i.e. another download's tagged temp.""" + dest = tmp_path / "model.safetensors" + x_part = self._tagged(dest, self.TAG_X, token="a1b2c3d4", data=b"X") + y_part = self._tagged(dest, self.TAG_Y, token="b2c3d4e5", data=b"Y") + + # An untagged sweep must not claim any tagged temp. + assert partial_paths_for(dest) == [] + assert cleanup_partials(dest) == 0 + assert x_part.exists() and y_part.exists() + + # A sweep scoped to X must not claim Y (whose token slice under X's own + # untagged legacy arm is `.<8>` = 21 chars). + under_untagged = y_part.name[len(f"{dest.name}.") : -len(".part")] + assert len(under_untagged) == 21 and "." in under_untagged + assert partial_paths_for(dest, tag=self.TAG_X) == [x_part] + + @patch("httpx.stream") + def test_a_long_name_with_a_tag_still_fits_and_both_shapes_are_found(self, mock_stream, tmp_path): + """Case 5: a destination name over the tagged stem budget truncates to a + different stem than the untagged prefix, the full tagged temp name still + fits in NAME_MAX, and the union matcher finds both shapes for that name.""" + # 250 bytes: past both the 241-byte untagged and 228-byte tagged budgets. + dest = tmp_path / ("m" * 238 + ".safetensors") + assert len(dest.name.encode()) > 228 + + tagged_prefix = file_utils._part_prefix(dest.name, self.TAG_X) + untagged_prefix = file_utils._part_prefix(dest.name) + # Different truncated stems — neither prefix extends the other, so a long + # name's two shapes can't be conflated. + assert tagged_prefix != untagged_prefix + assert not tagged_prefix.startswith(untagged_prefix) + assert not untagged_prefix.startswith(tagged_prefix) + # The full tagged temp name fits in one path component. + assert len((tagged_prefix + "a1b2c3d4" + file_utils._PART_SUFFIX).encode()) <= file_utils._NAME_MAX + + # A real tagged download of this over-long name still succeeds (mirrors the + # untagged ENAMETOOLONG regression test). + mock_stream.return_value = _make_ok_response(content=b"data") + download_file("http://example.com/model.safetensors", dest, part_tag=self.TAG_X) + assert dest.read_bytes() == b"data" + assert partial_paths_for(dest, tag=self.TAG_X) == [] + + # Both shapes for this long name are reclaimed by one tag-scoped query. + tagged_temp = dest.parent / (tagged_prefix + "a1b2c3d4" + file_utils._PART_SUFFIX) + untagged_temp = dest.parent / (untagged_prefix + "b2c3d4e5" + file_utils._PART_SUFFIX) + tagged_temp.write_bytes(b"tagged") + untagged_temp.write_bytes(b"untagged") + assert set(partial_paths_for(dest, tag=self.TAG_X)) == {tagged_temp, untagged_temp} + assert cleanup_partials(dest, tag=self.TAG_X) == 2 + + class TestDownloadHTTPStatusRetry: """Retry behavior for transient HTTP status codes (5xx, 429, 408)."""