fix: preserve causal delays after AgentX live warmup

crates/core/src/replay/loadgen/driver.rs

@@ -1548,6 +1548,23 @@ impl WorkloadDriver {
             .agentic_snapshots
             .as_ref()
             .context("missing live snapshots")?;
+        let states_by_request: FxHashMap<_, _> = self
+            .sessions
+            .iter()
+            .enumerate()
+            .map(|(index, session)| {
+                Ok((
+                    (
+                        state.node_to_play[index],
+                        session.turns[0]
+                            .request_id
+                            .as_deref()
+                            .context("live warmup requires source request IDs")?,
+                    ),
+                    state.node_states[index],
+                ))
+            })
+            .collect::<Result<_>>()?;
         let mut offsets = Vec::new();
         let mut first = if fresh_origin { 0.0 } else { f64::INFINITY };
         for (index, node) in state.node_states.iter().enumerate() {
@@ -1559,6 +1576,34 @@ impl WorkloadDriver {
             else {
                 bail!("live warmup requires immutable snapshot prompts");
             };
+            // Only the immediate conversation frontier carries a source-clock
+            // deadline across handoff. A pending join/barrier retains its gate,
+            // but future continuations and unspawned children must create their
+            // timers from actual dependency events, including after idle caps.
+            if *node == AgenticNodeState::Blocked
+                && context.graph.nodes[source_index]
+                    .dependencies
+                    .iter()
+                    .filter(|edge| {
+                        matches!(
+                            edge.relation,
+                            AgenticDependencyRelation::Sequence | AgenticDependencyRelation::Spawn
+                        )
+                    })
+                    .any(|edge| {
+                        states_by_request
+                            .get(&(state.node_to_play[index], edge.request_id.as_str()))
+                            .is_some_and(|source| {
+                                matches!(
+                                    source,
+                                    AgenticNodeState::Blocked | AgenticNodeState::Ready
+                                ) || (edge.trigger == AgenticDependencyTrigger::Completion
+                                    && *source == AgenticNodeState::Emitted)
+                            })
+                    })
+            {
+                continue;
+            }
             let offset = (context.recorded_start_ms(source_index)
                 - snapshots[state.node_to_play[index]].t_star_ms)
                 .max(0.0)

crates/core/src/replay/loadgen/phase.rs

@@ -1145,8 +1145,8 @@ mod tests {
         );
     }
 
-    fn linear_live_driver(nodes: usize, rootless: bool) -> WorkloadDriver {
-        let rows = (0..nodes)
+    fn linear_live_rows(nodes: usize, rootless: bool) -> Vec<AgenticMooncakeRow> {
+        (0..nodes)
             .map(|index| AgenticMooncakeRow {
                 request_id: format!("n{index}"),
                 play_id: "linear".into(),
@@ -1182,7 +1182,15 @@ mod tests {
                 },
                 ..Default::default()
             })
-            .collect();
+            .collect()
+    }
+
+    fn rows_live_driver(
+        rows: Vec<AgenticMooncakeRow>,
+        lanes: usize,
+        cut: f64,
+        speedup: f64,
+    ) -> WorkloadDriver {
         let graph = ValidatedAgenticGraph::from_agentic_mooncake_rows(
             AgenticMooncakeHeader {
                 schema: AGENTIC_MOONCAKE_SCHEMA.into(),
@@ -1198,17 +1206,16 @@ mod tests {
         )
         .unwrap();
         let prepared = graph
-            .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42))
-            .unwrap();
-        let play = prepared
-            .context()
-            .prepare_play(0, 0, Some(if rootless { 0.5 } else { 0.0 }))
+            .prepare_snapshots(lanes, AgenticSnapshotOptions::seeded(42))
             .unwrap();
+        let plays = (0..lanes)
+            .map(|lane| prepared.context().prepare_play(lane, 0, Some(cut)).unwrap())
+            .collect();
         let mut driver = WorkloadDriver::new_agentic_warmup_with_policy(
-            PreparedAgenticSnapshots::from_plays(vec![play]).unwrap(),
+            PreparedAgenticSnapshots::from_plays(plays).unwrap(),
             64,
             true,
-            1.0,
+            speedup,
             AgenticWarmupPolicy::LiveTrajectoryV1,
             10_000.0,
         )
@@ -1219,6 +1226,181 @@ mod tests {
         driver
     }
 
+    fn linear_live_driver(nodes: usize, rootless: bool) -> WorkloadDriver {
+        rows_live_driver(
+            linear_live_rows(nodes, rootless),
+            1,
+            if rootless { 0.5 } else { 0.0 },
+            1.0,
+        )
+    }
+
+    fn finish_ten_live_rounds(driver: &mut WorkloadDriver, lanes: usize) -> f64 {
+        driver.finish_agentic_preparation(0.0).unwrap();
+        for round in 0..10 {
+            let now = round as f64;
+            driver.advance_agentic_profile(now).unwrap();
+            let ready = driver.pop_ready(now, lanes);
+            assert_eq!(ready.len(), lanes);
+            for turn in ready {
+                driver.on_output_token(turn.request_uuid, 0).unwrap();
+                driver
+                    .on_terminal(
+                        turn.request_uuid,
+                        now + 1.0,
+                        ReplayTerminalStatus::Completed,
+                    )
+                    .unwrap();
+            }
+        }
+        assert_eq!(
+            driver.finish_agentic_preparation(10.0).unwrap(),
+            Some(AgenticPreparationTransition::OpenProfile)
+        );
+        10.0
+    }
+
+    #[test]
+    fn live_handoff_idle_cap_does_not_restore_future_continuation_timestamps() {
+        let mut rows = linear_live_rows(13, false);
+        rows[11].not_before_ms = 601_000.0;
+        rows[11].dependencies[0].delay_ms = 600_000.0;
+        rows[12].not_before_ms = 605_001.0;
+        rows[12].dependencies[0].delay_ms = 4_000.0;
+        let mut driver = rows_live_driver(rows, 2, 0.0, 1.0);
+        let start = finish_ten_live_rounds(&mut driver, 2);
+        let frontier = driver.pop_ready(start, 2);
+        assert_eq!(frontier.len(), 2);
+        assert!(
+            frontier
+                .iter()
+                .all(|turn| turn.authored_request_id.as_deref() == Some("n10"))
+        );
+        // Keep the second lane busy so only the 300-second per-tree guard fires.
+        driver
+            .on_terminal(
+                frontier[0].request_uuid,
+                start + 1.0,
+                ReplayTerminalStatus::Completed,
+            )
+            .unwrap();
+        let capped = start + 1.0 + 300_000.0;
+        assert_eq!(driver.next_ready_time_ms(), Some(capped));
+        driver.advance_agentic_profile(capped).unwrap();
+        let delayed = driver.pop_ready(capped, 1).pop().unwrap();
+        assert_eq!(delayed.authored_request_id.as_deref(), Some("n11"));
+        driver
+            .on_terminal(
+                delayed.request_uuid,
+                capped + 1.0,
+                ReplayTerminalStatus::Completed,
+            )
+            .unwrap();
+        let continuation = capped + 1.0 + 4_000.0;
+        assert_eq!(driver.next_ready_time_ms(), Some(continuation));
+        let ready = driver.pop_ready(continuation, 1).pop().unwrap();
+        assert_eq!(ready.authored_request_id.as_deref(), Some("n12"));
+        assert_eq!(
+            driver.agentic_profile_report().unwrap().idle_shifts.len(),
+            1
+        );
+    }
+
+    #[rstest::rstest]
+    #[case::join(AgenticDependencyRelation::Join)]
+    #[case::replay_barrier(AgenticDependencyRelation::ReplayBarrier)]
+    fn live_handoff_preserves_blocked_frontier_but_not_later_turns(
+        #[case] relation: AgenticDependencyRelation,
+    ) {
+        let mut rows = linear_live_rows(12, false);
+        rows[10].not_before_ms = 5_000.0;
+        rows[10].recorded_api_time_ms = Some(10_000.0);
+        rows[10].dependencies.push(AgenticDependency {
+            request_id: "child".into(),
+            delay_ms: 0.0,
+            relation,
+            trigger: AgenticDependencyTrigger::Completion,
+        });
+        rows[11].not_before_ms = 19_000.0;
+        rows[11].dependencies[0].delay_ms = 4_000.0;
+        let mut child = rows[0].clone();
+        child.request_id = "child".into();
+        child.session_id = "child".into();
+        child.not_before_ms = 1_000.0;
+        child.dependencies = vec![AgenticDependency {
+            request_id: "n9".into(),
+            delay_ms: 100.0,
+            relation: AgenticDependencyRelation::Spawn,
+            trigger: AgenticDependencyTrigger::Dispatch,
+        }];
+        rows.push(child.clone());
+        child.request_id = "future-child".into();
+        child.session_id = "future-child".into();
+        child.not_before_ms = 21_000.0;
+        child.dependencies[0].request_id = "n11".into();
+        child.dependencies[0].delay_ms = 2_000.0;
+        rows.push(child);
+
+        let mut driver = rows_live_driver(rows, 1, 0.0, 1.0);
+        let start = finish_ten_live_rounds(&mut driver, 1);
+        let ready = driver.pop_ready(start, 1).pop().unwrap();
+        assert_eq!(ready.authored_request_id.as_deref(), Some("child"));
+        driver
+            .on_terminal(
+                ready.request_uuid,
+                start + 1.0,
+                ReplayTerminalStatus::Completed,
+            )
+            .unwrap();
+        // The parent was already waiting at a join/barrier when warmup ended:
+        // keep its recorded deadline even if the remaining child finishes early.
+        let parent_at = start + 4_000.0;
+        assert_eq!(driver.next_ready_time_ms(), Some(parent_at));
+        let parent = driver.pop_ready(parent_at, 1).pop().unwrap();
+        assert_eq!(parent.authored_request_id.as_deref(), Some("n10"));
+        driver
+            .on_terminal(
+                parent.request_uuid,
+                parent_at + 1.0,
+                ReplayTerminalStatus::Completed,
+            )
+            .unwrap();
+        let next_at = parent_at + 1.0 + 4_000.0;
+        assert_eq!(driver.next_ready_time_ms(), Some(next_at));
+        let next = driver.pop_ready(next_at, 1).pop().unwrap();
+        assert_eq!(next.authored_request_id.as_deref(), Some("n11"));
+        // A child whose spawning turn was still blocked at handoff is not a
+        // pending handoff stream. Its timer starts only on actual dispatch.
+        assert_eq!(driver.next_ready_time_ms(), Some(next_at + 2_000.0));
+        let spawned = driver.pop_ready(next_at + 2_000.0, 1).pop().unwrap();
+        assert_eq!(spawned.authored_request_id.as_deref(), Some("future-child"));
+    }
+
+    #[test]
+    fn live_handoff_fresh_root_preserves_other_lanes_scaled_frontier_offset() {
+        let mut rows = linear_live_rows(10, false);
+        let mut other = linear_live_rows(12, false);
+        for row in &mut other {
+            row.play_id = "other".into();
+            row.request_id = format!("other-{}", row.request_id);
+            for edge in &mut row.dependencies {
+                edge.request_id = format!("other-{}", edge.request_id);
+            }
+        }
+        rows.extend(other);
+        let mut driver = rows_live_driver(rows, 2, 0.0, 2.0);
+        let start = finish_ten_live_rounds(&mut driver, 2);
+        // One lane finished exactly at its quota. Its new incarnation anchors
+        // the phase at zero, without also bursting the other lane's frontier.
+        let root = driver.pop_ready(start, 2);
+        assert_eq!(root.len(), 1);
+        assert_eq!(root[0].authored_request_id.as_deref(), Some("n0"));
+        assert_eq!(driver.next_ready_time_ms(), Some(start + 5.0));
+        let other = driver.pop_ready(start + 5.0, 2);
+        assert_eq!(other.len(), 1);
+        assert_eq!(other[0].authored_request_id.as_deref(), Some("other-n10"));
+    }
+
     #[test]
     fn zero_time_live_warmup_is_bounded_and_failed_reset_is_transactional() {
         let mut driver = linear_live_driver(1, false);