From 70d03d9a204fe3db3b9201f7015e9b618e59c973 Mon Sep 17 00:00:00 2001
From: ZY Liu <236001964+liuzhengyang699@users.noreply.github.com>
Date: Sat, 10 Oct 2026 23:44:45 +0800
Subject: [PATCH] fix: bound random edge sampling to available candidates

---
 .../edge_optimizer/parameterization.py        | 18 +++--
 test/swarm/graph/test_edge_sampling.py        | 67 +++++++++++++++++++
 2 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 test/swarm/graph/test_edge_sampling.py

diff --git a/swarm/optimizer/edge_optimizer/parameterization.py b/swarm/optimizer/edge_optimizer/parameterization.py
index b09cad4..877d773 100644
--- a/swarm/optimizer/edge_optimizer/parameterization.py
+++ b/swarm/optimizer/edge_optimizer/parameterization.py
@@ -44,10 +44,12 @@ def __init__(self,
 
     def random_sample_num_edges(self, graph: CompositeGraph, num_edges: int) -> CompositeGraph:
         _graph = deepcopy(graph)
-        while True:
-            if _graph.num_edges >= num_edges:
-                break
-            potential_connection = random.sample(self.potential_connections, 1)[0]
+        if _graph.num_edges >= num_edges:
+            return _graph
+
+        # Rejected edges stay invalid as the graph grows; try each candidate once.
+        candidates = random.sample(self.potential_connections, len(self.potential_connections))
+        for potential_connection in candidates:
             out_node = _graph.find_node(potential_connection[0])
             in_node = _graph.find_node(potential_connection[1])
 
@@ -57,7 +59,13 @@ def random_sample_num_edges(self, graph: CompositeGraph, num_edges: int) -> Comp
             if not _graph.check_cycle(in_node, {out_node}, set()):
                 out_node.add_successor(in_node)
                 in_node.add_predecessor(out_node)
-        return _graph
+                if _graph.num_edges >= num_edges:
+                    return _graph
+
+        raise ValueError(
+            f"Could only sample {_graph.num_edges} edges without a cycle; "
+            f"requested {num_edges} edges."
+        )
 
     def realize_ranks(self, graph, use_max: bool = False):
         log_probs = []
diff --git a/test/swarm/graph/test_edge_sampling.py b/test/swarm/graph/test_edge_sampling.py
new file mode 100644
index 0000000..bd2a092
--- /dev/null
+++ b/test/swarm/graph/test_edge_sampling.py
@@ -0,0 +1,67 @@
+import random
+
+import pytest
+
+from swarm.graph.composite_graph import CompositeGraph
+from swarm.graph.node import Node
+from swarm.optimizer.edge_optimizer.parameterization import EdgeWiseDistribution
+
+
+class DummyNode(Node):
+    def __init__(self, node_id):
+        super().__init__("test node", node_id, combine_inputs_as_one=False)
+
+    async def _execute(self, inputs, **kwargs):
+        return inputs
+
+
+@pytest.fixture
+def graph():
+    graph = CompositeGraph(DummyNode("output"), "test")
+    graph.add_node(DummyNode("a"))
+    graph.add_node(DummyNode("b"))
+    return graph
+
+
+@pytest.fixture
+def bounded_sampling(monkeypatch):
+    """Fail promptly if an impossible edge budget causes unbounded resampling."""
+    calls = 0
+    sample = random.Random(0).sample
+
+    def bounded_sample(population, k):
+        nonlocal calls
+        calls += 1
+        assert calls <= max(1, len(population)), "repeated sampling cannot satisfy this edge budget"
+        return sample(population, k)
+
+    monkeypatch.setattr(random, "sample", bounded_sample)
+
+
+@pytest.mark.mock_llm
+def test_sample_requested_edges_without_mutating_source(graph):
+    distribution = EdgeWiseDistribution([("a", "b"), ("b", "output")])
+
+    sampled = distribution.random_sample_num_edges(graph, 2)
+
+    assert sampled.num_edges == 2
+    assert graph.num_edges == 0
+
+
+@pytest.mark.mock_llm
+@pytest.mark.parametrize("connections", [[], [("a", "b")], [("a", "b"), ("b", "a")]])
+def test_unreachable_edge_budget_raises(graph, bounded_sampling, connections):
+    distribution = EdgeWiseDistribution(connections)
+
+    with pytest.raises(ValueError, match="edge"):
+        distribution.random_sample_num_edges(graph, 2)
+
+
+@pytest.mark.mock_llm
+def test_existing_edges_already_satisfy_budget(graph):
+    graph.nodes["a"].add_successor(graph.nodes["output"])
+    distribution = EdgeWiseDistribution([])
+
+    sampled = distribution.random_sample_num_edges(graph, 1)
+
+    assert sampled.num_edges == 1