fix(GCG): address review — schedule loss carry-over, stale state clearing, candidate loss pairing

- Replace RST roles with double-backtick references (check-no-rest-roles)
- ProgressiveMultiPromptAttack carries the inner loss on
  ``ProgressiveScheduleState.loss`` instead of a loose local, so
  ``last_schedule_state.loss`` reflects the final inner run; added a
  post-loop assertion guarding silent carry-over regressions
- ``last_run_state`` / ``last_schedule_state`` default to ``None`` and are
  cleared at the start of each run, so failed runs expose no stale data
- Annealing rejection no longer pairs the accepted suffix with the rejected
  candidate's loss: ``OptimizationRunState.candidate_loss`` tracks what was
  just evaluated while ``loss`` stays paired with ``state.control``

Signed-off-by: fei <204683769+feiiiiii5@users.noreply.github.com>

Author: feiiiiii5

pyrit/executor/promptgen/gcg/attack/base/attack_manager.py

@@ -68,14 +68,17 @@ class OptimizationRunState:
     Captures the current suffix, losses, best result, counters, and stop reason
     explicitly instead of leaving them as loose loop locals, so each phase of
     the optimization loop has a stable contract that can be asserted under
-    seeded tests. Exposed as ``MultiPromptAttack.last_run_state`` after a call
-    to :meth:`MultiPromptAttack.run`.
+    seeded tests. ``loss`` is the loss of the *active* control suffix;
+    ``candidate_loss`` is the loss of the most recently evaluated candidate,
+    which annealing may have rejected. Exposed as ``MultiPromptAttack.last_run_state`` after a call
+    to ``MultiPromptAttack.run``.
     """
 
     control: str
     best_control: str
     loss: float
     best_loss: float
+    candidate_loss: float | None = None
     steps_completed: int = 0
     runtime: float = 0.0
     stop_reason: StopReason | None = None
@@ -84,12 +87,12 @@ class OptimizationRunState:
 @dataclass
 class ProgressiveScheduleState:
     """
-    Typed schedule state for :class:`ProgressiveMultiPromptAttack`.
+    Typed schedule state for ``ProgressiveMultiPromptAttack``.
 
     Tracks how many goals and workers have been admitted so far, together with
     the shared step counter and the loss carried between progressive rounds.
     Exposed as ``ProgressiveMultiPromptAttack.last_schedule_state`` after a call
-    to :meth:`ProgressiveMultiPromptAttack.run`.
+    to ``ProgressiveMultiPromptAttack.run``.
     """
 
     goals_admitted: int
@@ -814,6 +817,10 @@ def disallowed_toks(self) -> torch.Tensor:
 class MultiPromptAttack:
     """A class used to manage multiple prompt-based attacks."""
 
+    #: State of the most recent `run` call; ``None`` until one completes
+    #: and cleared at the start of each run so failed runs expose no stale data.
+    last_run_state: OptimizationRunState | None = None
+
     def __init__(
         self,
         goals: list[str],
@@ -1041,6 +1048,10 @@ def control_weight_fn(_: int) -> float:
             def control_weight_fn(_: int) -> float:
                 return control_weight
 
+        # Clear eagerly: a run that raises mid-loop must not leave the previous
+        # run's state looking current.
+        self.last_run_state = None
+
         state = OptimizationRunState(
             control=self.control_str,
             best_control=self.control_str,
@@ -1082,9 +1093,13 @@ def control_weight_fn(_: int) -> float:
             if keep_control:
                 self.control_str = control
                 state.control = control
+                state.loss = loss
 
+            # ``candidate_loss`` tracks what was just evaluated even when
+            # annealing rejects it, so ``state.loss`` always describes the
+            # suffix in ``state.control``.
+            state.candidate_loss = loss
             prev_loss = loss
-            state.loss = loss
             if loss < state.best_loss:
                 state.best_loss = loss
                 state.best_control = control
@@ -1251,6 +1266,10 @@ def log(
 class ProgressiveMultiPromptAttack:
     """A class used to manage multiple progressive prompt-based attacks."""
 
+    #: State of the most recent `run` call; ``None`` until one completes
+    #: and cleared at the start of each run so failed runs expose no stale data.
+    last_schedule_state: ProgressiveScheduleState | None = None
+
     def __init__(
         self,
         goals: list[str],
@@ -1429,12 +1448,15 @@ def run(
             },
         )
 
+        # Clear eagerly: a run that raises mid-loop must not leave the previous
+        # run's state looking current.
+        self.last_schedule_state = None
+
         schedule = ProgressiveScheduleState(
             goals_admitted=1 if self.progressive_goals else len(self.goals),
             workers_admitted=1 if self.progressive_models else len(self.workers),
             stop_inner_on_success=self.progressive_goals,
         )
-        loss = np.inf
 
         while schedule.steps_completed < n_steps:
             attack = self.managers["MPA"](
@@ -1461,39 +1483,49 @@ def run(
                 control_weight=control_weight,
                 anneal=anneal,
                 anneal_from=schedule.steps_completed,
-                prev_loss=loss,
+                prev_loss=schedule.loss,
                 stop_on_success=schedule.stop_inner_on_success,
                 test_steps=test_steps,
                 filter_cand=filter_cand,
                 verbose=verbose,
             )
-            control, loss, inner_steps = inner_result
+            control, inner_loss, inner_steps = inner_result
+            schedule.loss = inner_loss
 
             schedule.steps_completed += inner_steps
             self.control = control
 
             if schedule.goals_admitted < len(self.goals):
                 schedule.goals_admitted += 1
-                loss = np.inf
+                schedule.loss = np.inf
             elif schedule.goals_admitted == len(self.goals):
                 if schedule.workers_admitted < len(self.workers):
                     schedule.workers_admitted += 1
-                    loss = np.inf
+                    schedule.loss = np.inf
                 elif schedule.workers_admitted == len(self.workers) and stop_on_success:
                     self._finalize_progressive_run(
-                        attack=attack, step=schedule.steps_completed, n_steps=n_steps, loss=loss, verbose=verbose
+                        attack=attack,
+                        step=schedule.steps_completed,
+                        n_steps=n_steps,
+                        loss=schedule.loss,
+                        verbose=verbose,
                     )
                     break
                 else:
                     if isinstance(control_weight, (int, float)) and incr_control:
                         if control_weight <= 0.09:
                             control_weight += 0.01
-                            loss = np.inf
+                            schedule.loss = np.inf
                             if verbose:
                                 logger.info(f"Control weight increased to {control_weight:.5}")
                         else:
                             schedule.stop_inner_on_success = False
 
+        # The inner run must have produced a measurable loss whenever any
+        # optimization happened; guards against silent carry-over regressions.
+        if schedule.steps_completed > 0:
+            assert not math.isinf(schedule.loss), "schedule.loss was never updated by the inner run"
+
         self.last_schedule_state = schedule
 
         return self.control, schedule.steps_completed

tests/unit/executor/promptgen/gcg/test_run_state.py

@@ -49,6 +49,11 @@ def test_counters_and_stop_reason_default(self) -> None:
         assert state.runtime == 0.0
         assert state.stop_reason is None
 
+    def test_candidate_loss_defaults_to_none(self) -> None:
+        state = OptimizationRunState(control="c", best_control="c", loss=1e6, best_loss=1e6)
+
+        assert state.candidate_loss is None
+
 
 class TestProgressiveScheduleState:
     def test_defaults(self) -> None:
@@ -86,22 +91,39 @@ def test_run_records_jailbroken_stop_reason_without_counting_final_check(self) -
         assert state.stop_reason == StopReason.ALL_PROMPTS_JAILBROKEN
         attack.step.assert_not_called()
 
-    def test_rejected_candidate_keeps_active_suffix_but_updates_loss(self) -> None:
+    def test_rejected_candidate_keeps_active_suffix_and_loss(self) -> None:
         attack = _bare_multi_prompt_attack([("better", 1.0), ("worse", 5.0)])
         random.seed(2026)
 
         control, loss, steps = attack.run(n_steps=2, prev_loss=2.0, stop_on_success=False, anneal=True)
 
         # The worse candidate must be rejected by annealing with overwhelming
-        # probability under this seed; the active suffix stays "better".
+        # probability under this seed; the active suffix stays "better" and the
+        # reported loss stays paired with it. The rejected candidate's loss is
+        # still observable through ``candidate_loss``.
         assert control == "better"
         assert steps == 2
         state: OptimizationRunState = attack.last_run_state
         assert state.best_control == "better"
         assert state.best_loss == 1.0
-        assert state.loss == 5.0
+        assert state.control == "better"
+        assert state.loss == 1.0
+        assert state.candidate_loss == 5.0
         assert state.stop_reason == StopReason.MAX_STEPS_REACHED
 
+    def test_failed_run_clears_stale_last_run_state(self) -> None:
+        attack = _bare_multi_prompt_attack([("better", 1.0)])
+        stale = OptimizationRunState(control="stale", best_control="stale", loss=0.1, best_loss=0.1)
+        attack.last_run_state = stale
+        attack.step = MagicMock(side_effect=RuntimeError("model exploded"))
+
+        with pytest.raises(RuntimeError, match="model exploded"):
+            attack.run(n_steps=3, stop_on_success=False)
+
+        # A run that raises mid-loop must not leave the previous run's state
+        # looking current.
+        assert attack.last_run_state is None
+
     def test_periodic_checkpoint_restores_active_suffix(self) -> None:
         attack = _bare_multi_prompt_attack([("better", 1.0), ("best-yet", 0.25)])
         attack.logfile = "unused-by-test.json"  # gate for periodic checkpoints; log/test_all are mocked
@@ -228,3 +250,17 @@ def test_schedule_exhaustion_continues_until_step_budget_spent(self) -> None:
             filter_cand=True,
             verbose=True,
         )
+
+    def test_schedule_loss_carried_on_schedule_object(self) -> None:
+        # The loss fed back between progressive rounds lives on the schedule
+        # state (not a loose local), so ``last_schedule_state.loss`` reflects
+        # the final inner run instead of staying at its ``inf`` default.
+        inner_attack = MagicMock()
+        inner_attack.run.return_value = ("ctrl", 0.75, 3)
+        progressive = self._bare_progressive_attack(inner_attack)
+
+        progressive.run(n_steps=6, stop_on_success=False)
+
+        schedule: ProgressiveScheduleState = progressive.last_schedule_state
+        assert schedule.steps_completed == 6
+        assert schedule.loss == 0.75