aisimulate_core/replay/loadgen/

trace.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};

use anyhow::{Context, Result, anyhow, bail, ensure};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use super::driver::WorkloadDriver;
use super::types::{
    AGENTIC_MOONCAKE_SCHEMA, AGENTIC_MOONCAKE_VERSION, AgenticDependency,
    AgenticDependencyRelation, AgenticDependencyTrigger, AgenticHashIdScope, AgenticMooncakeHeader,
    AgenticMooncakeRow, AgenticNode, AgenticPlay, AgenticSourceProvenance, AgenticTrace, DelaySpec,
    LengthSpec, MooncakeRow, ReplayRequestHashes, SessionPartitionSpec, SessionTrace,
    SyntheticTraceSpec, Trace, TraceFileFormat, TurnTrace, effective_replay_key,
};
use super::{SYNTHETIC_OUTPUT_SEED, planned_output_token_ids};
use crate::replay::protocol::DirectRequest;

#[derive(Debug, Deserialize)]
struct RawAppliedComputeAgenticRecord {
    num_turns: usize,
    input_prompt_length: usize,
    assistant_response_length: Vec<usize>,
    tool_call_output_length: Vec<usize>,
    tool_call_latency: Vec<f64>,
    final_assistant_response_length: usize,
}

#[derive(Debug, Deserialize)]
struct LegacyAgenticMooncakeRow {
    request_id: String,
    #[serde(default)]
    session_id: Option<String>,
    #[serde(default, alias = "input_tokens")]
    input_length: Option<usize>,
    #[serde(alias = "output_tokens")]
    output_length: usize,
    #[serde(default)]
    output_token_ids: Option<Vec<u32>>,
    hash_ids: Vec<u64>,
    #[serde(default, alias = "created_time")]
    timestamp: Option<f64>,
    #[serde(default)]
    delay: Option<f64>,
    #[serde(default)]
    delay_ms: Option<f64>,
    #[serde(default)]
    tool_wait_ms: f64,
    #[serde(default)]
    wait_for: Vec<String>,
    #[serde(default)]
    priority: Option<i32>,
    #[serde(default)]
    strict_priority: Option<u32>,
    #[serde(default)]
    policy_class: Option<String>,
}

/// Load either the versioned agentic-Mooncake graph or the legacy flat rows.
///
/// The versioned header owns its block size. ``legacy_trace_block_size`` is
/// consulted only when lowering headerless legacy rows.
pub fn load_agentic_mooncake(path: &Path, legacy_trace_block_size: usize) -> Result<AgenticTrace> {
    let file = File::open(path)
        .with_context(|| format!("failed to open trace file {}", path.display()))?;
    let mut lines = BufReader::new(file).lines();
    let first = loop {
        let line = lines
            .next()
            .transpose()
            .with_context(|| format!("failed to read trace file {}", path.display()))?
            .context("agentic trace file is empty")?;
        if !line.trim().is_empty() {
            break line;
        }
    };
    let first_json: serde_json::Value =
        serde_json::from_str(&first).context("failed to parse first agentic trace row")?;
    if first_json.get("schema").and_then(serde_json::Value::as_str) == Some(AGENTIC_MOONCAKE_SCHEMA)
    {
        return AgenticTrace::from_agentic_mooncake(path);
    }

    let mut raw_rows = vec![
        serde_json::from_value::<LegacyAgenticMooncakeRow>(first_json)
            .context("failed to parse first legacy agentic Mooncake row")?,
    ];
    for (line_index, line) in lines.enumerate() {
        let line = line.with_context(|| {
            format!(
                "failed to read legacy agentic trace line {}",
                line_index + 2
            )
        })?;
        if line.trim().is_empty() {
            continue;
        }
        raw_rows.push(serde_json::from_str(&line).with_context(|| {
            format!(
                "failed to parse legacy agentic trace line {}",
                line_index + 2
            )
        })?);
    }
    let mut rows = raw_rows
        .into_iter()
        .map(|raw| -> Result<AgenticMooncakeRow> {
            if raw.request_id.trim().is_empty() {
                bail!("request_id must be nonempty");
            }
            if raw.hash_ids.is_empty() {
                bail!("hash_ids must be nonempty");
            }
            if !raw.tool_wait_ms.is_finite() || raw.tool_wait_ms < 0.0 {
                bail!("tool_wait_ms must be finite and nonnegative");
            }
            if raw.delay.is_some() && raw.delay_ms.is_some() {
                bail!("delay and delay_ms cannot both be set");
            }
            let delay = raw.delay.or(raw.delay_ms).unwrap_or(0.0) + raw.tool_wait_ms;
            if !delay.is_finite() || delay < 0.0 {
                bail!("dependency delay must be finite and nonnegative");
            }
            let relation = if raw.wait_for.len() > 1 {
                AgenticDependencyRelation::Join
            } else {
                AgenticDependencyRelation::Sequence
            };
            let dependencies = raw
                .wait_for
                .into_iter()
                .map(|request_id| AgenticDependency {
                    request_id,
                    trigger: AgenticDependencyTrigger::Completion,
                    delay_ms: delay,
                    relation,
                })
                .collect::<Vec<_>>();
            Ok(AgenticMooncakeRow {
                request_id: raw.request_id,
                play_id: "agentic-play".to_string(),
                session_id: raw
                    .session_id
                    .unwrap_or_else(|| "agentic-session".to_string()),
                model: "unknown".to_string(),
                input_length: raw.input_length,
                output_length: Some(raw.output_length),
                output_token_ids: raw.output_token_ids,
                hash_ids: Some(raw.hash_ids),
                not_before_ms: if dependencies.is_empty() {
                    raw.timestamp.unwrap_or(0.0)
                } else {
                    0.0
                },
                priority: raw.priority,
                strict_priority: raw.strict_priority,
                policy_class: raw.policy_class,
                dependencies,
            })
        })
        .collect::<Result<Vec<_>>>()?;
    assign_dependency_component_play_ids(&mut rows, "legacy-play");
    AgenticTrace::from_agentic_mooncake_rows(
        AgenticMooncakeHeader {
            schema: AGENTIC_MOONCAKE_SCHEMA.to_string(),
            version: AGENTIC_MOONCAKE_VERSION,
            block_size: legacy_trace_block_size,
            hash_id_scope: AgenticHashIdScope::Local,
            source: AgenticSourceProvenance {
                format: "legacy_agentic_mooncake".to_string(),
                digest: format!("{}:{}", path.display(), rows.len()),
            },
        },
        rows,
    )
}

pub(super) fn assign_dependency_component_play_ids(rows: &mut [AgenticMooncakeRow], prefix: &str) {
    fn find(parent: &mut [usize], value: usize) -> usize {
        if parent[value] != value {
            parent[value] = find(parent, parent[value]);
        }
        parent[value]
    }

    let by_id = rows
        .iter()
        .enumerate()
        .map(|(index, row)| (row.request_id.clone(), index))
        .collect::<HashMap<_, _>>();
    let mut parent = (0..rows.len()).collect::<Vec<_>>();
    for (index, row) in rows.iter().enumerate() {
        for dependency in &row.dependencies {
            let Some(&source) = by_id.get(&dependency.request_id) else {
                continue;
            };
            let left = find(&mut parent, index);
            let right = find(&mut parent, source);
            if left != right {
                parent[left] = right;
            }
        }
    }
    let mut roots_by_component: HashMap<usize, Vec<usize>> = HashMap::new();
    for (index, row) in rows.iter().enumerate() {
        if row.dependencies.is_empty() {
            let component = find(&mut parent, index);
            roots_by_component.entry(component).or_default().push(index);
        }
    }
    let mut labels = HashMap::new();
    for (component, roots) in roots_by_component {
        let canonical = roots
            .iter()
            .copied()
            .min_by(|left, right| {
                rows[*left]
                    .not_before_ms
                    .total_cmp(&rows[*right].not_before_ms)
                    .then_with(|| rows[*left].request_id.cmp(&rows[*right].request_id))
            })
            .expect("a root component is nonempty");
        labels.insert(component, rows[canonical].request_id.clone());
    }
    for (index, row) in rows.iter_mut().enumerate() {
        let component = find(&mut parent, index);
        let label = labels
            .entry(component)
            .or_insert_with(|| row.request_id.clone());
        row.play_id = format!("{prefix}:{label}");
    }
}

#[derive(Debug, Default)]
struct HashIdInterner {
    canonical_ids: FxHashMap<u64, u32>,
}

impl HashIdInterner {
    fn intern_all(&mut self, hash_ids: Vec<u64>) -> Result<Vec<u32>> {
        hash_ids
            .into_iter()
            .map(|hash_id| self.intern(hash_id))
            .collect()
    }

    fn intern(&mut self, hash_id: u64) -> Result<u32> {
        let next_id = self.canonical_ids.len();
        match self.canonical_ids.entry(hash_id) {
            std::collections::hash_map::Entry::Occupied(entry) => Ok(*entry.get()),
            std::collections::hash_map::Entry::Vacant(entry) => {
                let canonical_id = u32::try_from(next_id)
                    .context("trace contains more unique hash IDs than u32 can represent")?;
                entry.insert(canonical_id);
                Ok(canonical_id)
            }
        }
    }
}

pub fn validate_trace_files(format: TraceFileFormat, paths: &[PathBuf]) -> Result<()> {
    if paths.is_empty() {
        bail!("trace replay requires at least one trace file");
    }
    if format != TraceFileFormat::Dynamo && paths.len() != 1 {
        bail!(
            "trace_format='{}' requires exactly one trace file, got {}",
            format.as_str(),
            paths.len()
        );
    }
    Ok(())
}

fn single_turn_request_uuid(_request_ordinal: usize) -> Uuid {
    Uuid::new_v4()
}

pub(super) fn validate_synthesizable_prompt(
    input_length: usize,
    hash_ids: &[u32],
    trace_block_size: usize,
) -> Result<()> {
    if trace_block_size == 0 {
        bail!("trace_block_size must be greater than 0");
    }
    let synthesizable_capacity = hash_ids
        .len()
        .checked_mul(trace_block_size)
        .context("synthesized prompt capacity overflow")?;
    let required_hash_ids = input_length.div_ceil(trace_block_size);
    if hash_ids.len() < required_hash_ids {
        bail!(
            "input_length {} exceeds synthesized capacity {}",
            input_length,
            synthesizable_capacity
        );
    }

    Ok(())
}

pub(super) fn synthesize_trace_tokens(
    input_length: usize,
    hash_ids: &[u32],
    trace_block_size: usize,
) -> Result<Vec<u32>> {
    validate_synthesizable_prompt(input_length, hash_ids, trace_block_size)?;

    Ok(synthesize_validated_trace_tokens(
        input_length,
        hash_ids,
        trace_block_size,
    ))
}

pub(super) fn synthesize_validated_trace_tokens(
    input_length: usize,
    hash_ids: &[u32],
    trace_block_size: usize,
) -> Vec<u32> {
    let mut tokens = Vec::with_capacity(input_length);
    for &hash_id in hash_ids {
        let remaining = input_length - tokens.len();
        tokens.extend(std::iter::repeat_n(
            hash_id,
            remaining.min(trace_block_size),
        ));
        if tokens.len() == input_length {
            break;
        }
    }

    debug_assert_eq!(tokens.len(), input_length);
    tokens
}

fn trace_to_replay_hashes(
    input_length: usize,
    hash_ids: &[u32],
    trace_block_size: usize,
    engine_block_size: usize,
) -> Result<ReplayRequestHashes> {
    if engine_block_size == 0 {
        bail!("engine_block_size must be greater than 0");
    }

    let tokens = synthesize_trace_tokens(input_length, hash_ids, trace_block_size)?;
    let engine_block_size =
        u32::try_from(engine_block_size).context("engine_block_size does not fit in u32")?;
    Ok(ReplayRequestHashes::from_tokens(&tokens, engine_block_size))
}

impl TurnTrace {
    pub fn synthesize_tokens(&self, trace_block_size: usize) -> Result<Vec<u32>> {
        synthesize_trace_tokens(self.input_length, &self.hash_ids, trace_block_size)
    }

    pub fn to_direct_request(
        &self,
        trace_block_size: usize,
        request_uuid: Uuid,
        arrival_timestamp_ms: Option<f64>,
    ) -> Result<DirectRequest> {
        let tokens = self.synthesize_tokens(trace_block_size)?;
        Ok(DirectRequest {
            tokens,
            max_output_tokens: self.max_output_tokens,
            output_token_ids: self.output_token_ids.clone(),
            uuid: Some(request_uuid),
            dp_rank: 0,
            preferred_dp_rank: None,
            arrival_timestamp_ms,
            priority: self.priority,
            strict_priority: self.strict_priority,
            policy_class: self.policy_class.clone(),
            replay_context: None,
        })
    }

    pub fn to_replay_hashes(
        &self,
        trace_block_size: usize,
        engine_block_size: usize,
    ) -> Result<ReplayRequestHashes> {
        trace_to_replay_hashes(
            self.input_length,
            &self.hash_ids,
            trace_block_size,
            engine_block_size,
        )
    }
}

struct MooncakeTraceBuilder {
    trace_block_size: usize,
    hash_id_interner: HashIdInterner,
    sessions: Vec<SessionTrace>,
    session_indices: HashMap<String, usize>,
    last_timestamps: Vec<Option<f64>>,
}

impl MooncakeTraceBuilder {
    fn new(trace_block_size: usize) -> Self {
        Self {
            trace_block_size,
            hash_id_interner: HashIdInterner::default(),
            sessions: Vec::new(),
            session_indices: HashMap::new(),
            last_timestamps: Vec::new(),
        }
    }

    fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }

    fn push(&mut self, line_idx: usize, raw: MooncakeRow) -> Result<()> {
        let request_id = raw.request_id;
        let raw_session_id = raw.session_id;
        let session_id = raw_session_id
            .clone()
            .unwrap_or_else(|| format!("request_{}", line_idx + 1));
        let hash_ids = raw
            .hash_ids
            .ok_or_else(|| anyhow!("trace line {} is missing hash_ids", line_idx + 1))?;
        let synthesizable_capacity = hash_ids
            .len()
            .checked_mul(self.trace_block_size)
            .ok_or_else(|| anyhow!("trace line {} synthesized capacity overflow", line_idx + 1))?;
        let input_length = raw.input_length.unwrap_or(synthesizable_capacity);
        ensure!(
            input_length <= synthesizable_capacity,
            "trace line {} input_length {} exceeds hash_ids capacity {}",
            line_idx + 1,
            input_length,
            synthesizable_capacity
        );
        let output_length = raw
            .output_length
            .ok_or_else(|| anyhow!("trace line {} is missing output_length", line_idx + 1))?;
        let output_token_ids = raw.output_token_ids;
        if let Some(output_token_ids) = output_token_ids.as_ref()
            && output_token_ids.len() != output_length
        {
            bail!(
                "trace line {} output_length {} does not match output_token_ids length {}",
                line_idx + 1,
                output_length,
                output_token_ids.len()
            );
        }
        let timestamp_ms = raw.timestamp;
        let explicit_delay_ms = raw.delay;
        let priority = raw.priority.unwrap_or(0);
        let strict_priority = raw.strict_priority.unwrap_or(0);
        let policy_class = raw.policy_class.clone();

        let session_index = *self
            .session_indices
            .entry(session_id.clone())
            .or_insert_with(|| {
                let idx = self.sessions.len();
                self.sessions.push(SessionTrace {
                    session_id: session_id.clone(),
                    first_arrival_timestamp_ms: timestamp_ms,
                    turns: Vec::new(),
                });
                self.last_timestamps.push(timestamp_ms);
                idx
            });

        let session = self
            .sessions
            .get_mut(session_index)
            .expect("newly inserted session must exist");
        let turn_idx = session.turns.len();
        let replay_key = output_token_ids.as_ref().map(|_| {
            effective_replay_key(
                request_id.as_deref(),
                raw_session_id.as_deref(),
                turn_idx,
                line_idx,
            )
        });
        let delay_after_previous_ms = if turn_idx == 0 {
            let delay = explicit_delay_ms.unwrap_or(0.0);
            if delay != 0.0 {
                bail!(
                    "trace line {} sets delay on the first turn of session {}",
                    line_idx + 1,
                    session.session_id
                );
            }
            0.0
        } else if let Some(delay_ms) = explicit_delay_ms {
            delay_ms
        } else if let Some(timestamp_ms) = timestamp_ms {
            let previous_timestamp_ms = self.last_timestamps[session_index].ok_or_else(|| {
                anyhow!(
                    "trace line {} for session {} cannot infer delay without a previous timestamp",
                    line_idx + 1,
                    session.session_id
                )
            })?;
            timestamp_ms - previous_timestamp_ms
        } else {
            0.0
        };

        if !delay_after_previous_ms.is_finite() || delay_after_previous_ms < 0.0 {
            bail!(
                "trace line {} has invalid delay {}",
                line_idx + 1,
                delay_after_previous_ms
            );
        }

        if hash_ids.len() * self.trace_block_size < input_length {
            bail!(
                "trace line {} input_length {} exceeds synthesized capacity {}",
                line_idx + 1,
                input_length,
                hash_ids.len() * self.trace_block_size
            );
        }

        let hash_ids = self.hash_id_interner.intern_all(hash_ids)?;

        session.turns.push(TurnTrace {
            input_length,
            max_output_tokens: output_length,
            output_token_ids,
            replay_key,
            hash_ids,
            delay_after_previous_ms,
            priority,
            strict_priority,
            policy_class,
        });
        if let Some(timestamp_ms) = timestamp_ms {
            self.last_timestamps[session_index] = Some(timestamp_ms);
        }
        Ok(())
    }

    fn finish(self) -> Trace {
        Trace {
            block_size: self.trace_block_size,
            sessions: self.sessions,
        }
    }
}

impl Trace {
    pub fn from_mooncake(path: &Path, trace_block_size: usize) -> Result<Self> {
        if trace_block_size == 0 {
            bail!("trace_block_size must be greater than 0");
        }

        let file = File::open(path)
            .with_context(|| format!("failed to open trace file {}", path.display()))?;
        let reader = BufReader::new(file);
        let mut builder = MooncakeTraceBuilder::new(trace_block_size);

        for (line_idx, line) in reader.lines().enumerate() {
            let line = line.with_context(|| {
                format!(
                    "failed to read line {} from {}",
                    line_idx + 1,
                    path.display()
                )
            })?;
            if line.trim().is_empty() {
                continue;
            }

            let row = serde_json::from_str(&line).with_context(|| {
                format!(
                    "failed to parse line {} from {} as JSON",
                    line_idx + 1,
                    path.display()
                )
            })?;
            builder.push(line_idx, row)?;
        }

        if builder.is_empty() {
            bail!("trace file {} did not contain any requests", path.display());
        }

        Ok(builder.finish())
    }

    pub fn from_mooncake_rows(rows: Vec<MooncakeRow>, trace_block_size: usize) -> Result<Self> {
        if trace_block_size == 0 {
            bail!("trace_block_size must be greater than 0");
        }
        let mut builder = MooncakeTraceBuilder::new(trace_block_size);
        for (line_idx, row) in rows.into_iter().enumerate() {
            builder.push(line_idx, row)?;
        }
        if builder.is_empty() {
            bail!("Mooncake rows did not contain any requests");
        }
        Ok(builder.finish())
    }

    pub fn from_applied_compute_agentic(
        path: &Path,
        trace_block_size: usize,
        shared_prefix_ratio: f64,
        num_prefix_groups: usize,
    ) -> Result<Self> {
        if trace_block_size == 0 {
            bail!("trace_block_size must be greater than 0");
        }
        if !(0.0..=1.0).contains(&shared_prefix_ratio) {
            bail!(
                "shared_prefix_ratio must be between 0.0 and 1.0, got {}",
                shared_prefix_ratio
            );
        }

        let file = File::open(path)
            .with_context(|| format!("failed to open trace file {}", path.display()))?;
        let reader = BufReader::new(file);
        let mut sessions = Vec::new();
        let mut hash_id_interner = HashIdInterner::default();
        let mut next_unique_hash = 1_u64;

        for (line_idx, line) in reader.lines().enumerate() {
            let line = line.with_context(|| {
                format!(
                    "failed to read line {} from {}",
                    line_idx + 1,
                    path.display()
                )
            })?;
            if line.trim().is_empty() {
                continue;
            }

            let raw: RawAppliedComputeAgenticRecord =
                serde_json::from_str(&line).with_context(|| {
                    format!(
                        "failed to parse line {} from {} as JSON",
                        line_idx + 1,
                        path.display()
                    )
                })?;

            for (name, values) in [
                (
                    "assistant_response_length",
                    raw.assistant_response_length.len(),
                ),
                ("tool_call_output_length", raw.tool_call_output_length.len()),
                ("tool_call_latency", raw.tool_call_latency.len()),
            ] {
                if values != raw.num_turns {
                    bail!(
                        "trace line {} field {} length {} does not match num_turns {}",
                        line_idx + 1,
                        name,
                        values,
                        raw.num_turns
                    );
                }
            }

            if raw.input_prompt_length == 0 {
                bail!(
                    "trace line {} input_prompt_length must be positive",
                    line_idx + 1
                );
            }

            let group_id = if shared_prefix_ratio > 0.0 && num_prefix_groups > 0 {
                Some(line_idx % num_prefix_groups)
            } else {
                None
            };
            let mut current_input_length = raw.input_prompt_length;
            let mut hash_ids = Vec::new();
            let shared_initial_blocks = ((current_input_length.div_ceil(trace_block_size) as f64)
                * shared_prefix_ratio)
                .round() as usize;
            extend_applied_compute_agentic_hash_ids(
                &mut hash_ids,
                current_input_length,
                trace_block_size,
                shared_initial_blocks,
                group_id,
                &mut next_unique_hash,
            )?;

            let mut turns = Vec::with_capacity(raw.num_turns + 1);
            let mut next_turn_delay_ms = 0.0;
            for turn_idx in 0..raw.num_turns {
                let tool_call_latency = raw.tool_call_latency[turn_idx];
                if !tool_call_latency.is_finite() || tool_call_latency < 0.0 {
                    bail!(
                        "trace line {} tool_call_latency[{}] must be a finite non-negative number",
                        line_idx + 1,
                        turn_idx
                    );
                }

                turns.push(TurnTrace {
                    input_length: current_input_length,
                    max_output_tokens: raw.assistant_response_length[turn_idx],
                    hash_ids: hash_id_interner.intern_all(hash_ids.clone())?,
                    delay_after_previous_ms: next_turn_delay_ms,
                    ..Default::default()
                });

                current_input_length = current_input_length
                    .checked_add(raw.assistant_response_length[turn_idx])
                    .and_then(|value| value.checked_add(raw.tool_call_output_length[turn_idx]))
                    .ok_or_else(|| {
                        anyhow!(
                            "trace line {} cumulative input length overflow",
                            line_idx + 1
                        )
                    })?;
                extend_applied_compute_agentic_hash_ids(
                    &mut hash_ids,
                    current_input_length,
                    trace_block_size,
                    shared_initial_blocks,
                    group_id,
                    &mut next_unique_hash,
                )?;
                next_turn_delay_ms = tool_call_latency * 1000.0;
            }

            turns.push(TurnTrace {
                input_length: current_input_length,
                max_output_tokens: raw.final_assistant_response_length,
                hash_ids: hash_id_interner.intern_all(hash_ids)?,
                delay_after_previous_ms: next_turn_delay_ms,
                ..Default::default()
            });

            sessions.push(SessionTrace {
                session_id: format!("applied_compute_agentic_session_{}", line_idx + 1),
                first_arrival_timestamp_ms: None,
                turns,
            });
        }

        if sessions.is_empty() {
            bail!("trace file {} did not contain any requests", path.display());
        }

        Ok(Self {
            block_size: trace_block_size,
            sessions,
        })
    }

    pub fn synthetic(spec: SyntheticTraceSpec) -> Result<Self> {
        if spec.block_size == 0 {
            bail!("block_size must be greater than 0");
        }
        if spec.num_sessions == 0 {
            bail!("num_sessions must be greater than 0");
        }
        if spec.turns_per_session == 0 {
            bail!("turns_per_session must be greater than 0");
        }
        if !(0.0..=1.0).contains(&spec.shared_prefix_ratio) {
            bail!(
                "shared_prefix_ratio must be between 0.0 and 1.0, got {}",
                spec.shared_prefix_ratio
            );
        }

        let mut rng = StdRng::seed_from_u64(spec.seed);
        let mut sessions = Vec::with_capacity(spec.num_sessions);
        let first_arrivals = spec
            .first_turn_arrivals
            .timestamps(spec.num_sessions, spec.arrival_seed)?;

        let mut next_unique_hash = 1_u64;
        let mut hash_id_interner = HashIdInterner::default();
        for (session_idx, first_arrival_timestamp_ms) in first_arrivals.into_iter().enumerate() {
            let group_id = if spec.num_prefix_groups > 0 && spec.shared_prefix_ratio > 0.0 {
                Some(rng.random_range(0..spec.num_prefix_groups) as u64)
            } else {
                None
            };
            let mut turns = Vec::with_capacity(spec.turns_per_session);
            for turn_idx in 0..spec.turns_per_session {
                let input_length = sample_length(&spec.input_tokens, 1, &mut rng);
                let max_output_tokens = sample_length(&spec.output_tokens, 1, &mut rng);
                let num_blocks = input_length.div_ceil(spec.block_size);
                let prefix_blocks =
                    ((num_blocks as f64) * spec.shared_prefix_ratio).round() as usize;
                let prefix_blocks = prefix_blocks.min(num_blocks);
                let mut hash_ids = Vec::with_capacity(num_blocks);

                for block_idx in 0..prefix_blocks {
                    if let Some(group_id) = group_id {
                        hash_ids.push(0xD00D_0000_0000_0000 | (group_id << 32) | block_idx as u64);
                    }
                }

                while hash_ids.len() < num_blocks {
                    hash_ids.push(next_unique_hash);
                    next_unique_hash = next_unique_hash
                        .checked_add(1)
                        .expect("synthetic hash id overflow");
                }

                turns.push(TurnTrace {
                    input_length,
                    max_output_tokens,
                    hash_ids: hash_id_interner.intern_all(hash_ids)?,
                    delay_after_previous_ms: if turn_idx == 0 {
                        0.0
                    } else {
                        sample_delay_ms(&spec.inter_turn_delays, &mut rng)?
                    },
                    ..Default::default()
                });
            }

            sessions.push(SessionTrace {
                session_id: format!("session_{session_idx}"),
                first_arrival_timestamp_ms: Some(first_arrival_timestamp_ms),
                turns,
            });
        }

        Ok(Self {
            block_size: spec.block_size,
            sessions,
        })
    }

    pub fn validate_for_trace_mode(&self) -> Result<()> {
        self.validate(false)
    }

    pub fn validate_for_concurrency_mode(&self) -> Result<()> {
        self.validate(true)
    }

    pub fn normalize_session_starts(mut self) -> Result<Self> {
        let Some(min_timestamp_ms) = self
            .sessions
            .iter()
            .filter_map(|session| session.first_arrival_timestamp_ms)
            .min_by(|left, right| left.total_cmp(right))
        else {
            return Ok(self);
        };

        for session in &mut self.sessions {
            if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
                *timestamp_ms -= min_timestamp_ms;
            }
        }
        Ok(self)
    }

    pub fn speed_up_timing(mut self, ratio: f64) -> Result<Self> {
        if !ratio.is_finite() || ratio <= 0.0 {
            bail!("ratio must be a finite positive number, got {ratio}");
        }

        for session in &mut self.sessions {
            if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
                *timestamp_ms /= ratio;
            }
            for turn in &mut session.turns {
                turn.delay_after_previous_ms /= ratio;
            }
        }
        Ok(self)
    }

    pub fn rescale_session_start_span(mut self, duration_ms: u64) -> Result<Self> {
        let Some(min_timestamp_ms) = self
            .sessions
            .iter()
            .filter_map(|session| session.first_arrival_timestamp_ms)
            .min_by(|left, right| left.total_cmp(right))
        else {
            return Ok(self);
        };
        let Some(max_timestamp_ms) = self
            .sessions
            .iter()
            .filter_map(|session| session.first_arrival_timestamp_ms)
            .max_by(|left, right| left.total_cmp(right))
        else {
            return Ok(self);
        };

        let target_span_ms = duration_ms as f64;
        let source_span_ms = max_timestamp_ms - min_timestamp_ms;
        for session in &mut self.sessions {
            if let Some(timestamp_ms) = session.first_arrival_timestamp_ms.as_mut() {
                *timestamp_ms = if source_span_ms == 0.0 {
                    0.0
                } else {
                    (*timestamp_ms - min_timestamp_ms) * target_span_ms / source_span_ms
                };
            }
        }
        Ok(self)
    }

    pub fn rescale_ready_span(mut self, duration_ms: u64) -> Result<Self> {
        let Some(min_start_ms) = self
            .sessions
            .iter()
            .map(|session| session.first_arrival_timestamp_ms.unwrap_or(0.0))
            .min_by(|left, right| left.total_cmp(right))
        else {
            return Ok(self);
        };

        let Some(max_ready_ms) = self
            .sessions
            .iter()
            .map(|session| {
                session.first_arrival_timestamp_ms.unwrap_or(0.0)
                    + session
                        .turns
                        .iter()
                        .enumerate()
                        .filter(|(turn_idx, _)| *turn_idx > 0)
                        .map(|(_, turn)| turn.delay_after_previous_ms)
                        .sum::<f64>()
            })
            .max_by(|left, right| left.total_cmp(right))
        else {
            return Ok(self);
        };

        let ratio = duration_ms as f64 / (max_ready_ms - min_start_ms).max(1.0);
        for session in &mut self.sessions {
            if let Some(start_ms) = session.first_arrival_timestamp_ms.as_mut() {
                *start_ms = (*start_ms - min_start_ms) * ratio;
            }
            for (turn_idx, turn) in session.turns.iter_mut().enumerate() {
                if turn_idx > 0 {
                    turn.delay_after_previous_ms *= ratio;
                }
            }
        }
        Ok(self)
    }

    pub fn expand_hash_prefix_depth(mut self, factor: usize) -> Self {
        if factor <= 1 {
            return self;
        }
        let factor = u32::try_from(factor).expect("hash prefix expansion factor exceeds u32");
        for session in &mut self.sessions {
            for turn in &mut session.turns {
                turn.input_length = turn
                    .input_length
                    .checked_mul(factor as usize)
                    .expect("input_length expansion overflow");
                turn.hash_ids = turn
                    .hash_ids
                    .iter()
                    .flat_map(|&hash_id| {
                        let base = hash_id
                            .checked_mul(factor)
                            .expect("hash prefix expansion overflow");
                        (0..factor).map(move |offset| {
                            base.checked_add(offset)
                                .expect("hash prefix expansion overflow")
                        })
                    })
                    .collect();
            }
        }
        self
    }

    pub fn duplicate_hash_space(mut self, copies: usize) -> Self {
        if copies <= 1 {
            return self;
        }

        let max_hash_id = self
            .sessions
            .iter()
            .flat_map(|session| session.turns.iter())
            .flat_map(|turn| turn.hash_ids.iter().copied())
            .max()
            .unwrap_or(0);
        let offset_base = max_hash_id
            .checked_add(1)
            .expect("hash duplication offset overflow");
        let original_sessions = self.sessions.clone();
        self.sessions.clear();

        for copy_idx in 0..copies {
            let copy_idx = u32::try_from(copy_idx).expect("hash copy index exceeds u32");
            let offset = offset_base
                .checked_mul(copy_idx)
                .expect("hash duplication offset overflow");
            for session in &original_sessions {
                let mut duplicated = session.clone();
                duplicated.session_id = format!("{}:copy_{copy_idx}", session.session_id);
                for turn in &mut duplicated.turns {
                    turn.hash_ids = turn
                        .hash_ids
                        .iter()
                        .map(|&hash_id| {
                            hash_id
                                .checked_add(offset)
                                .expect("hash duplication overflow")
                        })
                        .collect();
                }
                self.sessions.push(duplicated);
            }
        }
        self
    }

    pub fn partition_by_session(&self, spec: SessionPartitionSpec) -> Vec<Self> {
        let num_partitions = match spec {
            SessionPartitionSpec::Random { num_partitions, .. } => num_partitions,
            SessionPartitionSpec::RoundRobin { num_partitions } => num_partitions,
        }
        .max(1);
        let mut partitions = vec![
            Self {
                block_size: self.block_size,
                sessions: Vec::new(),
            };
            num_partitions
        ];

        let mut rng = match spec {
            SessionPartitionSpec::Random { seed, .. } => Some(StdRng::seed_from_u64(seed)),
            SessionPartitionSpec::RoundRobin { .. } => None,
        };

        for (session_idx, session) in self.sessions.iter().cloned().enumerate() {
            let partition_idx = match spec {
                SessionPartitionSpec::Random { .. } => rng
                    .as_mut()
                    .expect("random partitioner must exist")
                    .random_range(0..num_partitions),
                SessionPartitionSpec::RoundRobin { .. } => session_idx % num_partitions,
            };
            partitions[partition_idx].sessions.push(session);
        }

        partitions
    }

    pub fn to_single_turn_requests(&self) -> Result<Vec<DirectRequest>> {
        let mut requests = Vec::with_capacity(self.sessions.len());
        let mut output_rng = StdRng::seed_from_u64(SYNTHETIC_OUTPUT_SEED);
        for (request_ordinal, session) in self.sessions.iter().enumerate() {
            if session.turns.len() != 1 {
                bail!(
                    "to_single_turn_requests requires exactly one turn per session, but session {} has {} turns",
                    session.session_id,
                    session.turns.len()
                );
            }
            let request_uuid = single_turn_request_uuid(request_ordinal);
            let mut request = session.turns[0].to_direct_request(
                self.block_size,
                request_uuid,
                session.first_arrival_timestamp_ms,
            )?;
            request.output_token_ids = Some(planned_output_token_ids(
                request.output_token_ids,
                request.max_output_tokens,
                &mut output_rng,
            ));
            requests.push(request);
        }
        Ok(requests)
    }

    pub fn is_single_turn(&self) -> bool {
        self.sessions.iter().all(|session| session.turns.len() == 1)
    }

    pub fn into_trace_driver(self) -> Result<WorkloadDriver> {
        self.validate_for_trace_mode()?;
        let engine_block_size = self.block_size;
        WorkloadDriver::new_trace(self, engine_block_size)
    }

    pub fn into_concurrency_driver(self, max_in_flight: usize) -> Result<WorkloadDriver> {
        self.validate_for_concurrency_mode()?;
        let engine_block_size = self.block_size;
        WorkloadDriver::new_concurrency(self, engine_block_size, max_in_flight)
    }

    pub fn into_trace_driver_with_block_size(
        self,
        engine_block_size: usize,
    ) -> Result<WorkloadDriver> {
        self.validate_for_trace_mode()?;
        WorkloadDriver::new_trace(self, engine_block_size)
    }

    pub fn into_delta_accumulating_trace_driver_with_block_size(
        self,
        engine_block_size: usize,
    ) -> Result<WorkloadDriver> {
        self.validate_for_trace_mode()?;
        WorkloadDriver::new_trace_accumulating_deltas(self, engine_block_size)
    }

    pub fn into_concurrency_driver_with_block_size(
        self,
        engine_block_size: usize,
        max_in_flight: usize,
    ) -> Result<WorkloadDriver> {
        self.validate_for_concurrency_mode()?;
        WorkloadDriver::new_concurrency(self, engine_block_size, max_in_flight)
    }

    pub fn into_delta_accumulating_concurrency_driver_with_block_size(
        self,
        engine_block_size: usize,
        max_in_flight: usize,
    ) -> Result<WorkloadDriver> {
        self.validate_for_concurrency_mode()?;
        WorkloadDriver::new_concurrency_accumulating_deltas(self, engine_block_size, max_in_flight)
    }

    fn validate(&self, allow_missing_first_timestamp: bool) -> Result<()> {
        if self.block_size == 0 {
            bail!("block_size must be greater than 0");
        }
        if self.sessions.is_empty() {
            bail!("trace must contain at least one session");
        }

        for session in &self.sessions {
            if session.turns.is_empty() {
                bail!(
                    "session {} must contain at least one turn",
                    session.session_id
                );
            }
            if !allow_missing_first_timestamp {
                let timestamp_ms = session.first_arrival_timestamp_ms.ok_or_else(|| {
                    anyhow!(
                        "trace mode requires first_arrival_timestamp_ms for session {}",
                        session.session_id
                    )
                })?;
                if !timestamp_ms.is_finite() || timestamp_ms < 0.0 {
                    bail!(
                        "session {} has invalid first_arrival_timestamp_ms {}",
                        session.session_id,
                        timestamp_ms
                    );
                }
            } else if let Some(timestamp_ms) = session.first_arrival_timestamp_ms
                && (!timestamp_ms.is_finite() || timestamp_ms < 0.0)
            {
                bail!(
                    "session {} has invalid first_arrival_timestamp_ms {}",
                    session.session_id,
                    timestamp_ms
                );
            }

            for (turn_idx, turn) in session.turns.iter().enumerate() {
                if let Some(output_token_ids) = turn.output_token_ids.as_ref()
                    && output_token_ids.len() != turn.max_output_tokens
                {
                    bail!(
                        "session {} turn {} max_output_tokens {} does not match output_token_ids length {}",
                        session.session_id,
                        turn_idx,
                        turn.max_output_tokens,
                        output_token_ids.len()
                    );
                }
                if turn.input_length == 0 {
                    bail!(
                        "session {} turn {} must have a positive input_length",
                        session.session_id,
                        turn_idx
                    );
                }
                if turn.hash_ids.is_empty() {
                    bail!(
                        "session {} turn {} must contain at least one hash id",
                        session.session_id,
                        turn_idx
                    );
                }
                validate_synthesizable_prompt(turn.input_length, &turn.hash_ids, self.block_size)
                    .with_context(|| {
                    format!(
                        "session {} turn {} has invalid prompt",
                        session.session_id, turn_idx
                    )
                })?;
                if !turn.delay_after_previous_ms.is_finite() || turn.delay_after_previous_ms < 0.0 {
                    bail!(
                        "session {} turn {} has invalid delay {}",
                        session.session_id,
                        turn_idx,
                        turn.delay_after_previous_ms
                    );
                }
                if turn_idx == 0 && turn.delay_after_previous_ms != 0.0 {
                    bail!(
                        "session {} first turn must have delay_after_previous_ms == 0.0",
                        session.session_id
                    );
                }
            }
        }

        Ok(())
    }
}

struct AgenticTraceBuilder {
    header: AgenticMooncakeHeader,
    nodes: Vec<AgenticNode>,
    request_ids: std::collections::HashSet<String>,
}

/// Incremental builder for an already-versioned agentic request stream.
///
/// Producers can validate and compile rows without retaining a second full
/// row collection in memory. The returned graph owns the row payloads.
pub struct AgenticGraphBuilder {
    inner: AgenticTraceBuilder,
    row_count: usize,
}

impl AgenticGraphBuilder {
    pub fn new(header: AgenticMooncakeHeader) -> Result<Self> {
        Ok(Self {
            inner: AgenticTraceBuilder::new(header)?,
            row_count: 0,
        })
    }

    pub fn push(&mut self, row: AgenticMooncakeRow) -> Result<()> {
        self.inner.push(self.row_count + 1, row)?;
        self.row_count += 1;
        Ok(())
    }

    pub fn finish(self) -> Result<AgenticTrace> {
        if self.inner.is_empty() {
            bail!("agentic Mooncake rows did not contain any requests");
        }
        self.inner.finish()
    }
}

impl AgenticTraceBuilder {
    fn new(header: AgenticMooncakeHeader) -> Result<Self> {
        if header.schema != AGENTIC_MOONCAKE_SCHEMA {
            bail!(
                "unsupported agentic Mooncake schema {:?}; expected {:?}",
                header.schema,
                AGENTIC_MOONCAKE_SCHEMA
            );
        }
        if header.version != AGENTIC_MOONCAKE_VERSION {
            bail!(
                "unsupported agentic Mooncake version {}; expected {}",
                header.version,
                AGENTIC_MOONCAKE_VERSION
            );
        }
        if header.block_size == 0 {
            bail!("agentic Mooncake block_size must be greater than 0");
        }
        if header.source.format.trim().is_empty() || header.source.digest.trim().is_empty() {
            bail!("agentic Mooncake source provenance requires nonempty format and digest");
        }

        Ok(Self {
            header,
            nodes: Vec::new(),
            request_ids: std::collections::HashSet::new(),
        })
    }

    fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }

    fn push(&mut self, line_idx: usize, raw: AgenticMooncakeRow) -> Result<()> {
        if raw.request_id.trim().is_empty() {
            bail!("trace line {} has empty request_id", line_idx + 1);
        }
        if raw.play_id.trim().is_empty() {
            bail!("trace line {} has empty play_id", line_idx + 1);
        }
        if raw.session_id.trim().is_empty() {
            bail!("trace line {} has empty session_id", line_idx + 1);
        }
        if !self.request_ids.insert(raw.request_id.clone()) {
            bail!(
                "trace line {} duplicates request_id {}",
                line_idx + 1,
                raw.request_id
            );
        }

        let hash_ids = raw
            .hash_ids
            .ok_or_else(|| anyhow!("trace line {} is missing hash_ids", line_idx + 1))?;
        if hash_ids.is_empty() {
            bail!("trace line {} has empty hash_ids", line_idx + 1);
        }
        let input_length = raw
            .input_length
            .ok_or_else(|| anyhow!("trace line {} is missing input_length", line_idx + 1))?;
        if input_length == 0 {
            bail!("trace line {} has zero input_length", line_idx + 1);
        }
        let expected_hashes = input_length.div_ceil(self.header.block_size);
        if hash_ids.len() != expected_hashes {
            bail!(
                "trace line {} has input_length {} at block_size {} and requires exactly {} hash_ids, got {}",
                line_idx + 1,
                input_length,
                self.header.block_size,
                expected_hashes,
                hash_ids.len()
            );
        }
        if raw.model.trim().is_empty() {
            bail!("trace line {} has an empty model", line_idx + 1);
        }
        let output_length = raw
            .output_length
            .ok_or_else(|| anyhow!("trace line {} is missing output_length", line_idx + 1))?;
        let output_token_ids = raw.output_token_ids;
        if let Some(output_token_ids) = output_token_ids.as_ref()
            && output_token_ids.len() != output_length
        {
            bail!(
                "trace line {} output_length {} does not match output_token_ids length {}",
                line_idx + 1,
                output_length,
                output_token_ids.len()
            );
        }
        if !raw.not_before_ms.is_finite() || raw.not_before_ms < 0.0 {
            bail!(
                "trace line {} has invalid not_before_ms {}",
                line_idx + 1,
                raw.not_before_ms
            );
        }
        for dependency in &raw.dependencies {
            if dependency.request_id.trim().is_empty() {
                bail!(
                    "trace line {} has a dependency with an empty request_id",
                    line_idx + 1
                );
            }
            if !dependency.delay_ms.is_finite() || dependency.delay_ms < 0.0 {
                bail!(
                    "trace line {} has invalid dependency delay {}",
                    line_idx + 1,
                    dependency.delay_ms
                );
            }
            match (dependency.relation, dependency.trigger) {
                (AgenticDependencyRelation::Sequence, AgenticDependencyTrigger::Completion)
                | (AgenticDependencyRelation::Spawn, _)
                | (AgenticDependencyRelation::Join, AgenticDependencyTrigger::Completion)
                | (
                    AgenticDependencyRelation::ReplayBarrier,
                    AgenticDependencyTrigger::Completion,
                ) => {}
                (relation, trigger) => bail!(
                    "trace line {} has invalid {:?} dependency with {:?} trigger",
                    line_idx + 1,
                    relation,
                    trigger
                ),
            }
        }

        let replay_key = output_token_ids.as_ref().map(|_| {
            effective_replay_key(
                Some(raw.request_id.as_str()),
                Some(raw.session_id.as_str()),
                0,
                line_idx,
            )
        });
        self.nodes.push(AgenticNode {
            replay_key,
            request_id: raw.request_id,
            play_id: raw.play_id,
            session_id: raw.session_id,
            model: raw.model,
            input_length,
            max_output_tokens: output_length,
            output_token_ids,
            hash_ids,
            not_before_ms: raw.not_before_ms,
            priority: raw.priority.unwrap_or(0),
            strict_priority: raw.strict_priority.unwrap_or(0),
            policy_class: raw.policy_class,
            dependencies: raw.dependencies,
        });
        Ok(())
    }

    fn finish(mut self) -> Result<AgenticTrace> {
        self.nodes
            .sort_by(|left, right| left.request_id.cmp(&right.request_id));
        let index_by_id: HashMap<&str, usize> = self
            .nodes
            .iter()
            .enumerate()
            .map(|(index, node)| (node.request_id.as_str(), index))
            .collect();
        let mut nodes_by_play: HashMap<String, Vec<usize>> = HashMap::new();

        for (node_index, node) in self.nodes.iter().enumerate() {
            nodes_by_play
                .entry(node.play_id.clone())
                .or_default()
                .push(node_index);
            for dependency in &node.dependencies {
                let Some(&dependency_index) = index_by_id.get(dependency.request_id.as_str())
                else {
                    bail!(
                        "request {} references unknown request_id {}",
                        node.request_id,
                        dependency.request_id
                    );
                };
                let dependency_node = &self.nodes[dependency_index];
                if dependency_index == node_index {
                    bail!("request {} cannot depend on itself", node.request_id);
                }
                if dependency_node.play_id != node.play_id {
                    bail!(
                        "request {} in play {} depends on request {} in play {}",
                        node.request_id,
                        node.play_id,
                        dependency.request_id,
                        dependency_node.play_id
                    );
                }
            }
        }
        validate_agentic_trace_is_acyclic(&self.nodes, &index_by_id)?;

        let mut plays = Vec::with_capacity(nodes_by_play.len());
        for (play_id, mut node_indices) in nodes_by_play {
            node_indices.sort_unstable();
            let roots: Vec<_> = node_indices
                .iter()
                .copied()
                .filter(|node_index| self.nodes[*node_index].dependencies.is_empty())
                .collect();
            if roots.is_empty() {
                bail!(
                    "play {} must have at least one root request, found {}",
                    play_id,
                    roots.len()
                );
            }
            plays.push(AgenticPlay {
                play_id,
                root_nodes: roots,
                nodes: node_indices,
            });
        }
        plays.sort_by(|left, right| left.play_id.cmp(&right.play_id));
        let graph_digest = canonical_agentic_graph_digest(self.header.block_size, &mut self.nodes)?;

        Ok(AgenticTrace {
            block_size: self.header.block_size,
            source: self.header.source,
            graph_digest,
            nodes: self.nodes,
            plays,
        })
    }
}

impl AgenticTrace {
    pub fn from_agentic_mooncake(path: &Path) -> Result<Self> {
        let file = File::open(path)
            .with_context(|| format!("failed to open trace file {}", path.display()))?;
        let reader = BufReader::new(file);
        let mut builder = None;

        for (line_idx, line) in reader.lines().enumerate() {
            let line = line.with_context(|| {
                format!(
                    "failed to read line {} from {}",
                    line_idx + 1,
                    path.display()
                )
            })?;
            if line.trim().is_empty() {
                continue;
            }

            if builder.is_none() {
                let header = serde_json::from_str(&line).with_context(|| {
                    format!(
                        "failed to parse line {} from {} as the agentic Mooncake v2 header",
                        line_idx + 1,
                        path.display()
                    )
                })?;
                builder = Some(AgenticTraceBuilder::new(header)?);
                continue;
            }
            let row = serde_json::from_str(&line).with_context(|| {
                format!(
                    "failed to parse line {} from {} as an agentic Mooncake v2 request",
                    line_idx + 1,
                    path.display()
                )
            })?;
            builder
                .as_mut()
                .expect("builder was initialized from the v2 header")
                .push(line_idx, row)?;
        }

        let Some(builder) = builder else {
            bail!(
                "agentic trace file {} is missing its v2 header",
                path.display()
            );
        };
        if builder.is_empty() {
            bail!(
                "agentic trace file {} did not contain any requests",
                path.display()
            );
        }

        builder.finish()
    }

    pub fn from_agentic_mooncake_rows(
        header: AgenticMooncakeHeader,
        rows: Vec<AgenticMooncakeRow>,
    ) -> Result<Self> {
        let mut builder = AgenticTraceBuilder::new(header)?;
        for (line_idx, row) in rows.into_iter().enumerate() {
            builder.push(line_idx + 1, row)?;
        }
        if builder.is_empty() {
            bail!("agentic Mooncake rows did not contain any requests");
        }
        builder.finish()
    }

    pub fn normalize_starts(mut self) -> Self {
        let min_timestamp_ms = self
            .nodes
            .iter()
            .map(|node| node.not_before_ms)
            .min_by(|left, right| left.total_cmp(right))
            .unwrap_or(0.0);

        for node in &mut self.nodes {
            node.not_before_ms -= min_timestamp_ms;
        }
        self.graph_digest = canonical_agentic_graph_digest(self.block_size, &mut self.nodes)
            .expect("validated agentic graph remains serializable after normalization");
        self
    }

    pub fn speed_up_timing(mut self, ratio: f64) -> Result<Self> {
        if !ratio.is_finite() || ratio <= 0.0 {
            bail!("ratio must be a finite positive number, got {ratio}");
        }

        for node in &mut self.nodes {
            node.not_before_ms /= ratio;
            for dependency in &mut node.dependencies {
                dependency.delay_ms /= ratio;
            }
        }
        self.graph_digest = canonical_agentic_graph_digest(self.block_size, &mut self.nodes)?;
        Ok(self)
    }

    pub fn into_trace_driver_with_block_size(
        self,
        engine_block_size: usize,
    ) -> Result<WorkloadDriver> {
        WorkloadDriver::new_agentic_trace(self, engine_block_size)
    }

    pub fn into_trace_driver_with_options(
        self,
        engine_block_size: usize,
        include_replay_hashes: bool,
        agentic_lanes: Option<usize>,
    ) -> Result<WorkloadDriver> {
        WorkloadDriver::new_agentic_trace_with_options(
            self,
            engine_block_size,
            include_replay_hashes,
            agentic_lanes,
        )
    }
}

fn validate_agentic_trace_is_acyclic(
    nodes: &[AgenticNode],
    index_by_id: &HashMap<&str, usize>,
) -> Result<()> {
    let mut indegree = nodes
        .iter()
        .map(|node| node.dependencies.len())
        .collect::<Vec<_>>();
    let mut dependents = vec![Vec::new(); nodes.len()];
    for (node_index, node) in nodes.iter().enumerate() {
        for dependency in &node.dependencies {
            let dependency_index = *index_by_id
                .get(dependency.request_id.as_str())
                .expect("dependencies were prevalidated");
            dependents[dependency_index].push(node_index);
        }
    }
    let mut ready = std::collections::VecDeque::from_iter(
        indegree
            .iter()
            .enumerate()
            .filter_map(|(index, degree)| (*degree == 0).then_some(index)),
    );
    let mut visited = 0;
    while let Some(node_index) = ready.pop_front() {
        visited += 1;
        for dependent in &dependents[node_index] {
            indegree[*dependent] -= 1;
            if indegree[*dependent] == 0 {
                ready.push_back(*dependent);
            }
        }
    }
    if visited == nodes.len() {
        return Ok(());
    }
    let blocked = indegree
        .iter()
        .enumerate()
        .filter_map(|(index, degree)| (*degree > 0).then_some(nodes[index].request_id.as_str()))
        .take(8)
        .collect::<Vec<_>>();
    bail!(
        "cycle detected among {} agentic requests; blocked request IDs: {}",
        nodes.len() - visited,
        blocked.join(", ")
    )
}

fn canonical_agentic_graph_digest(block_size: usize, nodes: &mut [AgenticNode]) -> Result<String> {
    struct Blake3Writer<'a>(&'a mut blake3::Hasher);

    impl std::io::Write for Blake3Writer<'_> {
        fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
            self.0.update(bytes);
            Ok(bytes.len())
        }

        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    #[derive(Serialize)]
    struct CanonicalGraph<'a> {
        block_size: usize,
        nodes: &'a [AgenticNode],
    }

    for node in nodes.iter_mut() {
        node.dependencies.sort_by(|left, right| {
            left.request_id
                .cmp(&right.request_id)
                .then_with(|| trigger_rank(left.trigger).cmp(&trigger_rank(right.trigger)))
                .then_with(|| relation_rank(left.relation).cmp(&relation_rank(right.relation)))
                .then_with(|| left.delay_ms.total_cmp(&right.delay_ms))
        });
    }
    let mut hasher = blake3::Hasher::new();
    serde_json::to_writer(
        Blake3Writer(&mut hasher),
        &CanonicalGraph { block_size, nodes },
    )?;
    Ok(hasher.finalize().to_hex().to_string())
}

fn trigger_rank(trigger: AgenticDependencyTrigger) -> u8 {
    match trigger {
        AgenticDependencyTrigger::Dispatch => 0,
        AgenticDependencyTrigger::Completion => 1,
    }
}

fn relation_rank(relation: AgenticDependencyRelation) -> u8 {
    match relation {
        AgenticDependencyRelation::Sequence => 0,
        AgenticDependencyRelation::Spawn => 1,
        AgenticDependencyRelation::Join => 2,
        AgenticDependencyRelation::ReplayBarrier => 3,
    }
}

fn extend_applied_compute_agentic_hash_ids(
    hash_ids: &mut Vec<u64>,
    input_length: usize,
    trace_block_size: usize,
    shared_initial_blocks: usize,
    group_id: Option<usize>,
    next_unique_hash: &mut u64,
) -> Result<()> {
    let target_blocks = input_length.div_ceil(trace_block_size);
    while hash_ids.len() < target_blocks {
        let block_idx = hash_ids.len();
        if block_idx < shared_initial_blocks
            && let Some(group_id) = group_id
        {
            hash_ids.push(0xA63E_0000_0000_0000 | ((group_id as u64) << 32) | block_idx as u64);
            continue;
        }
        hash_ids.push(*next_unique_hash);
        *next_unique_hash = next_unique_hash
            .checked_add(1)
            .ok_or_else(|| anyhow!("synthetic hash id overflow"))?;
    }
    Ok(())
}

fn sample_delay_ms(spec: &DelaySpec, rng: &mut StdRng) -> Result<f64> {
    match spec {
        DelaySpec::None => Ok(0.0),
        DelaySpec::ConstantMs(delay_ms) => {
            if !delay_ms.is_finite() || *delay_ms < 0.0 {
                bail!("delay must be a finite non-negative number, got {delay_ms}");
            }
            Ok(*delay_ms)
        }
        DelaySpec::ExponentialMs { mean_ms } => {
            if !mean_ms.is_finite() || *mean_ms < 0.0 {
                bail!("mean_ms must be a finite non-negative number, got {mean_ms}");
            }
            Ok(sample_exponential_delay_ms(*mean_ms, rng))
        }
    }
}

fn sample_length(spec: &LengthSpec, min_value: usize, rng: &mut StdRng) -> usize {
    if spec.stddev == 0.0 {
        return spec.mean.max(min_value);
    }

    let stddev = spec.stddev.abs();
    let u1 = (1.0 - rng.random::<f64>()).clamp(f64::MIN_POSITIVE, 1.0);
    let u2 = rng.random::<f64>();
    let z0 = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos();
    let sample = spec.mean as f64 + z0 * stddev;
    sample.round().max(min_value as f64) as usize
}

fn sample_exponential_delay_ms(mean_ms: f64, rng: &mut StdRng) -> f64 {
    if mean_ms == 0.0 {
        return 0.0;
    }
    let u = (1.0 - rng.random::<f64>()).clamp(f64::MIN_POSITIVE, 1.0);
    -mean_ms * u.ln()
}