diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py index a01b66a..78fd595 100644 --- a/tests/test_prompt_graph.py +++ b/tests/test_prompt_graph.py @@ -361,6 +361,64 @@ class TestResolveText(unittest.TestCase): resolve_text(self._linked(source), "1"), source["class_type"] ) + def test_same_source_used_twice_resolves_both_branches(self): + # StringConcatenate(string_a=X, string_b=X): X is not a cycle, it is + # simply shared, and must resolve in both branches. + graph = { + "1": pm(["2", 0]), + "2": { + "class_type": "StringConcatenate", + "inputs": { + "string_a": ["3", 0], + "string_b": ["3", 0], + "delimiter": "+", + }, + }, + "3": {"class_type": "PrimitiveString", "inputs": {"value": "x"}}, + } + self.assertEqual(resolve_text(graph, "1"), "x+x") + + def test_diamond_graph_resolves(self): + # 1 <- 2 <- (3, 4) and both 3 and 4 read from 5 + graph = { + "1": pm(["2", 0]), + "2": { + "class_type": "StringConcatenate", + "inputs": { + "string_a": ["3", 0], + "string_b": ["4", 0], + "delimiter": " ", + }, + }, + "3": {"class_type": "ShowText|pysssss", "inputs": {"text": ["5", 0]}}, + "4": {"class_type": "Text Multiline", "inputs": {"text": ["5", 0]}}, + "5": {"class_type": "PrimitiveString", "inputs": {"value": "cat"}}, + } + self.assertEqual(resolve_text(graph, "1"), "cat cat") + + def test_indirect_cycle_returns_none(self): + graph = { + "1": pm(["2", 0]), + "2": {"class_type": "ShowText|pysssss", "inputs": {"text": ["3", 0]}}, + "3": {"class_type": "Text Multiline", "inputs": {"text": ["2", 0]}}, + } + self.assertIsNone(resolve_text(graph, "1")) + + def test_visit_budget_bounds_shared_fan_out(self): + # A chain of concatenations that each read the previous node twice + # expands exponentially without memoisation; MAX_VISITS must cap it. + graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": "a"}}} + for i in range(1, 400): + prev = [str(i - 1), 0] + graph[str(i)] = { + "class_type": "StringConcatenate", + "inputs": {"string_a": prev, "string_b": prev, "delimiter": ""}, + } + graph["pm"] = pm(["399", 0]) + start = time.monotonic() + self.assertIsNone(resolve_text(graph, "pm")) + self.assertLess(time.monotonic() - start, 2.0) + def test_cycles_and_missing_links(self): graph = { "1": pm(["2", 0]), diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index a22c4a7..2b96abc 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -181,24 +181,33 @@ def resolve_text(graph: Any, node_id: str) -> Optional[str]: """ if not isinstance(graph, dict): return None - seen = set() + # Cycle detection is scoped to the current path so a node read by two + # branches (a diamond, or StringConcatenate(string_a=X, string_b=X)) + # resolves in both; the visit counter bounds total work regardless. + path = {str(node_id)} + visits = [0] def resolve(value: Any) -> Optional[str]: if isinstance(value, str): return value - if not _is_link(value) or len(seen) >= MAX_VISITS: + if not _is_link(value): + return None + visits[0] += 1 + if visits[0] > MAX_VISITS: return None source_id = str(value[0]) - if source_id in seen: + if source_id in path: return None - seen.add(source_id) source = graph.get(source_id) if not isinstance(source, dict): return None - return _resolve_string_node( - str(source.get("class_type", "")), _inputs(source), resolve - ) + path.add(source_id) + try: + return _resolve_string_node( + str(source.get("class_type", "")), _inputs(source), resolve + ) + finally: + path.discard(source_id) - seen.add(str(node_id)) text = resolve(_inputs(graph.get(str(node_id))).get("text")) return text.strip() if isinstance(text, str) and text.strip() else None