From d6b02adaf57a7d0e3807610781d295e65f333c22 Mon Sep 17 00:00:00 2001
From: ZY Liu <236001964+liuzhengyang699@users.noreply.github.com>
Date: Sat, 10 Oct 2026 23:52:30 +0800
Subject: [PATCH] fix(mars): preserve full agent IDs when applying verification

---
 optillm/mars/workspace.py                 |  4 +--
 tests/test_mars_workspace_verification.py | 34 +++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_mars_workspace_verification.py

diff --git a/optillm/mars/workspace.py b/optillm/mars/workspace.py
index 874633f9..1a97daf1 100644
--- a/optillm/mars/workspace.py
+++ b/optillm/mars/workspace.py
@@ -69,7 +69,7 @@ def add_verification(self, verification: VerificationResult):
         # Extract agent_id from solution_id (format: "agent_X_iter_Y")
         if verification.solution_id.startswith("agent_"):
             try:
-                agent_id_str = verification.solution_id.split("_")[1]
+                agent_id_str = verification.solution_id.removeprefix("agent_").rsplit("_iter_", 1)[0]
 
                 # Handle both integer and string agent_ids for backward compatibility
                 for solution in self.solutions:
@@ -184,4 +184,4 @@ def get_summary(self) -> Dict[str, Any]:
             'has_consensus': self.has_consensus(),
             'final_solution': self.final_solution,
             'verification_summary': self._get_verification_summary()
-        }
\ No newline at end of file
+        }
diff --git a/tests/test_mars_workspace_verification.py b/tests/test_mars_workspace_verification.py
new file mode 100644
index 00000000..5baaf4bd
--- /dev/null
+++ b/tests/test_mars_workspace_verification.py
@@ -0,0 +1,34 @@
+"""Verification feedback must reach synthesized as well as original solutions."""
+
+from datetime import datetime
+
+import pytest
+
+from optillm.mars.workspace import AgentSolution, MARSWorkspace, VerificationResult
+
+
+@pytest.mark.parametrize("agent_id", [0, "1", "agg_123456", "strategy_iter_2"])
+def test_verification_updates_the_full_agent_id(agent_id):
+    workspace = MARSWorkspace("problem", {"verification_passes_required": 2, "consensus_threshold": 1})
+    solution = AgentSolution(agent_id, "answer", 0.8, 10, 20, 6)
+    workspace.add_solution(solution)
+
+    for _ in range(2):
+        workspace.add_verification(
+            VerificationResult(
+                verifier_id=2,
+                solution_id=f"agent_{agent_id}_iter_0",
+                assessment="CORRECT",
+                confidence=0.9,
+                issues=[],
+                suggestions=[],
+                detailed_report="Checked the answer",
+                timestamp=datetime.now(),
+            )
+        )
+
+    assert len(solution.verification_results) == 2
+    assert solution.verification_score == 1.0
+    assert workspace.get_verified_solutions() == [solution]
+    assert workspace.has_consensus()
+    assert not workspace.should_continue_iteration()