fix: make agentic warmup handoff robust

Signed-off-by: Cam Quilici <cjquilici@gmail.com>

Author: cquil11

src/aiperf/credit/callback_handler.py

@@ -492,23 +492,27 @@ async def on_credit_return(
         # behind ``can_send_child_turn`` instead — the phase-level
         # sending-complete flag is driven by root sampling exhaustion, not
         # by DAG work, but the global ``--request-count`` cap still
-        # applies. When the cap blocks a non-final child continuation, we
-        # notify the orchestrator (``on_child_stopped``) so the parent's
-        # join still drains instead of deadlocking on a child whose
-        # remaining turns will never be issued. Final-turn child returns
-        # are always passed through (the strategy is a no-op for them, but
-        # observer hooks still need to fire).
+        # applies. In terminal phases, a blocked non-final child notifies the
+        # orchestrator (``on_child_stopped``) so the parent's join can drain.
+        # Strategies requesting stopped returns (such as quota warmup) instead
+        # preserve resumable children across a phase handoff. Final-turn child
+        # returns are always passed through (the strategy is a no-op for them,
+        # but observer hooks still need to fire).
+        wants_stopped_returns = (
+            getattr(handler.strategy, "wants_returns_after_sending_complete", False)
+            is True
+        )
         is_child = credit.agent_depth > 0
         if not is_child:
-            wants_stopped_returns = (
-                getattr(handler.strategy, "wants_returns_after_sending_complete", False)
-                is True
-            )
             if handler.stop_checker.can_send_any_turn() or wants_stopped_returns:
                 await handler.strategy.handle_credit_return(
                     credit, error=credit_return.error
                 )
-        elif credit.is_final_turn or handler.stop_checker.can_send_child_turn():
+        elif (
+            credit.is_final_turn
+            or handler.stop_checker.can_send_child_turn()
+            or wants_stopped_returns
+        ):
             await handler.strategy.handle_credit_return(
                 credit, error=credit_return.error
             )

src/aiperf/dataset/loader/weka_synth_buf.py

@@ -431,7 +431,7 @@ def _assert_trailing_user(self) -> None:
     def turn_delta(self) -> TurnDelta:
         """Compute the raw_messages to emit for the just-completed turn.
 
-        Three cases:
+        Four cases:
           1. First call after ``init_turn_0`` (``_emitted_segment_count == 0``):
              emit ALL current segments, ``reset_context=False``. This is
              turn 0's baseline state.
@@ -441,6 +441,10 @@ def turn_delta(self) -> TurnDelta:
           3. Disturbance touched a previously-emitted segment (index
              ``< _emitted_segment_count``): emit ALL current segments,
              ``reset_context=True``.
+          4. An equal-context retry appended no segments: re-emit ALL current
+             segments with ``reset_context=True``. An empty delta would
+             otherwise render as an invalid empty user message instead of the
+             recorded repeated request.
 
         Updates ``_emitted_segment_count`` to ``len(self._segments)`` on
         return. Clears ``_last_disturbance_at`` to ``None``.
@@ -449,9 +453,16 @@ def turn_delta(self) -> TurnDelta:
             self._last_disturbance_at is not None
             and self._last_disturbance_at < self._emitted_segment_count
         )
-        if self._emitted_segment_count == 0 or disturbed_emitted:
+        unchanged_retry = (
+            self._emitted_segment_count > 0
+            and self._emitted_segment_count == len(self._segments)
+            and not disturbed_emitted
+        )
+        if self._emitted_segment_count == 0 or disturbed_emitted or unchanged_retry:
             source = self._segments
-            reset = self._emitted_segment_count != 0 and disturbed_emitted
+            reset = self._emitted_segment_count != 0 and (
+                disturbed_emitted or unchanged_retry
+            )
         else:
             source = self._segments[self._emitted_segment_count :]
             reset = False

src/aiperf/timing/strategies/agentic_replay.py

@@ -254,8 +254,9 @@ def __init__(
             self.scheduler.set_drain_observer(self.enforce_system_idle_cap)
         # Idle-gap cap (ms) for the t* boundary the load-time warp can't see (t*
         # is the sampling instant, not a request). Consumed two ways:
-        #   - WARMUP: clamps each warmup lead so priming doesn't start hours
-        #     early (``_capped_warmup_lead_ms``); priming spacing is meaningless.
+        #   - WARMUP: this cap and the global system-idle cap both clamp each
+        #     warmup lead so priming doesn't start hours early
+        #     (``_capped_warmup_lead_ms``); priming spacing is meaningless.
         #   - PROFILING: a single uniform shift caps the leading idle (t* ->
         #     earliest stream) while preserving recorded inter-stream spacing
         #     (``_leading_idle_shift_ms``); a per-stream clamp would collapse
@@ -548,9 +549,19 @@ def _capped_warmup_lead_ms(self, lead_ms: float) -> float:
         PROFILING dispatch offsets are NOT clamped this way -- see
         :meth:`_leading_idle_shift_ms`.
         """
-        if self._phase_offset_cap_ms is not None:
-            return min(lead_ms, self._phase_offset_cap_ms)
-        return lead_ms
+        caps_ms = [
+            cap
+            for cap in (
+                self._phase_offset_cap_ms,
+                (
+                    self._system_idle_gap_cap_seconds * MILLIS_PER_SECOND
+                    if self._system_idle_gap_cap_seconds is not None
+                    else None
+                ),
+            )
+            if cap is not None
+        ]
+        return min(lead_ms, *caps_ms) if caps_ms else lead_ms
 
     def _leading_idle_shift_ms(self, offsets: Iterable[float]) -> float:
         """Excess to subtract UNIFORMLY from every PROFILING dispatch offset so
@@ -878,15 +889,21 @@ async def _handle_accelerated_warmup_return(self, credit: Credit) -> None:
             await self.credit_issuer.issue_credit(turn)
 
     async def _issue_child_continuation_or_drain(self, turn: TurnToSend) -> None:
-        """Dispatch a DAG child continuation, draining the join on refusal.
+        """Dispatch a DAG child continuation, draining terminal refusals.
 
         ``dispatch_child_turn`` returns True iff the turn reached the wire; on
-        any refusal (e.g. the ``--request-count`` cap) notify the orchestrator
-        so the parent's join drains deterministically instead of deadlocking on
-        a child whose remaining turns will never be issued.
+        a terminal refusal notify the orchestrator so the parent's join drains
+        deterministically instead of deadlocking on a child whose remaining
+        turns will never be issued. Accelerated warmup refusals are different:
+        the remaining child and its active join are persisted for profiling, so
+        marking the child stopped here would release the parent prematurely.
         """
         on_wire = await self.credit_issuer.dispatch_child_turn(turn)
-        if not on_wire and self.branch_orchestrator is not None:
+        if (
+            not on_wire
+            and self.branch_orchestrator is not None
+            and not self.allows_pending_branch_handoff_after_sending_complete
+        ):
             await self.branch_orchestrator.on_child_stopped(turn.x_correlation_id)
 
     async def finalize_phase(self) -> None:
@@ -1460,11 +1477,12 @@ async def _dispatch_next_turn(self, credit: Credit) -> None:
 
         DAG child continuations (``agent_depth > 0``) go through the single
         child-issuance chokepoint (``_issue_child_continuation_or_drain``) so a
-        ``--request-count`` cap refusal is routed to ``on_child_stopped`` (drain
-        the parent join) instead of being silently swallowed by the discarded
-        ``issue_credit`` return -- including on the delayed (``delay_ms``) path,
-        where the refusal would otherwise fire long after the callback handler
-        decided the child could proceed. Root continuations keep ``issue_credit``.
+        terminal refusal is routed to ``on_child_stopped`` (drain the parent
+        join) instead of being silently swallowed by the discarded
+        ``issue_credit`` return. A warmup cutoff is preserved for profiling
+        handoff instead. This applies equally to delayed continuations, whose
+        refusal may happen after the callback handler decided the child could
+        proceed. Root continuations keep ``issue_credit``.
         """
         next_meta = self.conversation_source.get_next_turn_metadata(credit)
         turn = TurnToSend.from_previous_credit(credit, next_meta)

tests/unit/credit/test_callback_handler.py

@@ -790,6 +790,42 @@ async def test_cache_warmup_handoff_allows_paused_dag_work(
     mock_branch_orchestrator.has_pending_branch_work.assert_called_once_with()
 
 
+@pytest.mark.asyncio
+async def test_cache_warmup_handoff_preserves_non_final_child(
+    callback_handler,
+    mock_progress,
+    mock_lifecycle,
+    mock_stop_checker,
+    mock_strategy,
+    mock_branch_orchestrator,
+):
+    """A quota-stopped warmup child remains live for profiling handoff."""
+    mock_stop_checker.can_send_child_turn = MagicMock(return_value=False)
+    mock_strategy.wants_returns_after_sending_complete = True
+    mock_branch_orchestrator.has_pending_branch_work = MagicMock(return_value=True)
+    mock_branch_orchestrator.intercept = AsyncMock(return_value=False)
+    mock_branch_orchestrator.on_child_stopped = AsyncMock()
+    callback_handler.set_branch_orchestrator(mock_branch_orchestrator)
+    callback_handler.register_phase(
+        phase=CreditPhase.WARMUP,
+        progress=mock_progress,
+        lifecycle=mock_lifecycle,
+        stop_checker=mock_stop_checker,
+        strategy=mock_strategy,
+    )
+
+    credit = make_credit(
+        phase=CreditPhase.WARMUP,
+        turn_index=1,
+        num_turns=7,
+        agent_depth=1,
+    )
+    await callback_handler.on_credit_return("worker-1", make_credit_return(credit))
+
+    mock_strategy.handle_credit_return.assert_awaited_once_with(credit, error=None)
+    mock_branch_orchestrator.on_child_stopped.assert_not_awaited()
+
+
 # =============================================================================
 # Test: Credit Return - Unregistered/Complete Phase
 # =============================================================================

tests/unit/dataset/loader/test_weka_pathological.py

@@ -501,6 +501,29 @@ def test_duplicate_hash_ids_in_request_inflate_theoretical_hit_to_full(tmp_path)
     )
 
 
+def test_equal_context_retry_reemits_nonempty_full_prompt(tmp_path):
+    """An identical retry must not become an invalid empty user message."""
+    trace = _base_trace(
+        [
+            _normal(0.0, [1, 2], in_tokens=128),
+            _normal(1.0, [1, 2], in_tokens=128),
+        ],
+        trace_id="equal_retry",
+    )
+    path = tmp_path / "t.json"
+    path.write_text(json.dumps(trace))
+    loader = _make_loader(path, _mk_user_config())
+
+    convs = loader.convert_to_conversations(loader.load_dataset())
+    first, retry = convs[0].turns
+
+    assert first.raw_messages
+    assert retry.raw_messages == first.raw_messages
+    assert retry.reset_context is True
+    assert retry.raw_messages[-1]["role"] == "user"
+    assert retry.raw_messages[-1]["content"]
+
+
 def test_empty_requests_trace_reconstructs_empty_conversation(tmp_path):
     """A trace with zero requests yields a single empty Conversation, no crash."""
     trace = _base_trace([], trace_id="empty_trace")

tests/unit/dataset/loader/test_weka_synth_buf_turn_delta.py

@@ -119,6 +119,36 @@ def test_turn_delta_case_1_strict_append_emits_only_new_segments():
     assert r._last_disturbance_at is None
 
 
+def test_turn_delta_equal_context_retry_reemits_full_context():
+    """An unchanged retry resets to the full prompt instead of emitting ``[]``."""
+    r = _make_recon()
+    hash_ids = [1, 2]
+    in_tokens = 2 * BLOCK_SIZE
+    r.init_turn_0(
+        hash_ids=hash_ids,
+        in_tokens=in_tokens,
+        tool_tokens=0,
+        system_tokens=0,
+        seed="t:0",
+    )
+    first = r.turn_delta()
+
+    r.advance_turn(
+        prev_hash_ids=hash_ids,
+        prev_in_tokens=in_tokens,
+        prev_out_tokens=BLOCK_SIZE,
+        curr_hash_ids=hash_ids,
+        curr_in_tokens=in_tokens,
+        seed="t:1",
+    )
+    retry = r.turn_delta()
+
+    assert retry.reset_context is True
+    assert retry.delta_messages == first.delta_messages
+    assert retry.delta_messages
+    assert retry.delta_messages[-1]["role"] == "user"
+
+
 def test_turn_delta_case_1_strict_append_three_turns_chain():
     """Three sequential strict-append advances: each delta is incremental."""
     r = _make_recon()

tests/unit/timing/strategies/test_agentic_replay.py

@@ -818,8 +818,10 @@ async def capture(turn):
         credit_issuer=issuer,
         lifecycle=lifecycle,
     )
-    # Idle-gap cap of 60s (what the agentx scenario sets).
-    strategy._phase_offset_cap_ms = 60_000.0
+    # The AgentX scenario sets the global system-idle cap, not the per-trace
+    # timestamp-warp cap. Warmup priming must honor that real configuration.
+    strategy._phase_offset_cap_ms = None
+    strategy._system_idle_gap_cap_seconds = 60.0
 
     await strategy.setup_phase()
     await strategy.execute_phase()

tests/unit/timing/strategies/test_agentic_replay_child_continuation.py

@@ -19,9 +19,14 @@ def _make_strategy(
     branch_orchestrator: MagicMock | None = None,
     dispatch_result: bool = True,
     delay_ms: float | None = None,
+    phase: CreditPhase = CreditPhase.PROFILING,
+    accelerated_warmup: bool = False,
 ) -> tuple[AgenticReplayStrategy, MagicMock, MagicMock]:
     """Build a strategy with only the attributes ``_dispatch_next_turn`` reads."""
     strategy = AgenticReplayStrategy.__new__(AgenticReplayStrategy)
+    strategy.config = MagicMock(phase=phase)
+    strategy._cache_warmup_duration = None
+    strategy._cache_warmup_requests_per_lane = 10 if accelerated_warmup else None
 
     conversation_source = MagicMock()
     conversation_source.get_next_turn_metadata.return_value = TurnMetadata(
@@ -104,6 +109,25 @@ async def test_child_at_cap_routes_to_on_child_stopped() -> None:
     orch.on_child_stopped.assert_awaited_once_with("child-xcid")
 
 
+@pytest.mark.asyncio
+async def test_child_at_warmup_quota_is_preserved_for_profiling_handoff() -> None:
+    """A resumable warmup refusal must not prematurely satisfy the parent join."""
+    orch = MagicMock()
+    orch.on_child_stopped = AsyncMock()
+    strategy, issuer, _ = _make_strategy(
+        branch_orchestrator=orch,
+        dispatch_result=False,
+        phase=CreditPhase.WARMUP,
+        accelerated_warmup=True,
+    )
+
+    await strategy._dispatch_next_turn(_child_credit())
+
+    issuer.dispatch_child_turn.assert_awaited_once()
+    issuer.issue_credit.assert_not_called()
+    orch.on_child_stopped.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_child_at_cap_without_orchestrator_swallows_silently() -> None:
     """No orchestrator wired: refusal must not raise."""