fix: align balance description and preserve duplicates in sampling

Address Copilot review by matching schema wording with actual deterministic ordering and preventing unconditional signature-based de-duplication during final sampling, so enable_dedup=false semantics remain effective.

Made-with: Cursor

Author: SummerOneTwo

src/autocode_mcp/tools/problem.py

@@ -230,7 +230,7 @@ def input_schema(self) -> dict:
                 },
                 "enable_balance": {
                     "type": "boolean",
-                    "description": "启用平衡分布：在已满足「至少一半为 extreme/tle」后，将剩余名额在各非极限类型间尽量均衡分配；关闭时剩余名额按确定性签名顺序填充",
+                    "description": "启用平衡分布：在已满足「至少一半为 extreme/tle」后，将剩余名额在各非极限类型间尽量均衡分配；关闭时剩余名额按确定性的 (type_param, signature) 顺序填充",
                     "default": True,
                 },
                 "oversample_ratio": {
@@ -604,17 +604,18 @@ def _balance_and_sample(
         )
 
         result: list[CandidateTest] = []
-        used_sig: set[str] = set()
+        selected_ids: set[int] = set()
 
         for c in extreme_pool:
             if len(result) >= need_limit:
                 break
-            if c.signature in used_sig:
+            cid = id(c)
+            if cid in selected_ids:
                 continue
             result.append(c)
-            used_sig.add(c.signature)
+            selected_ids.add(cid)
 
-        remaining = [c for c in candidates if c.signature not in used_sig]
+        remaining = [c for c in candidates if id(c) not in selected_ids]
         need_more = target_count - len(result)
         if need_more <= 0:
             return result[:target_count]
@@ -637,10 +638,11 @@ def _balance_and_sample(
             for i, type_param in enumerate(type_order):
                 count = base_count + (1 if i < rem else 0)
                 for c in by_type[type_param][:count]:
-                    if c.signature in used_sig:
+                    cid = id(c)
+                    if cid in selected_ids:
                         continue
                     result.append(c)
-                    used_sig.add(c.signature)
+                    selected_ids.add(cid)
                     if len(result) >= target_count:
                         break
                 if len(result) >= target_count:
@@ -650,18 +652,20 @@ def _balance_and_sample(
                 for c in sorted(remaining, key=lambda c: (c.type_param, c.signature)):
                     if len(result) >= target_count:
                         break
-                    if c.signature in used_sig:
+                    cid = id(c)
+                    if cid in selected_ids:
                         continue
                     result.append(c)
-                    used_sig.add(c.signature)
+                    selected_ids.add(cid)
         else:
             for c in sorted(remaining, key=lambda c: (c.type_param, c.signature)):
                 if len(result) >= target_count:
                     break
-                if c.signature in used_sig:
+                cid = id(c)
+                if cid in selected_ids:
                     continue
                 result.append(c)
-                used_sig.add(c.signature)
+                selected_ids.add(cid)
 
         return result[:target_count]
 

tests/test_tools/test_problem.py

@@ -673,6 +673,22 @@ def mk(type_param: str, sig: str) -> CandidateTest:
     assert sum(1 for x in out11 if x.type_param in ("3", "4")) >= 6
 
 
+def test_balance_and_sample_keeps_duplicates_when_dedup_disabled():
+    """采样函数不应按 signature 强制去重（由 enable_dedup 控制前置候选）。"""
+    tool = ProblemGenerateTestsTool()
+
+    dup1 = CandidateTest("in-a", "out", "3", "same")
+    dup2 = CandidateTest("in-b", "out", "3", "same")
+    dup3 = CandidateTest("in-c", "out", "2", "same")
+    dup4 = CandidateTest("in-d", "out", "1", "same")
+    candidates = [dup1, dup2, dup3, dup4]
+
+    out = tool._balance_and_sample(candidates, 4, balance_remainder=False)
+    assert len(out) == 4
+    assert out.count(dup1) == 1
+    assert out.count(dup2) == 1
+
+
 @pytest.mark.asyncio
 async def test_problem_generate_tests_balance():
     """测试平衡分布功能。"""