From 8e2578ceb6093ded6d0fee5c26c83f401a3329f8 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 19:43:28 -0700 Subject: [PATCH] fix: bound resolve_text on diamond chains with memoisation, depth and length caps A per-path visited set alone makes a chain of diamonds (node i reads node i+1 on both inputs) exponential in depth. Each node's text is now memoised so it is expanded once, every expansion is charged to the shared MAX_VISITS budget, chains deeper than MAX_DEPTH (64) are unknowable so recursion stays bounded, and text longer than MAX_TEXT_LENGTH (100k) is unknowable so StringConcatenate doubling cannot exhaust memory. --- tests/test_prompt_graph.py | 75 ++++++++++++++++++++++++++++++++++---- utils/prompt_graph.py | 31 ++++++++++++---- 2 files changed, 90 insertions(+), 16 deletions(-) diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py index 78fd595..730e72d 100644 --- a/tests/test_prompt_graph.py +++ b/tests/test_prompt_graph.py @@ -4,10 +4,15 @@ import os import sys import time import unittest +from unittest import mock sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from utils import prompt_graph from utils.prompt_graph import ( + MAX_DEPTH, + MAX_TEXT_LENGTH, + MAX_VISITS, _upstream_links, follows_input, has_sampler, @@ -404,20 +409,74 @@ class TestResolveText(unittest.TestCase): } 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): + @staticmethod + def _diamond_chain(levels, leaf): + """Node i reads node i+1 on both inputs: 2**levels paths, levels nodes.""" + graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": leaf}}} + for i in range(1, levels + 1): 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() + graph["pm"] = pm([str(levels), 0]) + return graph + + def _count_expansions(self, graph): + real = prompt_graph._resolve_string_node + with mock.patch.object( + prompt_graph, "_resolve_string_node", side_effect=real + ) as spy: + start = time.monotonic() + text = resolve_text(graph, "pm") + elapsed = time.monotonic() - start + return text, spy.call_count, elapsed + + def test_forty_level_diamond_chain_expands_each_node_once(self): + text, expansions, elapsed = self._count_expansions(self._diamond_chain(40, "")) + self.assertIsNone(text) # the leaf is empty, so the prompt is empty + self.assertLessEqual(expansions, 41) + self.assertLess(elapsed, 0.5) + + def test_diamond_chain_with_text_stays_within_budget(self): + # Doubling "a" 400 times would be a 2**400-character string: the + # length cap turns it into None long before, and the visit budget + # holds regardless of how the graph is shaped. + text, expansions, elapsed = self._count_expansions( + self._diamond_chain(400, "a") + ) + self.assertIsNone(text) + self.assertLessEqual(expansions, MAX_VISITS) + self.assertLess(elapsed, 0.5) + + def test_shallow_diamond_chain_resolves_doubled_text(self): + text, expansions, _ = self._count_expansions(self._diamond_chain(4, "ab")) + self.assertEqual(text, "ab" * 16) + self.assertLessEqual(expansions, 5) + + def test_chain_deeper_than_max_depth_is_none_without_raising(self): + graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": "a"}}} + for i in range(1, 5000): + graph[str(i)] = { + "class_type": "ShowText|pysssss", + "inputs": {"text": [str(i - 1), 0]}, + } + graph["pm"] = pm(["4999", 0]) self.assertIsNone(resolve_text(graph, "pm")) - self.assertLess(time.monotonic() - start, 2.0) + graph["pm"] = pm([str(MAX_DEPTH - 1), 0]) + self.assertEqual(resolve_text(graph, "pm"), "a") + + def test_resolved_text_longer_than_cap_is_none(self): + graph = { + "1": pm(["2", 0]), + "2": { + "class_type": "PrimitiveString", + "inputs": {"value": "x" * (MAX_TEXT_LENGTH + 1)}, + }, + } + self.assertIsNone(resolve_text(graph, "1")) + graph["2"]["inputs"]["value"] = "x" * MAX_TEXT_LENGTH + self.assertEqual(len(resolve_text(graph, "1")), MAX_TEXT_LENGTH) def test_cycles_and_missing_links(self): graph = { diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index 2b96abc..a9b853d 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -15,6 +15,11 @@ PROMPT_MANAGER_TYPES = frozenset({"PromptManager", "PromptManagerText"}) # Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted MAX_VISITS = 500 +# Longest text resolve_text will build; StringConcatenate chains can double +# the text at every level, so anything larger is treated as unknowable +MAX_TEXT_LENGTH = 100_000 +# Deepest chain of string nodes resolve_text follows (recursion stays bounded) +MAX_DEPTH = 64 # Inputs that carry positive conditioning or prompt text towards a sampler _CONDITIONING_KEY = re.compile(r"^conditioning(_\w+)?$") @@ -183,31 +188,41 @@ def resolve_text(graph: Any, node_id: str) -> Optional[str]: return None # 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. + # resolves in both. Each node's text is computed once and memoised, so a + # chain of diamonds costs one expansion per node rather than one per path, + # and the shared visit budget bounds total work whatever the shape. A node + # that fails because of a cycle fails from every path (the cycle is in the + # graph, not in the walk), so memoising None is sound. path = {str(node_id)} - visits = [0] + memo: Dict[str, Optional[str]] = {} + budget = [MAX_VISITS] def resolve(value: Any) -> Optional[str]: if isinstance(value, str): - return value + return value if len(value) <= MAX_TEXT_LENGTH else None 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 path: + if source_id in memo: + return memo[source_id] + if source_id in path or budget[0] <= 0 or len(path) > MAX_DEPTH: return None + budget[0] -= 1 source = graph.get(source_id) if not isinstance(source, dict): + memo[source_id] = None return None path.add(source_id) try: - return _resolve_string_node( + text = _resolve_string_node( str(source.get("class_type", "")), _inputs(source), resolve ) finally: path.discard(source_id) + if isinstance(text, str) and len(text) > MAX_TEXT_LENGTH: + text = None + memo[source_id] = text + return text text = resolve(_inputs(graph.get(str(node_id))).get("text")) return text.strip() if isinstance(text, str) and text.strip() else None