fix(timing): carry has_forks onto continuation turns

src/aiperf/timing/strategies/fixed_schedule.py

@@ -153,7 +153,10 @@ async def handle_credit_return(
 
         # This contains the delay_ms or timestamp_ms for the next turn
         next_meta = self._conversation_source.get_next_turn_metadata(credit)
-        turn = TurnToSend.from_previous_credit(credit)
+        # Pass next_meta so has_forks rides onto the continuation turn: the
+        # sticky router defers parent-entry eviction until DAG children drain
+        # (dropping it premature-evicts a fork-bearing parent's later turns).
+        turn = TurnToSend.from_previous_credit(credit, next_meta)
 
         if next_meta.timestamp_ms is not None:
             self._scheduler.schedule_at_perf_sec(

src/aiperf/timing/strategies/user_centric_rate.py

@@ -340,7 +340,11 @@ async def handle_credit_return(
             raise ValueError(
                 f"User not found for x_correlation_id: {credit.x_correlation_id}"
             )
-        turn = TurnToSend.from_previous_credit(credit)
+        # Pass next-turn metadata so has_forks rides onto the continuation turn
+        # (the sticky router defers parent-entry eviction until DAG children
+        # drain); dropping it premature-evicts a fork-bearing parent's turns.
+        meta = self._conversation_source.get_next_turn_metadata(credit)
+        turn = TurnToSend.from_previous_credit(credit, meta)
 
         # If the next turn time already passed, the max() will
         # re-align their schedule to account for the delay.

tests/unit/timing/strategies/test_fixed_schedule.py

@@ -227,6 +227,33 @@ async def test_non_final_turn_schedules_next(self) -> None:
         await strategy.handle_credit_return(credit)
         scheduler.schedule_at_perf_sec.assert_called_once()
 
+    async def test_continuation_turn_carries_has_forks(self) -> None:
+        """has_forks from the NEXT turn's metadata must ride onto the
+        continuation turn (the sticky router defers parent-entry eviction until
+        DAG children drain). Regression: from_previous_credit(credit) was called
+        without next_meta, so every continuation turn got has_forks=False."""
+        strategy, scheduler, _ = make_strategy([(0, "c1"), (100, "c1")])
+        await strategy.setup_phase()
+        # Mark the next turn (index 1) as fork-bearing.
+        meta = strategy._conversation_source._metadata_lookup["c1"]
+        meta.turns[1].has_forks = True
+
+        captured: list = []
+        strategy._credit_issuer.issue_credit = (
+            lambda turn: captured.append(turn) or True
+        )
+        credit = Credit(
+            id=1,
+            phase=CreditPhase.PROFILING,
+            conversation_id="c1",
+            x_correlation_id="corr-c1",
+            turn_index=0,
+            num_turns=2,
+            issued_at_ns=1000,
+        )
+        await strategy.handle_credit_return(credit)
+        assert captured and captured[0].has_forks is True
+
 
 @pytest.mark.asyncio
 class TestFixedScheduleTimestampConversion:

tests/unit/timing/strategies/test_user_centric_rate.py

@@ -5,10 +5,13 @@
 import pytest
 
 from aiperf.common.enums import CreditPhase
-from aiperf.plugin.enums import TimingMode
+from aiperf.common.models import ConversationMetadata, DatasetMetadata, TurnMetadata
+from aiperf.credit.structs import Credit
+from aiperf.plugin.enums import DatasetSamplingStrategy, TimingMode
 from aiperf.timing.config import CreditPhaseConfig
+from aiperf.timing.conversation_source import ConversationSource
 from aiperf.timing.strategies.user_centric_rate import User, UserCentricStrategy
-from tests.unit.timing.conftest import OrchestratorHarness
+from tests.unit.timing.conftest import OrchestratorHarness, make_sampler
 
 TWO_TURN = [("c1", 2), ("c2", 2), ("c3", 2), ("c4", 2), ("c5", 2)]
 MULTI_TURN = [("c1", 3), ("c2", 3), ("c3", 3), ("c4", 3)]
@@ -143,6 +146,75 @@ async def test_stops_at_session_count(self, create_orchestrator_harness) -> None
         assert len([c for c in h.sent_credits if c.turn_index == 0]) == 10
 
 
+def make_strategy(num_turns: int = 2) -> UserCentricStrategy:
+    """Build a UserCentricStrategy over a real ConversationSource with one
+    num_turns-turn conversation and mocked scheduler/issuer/lifecycle."""
+    scheduler = MagicMock()
+    stop_checker = MagicMock()
+    issuer = MagicMock()
+    issuer.issue_credit = lambda *a, **k: True
+    lifecycle = MagicMock()
+    ds = DatasetMetadata(
+        conversations=[
+            ConversationMetadata(
+                conversation_id="c1",
+                turns=[TurnMetadata(delay_ms=None) for _ in range(num_turns)],
+            )
+        ],
+        sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL,
+    )
+    sampler = make_sampler(["c1"], DatasetSamplingStrategy.SEQUENTIAL)
+    src = ConversationSource(ds, sampler)
+    cfg = CreditPhaseConfig(
+        phase=CreditPhase.PROFILING,
+        timing_mode=TimingMode.USER_CENTRIC_RATE,
+        request_rate=10.0,
+        num_users=1,
+        total_expected_requests=num_turns,
+    )
+    return UserCentricStrategy(
+        config=cfg,
+        conversation_source=src,
+        scheduler=scheduler,
+        stop_checker=stop_checker,
+        credit_issuer=issuer,
+        lifecycle=lifecycle,
+    )
+
+
+@pytest.mark.asyncio
+class TestUserCentricCreditReturn:
+    async def test_continuation_turn_carries_has_forks(self) -> None:
+        """has_forks from the NEXT turn's metadata must ride onto the
+        continuation turn (the sticky router defers parent-entry eviction until
+        DAG children drain). Regression: from_previous_credit(credit) was called
+        without next-turn metadata, so every continuation turn got
+        has_forks=False."""
+        strategy = make_strategy()
+        await strategy.setup_phase()
+        # Mark the next turn (index 1) as fork-bearing.
+        meta = strategy._conversation_source._metadata_lookup["c1"]
+        meta.turns[1].has_forks = True
+
+        captured: list = []
+        strategy._credit_issuer.issue_credit = (
+            lambda turn: captured.append(turn) or True
+        )
+        # setup_phase registered the t=0 replacement user; use its session id.
+        x_correlation_id = next(iter(strategy._session_to_user))
+        credit = Credit(
+            id=1,
+            phase=CreditPhase.PROFILING,
+            conversation_id="c1",
+            x_correlation_id=x_correlation_id,
+            turn_index=0,
+            num_turns=2,
+            issued_at_ns=1000,
+        )
+        await strategy.handle_credit_return(credit)
+        assert captured and captured[0].has_forks is True
+
+
 @pytest.mark.asyncio
 class TestRealisticScenarios:
     async def test_chat_benchmark(self, create_orchestrator_harness) -> None: