From 926313240746537ccf0dd82bcec5894e7e0dd0e2 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 14:03:56 -0700 Subject: [PATCH 1/2] fix(run): refresh OAuth session in local partner-credential injection (BE-3361) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The local partner-node credential injector called resolve_cloud_credential with refresh=False, so a signed-in OAuth-only user whose short-lived access token had lapsed was skipped here and deterministically hit partner_node_requires_credential on any local run — the 'intermittent while signed in' failure from dogfooding. Resolve with refresh=True so the lapsed token is refreshed, and add an allow_clear kwarg to resolve_cloud_credential (forwarded through get_session into ensure_fresh_session) so this best-effort path passes allow_clear=False: a fatal refresh error never clears the shared session, and a transient failure returns the stale session that fails its own expiry check and falls through to env/stored key exactly as before. --- comfy_cli/command/run/credentials.py | 31 +++++-- comfy_cli/credentials.py | 10 ++- tests/comfy_cli/command/test_run.py | 75 +++++++++++++---- tests/comfy_cli/test_credentials.py | 121 +++++++++++++++++++++++++++ 4 files changed, 213 insertions(+), 24 deletions(-) diff --git a/comfy_cli/command/run/credentials.py b/comfy_cli/command/run/credentials.py index 57bc575b..c6d255f9 100644 --- a/comfy_cli/command/run/credentials.py +++ b/comfy_cli/command/run/credentials.py @@ -1,9 +1,17 @@ """Partner-API credential resolution for the local exec path. Cloud auto-injects credentials at submit time; local doesn't. This module -finds the first usable credential (active OAuth session, env var, or stored -API key) so the local exec can hand it to ``submit_prompt`` for nodes that -talk to partner APIs. +finds the first usable credential (OAuth session, env var, or stored API key) +so the local exec can hand it to ``submit_prompt`` for nodes that talk to +partner APIs. + +The OAuth session is refreshed when possible but never cleared from this +best-effort path (``refresh=True, allow_clear=False``): access tokens are +short-lived, so a signed-in user's token routinely lapses between commands; +refreshing it here keeps local runs working, while ``allow_clear=False`` +guarantees a fatal refresh error can't log the user out from under a +foreground command. A refresh that can't succeed falls through to the +env/stored-key tail exactly as before. """ from __future__ import annotations @@ -18,17 +26,24 @@ def _resolve_partner_credential() -> tuple[str, str] | None: OAuth-first chain (``comfy_cli.credentials``), same precedence as the cloud target resolution: - 1. Active (non-expired) OAuth session → ``auth_token_comfy_org`` + 1. Active OAuth session → ``auth_token_comfy_org`` 2. ``COMFY_CLOUD_API_KEY`` env var → ``api_key_comfy_org`` 3. Stored ``comfy-cloud-api-key`` provider record → ``api_key_comfy_org`` - ``refresh=False`` preserves this chain's historical behavior: the session - is read as-is, never refreshed here. Returns ``None`` when nothing usable - is configured. + The session is refreshed when possible (``refresh=True``): access tokens + are short-lived by design, so a signed-in user whose token has lapsed since + their last command would otherwise be skipped here and hit + ``partner_node_requires_credential`` on every local run. ``allow_clear=False`` + keeps this best-effort injector from ever destroying the shared login: a + fatal refresh error does not clear the stored session, and a transient + refresh failure returns the stale session (which fails its own expiry check + and falls through to env/stored key) — so behavior on a network flake is + unchanged, and concurrent ``comfy run`` fan-outs are safe under the OAuth + refresh lock. Returns ``None`` when nothing usable is configured. """ from comfy_cli.credentials import resolve_cloud_credential - cred = resolve_cloud_credential(purpose="cloud", refresh=False) + cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) if cred is None: return None if cred.kind == "oauth": diff --git a/comfy_cli/credentials.py b/comfy_cli/credentials.py index ace6908e..fa632986 100644 --- a/comfy_cli/credentials.py +++ b/comfy_cli/credentials.py @@ -123,6 +123,7 @@ def resolve_cloud_credential( explicit: str | None = None, base_url: str | None = None, refresh: bool = True, + allow_clear: bool = True, ) -> Credential | None: """Resolve the active credential for ``purpose``, or ``None``. @@ -137,12 +138,19 @@ def resolve_cloud_credential( authenticate against). 3. The purpose's env var (``COMFY_CLOUD_API_KEY`` / ``COMFY_API_KEY``). 4. The stored ``comfy-cloud-api-key`` key (``comfy cloud set-key``). + + ``allow_clear=False`` is forwarded into :func:`get_session` (and on to + ``ensure_fresh_session``) so a fatal refresh error on THIS call does not + wipe the shared stored session. Best-effort, read-mostly callers that must + never log the user off the shared session (e.g. the local partner-node + injector) pass this; a failed refresh then just falls through to the + env/stored-key tail instead of destroying the login. """ explicit_key = explicit.strip() if isinstance(explicit, str) else "" if explicit_key: return Credential(kind="api_key", value=explicit_key, source="flag") - session = get_session(refresh=refresh) + session = get_session(refresh=refresh, allow_clear=allow_clear) if ( session is not None and not session.is_expired() diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index d1729aa1..9d58430e 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/tests/comfy_cli/command/test_run.py @@ -618,13 +618,28 @@ def test_diff_detects_pruned_branch(self): class TestResolvePartnerCredential: """The credential the local submit can inject into ``extra_data`` so a - partner-API node finds it. Three sources, env > stored key > OAuth.""" + partner-API node finds it. Precedence session > env > stored key. + + The OAuth session is refreshed when possible (``refresh=True``) but never + cleared from this best-effort path (``allow_clear=False``): access tokens + are short-lived, so a signed-in user's token routinely lapses between + commands — refreshing keeps local runs working, without ever logging the + user off the shared session. The refresh happens inside + ``oauth.ensure_fresh_session`` (mocked here); its allow_clear semantics are + exercised end-to-end in ``tests/comfy_cli/test_credentials.py``. + """ + + def _no_session(self, monkeypatch: pytest.MonkeyPatch): + from comfy_cli.cloud import oauth + + monkeypatch.setattr(oauth, "ensure_fresh_session", lambda **kw: None) - def test_uses_env_var_first(self, monkeypatch: pytest.MonkeyPatch): + def test_uses_env_var_when_no_session(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("COMFY_CLOUD_API_KEY", "env-key-123") from comfy_cli.auth import store as auth_store monkeypatch.setattr(auth_store, "get", lambda _: None) + self._no_session(monkeypatch) assert _resolve_partner_credential() == ("api_key_comfy_org", "env-key-123") def test_falls_back_to_stored_provider_key(self, monkeypatch: pytest.MonkeyPatch): @@ -639,37 +654,67 @@ def test_falls_back_to_stored_provider_key(self, monkeypatch: pytest.MonkeyPatch "get", lambda name: record if name == CLOUD_API_KEY_PROVIDER else None, ) - monkeypatch.setattr(auth_store, "get_cloud_session", lambda: None) + self._no_session(monkeypatch) assert _resolve_partner_credential() == ("api_key_comfy_org", "stored-key-456") - def test_falls_back_to_oauth_token(self, monkeypatch: pytest.MonkeyPatch): + def test_refreshes_and_uses_oauth_token(self, monkeypatch: pytest.MonkeyPatch): + """A signed-in user whose access token lapsed gets a REFRESHED token + here — the whole point of BE-3361 — rather than being skipped and + hitting ``partner_node_requires_credential``.""" monkeypatch.delenv("COMFY_CLOUD_API_KEY", raising=False) from comfy_cli.auth import store as auth_store + from comfy_cli.cloud import oauth + + # ensure_fresh_session refreshes the lapsed token and returns a fresh, + # non-expired session carrying the NEW access token. + refreshed = MagicMock() + refreshed.is_expired.return_value = False + refreshed.access_token = "refreshed-bearer-789" + refreshed.base_url = "https://cloud.comfy.org" + monkeypatch.setattr(auth_store, "get", lambda _: None) + monkeypatch.setattr(oauth, "ensure_fresh_session", lambda **kw: refreshed) + assert _resolve_partner_credential() == ("auth_token_comfy_org", "refreshed-bearer-789") + + def test_passes_allow_clear_false_to_refresh(self, monkeypatch: pytest.MonkeyPatch): + """This best-effort injector must NEVER clear the shared session on a + fatal refresh: it refreshes with ``allow_clear=False``.""" + monkeypatch.delenv("COMFY_CLOUD_API_KEY", raising=False) + from comfy_cli.auth import store as auth_store + from comfy_cli.cloud import oauth + + seen: dict = {} + + def _refresh(**kw): + seen.update(kw) + return None - session = MagicMock() - session.is_expired.return_value = False - session.access_token = "oauth-bearer-789" monkeypatch.setattr(auth_store, "get", lambda _: None) - monkeypatch.setattr(auth_store, "get_cloud_session", lambda: session) - assert _resolve_partner_credential() == ("auth_token_comfy_org", "oauth-bearer-789") + monkeypatch.setattr(oauth, "ensure_fresh_session", _refresh) + _resolve_partner_credential() + assert seen.get("allow_clear") is False def test_returns_none_when_nothing_configured(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("COMFY_CLOUD_API_KEY", raising=False) from comfy_cli.auth import store as auth_store monkeypatch.setattr(auth_store, "get", lambda _: None) - monkeypatch.setattr(auth_store, "get_cloud_session", lambda: None) + self._no_session(monkeypatch) assert _resolve_partner_credential() is None - def test_treats_expired_session_as_no_creds(self, monkeypatch: pytest.MonkeyPatch): + def test_stale_session_from_transient_failure_falls_through(self, monkeypatch: pytest.MonkeyPatch): + """A transient refresh failure returns the STALE (expired) session; it + fails its own expiry check and the resolver falls through — unchanged + from the pre-BE-3361 behavior on a network flake.""" monkeypatch.delenv("COMFY_CLOUD_API_KEY", raising=False) from comfy_cli.auth import store as auth_store + from comfy_cli.cloud import oauth - session = MagicMock() - session.is_expired.return_value = True - session.access_token = "stale" + stale = MagicMock() + stale.is_expired.return_value = True + stale.access_token = "stale" + stale.base_url = "https://cloud.comfy.org" monkeypatch.setattr(auth_store, "get", lambda _: None) - monkeypatch.setattr(auth_store, "get_cloud_session", lambda: session) + monkeypatch.setattr(oauth, "ensure_fresh_session", lambda **kw: stale) assert _resolve_partner_credential() is None diff --git a/tests/comfy_cli/test_credentials.py b/tests/comfy_cli/test_credentials.py index 2b0d5b5c..5b37935e 100644 --- a/tests/comfy_cli/test_credentials.py +++ b/tests/comfy_cli/test_credentials.py @@ -255,6 +255,127 @@ def _refresh(**kw): assert cred is not None and cred.access_token == "forced" +class TestResolveAllowClear: + """``resolve_cloud_credential(allow_clear=...)`` plumbs the flag through + ``get_session`` into ``ensure_fresh_session`` so a best-effort caller (the + local partner-node injector, BE-3361) can REFRESH a lapsed session without + ever CLEARING it on a fatal refresh error. + + Exercised against the real secret store (tmp secrets file) so the whole + ``resolve_cloud_credential → get_session → ensure_fresh_session → + _locked_refresh`` stack — including the ``clear_cloud_session`` decision — is + covered end-to-end, not just the boundary. + """ + + @pytest.fixture + def persisted(self, tmp_path, monkeypatch: pytest.MonkeyPatch): + path = tmp_path / "secrets.json" + monkeypatch.setattr(auth_store, "secrets_path", lambda: path) + monkeypatch.delenv("COMFY_CLOUD_API_KEY", raising=False) + monkeypatch.delenv("COMFY_API_KEY", raising=False) + return path + + def _persist_expired(self): + import time + + auth_store.save_cloud_session( + base_url="https://cloud.comfy.org", + resource="https://cloud.comfy.org/api", + client_id="cid", + scope="s", + access_token="AT-stale", + refresh_token="RT0", + token_type="Bearer", + expires_at=int(time.time()) - 10, # already lapsed + ) + + def test_allow_clear_forwarded_to_ensure_fresh_session(self, clean_env): + """The flag reaches the refresher (cheap boundary check).""" + seen: dict = {} + + def _refresh(**kw): + seen.update(kw) + return None + + clean_env.setattr(oauth, "ensure_fresh_session", _refresh) + resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + assert seen.get("allow_clear") is False + + def test_allow_clear_defaults_true(self, clean_env): + seen: dict = {} + + def _refresh(**kw): + seen.update(kw) + return None + + clean_env.setattr(oauth, "ensure_fresh_session", _refresh) + resolve_cloud_credential(purpose="cloud", refresh=True) + assert seen.get("allow_clear") is True + + def test_expired_session_is_refreshed(self, persisted, monkeypatch: pytest.MonkeyPatch): + """Expired session + valid refresh token → a refresh happens and the + NEW access token is returned (the BE-3361 fix).""" + import time + + self._persist_expired() + + def _refreshed(**kw): + return oauth.TokenSet( + access_token="AT-fresh", + refresh_token="RT1", + token_type="Bearer", + expires_in=3600, + expires_at=int(time.time()) + 3600, + scope="s", + ) + + monkeypatch.setattr(oauth, "refresh_tokens", _refreshed) + cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + assert cred == Credential(kind="oauth", value="AT-fresh", source="session") + + def test_fatal_refresh_allow_clear_false_preserves_session_and_falls_through( + self, persisted, monkeypatch: pytest.MonkeyPatch + ): + """A fatal refresh error with ``allow_clear=False`` must NOT clear the + stored session; the resolver falls through to the env/stored key.""" + self._persist_expired() + monkeypatch.setenv("COMFY_CLOUD_API_KEY", "env-key") + + def _boom(**kw): + raise oauth.OAuthRefreshError( + "refresh failed: HTTP 400", + details={"status": 400, "body": "invalid_grant: refresh token reuse detected"}, + ) + + monkeypatch.setattr(oauth, "refresh_tokens", _boom) + cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + # fell through to the env key... + assert cred == Credential(kind="api_key", value="env-key", source="env:COMFY_CLOUD_API_KEY") + # ...and the shared login was NOT destroyed. + assert auth_store.get_cloud_session() is not None + assert auth_store.get_cloud_session().refresh_token == "RT0" + + def test_transient_refresh_failure_returns_stale_and_falls_through( + self, persisted, monkeypatch: pytest.MonkeyPatch + ): + """A transient (network) refresh failure returns the stale session, + which fails its own expiry check → resolver falls through. Same as the + pre-BE-3361 behavior on a flake; the session is preserved.""" + self._persist_expired() + monkeypatch.setenv("COMFY_CLOUD_API_KEY", "env-key") + + def _flake(**kw): + raise oauth.OAuthRefreshError( + "refresh failed: ", + details={"status": 0, "body": ""}, # transient, non-fatal + ) + + monkeypatch.setattr(oauth, "refresh_tokens", _flake) + cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + assert cred == Credential(kind="api_key", value="env-key", source="env:COMFY_CLOUD_API_KEY") + assert auth_store.get_cloud_session() is not None # stale session kept + + # --------------------------------------------------------------------------- # 2. no-direct-reads ratchet # --------------------------------------------------------------------------- From 64748c64ffcf76df72b567a7045400a027a1f797 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 14:41:13 -0700 Subject: [PATCH 2/2] fix(run): harden best-effort partner-cred refresh per cursor review (BE-3361) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the cursor-review panel findings on the refresh=False→True switch: - Skip _resolve_partner_credential() when an explicit --api-key already populated extra_data: the resolver can perform a synchronous OAuth refresh (network POST + refresh-lock wait), and an explicit key already satisfies the partner node — so an explicit-key run stays network-free. - Make the injector truly best-effort: ensure_fresh_session only swallows transient/timeout cases, so an unexpected OSError (acquiring the refresh lock or persisting the rotated token) previously propagated and aborted the run. Catch it and fall through to a network-free env/stored-key read, exactly the pre-BE-3361 behavior on this path. - Correct the docstring's leeway over-claim: a token inside the 30-60s window is returned as live rather than falling through. Co-Authored-By: Claude Opus 4.8 --- comfy_cli/command/run/__init__.py | 10 ++++++---- comfy_cli/command/run/credentials.py | 25 ++++++++++++++++++++----- tests/comfy_cli/command/test_run.py | 18 ++++++++++++++++++ 3 files changed, 44 insertions(+), 9 deletions(-) diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index fc230d00..482cf6d5 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -218,9 +218,12 @@ def execute( extra_data: dict | None = None if api_key: extra_data = {"api_key_comfy_org": api_key} - if partner_nodes: + # Only resolve an injected credential when an explicit --api-key hasn't + # already satisfied the partner node: the resolver may perform a network + # OAuth refresh, so skipping it here keeps an explicit-key run network-free. + if partner_nodes and not extra_data: cred = _resolve_partner_credential() - if cred is None and not extra_data: + if cred is None: msg = ( "Workflow uses partner-API node(s) that need an `api_key_comfy_org` " "credential the local server doesn't have: " + ", ".join(partner_nodes) + "." @@ -239,8 +242,7 @@ def execute( }, ) raise typer.Exit(code=1) - elif cred is not None and not extra_data: - extra_data = {cred[0]: cred[1]} + extra_data = {cred[0]: cred[1]} # Pre-submit validation via pure-Python CQL engine (checks class_types + input shapes). _preflight_validate(renderer, workflow, object_info, target_label="server") diff --git a/comfy_cli/command/run/credentials.py b/comfy_cli/command/run/credentials.py index c6d255f9..34826838 100644 --- a/comfy_cli/command/run/credentials.py +++ b/comfy_cli/command/run/credentials.py @@ -36,14 +36,29 @@ def _resolve_partner_credential() -> tuple[str, str] | None: ``partner_node_requires_credential`` on every local run. ``allow_clear=False`` keeps this best-effort injector from ever destroying the shared login: a fatal refresh error does not clear the stored session, and a transient - refresh failure returns the stale session (which fails its own expiry check - and falls through to env/stored key) — so behavior on a network flake is - unchanged, and concurrent ``comfy run`` fan-outs are safe under the OAuth - refresh lock. Returns ``None`` when nothing usable is configured. + refresh failure returns the stale session — which, once it is past the + resolver's own 30s expiry leeway, falls through to env/stored key. (A token + inside the 30–60s window is still returned as live; it is refreshed when the + network allows and otherwise good for the imminent submit.) Concurrent + ``comfy run`` fan-outs are safe under the OAuth refresh lock. + + Truly best-effort: the refresh path does network I/O plus a file-locked + persist, and ``ensure_fresh_session`` only swallows the transient/timeout + cases — an unexpected error (e.g. an ``OSError`` acquiring the lock or + saving the rotated token) would otherwise propagate and abort the run. We + catch it and fall through to a network-free read of env/stored keys, so this + injector never turns a refresh hiccup into a failed ``comfy run``. Returns + ``None`` when nothing usable is configured. """ from comfy_cli.credentials import resolve_cloud_credential - cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + try: + cred = resolve_cloud_credential(purpose="cloud", refresh=True, allow_clear=False) + except Exception: # noqa: BLE001 — best-effort: never abort the run on a refresh hiccup + # refresh=False reads the store as-is (no network, no lock): an expired + # session fails its own expiry check and the resolver returns env/stored + # key — exactly the pre-BE-3361 behavior on this path. + cred = resolve_cloud_credential(purpose="cloud", refresh=False, allow_clear=False) if cred is None: return None if cred.kind == "oauth": diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index 9d58430e..1ed22eed 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/tests/comfy_cli/command/test_run.py @@ -717,6 +717,24 @@ def test_stale_session_from_transient_failure_falls_through(self, monkeypatch: p monkeypatch.setattr(oauth, "ensure_fresh_session", lambda **kw: stale) assert _resolve_partner_credential() is None + def test_refresh_path_error_falls_through_to_env_key(self, monkeypatch: pytest.MonkeyPatch): + """The refresh leg does network + file-locked persist; ``ensure_fresh_session`` + only swallows transient/timeout cases, so an unexpected ``OSError`` (lock + acquire / token persist) would otherwise abort the run. This best-effort + injector must catch it and still return the env/stored key network-free.""" + monkeypatch.setenv("COMFY_CLOUD_API_KEY", "env-key-fallback") + from comfy_cli.auth import store as auth_store + from comfy_cli.cloud import oauth + + def _boom(**kw): + raise OSError("cannot acquire refresh lock") + + monkeypatch.setattr(auth_store, "get", lambda _: None) + # refresh=True raises; the network-free fallback reads the store as-is. + monkeypatch.setattr(oauth, "ensure_fresh_session", _boom) + monkeypatch.setattr(auth_store, "get_cloud_session", lambda: None) + assert _resolve_partner_credential() == ("api_key_comfy_org", "env-key-fallback") + class TestExecutePartnerNodePreflight: """Submitting a partner-API workflow to a local server with no