dynamo_data_gen/
mooncake.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Mooncake JSONL primitives.
5//!
6//! This module is producer- and consumer-agnostic: it defines the row schema,
7//! the block-hash-to-id mapping, the token-block hashing helper, and the JSONL
8//! writer. Workload-specific orchestration such as scheduling, tokenization,
9//! and parsing lives elsewhere.
10//!
11//! The [`MooncakeRow`] schema deliberately matches the externally-authored
12//! Mooncake trace format: `timestamp` and `delay` are `f64` milliseconds, and
13//! `input_length`/`output_length`/`timestamp`/`delay` accept the upstream
14//! aliases (`input_tokens`, `output_tokens`, `created_time`, `delay_ms`) on
15//! deserialization. Serialization emits the canonical names.
16
17use anyhow::{Context, Result, bail};
18use dynamo_kv_hashing::{Request, compute_hash_v2, compute_next_sequence_hash, compute_salt_hash};
19use rustc_hash::FxHashMap;
20use serde::{Deserialize, Serialize};
21use std::fs::File;
22use std::io::{BufWriter, Write};
23use std::path::Path;
24
25/// One row of a Mooncake replay trace.
26///
27/// `timestamp` is an absolute request arrival offset in milliseconds. Rows
28/// without a `session_id` are independent request arrivals. Rows that share a
29/// `session_id` are interpreted as closed-loop turns; later turns use `delay`
30/// or timestamp deltas relative to the previous row in that session.
31///
32/// The row type is `Serialize + Deserialize` so the same definition serves
33/// producers and consumers. Field-level aliases on deserialization accept the
34/// upstream Mooncake field names (`input_tokens`, `output_tokens`,
35/// `created_time`, `delay_ms`) without requiring producers to emit them.
36#[derive(Debug, Clone, Default, Serialize, Deserialize)]
37pub struct MooncakeRow {
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub request_id: Option<String>,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub session_id: Option<String>,
42    #[serde(default, alias = "input_tokens")]
43    pub input_length: Option<usize>,
44    #[serde(default, alias = "output_tokens")]
45    pub output_length: Option<usize>,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub output_token_ids: Option<Vec<u32>>,
48    #[serde(default)]
49    pub hash_ids: Option<Vec<u64>>,
50    #[serde(
51        default,
52        skip_serializing_if = "Option::is_none",
53        alias = "created_time"
54    )]
55    pub timestamp: Option<f64>,
56    #[serde(default, skip_serializing_if = "Option::is_none", alias = "delay_ms")]
57    pub delay: Option<f64>,
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub priority: Option<i32>,
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub strict_priority: Option<u32>,
62    #[serde(default, skip_serializing_if = "Option::is_none")]
63    pub policy_class: Option<String>,
64}
65
66pub const AGENTIC_MOONCAKE_SCHEMA: &str = "dynamo.agentic_mooncake";
67pub const AGENTIC_MOONCAKE_VERSION: u32 = 2;
68
69#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
70pub struct AgenticSourceProvenance {
71    pub format: String,
72    pub digest: String,
73}
74
75#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
76#[serde(rename_all = "snake_case")]
77pub enum AgenticHashIdScope {
78    Local,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
82pub struct AgenticMooncakeHeader {
83    pub schema: String,
84    pub version: u32,
85    pub block_size: usize,
86    pub hash_id_scope: AgenticHashIdScope,
87    pub source: AgenticSourceProvenance,
88}
89
90#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
91#[serde(rename_all = "snake_case")]
92pub enum AgenticDependencyTrigger {
93    Dispatch,
94    Completion,
95}
96
97#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
98#[serde(rename_all = "snake_case")]
99pub enum AgenticDependencyRelation {
100    Sequence,
101    Spawn,
102    Join,
103    ReplayBarrier,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
107pub struct AgenticDependency {
108    pub request_id: String,
109    pub trigger: AgenticDependencyTrigger,
110    pub delay_ms: f64,
111    pub relation: AgenticDependencyRelation,
112}
113
114/// One request row in the versioned Agentic Mooncake wire format.
115#[derive(Debug, Clone, Default, Serialize, Deserialize)]
116pub struct AgenticMooncakeRow {
117    pub request_id: String,
118    pub play_id: String,
119    pub session_id: String,
120    pub model: String,
121    #[serde(default, alias = "input_tokens")]
122    pub input_length: Option<usize>,
123    #[serde(default, alias = "output_tokens")]
124    pub output_length: Option<usize>,
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub output_token_ids: Option<Vec<u32>>,
127    #[serde(default)]
128    pub hash_ids: Option<Vec<u64>>,
129    pub not_before_ms: f64,
130    #[serde(default, skip_serializing_if = "Option::is_none")]
131    pub priority: Option<i32>,
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub strict_priority: Option<u32>,
134    #[serde(default, skip_serializing_if = "Option::is_none")]
135    pub policy_class: Option<String>,
136    #[serde(default, skip_serializing_if = "Vec::is_empty")]
137    pub dependencies: Vec<AgenticDependency>,
138}
139
140/// Maps sequence-aware block hashes to compact, stable `u64` ids.
141///
142/// The mapper is intentionally stateful and reusable across requests/turns: a
143/// block of tokens that appears at the same prefix position in two different
144/// requests will be assigned the same id. Equality of leading `hash_ids`
145/// between rows therefore signals shared prompt prefixes for replay purposes.
146///
147/// `hash_ids` here are workload identity labels, not literal Dynamo runtime
148/// KV-cache hashes. Producers should not try to reconcile them with a
149/// production cache.
150pub struct RollingHashIdMapper {
151    block_size: usize,
152    hash_to_id: FxHashMap<u64, u64>,
153    next_id: u64,
154}
155
156impl RollingHashIdMapper {
157    /// Create a new mapper for the given block size.
158    pub fn new(block_size: usize) -> Self {
159        Self {
160            block_size,
161            hash_to_id: FxHashMap::default(),
162            next_id: 0,
163        }
164    }
165
166    /// Block size that this mapper was constructed with.
167    pub fn block_size(&self) -> usize {
168        self.block_size
169    }
170
171    /// Hash a sequence of tokens into Mooncake `hash_ids`.
172    ///
173    /// Tokens are chunked by `block_size`; each complete block contributes one
174    /// compact id derived from Dynamo's shared KV-hashing contract. A trailing
175    /// partial block also contributes one compact id so replay capacity still
176    /// covers the full prompt length. Identical prefixes across requests
177    /// resolve to identical leading `hash_ids` once the mapper has seen them.
178    pub fn hash_token_blocks(&mut self, tokens: &[u32]) -> Vec<u64> {
179        hash_token_blocks(self, tokens)
180    }
181
182    /// Fallible variant of [`Self::hash_token_blocks`].
183    pub fn try_hash_token_blocks(&mut self, tokens: &[u32]) -> Result<Vec<u64>> {
184        try_hash_token_blocks(self, tokens)
185    }
186
187    /// Map precomputed sequence-aware block hashes into compact Mooncake IDs.
188    ///
189    /// This is useful for producers that record stable block hashes in the
190    /// serving path and only compact them during offline trace conversion.
191    pub fn ids_for_sequence_hashes(&mut self, sequence_hashes: &[u64]) -> Vec<u64> {
192        ids_for_sequence_hashes(self, sequence_hashes)
193    }
194}
195
196/// Token-block hashing helper for the Mooncake replay schema.
197///
198/// Splits `tokens` into chunks of `mapper.block_size()`, derives sequence-aware
199/// hashes for complete blocks through `dynamo-kv-hashing`, appends a sequence
200/// hash for a trailing partial block when present, and returns the
201/// compact ids assigned by `mapper`. Mirrors
202/// [`RollingHashIdMapper::hash_token_blocks`] as a free function so callers
203/// that already hold a mutable mapper reference can invoke it without
204/// re-borrowing.
205pub fn hash_token_blocks(mapper: &mut RollingHashIdMapper, tokens: &[u32]) -> Vec<u64> {
206    try_hash_token_blocks(mapper, tokens).expect("Mooncake token-block hashing failed")
207}
208
209/// Fallible token-block hashing helper for callers that want to surface
210/// invalid block-size or request-shape errors.
211pub fn try_hash_token_blocks(mapper: &mut RollingHashIdMapper, tokens: &[u32]) -> Result<Vec<u64>> {
212    let sequence_hashes = sequence_hashes_for_tokens(tokens, mapper.block_size)?;
213    Ok(ids_for_sequence_hashes(mapper, &sequence_hashes))
214}
215
216/// Derive the sequence-aware block hashes recorded by Dynamo request traces.
217pub fn sequence_hashes_for_tokens(tokens: &[u32], block_size: usize) -> Result<Vec<u64>> {
218    require_positive("block size", block_size)?;
219    let block_size_u32: u32 = block_size
220        .try_into()
221        .context("block_size does not fit u32")?;
222    let request = Request::builder().tokens(tokens.to_vec()).build()?;
223    let salt_hash = request.salt_hash()?;
224    let mut sequence_hashes = request.into_sequence_hashes(block_size_u32)?;
225    if let Some(partial_hash) =
226        trailing_partial_sequence_hash(salt_hash, block_size, tokens, &sequence_hashes)
227    {
228        sequence_hashes.push(partial_hash);
229    }
230    Ok(sequence_hashes)
231}
232
233/// Extends reusable full-block sequence hashes to cover `tokens`.
234///
235/// `prefix_hashes` must be the leading full-block hashes that [`sequence_hashes_for_tokens`]
236/// returns for `tokens`. Only the remaining blocks are hashed, and the result equals
237/// `sequence_hashes_for_tokens(tokens, block_size)`.
238pub fn extend_sequence_hashes(
239    prefix_hashes: &[u64],
240    tokens: &[u32],
241    block_size: usize,
242) -> Result<Vec<u64>> {
243    require_positive("block size", block_size)?;
244    let full_blocks = tokens.len() / block_size;
245    if prefix_hashes.len() > full_blocks {
246        bail!(
247            "{} reusable sequence hashes exceed {full_blocks} full blocks",
248            prefix_hashes.len()
249        );
250    }
251    let salt_hash = compute_salt_hash(None, None)?;
252    let mut sequence_hashes = Vec::with_capacity(full_blocks + 1);
253    sequence_hashes.extend_from_slice(prefix_hashes);
254    let mut block_bytes = Vec::with_capacity(block_size * std::mem::size_of::<u32>());
255    for block in
256        tokens[prefix_hashes.len() * block_size..full_blocks * block_size].chunks_exact(block_size)
257    {
258        block_bytes.clear();
259        for token in block {
260            block_bytes.extend_from_slice(&token.to_ne_bytes());
261        }
262        let block_hash = compute_hash_v2(&block_bytes, salt_hash);
263        let sequence_hash = match sequence_hashes.last() {
264            Some(parent) => compute_next_sequence_hash(*parent, block_hash),
265            None => block_hash,
266        };
267        sequence_hashes.push(sequence_hash);
268    }
269    if let Some(partial_hash) =
270        trailing_partial_sequence_hash(salt_hash, block_size, tokens, &sequence_hashes)
271    {
272        sequence_hashes.push(partial_hash);
273    }
274    Ok(sequence_hashes)
275}
276
277fn trailing_partial_sequence_hash(
278    salt_hash: u64,
279    block_size: usize,
280    tokens: &[u32],
281    complete_sequence_hashes: &[u64],
282) -> Option<u64> {
283    let tail_len = tokens.len() % block_size;
284    if tail_len == 0 {
285        return None;
286    }
287
288    let tail = &tokens[tokens.len() - tail_len..];
289    let mut tail_bytes = Vec::with_capacity(std::mem::size_of_val(tail));
290    for token in tail {
291        tail_bytes.extend_from_slice(&token.to_ne_bytes());
292    }
293    let tail_block_hash = compute_hash_v2(&tail_bytes, salt_hash);
294    Some(match complete_sequence_hashes.last().copied() {
295        Some(parent) => compute_next_sequence_hash(parent, tail_block_hash),
296        None => tail_block_hash,
297    })
298}
299
300/// Map stable sequence hashes to compact Mooncake IDs with a shared mapper.
301pub fn ids_for_sequence_hashes(
302    mapper: &mut RollingHashIdMapper,
303    sequence_hashes: &[u64],
304) -> Vec<u64> {
305    sequence_hashes
306        .iter()
307        .map(|sequence_hash| {
308            *mapper.hash_to_id.entry(*sequence_hash).or_insert_with(|| {
309                let next_id = mapper.next_id;
310                mapper.next_id += 1;
311                next_id
312            })
313        })
314        .collect()
315}
316
317/// Counters for what a [`MooncakeJsonlWriter`] has emitted.
318#[derive(Debug, Clone, Copy, Default)]
319pub struct WriterStats {
320    pub row_count: usize,
321    pub sidecar_count: usize,
322}
323
324/// JSONL writer for Mooncake rows plus an optional sidecar stream.
325///
326/// The sidecar stream is configured at construction time. Producers that do
327/// not emit sidecar metadata pass `None` for `sidecar_path` and never call
328/// [`Self::write_sidecar`]. When a sidecar path is configured, callers are
329/// responsible for choosing the path -- this writer does not enforce a naming
330/// convention.
331pub struct MooncakeJsonlWriter {
332    output: BufWriter<File>,
333    sidecar: Option<BufWriter<File>>,
334    stats: WriterStats,
335    agentic_header_written: bool,
336}
337
338impl MooncakeJsonlWriter {
339    /// Create a writer at `output_path`, optionally with a paired sidecar
340    /// JSONL file at `sidecar_path`. Parent directories are created as needed.
341    pub fn create(output_path: &Path, sidecar_path: Option<&Path>) -> Result<Self> {
342        if let Some(parent) = output_path.parent() {
343            std::fs::create_dir_all(parent)?;
344        }
345        let output = BufWriter::new(
346            File::create(output_path)
347                .with_context(|| format!("failed to create {}", output_path.display()))?,
348        );
349        let sidecar = if let Some(path) = sidecar_path {
350            if let Some(parent) = path.parent() {
351                std::fs::create_dir_all(parent)?;
352            }
353            Some(BufWriter::new(File::create(path).with_context(|| {
354                format!("failed to create {}", path.display())
355            })?))
356        } else {
357            None
358        };
359        Ok(Self {
360            output,
361            sidecar,
362            stats: WriterStats::default(),
363            agentic_header_written: false,
364        })
365    }
366
367    /// Append one Mooncake row.
368    pub fn write_row(&mut self, row: &MooncakeRow) -> Result<()> {
369        if self.agentic_header_written {
370            bail!("standard Mooncake rows cannot follow an agentic Mooncake v2 header");
371        }
372        serde_json::to_writer(&mut self.output, row)?;
373        self.output.write_all(b"\n")?;
374        self.stats.row_count += 1;
375        Ok(())
376    }
377
378    pub fn write_agentic_header(&mut self, header: &AgenticMooncakeHeader) -> Result<()> {
379        if self.agentic_header_written || self.stats.row_count != 0 {
380            bail!("agentic Mooncake header must be the first JSONL record");
381        }
382        serde_json::to_writer(&mut self.output, header)?;
383        self.output.write_all(b"\n")?;
384        self.agentic_header_written = true;
385        Ok(())
386    }
387
388    /// Append one agentic Mooncake row after its required v2 header.
389    pub fn write_agentic_row(&mut self, row: &AgenticMooncakeRow) -> Result<()> {
390        if !self.agentic_header_written {
391            bail!("agentic Mooncake row requires a preceding v2 header");
392        }
393        serde_json::to_writer(&mut self.output, row)?;
394        self.output.write_all(b"\n")?;
395        self.stats.row_count += 1;
396        Ok(())
397    }
398
399    /// Append one sidecar entry. Errors if no sidecar was configured.
400    pub fn write_sidecar<S: Serialize>(&mut self, sidecar: &S) -> Result<()> {
401        let writer = self
402            .sidecar
403            .as_mut()
404            .ok_or_else(|| anyhow::anyhow!("sidecar was not configured for this writer"))?;
405        serde_json::to_writer(writer, sidecar)?;
406        let writer = self.sidecar.as_mut().unwrap();
407        writer.write_all(b"\n")?;
408        self.stats.sidecar_count += 1;
409        Ok(())
410    }
411
412    /// True if a sidecar stream is configured.
413    pub fn has_sidecar(&self) -> bool {
414        self.sidecar.is_some()
415    }
416
417    /// Snapshot of how many rows and sidecar entries have been written so far.
418    pub fn stats(&self) -> WriterStats {
419        self.stats
420    }
421
422    /// Flush both streams and return the final stats.
423    pub fn finish(mut self) -> Result<WriterStats> {
424        self.output.flush()?;
425        if let Some(sidecar) = self.sidecar.as_mut() {
426            sidecar.flush()?;
427        }
428        Ok(self.stats)
429    }
430}
431
432/// Create both files empty (touch-equivalent), preserving directory creation
433/// semantics for callers that want a "no rows produced" outcome to still emit
434/// well-formed (empty) JSONL files.
435pub fn write_empty_files(output_path: &Path, sidecar_path: Option<&Path>) -> Result<()> {
436    if let Some(parent) = output_path.parent() {
437        std::fs::create_dir_all(parent)?;
438    }
439    File::create(output_path)
440        .with_context(|| format!("failed to create {}", output_path.display()))?;
441    if let Some(path) = sidecar_path {
442        if let Some(parent) = path.parent() {
443            std::fs::create_dir_all(parent)?;
444        }
445        File::create(path).with_context(|| format!("failed to create {}", path.display()))?;
446    }
447    Ok(())
448}
449
450/// Sentinel used by callers that want to bail when neither block_size nor
451/// worker count is allowed to be zero. Producers may also enforce this on
452/// their own configuration types.
453pub fn require_positive(name: &str, value: usize) -> Result<()> {
454    if value == 0 {
455        bail!("{name} must be greater than 0");
456    }
457    Ok(())
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463    use serde_json::{Value, json};
464    use tempfile::TempDir;
465
466    #[test]
467    fn extended_sequence_hashes_match_full_hashing() {
468        let mut state = 0x9e37_79b9_u32;
469        let tokens = (0..300)
470            .map(|_| {
471                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
472                state
473            })
474            .collect::<Vec<_>>();
475        for block_size in [1, 3, 64] {
476            for length in [0, 1, 2, 63, 64, 65, 128, 191, 300] {
477                let tokens = &tokens[..length];
478                let expected = sequence_hashes_for_tokens(tokens, block_size).unwrap();
479                for reused in 0..=length / block_size {
480                    let extended =
481                        extend_sequence_hashes(&expected[..reused], tokens, block_size).unwrap();
482                    assert_eq!(
483                        extended, expected,
484                        "block={block_size} len={length} reused={reused}"
485                    );
486                }
487            }
488        }
489        assert!(extend_sequence_hashes(&[1, 2], &tokens[..3], 2).is_err());
490    }
491
492    #[test]
493    fn shared_prefix_yields_shared_leading_hash_ids() {
494        let mut mapper = RollingHashIdMapper::new(2);
495        let prefix = vec![1u32, 2, 3, 4];
496        let extended = vec![1u32, 2, 3, 4, 5, 6];
497
498        let prefix_ids = mapper.hash_token_blocks(&prefix);
499        let extended_ids = mapper.hash_token_blocks(&extended);
500
501        assert_eq!(prefix_ids.len(), 2);
502        assert_eq!(extended_ids.len(), 3);
503        assert_eq!(extended_ids[..2], prefix_ids[..]);
504    }
505
506    #[test]
507    fn mapper_state_is_reused_across_requests() {
508        let mut mapper = RollingHashIdMapper::new(4);
509        let request_a = vec![10u32, 20, 30, 40, 50, 60, 70, 80];
510        let request_b = vec![10u32, 20, 30, 40, 50, 60, 70, 80];
511        let request_c = vec![10u32, 20, 30, 40, 99, 99, 99, 99];
512
513        let ids_a = mapper.hash_token_blocks(&request_a);
514        let ids_b = mapper.hash_token_blocks(&request_b);
515        let ids_c = mapper.hash_token_blocks(&request_c);
516
517        assert_eq!(ids_a, ids_b);
518        assert_eq!(ids_c[0], ids_a[0], "shared first block should keep its id");
519        assert_ne!(
520            ids_c[1], ids_a[1],
521            "diverging tail block must get a fresh id"
522        );
523    }
524
525    #[test]
526    fn free_function_and_method_agree() {
527        let mut mapper_a = RollingHashIdMapper::new(2);
528        let mut mapper_b = RollingHashIdMapper::new(2);
529        let tokens = vec![7u32, 8, 9, 10, 11];
530
531        let via_method = mapper_a.hash_token_blocks(&tokens);
532        let via_function = hash_token_blocks(&mut mapper_b, &tokens);
533
534        assert_eq!(via_method, via_function);
535    }
536
537    #[test]
538    fn exact_token_blocks_match_shared_kv_hashing_contract() {
539        let tokens = vec![7u32, 8, 9, 10, 11, 12, 13, 14];
540        let request = Request::builder().tokens(tokens.clone()).build().unwrap();
541        let expected = request.into_sequence_hashes(4).unwrap();
542
543        assert_eq!(sequence_hashes_for_tokens(&tokens, 4).unwrap(), expected);
544    }
545
546    #[test]
547    fn empty_token_input_yields_empty_hash_ids() {
548        let mut mapper = RollingHashIdMapper::new(4);
549        assert!(mapper.hash_token_blocks(&[]).is_empty());
550    }
551
552    #[test]
553    fn trailing_partial_block_preserves_replay_capacity() {
554        let mut mapper = RollingHashIdMapper::new(4);
555
556        assert_eq!(mapper.hash_token_blocks(&[1, 2, 3]), vec![0]);
557        assert_eq!(mapper.hash_token_blocks(&[1, 2, 3, 4, 5, 6]), vec![1, 2]);
558    }
559
560    #[test]
561    fn trailing_partial_block_uses_shared_chain_contract() {
562        let tokens = vec![1u32, 2, 3, 4, 5, 6];
563        let request = Request::builder().tokens(tokens.clone()).build().unwrap();
564        let salt_hash = request.salt_hash().unwrap();
565        let mut expected = request.into_sequence_hashes(4).unwrap();
566        let mut tail_bytes = Vec::new();
567        for token in &tokens[4..] {
568            tail_bytes.extend_from_slice(&token.to_ne_bytes());
569        }
570        let tail_block_hash = compute_hash_v2(&tail_bytes, salt_hash);
571        expected.push(compute_next_sequence_hash(expected[0], tail_block_hash));
572
573        assert_eq!(sequence_hashes_for_tokens(&tokens, 4).unwrap(), expected);
574    }
575
576    #[test]
577    fn exact_block_boundary_does_not_add_partial_hash_id() {
578        let mut mapper = RollingHashIdMapper::new(4);
579
580        assert_eq!(mapper.hash_token_blocks(&[1, 2, 3, 4]), vec![0]);
581        assert_eq!(
582            mapper.hash_token_blocks(&[1, 2, 3, 4, 5, 6, 7, 8]),
583            vec![0, 1]
584        );
585    }
586
587    #[test]
588    fn try_hash_token_blocks_rejects_zero_block_size() {
589        let mut mapper = RollingHashIdMapper::new(0);
590        let err = mapper.try_hash_token_blocks(&[1, 2, 3]).unwrap_err();
591
592        assert!(err.to_string().contains("block size"));
593    }
594
595    #[test]
596    fn precomputed_sequence_hashes_map_to_stable_ids() {
597        let mut mapper = RollingHashIdMapper::new(64);
598
599        let first = mapper.ids_for_sequence_hashes(&[101, 202, 303]);
600        let second = mapper.ids_for_sequence_hashes(&[101, 202, 404]);
601
602        assert_eq!(first[..2], second[..2]);
603        assert_ne!(first[2], second[2]);
604    }
605
606    #[test]
607    fn row_omits_timestamp_and_delay_when_absent() {
608        let row = MooncakeRow {
609            session_id: Some("s".to_string()),
610            input_length: Some(4),
611            output_length: Some(1),
612            hash_ids: Some(vec![0, 1]),
613            timestamp: None,
614            delay: None,
615            ..Default::default()
616        };
617        let rendered: Value = serde_json::to_value(&row).unwrap();
618        assert!(rendered.get("timestamp").is_none());
619        assert!(rendered.get("delay").is_none());
620        assert_eq!(rendered["hash_ids"], json!([0, 1]));
621    }
622
623    #[test]
624    fn row_serializes_optional_fields_when_set() {
625        let with_timestamp = MooncakeRow {
626            session_id: Some("s".to_string()),
627            input_length: Some(4),
628            output_length: Some(1),
629            hash_ids: Some(vec![]),
630            timestamp: Some(0.0),
631            delay: None,
632            ..Default::default()
633        };
634        let with_delay = MooncakeRow {
635            session_id: Some("s".to_string()),
636            input_length: Some(4),
637            output_length: Some(1),
638            hash_ids: Some(vec![]),
639            timestamp: None,
640            delay: Some(123.0),
641            ..Default::default()
642        };
643        let v_ts: Value = serde_json::to_value(&with_timestamp).unwrap();
644        let v_dl: Value = serde_json::to_value(&with_delay).unwrap();
645        assert_eq!(v_ts["timestamp"], json!(0.0));
646        assert!(v_ts.get("delay").is_none());
647        assert_eq!(v_dl["delay"], json!(123.0));
648        assert!(v_dl.get("timestamp").is_none());
649    }
650
651    #[test]
652    fn row_deserializes_canonical_field_names() {
653        let raw = r#"{"session_id":"s","input_length":4,"output_length":1,"hash_ids":[0,1],"timestamp":12.5,"delay":3.0}"#;
654        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
655        assert_eq!(row.session_id.as_deref(), Some("s"));
656        assert_eq!(row.input_length, Some(4));
657        assert_eq!(row.output_length, Some(1));
658        assert_eq!(row.hash_ids, Some(vec![0, 1]));
659        assert_eq!(row.timestamp, Some(12.5));
660        assert_eq!(row.delay, Some(3.0));
661    }
662
663    #[test]
664    fn row_deserializes_upstream_mooncake_aliases() {
665        let raw = r#"{"input_tokens":4,"output_tokens":1,"hash_ids":[0,1],"created_time":12.5,"delay_ms":3.0}"#;
666        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
667        assert_eq!(row.input_length, Some(4));
668        assert_eq!(row.output_length, Some(1));
669        assert_eq!(row.timestamp, Some(12.5));
670        assert_eq!(row.delay, Some(3.0));
671    }
672
673    #[test]
674    fn row_alias_input_round_trips_to_canonical_fields() {
675        let raw = r#"{"input_tokens":8,"output_tokens":2,"created_time":12.5,"delay_ms":3.0}"#;
676        let mut row: MooncakeRow = serde_json::from_str(raw).unwrap();
677
678        let tokens: Vec<u32> = (0..row.input_length.unwrap() as u32).collect();
679        let mut mapper = RollingHashIdMapper::new(4);
680        row.hash_ids = Some(mapper.hash_token_blocks(&tokens));
681
682        let rendered: Value = serde_json::to_value(&row).unwrap();
683        assert_eq!(rendered["input_length"], json!(8));
684        assert_eq!(rendered["output_length"], json!(2));
685        assert_eq!(rendered["timestamp"], json!(12.5));
686        assert_eq!(rendered["delay"], json!(3.0));
687        assert_eq!(rendered["hash_ids"], json!([0, 1]));
688        assert!(rendered.get("input_tokens").is_none());
689        assert!(rendered.get("output_tokens").is_none());
690        assert!(rendered.get("created_time").is_none());
691        assert!(rendered.get("delay_ms").is_none());
692    }
693
694    #[test]
695    fn row_replay_fields_round_trip_canonical_and_alias_inputs() {
696        let canonical = r#"{"request_id":"r1","session_id":"s","input_length":8,"output_length":3,"output_token_ids":[101,102,103],"hash_ids":[0,1],"timestamp":12.5,"delay":3.0}"#;
697        let row: MooncakeRow = serde_json::from_str(canonical).unwrap();
698        assert_eq!(row.request_id.as_deref(), Some("r1"));
699        assert_eq!(row.output_length, Some(3));
700        assert_eq!(row.output_token_ids, Some(vec![101, 102, 103]));
701
702        let rendered: Value = serde_json::to_value(&row).unwrap();
703        assert_eq!(rendered["request_id"], json!("r1"));
704        assert_eq!(rendered["output_token_ids"], json!([101, 102, 103]));
705        let decoded: MooncakeRow = serde_json::from_value(rendered).unwrap();
706        assert_eq!(decoded.request_id.as_deref(), Some("r1"));
707        assert_eq!(decoded.output_token_ids, Some(vec![101, 102, 103]));
708
709        let aliased = r#"{"request_id":"r2","input_tokens":4,"output_tokens":2,"output_token_ids":[201,202],"hash_ids":[7],"created_time":1.5,"delay_ms":0.5}"#;
710        let row: MooncakeRow = serde_json::from_str(aliased).unwrap();
711        assert_eq!(row.request_id.as_deref(), Some("r2"));
712        assert_eq!(row.input_length, Some(4));
713        assert_eq!(row.output_length, Some(2));
714        assert_eq!(row.output_token_ids, Some(vec![201, 202]));
715
716        let rendered: Value = serde_json::to_value(&row).unwrap();
717        assert_eq!(rendered["input_length"], json!(4));
718        assert_eq!(rendered["output_length"], json!(2));
719        assert_eq!(rendered["output_token_ids"], json!([201, 202]));
720        assert!(rendered.get("input_tokens").is_none());
721        assert!(rendered.get("output_tokens").is_none());
722    }
723
724    #[test]
725    fn row_canonical_input_round_trips_without_renaming() {
726        let raw = r#"{"input_length":8,"output_length":2,"timestamp":12.5,"delay":3.0}"#;
727        let mut row: MooncakeRow = serde_json::from_str(raw).unwrap();
728
729        let tokens: Vec<u32> = (0..row.input_length.unwrap() as u32).collect();
730        let mut mapper = RollingHashIdMapper::new(4);
731        row.hash_ids = Some(mapper.hash_token_blocks(&tokens));
732
733        let rendered: Value = serde_json::to_value(&row).unwrap();
734        assert_eq!(rendered["input_length"], json!(8));
735        assert_eq!(rendered["output_length"], json!(2));
736        assert_eq!(rendered["timestamp"], json!(12.5));
737        assert_eq!(rendered["delay"], json!(3.0));
738        assert_eq!(rendered["hash_ids"], json!([0, 1]));
739        assert!(rendered.get("input_tokens").is_none());
740        assert!(rendered.get("output_tokens").is_none());
741        assert!(rendered.get("created_time").is_none());
742        assert!(rendered.get("delay_ms").is_none());
743    }
744
745    #[test]
746    fn row_deserializes_with_missing_optional_fields() {
747        let raw = r#"{"output_length":2}"#;
748        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
749        assert_eq!(row.session_id, None);
750        assert_eq!(row.input_length, None);
751        assert_eq!(row.output_length, Some(2));
752        assert_eq!(row.hash_ids, None);
753        assert_eq!(row.timestamp, None);
754        assert_eq!(row.delay, None);
755        assert_eq!(row.priority, None);
756        assert_eq!(row.strict_priority, None);
757        assert_eq!(row.policy_class, None);
758        let rendered: Value = serde_json::to_value(&row).unwrap();
759        assert!(rendered.get("priority").is_none());
760        assert!(rendered.get("strict_priority").is_none());
761        assert!(rendered.get("policy_class").is_none());
762        assert!(rendered.get("request_id").is_none());
763        assert!(rendered.get("output_token_ids").is_none());
764    }
765
766    #[test]
767    fn row_round_trips_priorities() {
768        for priority in [Some(7), Some(0), Some(-3)] {
769            let raw = json!({
770                "output_length": 2,
771                "priority": priority,
772                "strict_priority": 9,
773                "policy_class": "latency"
774            });
775            let row: MooncakeRow = serde_json::from_value(raw).unwrap();
776            assert_eq!(row.priority, priority);
777            assert_eq!(row.strict_priority, Some(9));
778            assert_eq!(row.policy_class.as_deref(), Some("latency"));
779
780            let rendered: Value = serde_json::to_value(&row).unwrap();
781            assert_eq!(rendered["priority"], json!(priority.unwrap()));
782            assert_eq!(rendered["strict_priority"], json!(9));
783            assert_eq!(rendered["policy_class"], json!("latency"));
784        }
785    }
786
787    #[test]
788    fn agentic_v2_row_round_trips_typed_dependencies() {
789        let raw = r#"{"request_id":"r2","play_id":"play","session_id":"child","model":"model","input_length":4,"output_length":1,"hash_ids":[7],"not_before_ms":10.0,"dependencies":[{"request_id":"r1","trigger":"dispatch","delay_ms":3.0,"relation":"spawn"}]}"#;
790        let row: AgenticMooncakeRow = serde_json::from_str(raw).unwrap();
791
792        assert_eq!(row.play_id, "play");
793        assert_eq!(row.model, "model");
794        assert_eq!(row.dependencies.len(), 1);
795        assert_eq!(
796            row.dependencies[0].trigger,
797            AgenticDependencyTrigger::Dispatch
798        );
799        assert_eq!(
800            row.dependencies[0].relation,
801            AgenticDependencyRelation::Spawn
802        );
803        assert_eq!(row.dependencies[0].delay_ms, 3.0);
804        let rendered: Value = serde_json::to_value(&row).unwrap();
805        assert_eq!(rendered["not_before_ms"], json!(10.0));
806        assert_eq!(rendered["dependencies"][0]["trigger"], json!("dispatch"));
807    }
808
809    #[test]
810    fn agentic_row_rejects_header_fields_by_shape() {
811        let raw = r#"{"request_id":"r1","input_length":4,"output_length":1,"hash_ids":[0],"timestamp":10.0}"#;
812        assert!(serde_json::from_str::<AgenticMooncakeRow>(raw).is_err());
813    }
814
815    #[test]
816    fn writer_writes_rows_and_sidecar_jsonl() {
817        let temp = TempDir::new().unwrap();
818        let output = temp.path().join("trace.jsonl");
819        let sidecar = temp.path().join("trace.sidecar.jsonl");
820
821        let mut writer = MooncakeJsonlWriter::create(&output, Some(&sidecar)).unwrap();
822        writer
823            .write_row(&MooncakeRow {
824                session_id: Some("s".to_string()),
825                input_length: Some(2),
826                output_length: Some(1),
827                hash_ids: Some(vec![0]),
828                timestamp: Some(0.0),
829                delay: None,
830                ..Default::default()
831            })
832            .unwrap();
833        writer.write_sidecar(&json!({"k": "v"})).unwrap();
834        let stats = writer.finish().unwrap();
835
836        assert_eq!(stats.row_count, 1);
837        assert_eq!(stats.sidecar_count, 1);
838
839        let row_lines: Vec<Value> = std::fs::read_to_string(&output)
840            .unwrap()
841            .lines()
842            .map(|line| serde_json::from_str(line).unwrap())
843            .collect();
844        let sidecar_lines: Vec<Value> = std::fs::read_to_string(&sidecar)
845            .unwrap()
846            .lines()
847            .map(|line| serde_json::from_str(line).unwrap())
848            .collect();
849        assert_eq!(row_lines.len(), 1);
850        assert_eq!(sidecar_lines, vec![json!({"k": "v"})]);
851        assert_eq!(row_lines[0]["session_id"], json!("s"));
852        assert!(row_lines[0].get("delay").is_none());
853    }
854
855    #[test]
856    fn writer_writes_agentic_rows() {
857        let temp = TempDir::new().unwrap();
858        let output = temp.path().join("agentic.jsonl");
859        let mut writer = MooncakeJsonlWriter::create(&output, None).unwrap();
860        writer
861            .write_agentic_header(&AgenticMooncakeHeader {
862                schema: AGENTIC_MOONCAKE_SCHEMA.to_string(),
863                version: AGENTIC_MOONCAKE_VERSION,
864                block_size: 2,
865                hash_id_scope: AgenticHashIdScope::Local,
866                source: AgenticSourceProvenance {
867                    format: "test".to_string(),
868                    digest: "digest".to_string(),
869                },
870            })
871            .unwrap();
872        writer
873            .write_agentic_row(&AgenticMooncakeRow {
874                request_id: "r1".to_string(),
875                play_id: "play".to_string(),
876                session_id: "session".to_string(),
877                model: "model".to_string(),
878                input_length: Some(2),
879                output_length: Some(1),
880                hash_ids: Some(vec![0]),
881                not_before_ms: 0.0,
882                ..Default::default()
883            })
884            .unwrap();
885        let stats = writer.finish().unwrap();
886
887        assert_eq!(stats.row_count, 1);
888        let row_lines: Vec<Value> = std::fs::read_to_string(&output)
889            .unwrap()
890            .lines()
891            .map(|line| serde_json::from_str(line).unwrap())
892            .collect();
893        assert_eq!(row_lines.len(), 2);
894        assert_eq!(row_lines[0]["version"], json!(2));
895        assert_eq!(row_lines[1]["request_id"], json!("r1"));
896    }
897
898    #[test]
899    fn writer_without_sidecar_rejects_sidecar_writes() {
900        let temp = TempDir::new().unwrap();
901        let output = temp.path().join("trace.jsonl");
902        let mut writer = MooncakeJsonlWriter::create(&output, None).unwrap();
903        assert!(!writer.has_sidecar());
904        let err = writer.write_sidecar(&json!({})).unwrap_err();
905        assert!(err.to_string().contains("sidecar was not configured"));
906    }
907}
dynamo_data_gen/mooncake.rs

dynamo_data_gen/
mooncake.rs
