Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 3 additions & 18 deletions comfy_cli/command/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import typer

from comfy_cli import tracking
from comfy_cli.file_utils import atomic_write_text
from comfy_cli.output import get_renderer, rprint

app = typer.Typer(no_args_is_help=True, help="Slot-based editing of frontend-format ComfyUI workflows.")
Expand Down Expand Up @@ -115,22 +116,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:
Expand Down Expand Up @@ -275,7 +260,7 @@ def set_slot_cmd(
sys.stdout.write("\n")
return

_atomic_write_text(p, serialized)
atomic_write_text(p, serialized)

payload = {
"workflow": str(p),
Expand Down Expand Up @@ -372,7 +357,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))
else:
import sys
Expand Down
12 changes: 4 additions & 8 deletions comfy_cli/cql/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

# Cap raw bytes read from disk or the network. Real `object_info` dumps are a
Expand Down Expand Up @@ -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:
Expand Down
77 changes: 77 additions & 0 deletions comfy_cli/file_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import json
import os
import pathlib
import stat
import subprocess
import tempfile
import time
import zipfile
from http import HTTPStatus
Expand All @@ -13,6 +15,81 @@
from comfy_cli import constants, ui


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.
"""
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, "w", encoding="utf-8") as f:
f.write(content)
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)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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

Expand Down
17 changes: 3 additions & 14 deletions comfy_cli/jobs_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,14 @@
from __future__ import annotations

import json
import os
import re
import secrets as _secrets
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from comfy_cli import constants, locking
from comfy_cli.file_utils import atomic_write_text
from comfy_cli.utils import get_os

TERMINAL_STATUSES = frozenset({"completed", "error", "cancelled"})
Expand Down Expand Up @@ -123,18 +122,8 @@ def write(state: JobState) -> Path | None:
# Lock per-file so a watcher and a foreground update can't tear each
# other's writes.
with locking.file_lock(path.with_suffix(".lock")):
tmp = path.with_suffix(f".{os.getpid()}.{_secrets.token_hex(4)}.tmp")
tmp.write_text(json.dumps(state.to_dict(), indent=2, default=str), encoding="utf-8")
# fsync for durability before atomic rename
try:
fd = os.open(str(tmp), os.O_RDONLY)
try:
os.fsync(fd)
finally:
os.close(fd)
except OSError:
pass
os.replace(tmp, path)
# fsync=True: durability against power loss before the atomic rename.
atomic_write_text(path, json.dumps(state.to_dict(), indent=2, default=str), fsync=True)
return path


Expand Down
31 changes: 7 additions & 24 deletions comfy_cli/skills/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@
from pathlib import Path
from typing import Literal

from comfy_cli.file_utils import atomic_write_text

# Where the bundled skills live. Each tuple is (skill_name, package_subdir).
# ``skill_name`` is the public identifier used in AGENTS.md fences and as the
# subdir name in installed targets. ``package_subdir`` is the local resource
Expand Down Expand Up @@ -199,11 +201,7 @@ def read_manifest() -> dict:

def write_manifest(manifest: dict) -> None:
"""Atomically write the manifest (tmp + rename so a SIGINT can't corrupt it)."""
path = manifest_path()
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(f".{os.getpid()}.tmp")
tmp.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
os.replace(tmp, path)
atomic_write_text(manifest_path(), json.dumps(manifest, indent=2))


def _sha256(text: str) -> str:
Expand Down Expand Up @@ -577,24 +575,9 @@ def _backup_if_user_edited(path: Path, expected_content: str) -> Path | None:
return bak


def _atomic_write_text(path: Path, content: str) -> None:
"""Write via tmp + rename so a SIGINT mid-write can't leave the file empty."""
path.parent.mkdir(parents=True, exist_ok=True)
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 _write_claude_skill(path: Path, content: str) -> None:
_backup_if_user_edited(path, content)
_atomic_write_text(path, content)
atomic_write_text(path, content)


def _cursor_description_for(skill_name: str) -> str:
Expand All @@ -610,14 +593,14 @@ def _write_cursor_rule(path: Path, content: str, *, skill_name: str) -> None:
body = _strip_frontmatter(content)
rule = f'---\ndescription: {_cursor_description_for(skill_name)}\nglobs: "**/*"\nalwaysApply: false\n---\n\n{body}'
_backup_if_user_edited(path, rule)
_atomic_write_text(path, rule)
atomic_write_text(path, rule)


def _upsert_agents_md_block(path: Path, content: str, *, skill_name: str) -> None:
start, end = _agents_fence(skill_name)
block = f"\n{start}\n{content}\n{end}\n"
if not path.exists():
_atomic_write_text(path, block.lstrip("\n"))
atomic_write_text(path, block.lstrip("\n"))
return
existing = path.read_text(encoding="utf-8")
if start in existing and end in existing:
Expand All @@ -626,7 +609,7 @@ def _upsert_agents_md_block(path: Path, content: str, *, skill_name: str) -> Non
new = before.rstrip() + "\n\n" + block.lstrip("\n") + after.lstrip("\n")
else:
new = existing.rstrip() + "\n" + block
_atomic_write_text(path, new)
atomic_write_text(path, new)


def _remove_agents_md_block(path: Path, *, skill_name: str) -> bool:
Expand Down
Loading
Loading