fix(run): refresh OAuth session in local partner-credential injection (BE-3361)

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")

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,39 @@ 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, 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=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":

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()

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,39 +654,87 @@ 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
 
+    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

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: <urlopen error>",
+                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
 # ---------------------------------------------------------------------------