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.
2026-10-08 18:01:34 -07:00
parent 5b7fa1476a
commit 8e2578ceb6
2 changed files with 90 additions and 16 deletions
@@ -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 = {
@@ -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