refactor: finish the atomic-writer consolidation and document the write policy

comfy_cli/auth/store.py

@@ -112,6 +112,9 @@ def _read_all(path: Path) -> dict[str, Any]:
 def _write_all(path: Path, payload: dict[str, Any]) -> None:
     """Atomic write with 0600 mode from inception.
 
+    Deliberately not ``file_utils.atomic_write_text`` — see the write-policy note
+    in ``comfy_cli/file_utils.py`` (tier 1, secrets).
+
     The tmp file is opened with ``O_CREAT|O_EXCL`` and explicit mode 0o600 so
     the secrets are never world-readable, even briefly, even on systems with a
     permissive umask. A unique tmp name (`<basename>.<pid>.<rand>.tmp`) avoids

comfy_cli/command/outdated.py

@@ -36,6 +36,7 @@
 
 from comfy_cli.command.pack_scan import iter_pack_dirs as _iter_pack_dirs
 from comfy_cli.command.pack_scan import read_pyproject as _read_pyproject
+from comfy_cli.file_utils import atomic_write_text
 from comfy_cli.registry import RegistryAPI
 
 CACHE_TTL_SECONDS = 3600  # 1 hour
@@ -77,24 +78,13 @@ def _load_cache() -> dict[str, Any]:
 def _save_cache(cache: dict[str, Any]) -> None:
     path = _cache_path()
     try:
-        path.parent.mkdir(parents=True, exist_ok=True)
-        # Write-then-rename: an interrupt mid-write must not leave a truncated
-        # file that the next `_load_cache` silently resets to `{}`. `os.replace`
-        # is atomic within a filesystem, and the temp file is a sibling so the
-        # rename never crosses one. Unique per process — two concurrent writers
-        # must not share (and truncate) one temp path.
-        tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp")
-        try:
-            # `encoding` pinned rather than left to the platform locale:
-            # `_load_cache` decodes bytes as JSON, which only accepts
-            # UTF-8/16/32. Today `json.dumps` defaults to `ensure_ascii=True`,
-            # so the payload is pure ASCII and any locale would round-trip —
-            # this keeps that a property of the writer, not a lucky default, if
-            # a non-ASCII pack id ever reaches the file verbatim.
-            tmp.write_text(json.dumps(cache), encoding="utf-8")
-            os.replace(tmp, path)
-        finally:
-            tmp.unlink(missing_ok=True)
+        # Write-then-rename (tier 4, regenerable cache — see the write policy in
+        # comfy_cli/file_utils.py): an interrupt mid-write must not leave a
+        # truncated file that the next `_load_cache` silently resets to `{}`.
+        # The helper creates the parent dir, writes UTF-8 (`_load_cache` decodes
+        # bytes as JSON, which only accepts UTF-8/16/32), and cleans up its own
+        # temp file on failure.
+        atomic_write_text(path, json.dumps(cache), fsync=False)
     except OSError:
         # A read-only cache dir must never break a read-only report.
         pass

comfy_cli/command/project.py

@@ -15,7 +15,6 @@
 
 import hashlib
 import json
-import os
 import urllib.error
 from datetime import datetime, timezone
 from pathlib import Path
@@ -25,6 +24,7 @@
 
 from comfy_cli import tracking
 from comfy_cli.command.transfer import _upload_file
+from comfy_cli.file_utils import atomic_write_text
 from comfy_cli.output import get_renderer, rprint
 from comfy_cli.project import (
     ASSETS_LOCK_SCHEMA,
@@ -282,18 +282,11 @@ def assets_push_cmd(
 
 
 def _write_assets_lock(path: Path, assets: dict) -> None:
-    """Atomically rewrite the lock (tmp + fsync + rename, the jobs_state
-    pattern) so a crash mid-push can't leave a torn JSON file."""
-    path.parent.mkdir(parents=True, exist_ok=True)
+    """Atomically rewrite the lock so a crash mid-push can't leave a torn JSON file."""
     doc = {"schema": ASSETS_LOCK_SCHEMA, "assets": assets}
-    tmp = path.with_suffix(f".{os.getpid()}.tmp")
-    tmp.write_text(json.dumps(doc, indent=2, sort_keys=True), encoding="utf-8")
-    fd = os.open(str(tmp), os.O_RDONLY)
-    try:
-        os.fsync(fd)
-    finally:
-        os.close(fd)
-    os.replace(tmp, path)
+    # Tier 3 (cross-process state) per the write policy in comfy_cli/file_utils.py:
+    # fsync=True, because a lock lost to power failure means re-pushing assets.
+    atomic_write_text(path, json.dumps(doc, indent=2, sort_keys=True), fsync=True)
 
 
 # ---------------------------------------------------------------------------

comfy_cli/command/templates.py

@@ -19,7 +19,6 @@
 
 import json
 import os
-import tempfile
 import time
 import urllib.error
 import urllib.parse
@@ -30,6 +29,7 @@
 import typer
 
 from comfy_cli import tracking
+from comfy_cli.file_utils import atomic_write_bytes
 from comfy_cli.http import plain_urlopen
 from comfy_cli.output import get_renderer, rprint
 
@@ -138,18 +138,14 @@ def _persist_cache(cache: Path, data: bytes) -> None:
       propagated.
     """
     try:
-        cache.parent.mkdir(parents=True, exist_ok=True)
-        fd, tmp = tempfile.mkstemp(dir=str(cache.parent), prefix=".index-", suffix=".tmp")
-        try:
-            with os.fdopen(fd, "wb") as f:
-                f.write(data)
-            os.replace(tmp, cache)
-        except OSError:
-            try:
-                os.unlink(tmp)
-            except OSError:
-                pass
-            raise
+        # Tier 4 (regenerable cache) per the write policy in
+        # comfy_cli/file_utils.py — fsync=False, wrapped best-effort below.
+        # The previous inline `tempfile.mkstemp` left the cache at mode 0600
+        # (mkstemp hardcodes that and `os.replace` carries it onto the
+        # destination); the helper restores the umask-derived mode instead, and
+        # that relaxation is intended: the payload is the public
+        # template-gallery index, not user data.
+        atomic_write_bytes(cache, data, fsync=False)
     except OSError:
         # Couldn't persist (read-only dir, disk full, …). We still have valid
         # data in hand, so proceed without caching rather than failing the run.

comfy_cli/command/workflow.py

@@ -28,6 +28,7 @@
 import typer
 
 from comfy_cli import tracking
+from comfy_cli.file_utils import atomic_write_text
 
 # Aliased at module scope rather than lazy-imported: a class used in ``except``
 # clauses at module scope cannot be resolved lazily. ``comfy_cli.http`` is
@@ -126,22 +127,6 @@ def _get_graph(input_path: str | None, host: str | None, port: int | None, on_st
         raise typer.Exit(code=1) from e
 
 
-def _atomic_write_text(path: Path, content: str) -> None:
-    """Write via tmp + rename so SIGINT mid-write can't leave a half-written file."""
-    import os
-
-    tmp = path.with_suffix(path.suffix + f".{os.getpid()}.tmp")
-    try:
-        tmp.write_text(content, encoding="utf-8")
-        os.replace(tmp, path)
-    except Exception:
-        try:
-            tmp.unlink()
-        except OSError:
-            pass
-        raise
-
-
 def _parse_value(raw: str) -> Any:
     """Parse a CLI-supplied value as JSON; fall back to the literal string."""
     try:
@@ -306,7 +291,7 @@ def set_slot_cmd(
         return
 
     if not stdout:
-        _atomic_write_text(p, json.dumps(new_workflow, indent=2))
+        atomic_write_text(p, json.dumps(new_workflow, indent=2))
 
     payload: dict[str, Any] = {
         "workflow": str(p),
@@ -421,7 +406,7 @@ def vary_cmd(
         out.mkdir(parents=True, exist_ok=True)
         for i, wf in enumerate(workflows):
             target = out / f"{p.stem}_{i:03d}.json"
-            _atomic_write_text(target, json.dumps(wf, indent=2))
+            atomic_write_text(target, json.dumps(wf, indent=2))
             written.append(str(target))
     elif renderer.is_pretty():
         import sys

comfy_cli/cql/loader.py

@@ -30,6 +30,7 @@
 
 from comfy_cli.cql._net import is_loopback_host
 from comfy_cli.cql.errors import CQLRuntimeError
+from comfy_cli.file_utils import atomic_write_text
 from comfy_cli.http import NoRedirectHandler, build_http_only_opener
 
 # Cap raw bytes read from disk or the network. Real `object_info` dumps are a
@@ -310,16 +311,11 @@ def write_object_info_cache(host_key: str, data: dict[str, Any]) -> None:
     """
     path = object_info_cache_path(host_key)
     try:
-        path.parent.mkdir(parents=True, exist_ok=True)
-        tmp = path.with_suffix(path.suffix + f".{os.getpid()}.tmp")
-        tmp.write_text(json.dumps(data), encoding="utf-8")
-        os.replace(tmp, path)
+        atomic_write_text(path, json.dumps(data))
     except OSError:
         # A cache we can't write is not worth failing the command over.
-        try:
-            tmp.unlink()  # type: ignore[possibly-undefined]
-        except (OSError, NameError, UnboundLocalError):
-            pass
+        # atomic_write_text already cleaned up its own tmp file on failure.
+        pass
 
 
 def read_object_info_cache(host_key: str) -> dict[str, Any] | None:

comfy_cli/download_state.py

@@ -215,6 +215,9 @@ def write(workspace: Path, state: DownloadState) -> Path:
 def write_path(path: Path, state: DownloadState) -> Path:
     """Atomically persist ``state`` at ``path`` (write a tmp file, ``os.replace``).
 
+    Deliberately not ``file_utils.atomic_write_text`` — see the write-policy note
+    in ``comfy_cli/file_utils.py`` (tier 2, secret-adjacent state).
+
     The worker addresses its state file by path rather than by workspace: it is
     handed ``--state <file>`` and must not have to re-resolve a workspace that
     the foreground already resolved.

comfy_cli/file_utils.py

@@ -1,7 +1,9 @@
 import json
 import os
 import pathlib
+import stat
 import subprocess
+import tempfile
 import time
 import zipfile
 from collections.abc import Callable
@@ -14,6 +16,150 @@
 from comfy_cli import constants, ui
 from comfy_cli.output.sanitize import sanitize_value
 
+# ---------------------------------------------------------------------------
+# Atomic writes — the write policy
+# ---------------------------------------------------------------------------
+#
+# Every hand-rolled tmp+``os.replace`` writer in comfy-cli falls into one of
+# four tiers. The tier decides whether the writer uses the shared helpers below
+# or stays bespoke; it is decided by *what the file can contain*, never by
+# convenience. The helpers deliberately take **no** ``mode=`` parameter —
+# parameterizing permissions was considered and rejected, because a writer whose
+# payload needs 0600 needs more than a mode (an ``O_EXCL`` open at that mode, a
+# 0700 parent, an fsync) and is clearer written out in full at its call site.
+#
+# Tier 1 — secrets. ``comfy_cli/auth/store.py`` ``_write_all`` stays bespoke:
+#   the tmp file is opened ``O_CREAT|O_EXCL`` with mode 0600 *at open*, so the
+#   OAuth tokens are never briefly world-readable under a permissive umask; the
+#   parent directory is forced to 0700; the write is fsynced.
+#
+# Tier 2 — secret-adjacent state. ``comfy_cli/download_state.py`` ``write_path``
+#   stays bespoke: chmod 0600 + 0700 parent + fsync, because the resolved
+#   download URL persisted in the state file can embed a presigned/SAS token,
+#   which is bearer credential material even though the file itself is nominally
+#   bookkeeping.
+#
+# Tier 3 — cross-process state. ``atomic_write_text(..., fsync=True)``:
+#   umask mode, durable. A second process (a watcher, a later CLI invocation, an
+#   agent shelling out) reads these, and losing one to power failure loses real
+#   work. Members: ``jobs_state.write``, ``command/project.py``
+#   ``_write_assets_lock``.
+#
+# Tier 4 — regenerable caches and manifests. ``fsync=False``: the content can be
+#   recomputed or refetched, so paying a sync per write buys nothing. Callers may
+#   additionally wrap the call best-effort (``except OSError: pass``) when a
+#   read-only or full cache directory must not break the command. Members:
+#   ``cql/loader``, ``command/outdated.py`` ``_save_cache``,
+#   ``command/templates.py`` ``_persist_cache``, the skills manifest writers,
+#   ``command/workflow.py``.
+#
+# Two invariants recorded here because they are the *reason* for a tier
+# assignment, and a future change to either moves a writer between tiers:
+#
+# (a) ``JobState`` files are tier 3 (umask mode), not tier 2, because their
+#     ``outputs`` are plain ``/view?filename=...`` URLs — they carry no embedded
+#     credential, and the CLI never asks the cloud jobs API for ``?short_link=``
+#     responses. If either changes (outputs start carrying signed URLs, or the
+#     CLI starts requesting short links), jobs_state moves to tier 2 and needs
+#     0600 like download_state.
+#
+# (b) No writer here fsyncs the parent directory *before* ``os.replace`` returns
+#     durably — ``fsync=True`` syncs the file's contents and then the parent
+#     directory, but on a power failure between the two an already-``replace``d
+#     rename can still be lost. This is accepted: the failure mode is losing the
+#     *newest* write, never observing a torn or half-written file. Contents are
+#     never torn; renames may be lost.
+
+
+def atomic_write_text(path: pathlib.Path, content: str, *, fsync: bool = False) -> None:
+    """Atomically write ``content`` to ``path`` via a sibling tmp file + ``os.replace``.
+
+    The write goes to a uniquely-named tmp file in the same directory (so the rename
+    stays on one filesystem and is atomic), which is then renamed over ``path``.
+    A SIGINT or crash mid-write therefore never leaves a half-written or empty
+    file at the destination — readers see either the old contents or the new.
+    On any failure the tmp file is cleaned up and the exception re-raised.
+
+    The tmp file is created with ``tempfile.mkstemp`` (``O_CREAT | O_EXCL``, no
+    symlink following) in the destination directory, so concurrent writers never
+    collide on a shared name and a pre-planted symlink can't redirect the write
+    (CWE-377).
+
+    Args:
+        path: destination file. Parent directories are created if missing.
+        content: text to write (UTF-8).
+        fsync: if True, flush the tmp file's contents to disk before the rename
+            and fsync the destination directory afterwards so both the data and
+            the rename survive power loss, at the cost of a sync. Best-effort:
+            a failing fsync is ignored, matching the prior per-site behavior.
+    """
+    _atomic_write(path, content.encode("utf-8"), fsync=fsync)
+
+
+def atomic_write_bytes(path: pathlib.Path, data: bytes, *, fsync: bool = False) -> None:
+    """Atomically write ``data`` to ``path`` — the bytes twin of :func:`atomic_write_text`.
+
+    Identical semantics (same tmp-file creation, same permission handling, same
+    cleanup-on-failure); it just skips the UTF-8 encode for a caller that already
+    holds bytes. See :func:`atomic_write_text` for the full contract.
+    """
+    _atomic_write(path, data, fsync=fsync)
+
+
+def _atomic_write(path: pathlib.Path, data: bytes, *, fsync: bool) -> None:
+    """Shared implementation behind :func:`atomic_write_text` / :func:`atomic_write_bytes`."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # mkstemp gives a unique, O_EXCL, non-symlink-following fd opened O_RDWR in the
+    # destination directory — same filesystem, so the os.replace below is atomic.
+    fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=path.name + ".", suffix=".tmp")
+    try:
+        with os.fdopen(fd, "wb") as f:
+            f.write(data)
+            if fsync:
+                f.flush()
+                try:
+                    os.fsync(f.fileno())  # O_RDWR fd, so fsync works on Windows too
+                except OSError:
+                    pass
+        # mkstemp hardcodes the tmp file to 0600, and os.replace carries that mode
+        # onto the destination — so without this a first atomic write would quietly
+        # strip the group/other read access a shared output is meant to have. Restore
+        # the intended mode before the rename: reuse the existing destination's bits,
+        # else fall back to the umask-derived default (0666 & ~umask) for a new file.
+        try:
+            dest_mode = stat.S_IMODE(os.stat(path).st_mode)
+        except OSError:
+            umask = os.umask(0)
+            os.umask(umask)
+            dest_mode = 0o666 & ~umask
+        try:
+            os.chmod(tmp_name, dest_mode)
+        except OSError:
+            # Windows / filesystems without POSIX perms: best-effort, matching fsync.
+            pass
+        os.replace(tmp_name, path)
+        if fsync:
+            # Also fsync the parent directory so the rename itself is durable.
+            try:
+                dir_fd = os.open(str(path.parent), os.O_RDONLY)
+                try:
+                    os.fsync(dir_fd)
+                except OSError:
+                    pass
+                finally:
+                    os.close(dir_fd)
+            except OSError:
+                # e.g. Windows can't open a directory for fsync; best-effort.
+                pass
+    except BaseException:
+        # BaseException (not just Exception) so a KeyboardInterrupt mid-write
+        # still cleans up the tmp file, per the docstring.
+        try:
+            os.unlink(tmp_name)
+        except OSError:
+            pass
+        raise
+
 
 class DownloadException(Exception):
     pass