harden branch collector local-slot mapping

nodes/branch_collector.py

@@ -79,6 +79,14 @@ class DistributedBranchCollector:
             if assigned_value is not None:
                 return assigned_value, assigned_idx
 
+            # In participant-pruned prompts, each participant should typically have
+            # exactly one local branch input. If socket names drifted (e.g. branch_2
+            # connected for a participant assigned to slot 0), remap that sole value
+            # to the assigned slot so the participant still contributes correctly.
+            non_none_indices = [idx for idx, value in enumerate(branch_values) if value is not None]
+            if len(non_none_indices) == 1:
+                return branch_values[non_none_indices[0]], assigned_idx
+
         for idx, value in enumerate(branch_values):
             if value is not None:
                 return value, idx
@@ -256,6 +264,8 @@ class DistributedBranchCollector:
                 outputs[idx] = branch_values[idx]
 
         local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, assigned_branch)
+        if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
+            outputs[local_branch_idx] = local_value
 
         if is_worker:
             if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:

tests/test_join.py

@@ -223,6 +223,26 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(float(outputs[1].abs().sum().item()), 0.0)
         self.assertEqual(float(outputs[2].abs().sum().item()), 0.0)
 
+    async def test_assigned_slot_can_be_filled_from_single_mismatched_local_input(self):
+        fake_server = _FakePromptServer()
+        branch_collector_module._get_prompt_server_instance = lambda: fake_server
+
+        # No worker results expected for this check.
+        fake_server.distributed_pending_jobs["join-job"] = asyncio.Queue()
+
+        local_tensor = torch.full((1, 4, 4, 3), 0.33, dtype=torch.float32)
+        outputs = await self.node.execute(
+            # Local value arrives on slot 1 even though participant is assigned slot 0.
+            [None, local_tensor, None, None, None, None, None, None, None, None],
+            num_branches=3,
+            multi_job_id="join-job",
+            is_worker=False,
+            enabled_worker_ids=json.dumps([]),
+            assigned_branch=0,
+        )
+
+        self.assertTrue(torch.equal(outputs[0], local_tensor))
+
 
 if __name__ == "__main__":
     unittest.main()