From 3a379943a498342dc1460d48a5b811f9a98458ba Mon Sep 17 00:00:00 2001
From: ZY Liu <236001964+liuzhengyang699@users.noreply.github.com>
Date: Sun, 11 Oct 2026 00:21:47 +0800
Subject: [PATCH] fix: honor a zero threshold when realizing agent graphs

---
 .../edge_optimizer/parameterization.py        |  2 +-
 test/swarm/optimizer/test_edge_threshold.py   | 37 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 test/swarm/optimizer/test_edge_threshold.py

diff --git a/swarm/optimizer/edge_optimizer/parameterization.py b/swarm/optimizer/edge_optimizer/parameterization.py
index b09cad4..ef735be 100644
--- a/swarm/optimizer/edge_optimizer/parameterization.py
+++ b/swarm/optimizer/edge_optimizer/parameterization.py
@@ -105,7 +105,7 @@ def realize(self,
             addable_if_not_used_learned_order = (not use_learned_order) and (not _graph.check_cycle(in_node, {out_node}, set()))
             if addable_if_not_used_learned_order or addable_if_use_learned_order:
                 edge_prob = torch.sigmoid(edge_logit / temperature)
-                if threshold:
+                if threshold is not None:
                     edge_prob = torch.tensor(1 if edge_prob > threshold else 0)
                 if torch.rand(1) < edge_prob:
                     out_node.add_successor(in_node)
diff --git a/test/swarm/optimizer/test_edge_threshold.py b/test/swarm/optimizer/test_edge_threshold.py
new file mode 100644
index 0000000..79d9ddc
--- /dev/null
+++ b/test/swarm/optimizer/test_edge_threshold.py
@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from swarm.graph.composite_graph import CompositeGraph
+from swarm.graph.node import Node
+from swarm.optimizer.edge_optimizer.parameterization import EdgeWiseDistribution
+
+pytestmark = pytest.mark.mock_llm
+
+
+class EmptyNode(Node):
+    async def _execute(self, inputs, **kwargs):
+        return {"output": "test"}
+
+
+@pytest.mark.parametrize("learned_order", [False, True])
+@pytest.mark.parametrize("threshold, expected_edges", [(0.0, 1), (1.0, 0), (None, 0)])
+def test_threshold_zero_is_deterministic(monkeypatch, learned_order, threshold, expected_edges):
+    source = EmptyNode("source", "source", True)
+    target = EmptyNode("target", "target", True)
+    graph = CompositeGraph(target, "test")
+    graph.add_node(source)
+    distribution = EdgeWiseDistribution([("source", "target")], initial_probability=0.25)
+    monkeypatch.setattr(torch, "rand", lambda *args: torch.tensor([0.9]))
+    monkeypatch.setattr(
+        distribution,
+        "realize_ranks",
+        lambda *args: ({"source": 0, "target": 1}, torch.tensor(0.0)),
+    )
+
+    realized, log_prob = distribution.realize(
+        graph, threshold=threshold, use_learned_order=learned_order
+    )
+
+    assert realized.num_edges == expected_edges
+    assert torch.isfinite(log_prob)
+    assert graph.num_edges == 0