From 8241bb95851c50be505a3f2a6078b9c250ee91ca Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 2 Mar 2026 07:54:44 +0000 Subject: [PATCH] prune worker nodes downstream of branch collector --- api/orchestration/prompt_transform.py | 3 ++- tests/test_prompt_transform.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 61051f2..9e8d02c 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -317,7 +317,8 @@ def prune_prompt_for_worker(prompt_obj): if branch_ids: connected.update(_find_downstream_nodes(prompt_obj, branch_ids)) if branch_collector_ids: - connected.update(_find_downstream_nodes(prompt_obj, branch_collector_ids)) + downstream_of_collectors = _find_downstream_nodes(prompt_obj, branch_collector_ids) + connected -= (downstream_of_collectors - set(branch_collector_ids)) pruned_prompt = {} for node_id in connected: diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 8c0d23d..c56a793 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -302,11 +302,11 @@ class PrunePromptForWorkerTests(unittest.TestCase): self.assertIn("5", result) self.assertIn("6", result) - def test_branch_collector_anchor_keeps_collector_and_downstream(self): + def test_branch_collector_anchor_keeps_collector_but_prunes_downstream(self): prompt = _branch_collector_prompt() result = pt.prune_prompt_for_worker(prompt) self.assertIn("5", result) - self.assertIn("6", result) + self.assertNotIn("6", result) # ---------------------------------------------------------------------------