diff --git a/comfy_cli/command/generate/app.py b/comfy_cli/command/generate/app.py index 500b8e90..221ae4c4 100644 --- a/comfy_cli/command/generate/app.py +++ b/comfy_cli/command/generate/app.py @@ -2,19 +2,28 @@ UX shape, modeled on fal-ai's genmedia but creative-user-first: - comfy generate [-- value]... [--download P] [--async] + comfy generate [-- value]... [--download P] [--async] [--yes] comfy generate list [--partner P] [--style S] comfy generate schema comfy generate refresh comfy generate resume [--download P] + comfy generate consent [show|always|ask] The first positional is either a reserved action (``list``/``schema``/ -``refresh``/``resume``) or a model alias (``flux-pro``, ``ideogram-edit``, …). -Anything not in the reserved set falls through to the generate path. +``refresh``/``resume``/``consent``) or a model alias (``flux-pro``, +``ideogram-edit``, …). Anything not in the reserved set falls through to the +generate path. + +Spend gate: a generation call spends Comfy credits, so the proxy call sits +behind a consent interlock (``_confirm_spend``) — an interactive TTY prompt, +bypassed by ``--yes`` or the persisted ``spend.auto_confirm`` config, and +fail-closed (error, no spend) when neither is present and no prompt is +possible (``--json`` or no TTY). """ from __future__ import annotations +import sys import uuid from pathlib import Path from typing import Annotated @@ -32,8 +41,9 @@ from rich import print as rprint from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn -from comfy_cli import tracking, ui +from comfy_cli import constants, tracking, ui from comfy_cli.command.generate import adapters, client, emit, output, poll, schema, spec, upload +from comfy_cli.config_manager import ConfigManager from comfy_cli.output.renderer import get_renderer _HELP = "Generate images via ComfyUI partner nodes (Flux, Ideogram, DALL·E, Recraft, Stability, …)." @@ -58,7 +68,7 @@ def _generate_entry( str | None, typer.Argument( help="A model alias (e.g. flux-pro, ideogram-edit, dalle) " - "or one of: list, schema, refresh, upload, resume.", + "or one of: list, schema, refresh, upload, resume, consent.", ), ] = None, ) -> None: @@ -87,12 +97,16 @@ def _generate_entry( {"model": resume_model, "job_id": resume_job_id}, ) return _resume(extra) + if target == "consent": + consent_action = extra[0] if extra and not extra[0].startswith("-") else None + tracking.track_event("generate:consent", {"action": consent_action}) + return _consent(extra) _generate(target, extra) def _separate_meta_flags(extra_args: list[str]) -> tuple[list[str], dict[str, str | bool]]: """Pull run-level flags out of the user's argv tail.""" - meta_names = {"download", "async", "json", "timeout", "api-key", "emit-workflow", "output-prefix"} + meta_names = {"download", "async", "json", "timeout", "api-key", "emit-workflow", "output-prefix", "yes"} meta: dict[str, str | bool] = {} remaining: list[str] = [] i = 0 @@ -104,7 +118,7 @@ def _separate_meta_flags(extra_args: list[str]) -> tuple[list[str], dict[str, st if "=" in body: body, raw = body.split("=", 1) if body in meta_names: - if body in {"async", "json"}: + if body in {"async", "json", "yes"}: meta[body] = True if raw is None else raw.lower() not in {"false", "0", "no"} i += 1 continue @@ -173,6 +187,90 @@ def _emit_result(result: poll.PollResult, *, request_id: str, download: str | No rprint("[yellow]--download requested but no image URLs found in response.[/yellow]") +class SpendNotConfirmed(RuntimeError): + """A credit-spending call lacked consent — declined at the prompt, or + fail-closed because no prompt was possible and nothing pre-authorized it.""" + + +def _spend_auto_confirmed() -> bool: + """True only when the persisted ``spend.auto_confirm`` config is an + affirmative boolean. A missing key or a garbage value never authorizes + spending — the gate fails closed.""" + try: + return bool(ConfigManager().get_bool(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM)) + except ValueError: + return False + + +def _stdin_is_tty() -> bool: + try: + return sys.stdin is not None and sys.stdin.isatty() + except ValueError: + # stdin already closed (e.g. daemonized caller) — no prompt possible. + return False + + +def _confirm_spend(*, model_name: str, assume_yes: bool, as_json: bool) -> None: + """The money interlock: explicit consent before a credit-spending proxy call. + + - ``--yes`` or persisted ``spend.auto_confirm=true`` → proceed. + - Interactive TTY → prompt, default No. + - ``--json`` or no TTY with neither → fail closed: raise, spend nothing. + Never hang on a prompt a machine caller can't answer. + """ + if assume_yes or _spend_auto_confirmed(): + return + if as_json or not _stdin_is_tty(): + msg = ( + f"`comfy generate {model_name}` spends Comfy credits and no consent was given. " + "Re-run with --yes, or persist consent with `comfy generate consent always`." + ) + if as_json: + output.print_json({"error": msg, "code": "spend_consent_required"}) + else: + rprint(f"[bold red]{msg}[/bold red]") + raise SpendNotConfirmed(msg) + rprint(f"[bold]{model_name}[/bold] runs via the partner API and [bold]spends Comfy credits[/bold].") + rprint("[dim]Skip this prompt with --yes; persist always-proceed with `comfy generate consent always`.[/dim]") + if not typer.confirm("Proceed?", default=False): + rprint("Canceled — no credits were spent.") + raise SpendNotConfirmed("user declined the spend confirmation prompt") + + +def _consent(extra_args: list[str]) -> None: + """``comfy generate consent [show|always|ask]`` — inspect or persist the + spend gate's always-proceed setting (``spend.auto_confirm`` in config.ini, + the same store that backs the CLI's other persisted settings).""" + try: + clean, meta = _separate_meta_flags(extra_args) + except schema.SchemaError as e: + rprint(f"[bold red]{e}[/bold red]") + raise typer.Exit(code=1) + as_json = bool(meta.get("json", False)) + action = clean[0] if clean and not clean[0].startswith("-") else "show" + if action not in {"show", "always", "ask"}: + msg = f"Unknown consent action {action!r}. Usage: comfy generate consent [show|always|ask]" + if as_json: + output.print_json({"error": msg}) + else: + rprint(f"[bold red]{msg}[/bold red]") + raise typer.Exit(code=1) + if action == "always": + ConfigManager().set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "true") + elif action == "ask": + ConfigManager().set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "false") + auto = _spend_auto_confirmed() + if as_json: + output.print_json({"spend_auto_confirm": auto, "action": action}) + return + if auto: + rprint("[bold]spend.auto_confirm: true[/bold] — `comfy generate` spends credits without prompting.") + rprint("[dim]Revert with `comfy generate consent ask`.[/dim]") + else: + rprint("[bold]spend.auto_confirm: false[/bold] — credit-spending calls prompt first (or need --yes).") + rprint("[dim]Persist always-proceed with `comfy generate consent always`.[/dim]") + + def _generate(model: str, extra_args: list[str]) -> None: # --help short-circuits before tracking — it's a help-display action, not an execution attempt. # If the model is unknown, fall through so the tracking path records the schema error. @@ -272,6 +370,20 @@ def _track_error(error_kind: str, exc: BaseException) -> None: ) return + # Spend gate — a proxy call spends Comfy credits, so consent comes + # BEFORE any network side effect (auth refresh, asset uploads, the + # generation request itself). Everything above this line is local: + # spec lookup, arg parsing, emit-workflow. + try: + _confirm_spend( + model_name=str(gen_props["model_alias"] or ep.id), + assume_yes=bool(meta.get("yes", False)), + as_json=as_json, + ) + except SpendNotConfirmed as e: + _track_error("consent", e) + raise typer.Exit(code=1) from e + try: api_key = client.resolve_api_key(meta.get("api-key") if isinstance(meta.get("api-key"), str) else None) except client.ApiError as e: @@ -683,7 +795,7 @@ def _print_top_help() -> None: rprint("[bold]comfy generate[/bold] — call ComfyUI partner nodes") rprint("") rprint("[bold]Usage:[/bold]") - rprint(" comfy generate [-- value]... [--download PATH] [--async] [--api-key KEY]") + rprint(" comfy generate [-- value]... [--download PATH] [--async] [--yes] [--api-key KEY]") rprint("") rprint("[bold]Examples:[/bold]") rprint(' comfy generate flux-pro --prompt "a cat on the moon" --width 1024 --height 1024 --download cat.png') @@ -702,7 +814,12 @@ def _print_top_help() -> None: rprint(" comfy generate refresh Refresh the model catalog") rprint(" comfy generate upload Host a local file or remote URL and print its signed URL") rprint(" comfy generate resume Resume an async job") + rprint(" comfy generate consent [show|always|ask] Inspect/persist the credit-spend confirmation") rprint("") + rprint( + "[dim]Generation spends Comfy credits: interactive runs confirm first; pass --yes (or " + "`comfy generate consent always`) to skip, required for --json / non-TTY runs.[/dim]" + ) rprint( "[dim]Auth: run `comfy cloud login` (session outranks env var), set COMFY_API_KEY, or pass --api-key. Get one at https://platform.comfy.org.[/dim]" ) diff --git a/comfy_cli/command/generate/schema.py b/comfy_cli/command/generate/schema.py index aaa67d61..03dd6fc3 100644 --- a/comfy_cli/command/generate/schema.py +++ b/comfy_cli/command/generate/schema.py @@ -259,6 +259,8 @@ def help_text(endpoint: Endpoint, flags: list[FlagDef]) -> str: lines.append(" --download Save outputs locally. Supports {request_id}, {index}, {ext}.") lines.append(" --async Submit and return job id without waiting.") lines.append(" --json Emit raw JSON response instead of pretty output.") + lines.append(" --yes Skip the credit-spend confirmation (required for --json/non-TTY runs") + lines.append(" unless `comfy generate consent always` was set).") lines.append(" --timeout Override sync-poll timeout (default 300).") lines.append(" --api-key Per-call override. Auth: --api-key > login session > COMFY_API_KEY env.") lines.append(" --emit-workflow Write a runnable partner-node workflow instead of calling the proxy.") diff --git a/comfy_cli/constants.py b/comfy_cli/constants.py index 6e4173f9..16f814e7 100644 --- a/comfy_cli/constants.py +++ b/comfy_cli/constants.py @@ -45,6 +45,9 @@ class PROC(str, Enum): CONFIG_KEY_MANAGER_GUI_ENABLED = "manager_gui_enabled" # Legacy, kept for backward compatibility CONFIG_KEY_MANAGER_GUI_MODE = "manager_gui_mode" # Valid: "disable", "enable-gui", "disable-gui", "enable-legacy-gui" CONFIG_KEY_UV_COMPILE_DEFAULT = "uv_compile_default" +# Persistent "always proceed" for credit-spending `comfy generate` calls +# (the spend gate). Set via `comfy generate consent always|ask`. +CONFIG_KEY_SPEND_AUTO_CONFIRM = "spend.auto_confirm" CIVITAI_API_TOKEN_KEY = "civitai_api_token" CIVITAI_API_TOKEN_ENV_KEY = "CIVITAI_API_TOKEN" diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index 4cba21d3..16c4996f 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -577,6 +577,12 @@ class ErrorCode: "`generate --emit-workflow` could not build the partner-node workflow.", "check the model name and that all required inputs are provided", ), + ErrorCode( + "spend_consent_required", + "A credit-spending `comfy generate` call ran non-interactively (--json / no TTY) " + "with no consent — the spend gate fails closed and nothing was spent.", + "re-run with --yes, or persist consent with `comfy generate consent always`", + ), # --- feedback ------------------------------------------------------------ ErrorCode( "feedback_message_required", diff --git a/comfy_cli/skills/comfy/SKILL.md b/comfy_cli/skills/comfy/SKILL.md index 6f9a3649..049d658a 100644 --- a/comfy_cli/skills/comfy/SKILL.md +++ b/comfy_cli/skills/comfy/SKILL.md @@ -646,6 +646,15 @@ comfy generate resume --download outputs/x.mp4 Mechanical contracts that bite agents — encode them, don't rediscover: +- **Spend gate — generation spends the user's Comfy credits and requires + consent.** Interactive TTY runs confirm before spending; `--json` / non-TTY + runs **fail closed** with error code `spend_consent_required` (exit 1, + nothing spent) unless consent is supplied. Pass `--yes` only when the human + has actually approved the spend — do not reflexively add it to make the + error go away. A human can persist always-proceed with + `comfy generate consent always` (revert: `consent ask`; inspect: + `consent show`). `list` / `schema` / `refresh` / `upload` / `resume` / + `--emit-workflow` spend nothing and are not gated. - **Machine-readable output:** `generate --json` prints the raw API response as JSON; `generate upload --json` emits structured `{url, expires_at, …}`; `generate --emit-workflow out.json` goes diff --git a/tests/comfy_cli/command/generate/conftest.py b/tests/comfy_cli/command/generate/conftest.py new file mode 100644 index 00000000..cdf67649 --- /dev/null +++ b/tests/comfy_cli/command/generate/conftest.py @@ -0,0 +1,26 @@ +"""Shared fixtures for the ``comfy generate`` test package.""" + +import pytest + + +@pytest.fixture(autouse=True) +def _auto_confirm_spend(): + """Pre-authorize the credit-spend gate (BE-4103) for generate tests. + + Most tests in this package exercise behavior *downstream* of the consent + interlock, and under CliRunner there is no TTY — without this, every + execution-path invocation would fail closed before reaching the code + under test. The gate's own tests override this fixture by name in + ``test_spend_gate.py`` to put the gate back in force. + + ConfigManager is a process-wide singleton, so the teardown removes the + in-memory key rather than relying on the per-test config-dir isolation + alone (the file is isolated; the loaded configparser is not). + """ + from comfy_cli import constants + from comfy_cli.config_manager import ConfigManager + + cm = ConfigManager() + cm.set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "true") + yield + cm.config.remove_option("DEFAULT", constants.CONFIG_KEY_SPEND_AUTO_CONFIRM) diff --git a/tests/comfy_cli/command/generate/test_spend_gate.py b/tests/comfy_cli/command/generate/test_spend_gate.py new file mode 100644 index 00000000..7dcb952e --- /dev/null +++ b/tests/comfy_cli/command/generate/test_spend_gate.py @@ -0,0 +1,252 @@ +"""Spend-gate tests for ``comfy generate`` (BE-4103). + +A generation call spends Comfy credits, so the proxy call sits behind a +consent interlock: interactive TTY runs prompt first, ``--yes`` or the +persisted ``spend.auto_confirm`` config bypass the prompt, and ``--json`` / +non-TTY runs with neither **fail closed** — error out, spend nothing, never +hang on a prompt no one can answer. +""" + +import json + +import httpx +import pytest +from typer.testing import CliRunner + +from comfy_cli import constants +from comfy_cli.cmdline import app as cli_app +from comfy_cli.command.generate import app as gen_app +from comfy_cli.config_manager import ConfigManager + + +@pytest.fixture(autouse=True) +def _auto_confirm_spend(): + """Override the package conftest's pre-authorization: these tests + exercise the gate itself, so it starts (and ends) cleared.""" + cm = ConfigManager() + cm.config.remove_option("DEFAULT", constants.CONFIG_KEY_SPEND_AUTO_CONFIRM) + yield + cm.config.remove_option("DEFAULT", constants.CONFIG_KEY_SPEND_AUTO_CONFIRM) + + +@pytest.fixture(autouse=True) +def disable_tracking_prompt(monkeypatch): + monkeypatch.setattr("comfy_cli.tracking.prompt_tracking_consent", lambda *a, **kw: None) + monkeypatch.setattr("comfy_cli.tracking.track_event", lambda *a, **kw: None) + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def api_key(monkeypatch): + monkeypatch.setenv("COMFY_API_KEY", "comfyui-test") + return "comfyui-test" + + +@pytest.fixture +def post_spy(monkeypatch): + """Record every proxy POST; each recorded call is a would-be credit spend.""" + calls: list = [] + + def _post(*a, **kw): + calls.append((a, kw)) + return httpx.Response(200, json={"data": [{"url": "https://cdn.example/a.png"}]}) + + monkeypatch.setattr(gen_app.client.httpx, "post", _post) + return calls + + +@pytest.fixture +def interactive_tty(monkeypatch): + """Simulate a human at a terminal — CliRunner's piped stdin is never a TTY.""" + monkeypatch.setattr(gen_app, "_stdin_is_tty", lambda: True) + + +# ─── Fail-closed: no consent, no TTY → error, nothing spent ─────────────── + + +def test_json_without_consent_fails_closed_and_spends_nothing(runner, api_key, post_spy): + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 1 + assert post_spy == [] + payload = json.loads(r.stdout) + assert payload["code"] == "spend_consent_required" + assert "--yes" in payload["error"] + + +def test_pretty_non_tty_without_consent_fails_closed(runner, api_key, post_spy): + # No --json, but stdin is a pipe (CliRunner): never hang on a prompt. + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x"]) + assert r.exit_code == 1 + assert post_spy == [] + assert "spends Comfy credits" in r.stdout + assert "--yes" in r.stdout + + +def test_async_submit_is_gated_too(runner, api_key, post_spy): + # --async still spends on submit — the gate must cover it. + r = runner.invoke( + cli_app, + ["generate", "flux-pro", "--prompt", "x", "--width", "1", "--height", "1", "--async", "--json"], + ) + assert r.exit_code == 1 + assert post_spy == [] + + +def test_gate_runs_before_auth_resolution(runner, post_spy): + # No COMFY_API_KEY in env: consent still decides first, so an + # unauthenticated machine caller sees the consent error, not an auth one, + # and no OAuth refresh / network happens without consent. + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 1 + assert post_spy == [] + assert json.loads(r.stdout)["code"] == "spend_consent_required" + + +# ─── Bypasses: --yes flag and spend.auto_confirm config ─────────────────── + + +def test_json_with_yes_proceeds(runner, api_key, post_spy): + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json", "--yes"]) + assert r.exit_code == 0, r.stdout + assert len(post_spy) == 1 + + +def test_json_with_auto_confirm_config_proceeds(runner, api_key, post_spy): + ConfigManager().set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "true") + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 0, r.stdout + assert len(post_spy) == 1 + + +def test_auto_confirm_false_still_fails_closed(runner, api_key, post_spy): + ConfigManager().set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "false") + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 1 + assert post_spy == [] + + +def test_auto_confirm_garbage_value_fails_closed(runner, api_key, post_spy): + # A corrupt config value must never silently authorize spending. + ConfigManager().set(constants.CONFIG_KEY_SPEND_AUTO_CONFIRM, "banana") + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 1 + assert post_spy == [] + + +# ─── Interactive TTY: prompt before spending ────────────────────────────── + + +def test_interactive_prompt_accept_proceeds(runner, api_key, post_spy, interactive_tty): + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x"], input="y\n") + assert r.exit_code == 0, r.stdout + assert len(post_spy) == 1 + assert "spends Comfy credits" in r.stdout + + +def test_interactive_prompt_decline_spends_nothing(runner, api_key, post_spy, interactive_tty): + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x"], input="n\n") + assert r.exit_code == 1 + assert post_spy == [] + assert "no credits were spent" in r.stdout + + +def test_interactive_prompt_default_is_no(runner, api_key, post_spy, interactive_tty): + # Bare Enter must not spend. + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x"], input="\n") + assert r.exit_code == 1 + assert post_spy == [] + + +# ─── Ungated paths: nothing that spends ─────────────────────────────────── + + +def test_emit_workflow_is_not_gated(runner, post_spy, tmp_path): + # --emit-workflow writes a local artifact, calls no proxy, spends nothing. + out = tmp_path / "wf.json" + r = runner.invoke(cli_app, ["generate", "flux-pro", "--prompt", "x", "--emit-workflow", str(out)]) + assert r.exit_code == 0, r.stdout + assert out.exists() + assert post_spy == [] + + +def test_schema_error_surfaces_before_consent(runner, api_key, post_spy): + # Bad args fail on validation, not on consent — no prompt, no spend. + r = runner.invoke(cli_app, ["generate", "flux-pro", "--prompt", "x", "--width", "abc", "--height", "1"]) + assert r.exit_code == 1 + assert post_spy == [] + assert "spend" not in r.stdout.lower() + + +# ─── Lifecycle: consent failures emit generate:error(kind=consent) ──────── + + +def test_consent_failure_emits_error_kind_consent(runner, api_key, post_spy, monkeypatch): + events: list[tuple[str, dict]] = [] + + def _record(event_name, properties=None, *, mixpanel_name=None): + events.append((event_name, dict(properties or {}))) + + monkeypatch.setattr("comfy_cli.tracking.track_event", _record) + monkeypatch.setattr("comfy_cli.command.generate.app.tracking.track_event", _record) + monkeypatch.setattr("comfy_cli.cmdline.tracking.track_event", _record) + + r = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r.exit_code == 1 + names = [n for n, _ in events] + assert names.count("generate:start") == 1 + err_props = [p for n, p in events if n == "generate:error"] + assert len(err_props) == 1 + assert err_props[0]["error_kind"] == "consent" + assert "generate:success" not in names + + +# ─── `comfy generate consent` action ────────────────────────────────────── + + +def test_consent_show_defaults_to_false(runner): + r = runner.invoke(cli_app, ["generate", "consent", "--json"]) + assert r.exit_code == 0 + payload = json.loads(r.stdout) + assert payload == {"spend_auto_confirm": False, "action": "show"} + + +def test_consent_always_persists_and_unlocks_generate(runner, api_key, post_spy): + r = runner.invoke(cli_app, ["generate", "consent", "always", "--json"]) + assert r.exit_code == 0 + assert json.loads(r.stdout)["spend_auto_confirm"] is True + # Persisted to the config file on disk, not just in memory. + from pathlib import Path + + cfg_text = Path(ConfigManager().get_config_file_path()).read_text() + assert "spend.auto_confirm = true" in cfg_text + + r2 = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r2.exit_code == 0, r2.stdout + assert len(post_spy) == 1 + + +def test_consent_ask_reverts(runner, api_key, post_spy): + runner.invoke(cli_app, ["generate", "consent", "always"]) + r = runner.invoke(cli_app, ["generate", "consent", "ask", "--json"]) + assert r.exit_code == 0 + assert json.loads(r.stdout)["spend_auto_confirm"] is False + + r2 = runner.invoke(cli_app, ["generate", "dalle", "--prompt", "x", "--json"]) + assert r2.exit_code == 1 + assert post_spy == [] + + +def test_consent_show_pretty_output(runner): + r = runner.invoke(cli_app, ["generate", "consent"]) + assert r.exit_code == 0 + assert "spend.auto_confirm: false" in r.stdout + + +def test_consent_unknown_action_errors(runner): + r = runner.invoke(cli_app, ["generate", "consent", "sometimes"]) + assert r.exit_code == 1 + assert "Unknown consent action" in r.stdout