fix: bound random edge sampling to available candidates

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 = []

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