diff --git a/comfy_cli/command/workflow.py b/comfy_cli/command/workflow.py index d953c10e..c11f6413 100644 --- a/comfy_cli/command/workflow.py +++ b/comfy_cli/command/workflow.py @@ -433,9 +433,24 @@ def vary_cmd( # silently writing a partial workflow and reporting success. _USERDATA_MAX_BYTES = 64 * 1024 * 1024 +# Same cap for a single cloud API response. Kept separate from +# ``_USERDATA_MAX_BYTES`` so the two surfaces can diverge without surprise. +_HTTP_MAX_BYTES = 64 * 1024 * 1024 + class _ResponseTooLarge(Exception): - """A ``/userdata`` response exceeded ``_USERDATA_MAX_BYTES`` — refuse to truncate.""" + """A response exceeded the surface's byte cap — refuse to truncate.""" + + +# Per-operation guidance for an oversize cloud response. ``save``/``delete`` +# have already sent their request by the time the response is read, so the +# server-side write may well have landed — say so rather than implying it did not. +_TOO_LARGE_HINTS = { + "list": "narrow the result set with `--limit` or `--name`", + "get": "the saved workflow is unexpectedly large; inspect it directly in the cloud UI", + "save": "the workflow may still have been saved; confirm with `comfy --json workflow list`", + "delete": "the workflow may still have been deleted; confirm with `comfy --json workflow list`", +} # Map the cloud ``--sort`` fields onto local FileInfo keys (client-side sort; @@ -606,7 +621,9 @@ def _http_request( url: str, target, *, method: str = "GET", body: dict | None = None, timeout: float = 30.0 ) -> tuple[int, dict | None]: """Authed HTTP call returning (status, parsed_json_or_none). Raises - urllib errors verbatim so callers can surface the right error code.""" + urllib errors verbatim so callers can surface the right error code, and + ``_ResponseTooLarge`` when the body exceeds ``_HTTP_MAX_BYTES`` — an + oversize body must not masquerade as an unparseable one.""" import urllib.request data = json.dumps(body).encode("utf-8") if body is not None else None @@ -614,12 +631,17 @@ def _http_request( req = _authed_request(url, target, method=method, data=data, content_type=ct) with urllib.request.urlopen(req, timeout=timeout) as resp: status = resp.status - raw = resp.read(64 * 1024 * 1024) # 64 MiB cap + # Read one byte past the cap so we can tell a full body from a truncated one. + raw = resp.read(_HTTP_MAX_BYTES + 1) + if len(raw) > _HTTP_MAX_BYTES: + raise _ResponseTooLarge() if not raw: return status, None try: return status, json.loads(raw) - except json.JSONDecodeError: + except (json.JSONDecodeError, UnicodeDecodeError): + # UnicodeDecodeError is a ValueError but *not* a JSONDecodeError, so a + # body that isn't valid UTF-8 needs naming here or it escapes uncaught. return status, None @@ -627,8 +649,17 @@ def _handle_cloud_http_error(renderer, e, *, operation: str, workflow_id: str | """Map HTTP failures to envelope codes. Returns an Exit to ``raise from``.""" import urllib.error - if isinstance(e, urllib.error.HTTPError): - body = (e.read() or b"")[:1000].decode("utf-8", "replace") + if isinstance(e, _ResponseTooLarge): + renderer.error( + code="workflow_too_large", + message=f"cloud API response during {operation} exceeded the {_HTTP_MAX_BYTES // (1024 * 1024)} MiB cap", + hint=_TOO_LARGE_HINTS.get(operation, "the cloud response was unexpectedly large"), + details={"operation": operation, "workflow_id": workflow_id, "limit_bytes": _HTTP_MAX_BYTES}, + ) + elif isinstance(e, urllib.error.HTTPError): + # Bound the read itself; slicing after an unbounded read would still + # have pulled an arbitrarily large error body into memory first. + body = (e.read(1000) or b"").decode("utf-8", "replace") if e.code == 404: renderer.error( code="workflow_not_found", @@ -935,7 +966,7 @@ def list_cmd( try: _, body = _http_request(url, target) - except (urllib.error.HTTPError, urllib.error.URLError, OSError) as e: + except (urllib.error.HTTPError, urllib.error.URLError, OSError, _ResponseTooLarge) as e: raise _handle_cloud_http_error(renderer, e, operation="list") from e rows = (body or {}).get("data") or [] @@ -1006,7 +1037,7 @@ def get_cmd( url = target.url("workflows", _up.quote(workflow_id, safe=""), "content") try: _, body = _http_request(url, target) - except (urllib.error.HTTPError, urllib.error.URLError, OSError) as e: + except (urllib.error.HTTPError, urllib.error.URLError, OSError, _ResponseTooLarge) as e: raise _handle_cloud_http_error(renderer, e, operation="get", workflow_id=workflow_id) from e if not isinstance(body, dict) or "workflow_json" not in body: @@ -1101,7 +1132,7 @@ def save_cmd( url = target.url("workflows") try: _, resp = _http_request(url, target, method="POST", body=body) - except (urllib.error.HTTPError, urllib.error.URLError, OSError) as e: + except (urllib.error.HTTPError, urllib.error.URLError, OSError, _ResponseTooLarge) as e: raise _handle_cloud_http_error(renderer, e, operation="save") from e workflow_id = (resp or {}).get("id") if isinstance(resp, dict) else None @@ -1137,7 +1168,7 @@ def delete_cmd( url = target.url("workflows", _up.quote(workflow_id, safe="")) try: _, _body = _http_request(url, target, method="DELETE") - except (urllib.error.HTTPError, urllib.error.URLError, OSError) as e: + except (urllib.error.HTTPError, urllib.error.URLError, OSError, _ResponseTooLarge) as e: raise _handle_cloud_http_error(renderer, e, operation="delete", workflow_id=workflow_id) from e payload = {"workflow_id": workflow_id, "deleted": True} diff --git a/tests/comfy_cli/command/test_workflow_saved.py b/tests/comfy_cli/command/test_workflow_saved.py index 2328480b..2d71cef9 100644 --- a/tests/comfy_cli/command/test_workflow_saved.py +++ b/tests/comfy_cli/command/test_workflow_saved.py @@ -466,6 +466,85 @@ def test_404_surfaces_workflow_not_found(self, cloud_target, monkeypatch, capsys assert env["ok"] is False assert env["error"]["code"] == "workflow_not_found" + def test_response_over_cap_refuses_to_truncate(self, cloud_target, monkeypatch, capsys): + # Shrink the cap so the mocked body exceeds it. An oversize body used to + # be truncated, fail to parse, and get swallowed into (status, None) — + # indistinguishable from an empty body. It must fail loudly instead. + monkeypatch.setattr(workflow_cmd, "_HTTP_MAX_BYTES", 4) + _patch_urlopen(monkeypatch, {"/api/workflows/wf-uuid/content": _WORKFLOW_CONTENT_RESPONSE}) + env = _run(["get", "wf-uuid", "--where", "cloud"], capsys) + assert env["ok"] is False + assert env["error"]["code"] == "workflow_too_large" + assert env["error"]["details"]["limit_bytes"] == 4 + + +class TestHttpRequestCap: + """``_http_request`` must distinguish an oversize body from an empty one.""" + + def test_oversize_body_raises_rather_than_returning_none(self, cloud_target, monkeypatch): + monkeypatch.setattr(workflow_cmd, "_HTTP_MAX_BYTES", 4) + _patch_urlopen(monkeypatch, {"/api/workflows": (b'{"data": []}', 200)}) + target = workflow_cmd._resolve_where_target("cloud") + with pytest.raises(workflow_cmd._ResponseTooLarge): + workflow_cmd._http_request(target.url("workflows"), target) + + def test_body_exactly_at_cap_still_parses(self, cloud_target, monkeypatch): + # Boundary: len(raw) == cap is a *complete* body, not a truncated one. + body = b'{"a": 1}' + monkeypatch.setattr(workflow_cmd, "_HTTP_MAX_BYTES", len(body)) + _patch_urlopen(monkeypatch, {"/api/workflows": (body, 200)}) + target = workflow_cmd._resolve_where_target("cloud") + assert workflow_cmd._http_request(target.url("workflows"), target) == (200, {"a": 1}) + + def test_empty_body_still_returns_none(self, cloud_target, monkeypatch): + # The (status, None) contract for a genuinely empty body is unchanged. + _patch_urlopen(monkeypatch, {"/api/workflows": (b"", 200)}) + target = workflow_cmd._resolve_where_target("cloud") + assert workflow_cmd._http_request(target.url("workflows"), target) == (200, None) + + def test_non_utf8_body_returns_none_rather_than_raising(self, cloud_target, monkeypatch): + # UnicodeDecodeError is a ValueError but not a JSONDecodeError; if it is + # not caught it escapes _http_request and no call site handles it. + _patch_urlopen(monkeypatch, {"/api/workflows": (b"\xff\xfe\x00not json", 200)}) + target = workflow_cmd._resolve_where_target("cloud") + assert workflow_cmd._http_request(target.url("workflows"), target) == (200, None) + + def test_non_utf8_body_surfaces_envelope_not_traceback(self, cloud_target, monkeypatch, capsys): + # End-to-end: the undecodable body must reach the user as an envelope. + _patch_urlopen(monkeypatch, {"/api/workflows/wf-uuid/content": (b"\xff\xfe\x00not json", 200)}) + env = _run(["get", "wf-uuid", "--where", "cloud"], capsys) + assert env["ok"] is False + + @pytest.mark.parametrize("verb", ["list", "get", "save", "delete"]) + def test_every_call_site_routes_oversize_to_envelope(self, verb, cloud_target, tmp_path, monkeypatch, capsys): + # Each of the four _http_request call sites must catch _ResponseTooLarge; + # an uncaught one would surface as a traceback rather than an envelope. + monkeypatch.setattr(workflow_cmd, "_HTTP_MAX_BYTES", 4) + _patch_urlopen(monkeypatch, {"/api/workflows": (b'{"data": []}', 200)}) + args = { + "list": ["list"], + "get": ["get", "wf-uuid"], + "delete": ["delete", "wf-uuid"], + }.get(verb) + if args is None: + wf_path = tmp_path / "wf.json" + wf_path.write_text(json.dumps({"1": {"class_type": "KSampler", "inputs": {}}})) + args = ["save", str(wf_path), "--name", "x"] + env = _run([*args, "--where", "cloud"], capsys) + assert env["ok"] is False + assert env["error"]["code"] == "workflow_too_large" + assert env["error"]["details"]["operation"] == verb + # The hint must speak to the operation that actually failed; a single + # hardcoded "saved workflow is too large" is wrong for list/save/delete. + assert env["error"]["hint"] == workflow_cmd._TOO_LARGE_HINTS[verb] + + @pytest.mark.parametrize("verb", ["save", "delete"]) + def test_mutating_verbs_do_not_claim_the_write_was_rejected(self, verb): + # save/delete have already sent their request by the time the response + # is read, so the server-side change may have landed. The hint must not + # imply otherwise, or users will wrongly retry a completed mutation. + assert "may still have been" in workflow_cmd._TOO_LARGE_HINTS[verb] + # --------------------------------------------------------------------------- # save