From e7fbb3c2db9bdcfae8d1ad0f4f0b08dcdf30124d Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 15 Apr 2026 17:30:20 -0700 Subject: [PATCH 1/5] Add deploy environment header to partner node API calls Read a .comfy_environment file from the ComfyUI base directory to determine the deployment environment (e.g. standalone, portable, desktop). Defaults to 'local_git' when the file is absent. The value is sent as an X-Comfy-Deploy-Env header on all requests to api.comfy.org, allowing the API to differentiate between environment types. The .comfy_environment file is gitignored so launchers/installers can write it without affecting the repository. Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-019d939e-6b4d-738b-8d1a-ac7cbf6736a4 --- .gitignore | 1 + comfy/deploy_environment.py | 33 +++++++++++++++++++++++++++++++++ comfy_api_nodes/util/client.py | 3 +++ 3 files changed, 37 insertions(+) create mode 100644 comfy/deploy_environment.py diff --git a/.gitignore b/.gitignore index 2700ad5c2..b7a7398ac 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,4 @@ web_custom_versions/ openapi.yaml filtered-openapi.yaml uv.lock +.comfy_environment diff --git a/comfy/deploy_environment.py b/comfy/deploy_environment.py new file mode 100644 index 000000000..b4684b363 --- /dev/null +++ b/comfy/deploy_environment.py @@ -0,0 +1,33 @@ +import logging +import os + +import folder_paths + +logger = logging.getLogger(__name__) + +_DEFAULT_DEPLOY_ENV = "local_git" +_ENV_FILENAME = ".comfy_environment" + +_cached_value: str | None = None + + +def get_deploy_environment() -> str: + global _cached_value + if _cached_value is not None: + return _cached_value + + env_file = os.path.join(folder_paths.base_path, _ENV_FILENAME) + try: + with open(env_file, encoding="utf-8") as f: + first_line = f.readline().strip() + value = "".join(c for c in first_line if 32 <= ord(c) < 127) + if value: + _cached_value = value + return _cached_value + except FileNotFoundError: + pass + except Exception as e: + logger.warning("Failed to read %s: %s", env_file, e) + + _cached_value = _DEFAULT_DEPLOY_ENV + return _cached_value diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index 9d730b81a..f09383dc1 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -19,6 +19,8 @@ from comfy import utils from comfy_api.latest import IO from server import PromptServer +from comfy.deploy_environment import get_deploy_environment + from . import request_logger from ._helpers import ( default_base_url, @@ -617,6 +619,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool): payload_headers = {"Accept": "*/*"} if expect_binary else {"Accept": "application/json"} if not parsed_url.scheme and not parsed_url.netloc: # is URL relative? payload_headers.update(get_auth_header(cfg.node_cls)) + payload_headers["X-Comfy-Env"] = get_deploy_environment() if cfg.endpoint.headers: payload_headers.update(cfg.endpoint.headers) From f350a175c315a5a4fa1ab4a00fb9d3d67dc5c1c3 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 21 Apr 2026 14:51:51 -0700 Subject: [PATCH 2/5] Change deploy environment read failure log from warning to error Amp-Thread-ID: https://ampcode.com/threads/T-019db205-95da-7654-ace4-40f12a5f6e69 Co-authored-by: Amp --- comfy/deploy_environment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/deploy_environment.py b/comfy/deploy_environment.py index b4684b363..8085ddc9c 100644 --- a/comfy/deploy_environment.py +++ b/comfy/deploy_environment.py @@ -27,7 +27,7 @@ def get_deploy_environment() -> str: except FileNotFoundError: pass except Exception as e: - logger.warning("Failed to read %s: %s", env_file, e) + logger.error("Failed to read %s: %s", env_file, e) _cached_value = _DEFAULT_DEPLOY_ENV return _cached_value From 06e416bd0dcf34f674efa59527a4b5e818029c47 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 4 May 2026 06:53:40 -0700 Subject: [PATCH 3/5] Bound .comfy_environment read at 128 bytes; add unit tests Defense-in-depth: cap readline() so a malformed or maliciously-large single-line file cannot blow up memory before the value is sanitized. Adds tests-unit/deploy_environment_test.py covering: missing file fallback, basic read, whitespace strip, multi-line (only first line used), empty + whitespace-only files, control-char stripping (header-injection protection), non-ASCII stripping, 128-byte read cap, cache stickiness, and OSError fallback. Amp-Thread-ID: https://ampcode.com/threads/T-019df26e-96f4-7518-94da-0e4263680e3c Co-authored-by: Amp --- comfy/deploy_environment.py | 4 +- tests-unit/deploy_environment_test.py | 84 +++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) create mode 100644 tests-unit/deploy_environment_test.py diff --git a/comfy/deploy_environment.py b/comfy/deploy_environment.py index 8085ddc9c..fb6ee6b47 100644 --- a/comfy/deploy_environment.py +++ b/comfy/deploy_environment.py @@ -19,7 +19,9 @@ def get_deploy_environment() -> str: env_file = os.path.join(folder_paths.base_path, _ENV_FILENAME) try: with open(env_file, encoding="utf-8") as f: - first_line = f.readline().strip() + # Cap the read so a malformed or maliciously crafted file (e.g. + # a single huge line with no newline) can't blow up memory. + first_line = f.readline(128).strip() value = "".join(c for c in first_line if 32 <= ord(c) < 127) if value: _cached_value = value diff --git a/tests-unit/deploy_environment_test.py b/tests-unit/deploy_environment_test.py new file mode 100644 index 000000000..77f7ad35f --- /dev/null +++ b/tests-unit/deploy_environment_test.py @@ -0,0 +1,84 @@ +"""Tests for comfy.deploy_environment.""" + +import os + +import pytest + +from comfy import deploy_environment +from comfy.deploy_environment import get_deploy_environment + + +@pytest.fixture(autouse=True) +def _reset_cache_and_base_path(tmp_path, monkeypatch): + """Reset the module cache and point folder_paths.base_path at a tmp dir for each test.""" + monkeypatch.setattr(deploy_environment, "_cached_value", None) + import folder_paths + monkeypatch.setattr(folder_paths, "base_path", str(tmp_path)) + yield + monkeypatch.setattr(deploy_environment, "_cached_value", None) + + +def _write_env_file(tmp_path, content: str) -> str: + path = os.path.join(str(tmp_path), ".comfy_environment") + with open(path, "w", encoding="utf-8") as f: + f.write(content) + return path + + +class TestGetDeployEnvironment: + def test_returns_local_git_when_file_missing(self): + assert get_deploy_environment() == "local_git" + + def test_reads_value_from_file(self, tmp_path): + _write_env_file(tmp_path, "local_desktop2_standalone\n") + assert get_deploy_environment() == "local_desktop2_standalone" + + def test_strips_trailing_whitespace_and_newline(self, tmp_path): + _write_env_file(tmp_path, " local_desktop2_standalone \n") + assert get_deploy_environment() == "local_desktop2_standalone" + + def test_only_first_line_is_used(self, tmp_path): + _write_env_file(tmp_path, "first_line\nsecond_line\n") + assert get_deploy_environment() == "first_line" + + def test_empty_file_falls_back_to_default(self, tmp_path): + _write_env_file(tmp_path, "") + assert get_deploy_environment() == "local_git" + + def test_empty_after_whitespace_strip_falls_back_to_default(self, tmp_path): + _write_env_file(tmp_path, " \n") + assert get_deploy_environment() == "local_git" + + def test_strips_control_chars_within_first_line(self, tmp_path): + # Embedded NUL/control chars in the value should be stripped + # (header-injection / smuggling protection). + _write_env_file(tmp_path, "abc\x00\x07xyz\n") + assert get_deploy_environment() == "abcxyz" + + def test_strips_non_ascii_characters(self, tmp_path): + _write_env_file(tmp_path, "café-é\n") + assert get_deploy_environment() == "caf-" + + def test_caps_read_at_128_bytes(self, tmp_path): + # A single huge line with no newline must not be fully read into memory. + huge = "x" * 10_000 + _write_env_file(tmp_path, huge) + result = get_deploy_environment() + assert result == "x" * 128 + + def test_result_is_cached_across_calls(self, tmp_path): + path = _write_env_file(tmp_path, "first_value\n") + assert get_deploy_environment() == "first_value" + # Overwrite the file — cached value should still be returned. + with open(path, "w", encoding="utf-8") as f: + f.write("second_value\n") + assert get_deploy_environment() == "first_value" + + def test_unreadable_file_falls_back_to_default(self, tmp_path, monkeypatch): + _write_env_file(tmp_path, "should_not_be_used\n") + + def _boom(*args, **kwargs): + raise OSError("simulated read failure") + + monkeypatch.setattr("builtins.open", _boom) + assert get_deploy_environment() == "local_git" From 22186b3daed7c51f9833033a1f29efe59601504e Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 4 May 2026 06:55:54 -0700 Subject: [PATCH 4/5] Use functools.cache instead of manual global cache Replaces the hand-rolled '_cached_value' module global with @functools.cache, which is the standard Python idiom for memoization. Tests now use the built-in get_deploy_environment.cache_clear() to reset between cases. Amp-Thread-ID: https://ampcode.com/threads/T-019df26e-96f4-7518-94da-0e4263680e3c Co-authored-by: Amp --- comfy/deploy_environment.py | 14 ++++---------- tests-unit/deploy_environment_test.py | 7 +++---- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/comfy/deploy_environment.py b/comfy/deploy_environment.py index fb6ee6b47..ba50ee213 100644 --- a/comfy/deploy_environment.py +++ b/comfy/deploy_environment.py @@ -1,3 +1,4 @@ +import functools import logging import os @@ -8,14 +9,9 @@ logger = logging.getLogger(__name__) _DEFAULT_DEPLOY_ENV = "local_git" _ENV_FILENAME = ".comfy_environment" -_cached_value: str | None = None - +@functools.cache def get_deploy_environment() -> str: - global _cached_value - if _cached_value is not None: - return _cached_value - env_file = os.path.join(folder_paths.base_path, _ENV_FILENAME) try: with open(env_file, encoding="utf-8") as f: @@ -24,12 +20,10 @@ def get_deploy_environment() -> str: first_line = f.readline(128).strip() value = "".join(c for c in first_line if 32 <= ord(c) < 127) if value: - _cached_value = value - return _cached_value + return value except FileNotFoundError: pass except Exception as e: logger.error("Failed to read %s: %s", env_file, e) - _cached_value = _DEFAULT_DEPLOY_ENV - return _cached_value + return _DEFAULT_DEPLOY_ENV diff --git a/tests-unit/deploy_environment_test.py b/tests-unit/deploy_environment_test.py index 77f7ad35f..f2b42f350 100644 --- a/tests-unit/deploy_environment_test.py +++ b/tests-unit/deploy_environment_test.py @@ -4,18 +4,17 @@ import os import pytest -from comfy import deploy_environment from comfy.deploy_environment import get_deploy_environment @pytest.fixture(autouse=True) def _reset_cache_and_base_path(tmp_path, monkeypatch): - """Reset the module cache and point folder_paths.base_path at a tmp dir for each test.""" - monkeypatch.setattr(deploy_environment, "_cached_value", None) + """Reset the functools cache and point folder_paths.base_path at a tmp dir for each test.""" + get_deploy_environment.cache_clear() import folder_paths monkeypatch.setattr(folder_paths, "base_path", str(tmp_path)) yield - monkeypatch.setattr(deploy_environment, "_cached_value", None) + get_deploy_environment.cache_clear() def _write_env_file(tmp_path, content: str) -> str: From 2001646f786fcedcedd91aa7b244152a42691f20 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 4 May 2026 07:20:22 -0700 Subject: [PATCH 5/5] Switch deploy-environment value convention from underscores to dashes Default value is now 'local-git' (was 'local_git'). Dashes are easier to type and more conventional in HTTP-header-adjacent identifiers. Tests updated accordingly. Amp-Thread-ID: https://ampcode.com/threads/T-019df26e-96f4-7518-94da-0e4263680e3c Co-authored-by: Amp --- comfy/deploy_environment.py | 2 +- tests-unit/deploy_environment_test.py | 20 ++++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/comfy/deploy_environment.py b/comfy/deploy_environment.py index ba50ee213..cd45d9576 100644 --- a/comfy/deploy_environment.py +++ b/comfy/deploy_environment.py @@ -6,7 +6,7 @@ import folder_paths logger = logging.getLogger(__name__) -_DEFAULT_DEPLOY_ENV = "local_git" +_DEFAULT_DEPLOY_ENV = "local-git" _ENV_FILENAME = ".comfy_environment" diff --git a/tests-unit/deploy_environment_test.py b/tests-unit/deploy_environment_test.py index f2b42f350..ef803f737 100644 --- a/tests-unit/deploy_environment_test.py +++ b/tests-unit/deploy_environment_test.py @@ -26,27 +26,27 @@ def _write_env_file(tmp_path, content: str) -> str: class TestGetDeployEnvironment: def test_returns_local_git_when_file_missing(self): - assert get_deploy_environment() == "local_git" + assert get_deploy_environment() == "local-git" def test_reads_value_from_file(self, tmp_path): - _write_env_file(tmp_path, "local_desktop2_standalone\n") - assert get_deploy_environment() == "local_desktop2_standalone" + _write_env_file(tmp_path, "local-desktop2-standalone\n") + assert get_deploy_environment() == "local-desktop2-standalone" def test_strips_trailing_whitespace_and_newline(self, tmp_path): - _write_env_file(tmp_path, " local_desktop2_standalone \n") - assert get_deploy_environment() == "local_desktop2_standalone" + _write_env_file(tmp_path, " local-desktop2-standalone \n") + assert get_deploy_environment() == "local-desktop2-standalone" def test_only_first_line_is_used(self, tmp_path): - _write_env_file(tmp_path, "first_line\nsecond_line\n") - assert get_deploy_environment() == "first_line" + _write_env_file(tmp_path, "first-line\nsecond-line\n") + assert get_deploy_environment() == "first-line" def test_empty_file_falls_back_to_default(self, tmp_path): _write_env_file(tmp_path, "") - assert get_deploy_environment() == "local_git" + assert get_deploy_environment() == "local-git" def test_empty_after_whitespace_strip_falls_back_to_default(self, tmp_path): _write_env_file(tmp_path, " \n") - assert get_deploy_environment() == "local_git" + assert get_deploy_environment() == "local-git" def test_strips_control_chars_within_first_line(self, tmp_path): # Embedded NUL/control chars in the value should be stripped @@ -80,4 +80,4 @@ class TestGetDeployEnvironment: raise OSError("simulated read failure") monkeypatch.setattr("builtins.open", _boom) - assert get_deploy_environment() == "local_git" + assert get_deploy_environment() == "local-git"