From 248d5880e6beb161dfc3d0001a340b50e84f4743 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 18:31:05 -0700 Subject: [PATCH] fix: never treat negative-only PromptManager nodes as the run prompt run_prompt_nodes fell back to every PromptManager node whenever no positive node was traced, including graphs that do have a sampler and whose only PromptManager node feeds its negative input. Fall back only when the graph has no recognisable sampler (new has_sampler helper). Samplers named like one are recognised even without a model link. Adversarial graphs (cycles, 10k-node fan-out, non-dict inputs) are covered. --- tests/test_prompt_graph.py | 101 +++++++++++++++++++++++++++++++++++++ utils/prompt_graph.py | 40 ++++++++++----- 2 files changed, 128 insertions(+), 13 deletions(-) diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py index b0b6185..9280f6d 100644 --- a/tests/test_prompt_graph.py +++ b/tests/test_prompt_graph.py @@ -2,11 +2,13 @@ import os import sys +import time import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.prompt_graph import ( + has_sampler, positive_prompt_nodes, resolve_text, run_prompt_nodes, @@ -167,6 +169,105 @@ class TestRunPromptNodes(unittest.TestCase): def test_malformed_graph(self): self.assertEqual(run_prompt_nodes(None), []) + def test_run_prompt_nodes_negative_only_graph_is_empty(self): + graph = { + "1": {"class_type": "PromptManager", "inputs": {"text": "bad hands"}}, + "2": { + "class_type": "KSampler", + "inputs": {"negative": ["1", 0], "positive": ["9", 0]}, + }, + "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}, + } + self.assertEqual(run_prompt_nodes(graph), []) + + def test_run_prompt_nodes_falls_back_when_no_sampler(self): + graph = {"1": {"class_type": "PromptManagerText", "inputs": {"text": "x"}}} + self.assertEqual(run_prompt_nodes(graph), ["1"]) + + +class TestHasSampler(unittest.TestCase): + def test_structural_sampler_and_guider_are_detected(self): + self.assertTrue(has_sampler({"5": ksampler("1", "2"), "20": LOADER})) + guider = { + "class_type": "BasicGuider", + "inputs": {"model": ["20", 0], "conditioning": ["1", 0]}, + } + self.assertTrue(has_sampler({"6": guider, "20": LOADER})) + + def test_sampler_without_model_link_is_detected_by_name(self): + sampler = { + "class_type": "KSamplerAdvanced", + "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, + } + self.assertTrue(has_sampler({"5": sampler})) + + def test_non_samplers_and_malformed_graphs(self): + controlnet = { + "class_type": "ControlNetApplyAdvanced", + "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, + } + self.assertFalse(has_sampler({"7": controlnet})) + exotic = {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}} + self.assertFalse(has_sampler({"9": exotic})) + for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}, {}): + self.assertFalse(has_sampler(bad), repr(bad)) + + +class TestAdversarialGraphs(unittest.TestCase): + """Queued and embedded graphs are untrusted: never hang, never raise.""" + + def _assert_fast(self, graph): + started = time.monotonic() + for fn in (positive_prompt_nodes, run_prompt_nodes, has_sampler): + fn(graph) + resolve_text(graph, "1") + self.assertLess(time.monotonic() - started, 1.0) + + def test_cyclic_graph_terminates(self): + graph = { + "1": pm(["2", 0]), + "2": {"class_type": "StringConcatenate", "inputs": {"string_a": ["1", 0]}}, + "3": { + "class_type": "ConditioningCombine", + "inputs": {"conditioning_1": ["3", 0], "conditioning_2": ["1", 0]}, + }, + "5": ksampler("3", "3"), + "20": LOADER, + } + self._assert_fast(graph) + self.assertEqual(run_prompt_nodes(graph), ["1"]) + + def test_ten_thousand_node_fan_out_terminates(self): + graph = {"20": LOADER} + combine_inputs = {} + for i in range(10_000): + graph[str(100 + i)] = pm(f"prompt {i}") + combine_inputs[f"conditioning_{i}"] = [str(100 + i), 0] + graph["3"] = {"class_type": "ConditioningCombine", "inputs": combine_inputs} + graph["5"] = ksampler("3", "3") + self._assert_fast(graph) + self.assertTrue(run_prompt_nodes(graph)) + + def test_non_dict_inputs_and_odd_values(self): + graph = { + "1": {"class_type": "PromptManager", "inputs": ["not", "a", "dict"]}, + "2": {"class_type": "KSampler", "inputs": 42}, + "3": {"class_type": None, "inputs": {"positive": "1", "negative": [1]}}, + "4": "just a string", + "5": { + "class_type": "KSampler", + "inputs": { + "positive": [None, 0], + "negative": [{"x": 1}, 0], + "model": [[], 0], + }, + }, + 6: {"class_type": "PromptManager", "inputs": {"text": 7}}, + } + self._assert_fast(graph) + for fn in (positive_prompt_nodes, run_prompt_nodes): + self.assertIsInstance(fn(graph), list) + class TestResolveText(unittest.TestCase): """Text known at queue time: typed, or from known pure string nodes.""" diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index fbb0df0..0e01bca 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -36,19 +36,26 @@ def _inputs(node: Any) -> Dict[str, Any]: def _positive_roots(graph: Dict[str, Any]) -> List[Any]: - """Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...).""" + """Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...). + + A sampler is a node taking a model plus positive/negative conditioning, or a + guider taking a model plus conditioning. Nodes named like a sampler also count + without a model link, so a negative-only graph is never mistaken for one with + no sampler at all. + """ roots = [] for node in graph.values(): - inputs = _inputs(node) - if not _is_link(inputs.get("model")): + if not isinstance(node, dict): continue + inputs = _inputs(node) + class_name = str(node.get("class_type", "")).lower() + has_model = _is_link(inputs.get("model")) if _is_link(inputs.get("positive")) and _is_link(inputs.get("negative")): - roots.append(inputs["positive"]) - elif ( - _is_link(inputs.get("conditioning")) - and "guider" in str(node.get("class_type", "")).lower() - ): - roots.append(inputs["conditioning"]) + if has_model or "sampler" in class_name: + roots.append(inputs["positive"]) + elif has_model and _is_link(inputs.get("conditioning")): + if "guider" in class_name: + roots.append(inputs["conditioning"]) return roots @@ -63,6 +70,11 @@ def _upstream_links(inputs: Dict[str, Any]) -> List[Any]: return links +def has_sampler(graph: Any) -> bool: + """Whether the graph contains a node recognised as a sampler or guider.""" + return isinstance(graph, dict) and bool(_positive_roots(graph)) + + def positive_prompt_nodes(graph: Any) -> List[str]: """Ids of PromptManager/PromptManagerText nodes that feed a positive input. @@ -92,12 +104,14 @@ def positive_prompt_nodes(graph: Any) -> List[str]: def run_prompt_nodes(graph: Any) -> List[str]: """Prompt nodes that represent a run: the positive ones, if any can be found. - When no PromptManager node can be traced to a positive input (an unrecognised - custom sampler, say), every PromptManager node counts, as before 3.2.4, so an - exotic workflow never silently loses usage counting or image linking. + When the graph has no recognisable sampler (an unrecognised custom sampler, + say), every PromptManager node counts, as before 3.2.4, so an exotic workflow + never silently loses usage counting or image linking. When a sampler exists + and no PromptManager node feeds its positive input, nothing counts: the only + PromptManager nodes are negative prompts. """ positive = positive_prompt_nodes(graph) - if positive or not isinstance(graph, dict): + if positive or not isinstance(graph, dict) or has_sampler(graph): return positive return [ str(node_id)