fix(mars): populate diverse aggregation groups

optillm/mars/aggregator.py

@@ -148,7 +148,7 @@ def _sample_solutions_for_aggregation(self, workspace: MARSWorkspace) -> List[Li
 
                 sample = []
                 for agent_solutions in by_agent.values():
-                    if sample and len(sample) < self.aggregation_size:
+                    if len(sample) < self.aggregation_size:
                         sample.append(max(agent_solutions, key=lambda s: s.confidence))
                     if len(sample) >= self.aggregation_size:
                         break
@@ -331,4 +331,4 @@ def _ensure_population_size(self, workspace: MARSWorkspace) -> None:
             logger.warning(f"Population size ({current_size}) < aggregation size ({self.aggregation_size})")
             logger.warning("Aggregation may be less effective with limited diversity")
 
-        logger.info(f"Population ready: {current_size} solutions available for aggregation")
+        logger.info(f"Population ready: {current_size} solutions available for aggregation")

tests/test_mars_aggregation_sampling.py

@@ -0,0 +1,41 @@
+"""Check that aggregation retains candidates from different reasoning agents."""
+
+from optillm.mars.aggregator import MARSAggregator
+from optillm.mars.workspace import AgentSolution, MARSWorkspace
+
+
+def solution(agent_id, confidence, score):
+    return AgentSolution(
+        agent_id=agent_id,
+        solution=f"{agent_id}: {confidence}",
+        confidence=confidence,
+        reasoning_tokens=0,
+        total_tokens=0,
+        solution_length=0,
+        verification_score=score,
+    )
+
+
+def test_diverse_group_selects_best_confidence_from_each_agent():
+    workspace = MARSWorkspace("problem", {})
+    strongest = [solution("a", 0.4, 0.9), solution("a", 0.8, 0.9), solution("a", 0.6, 0.9)]
+    other_agents = [solution("b", 0.7, 0.5), solution("c", 0.6, 0.4)]
+    workspace.solutions = strongest + other_agents
+    aggregator = MARSAggregator(None, "test-model", {"population_size": 6, "aggregation_size": 3})
+
+    groups = aggregator._sample_solutions_for_aggregation(workspace)
+
+    assert groups[0] == strongest
+    assert groups[1] == [strongest[1], *other_agents]
+
+
+def test_diverse_group_fills_remaining_slots_without_duplicates():
+    workspace = MARSWorkspace("problem", {})
+    candidates = [solution("a", 0.8, 0.9), solution("a", 0.5, 0.8), solution("b", 0.7, 0.5)]
+    workspace.solutions = candidates
+    aggregator = MARSAggregator(None, "test-model", {"population_size": 6, "aggregation_size": 3})
+
+    groups = aggregator._sample_solutions_for_aggregation(workspace)
+
+    assert all(candidate in groups[1] for candidate in candidates)
+    assert len({id(candidate) for candidate in groups[1]}) == 3