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 (session scheduling, tokenization,
9//! parsing) lives elsewhere -- the Claude exporter in `dynamo-bench` is one
10//! such producer; the `dynamo-mocker` load generator is one such consumer.
11//!
12//! The [`MooncakeRow`] schema deliberately matches the externally-authored
13//! Mooncake trace format: `timestamp` and `delay` are `f64` milliseconds, and
14//! `input_length`/`output_length`/`timestamp`/`delay` accept the upstream
15//! aliases (`input_tokens`, `output_tokens`, `created_time`, `delay_ms`) on
16//! deserialization. Dynamo-produced traces always emit the canonical names.
17
18use anyhow::{Context, Result, bail};
19use bytemuck::cast_slice;
20use dynamo_tokens::compute_hash_v2;
21use rustc_hash::FxHashMap;
22use serde::{Deserialize, Serialize};
23use std::fs::File;
24use std::io::{BufWriter, Write};
25use std::path::Path;
26
27/// One row of a Mooncake replay trace.
28///
29/// `timestamp` is an absolute request arrival offset in milliseconds. Rows
30/// without a `session_id` are independent request arrivals. Rows that share a
31/// `session_id` are interpreted as closed-loop turns; later turns use `delay`
32/// or timestamp deltas relative to the previous row in that session.
33///
34/// The row type is `Serialize + Deserialize` so the same definition serves
35/// producers and consumers. Field-level aliases on deserialization accept the
36/// upstream Mooncake field names (`input_tokens`, `output_tokens`,
37/// `created_time`, `delay_ms`) without requiring producers to emit them.
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct MooncakeRow {
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)]
47    pub hash_ids: Option<Vec<u64>>,
48    #[serde(
49        default,
50        skip_serializing_if = "Option::is_none",
51        alias = "created_time"
52    )]
53    pub timestamp: Option<f64>,
54    #[serde(default, skip_serializing_if = "Option::is_none", alias = "delay_ms")]
55    pub delay: Option<f64>,
56}
57
58/// One row of an agentic Mooncake replay trace.
59///
60/// This format keeps the request/cache fields from [`MooncakeRow`] and adds a
61/// tiny workflow layer above them. `request_id` names the row. `wait_for` names
62/// request ids whose simulated completions must arrive before this row becomes
63/// eligible. Once all dependencies are satisfied, replay waits `delay` plus
64/// `tool_wait_ms` before dispatching the request. Rows with no dependencies
65/// use `timestamp` as their open-loop start time.
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct AgenticMooncakeRow {
68    pub request_id: String,
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub session_id: Option<String>,
71    #[serde(default, alias = "input_tokens")]
72    pub input_length: Option<usize>,
73    #[serde(default, alias = "output_tokens")]
74    pub output_length: Option<usize>,
75    #[serde(default)]
76    pub hash_ids: Option<Vec<u64>>,
77    #[serde(
78        default,
79        skip_serializing_if = "Option::is_none",
80        alias = "created_time"
81    )]
82    pub timestamp: Option<f64>,
83    #[serde(default, skip_serializing_if = "Option::is_none", alias = "delay_ms")]
84    pub delay: Option<f64>,
85    #[serde(default, skip_serializing_if = "Vec::is_empty")]
86    pub wait_for: Vec<String>,
87    #[serde(default, skip_serializing_if = "Vec::is_empty")]
88    pub branches: Vec<String>,
89    #[serde(default, skip_serializing_if = "Option::is_none")]
90    pub prefix_reset: Option<bool>,
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub tool_wait_ms: Option<f64>,
93    #[serde(default, skip_serializing_if = "Vec::is_empty")]
94    pub tool_events: Vec<AgenticToolEvent>,
95}
96
97impl AgenticMooncakeRow {
98    /// Return the total wait after all dependencies complete.
99    pub fn dependency_delay_ms(&self) -> f64 {
100        self.delay.unwrap_or(0.0) + self.tool_wait_ms.unwrap_or(0.0)
101    }
102}
103
104/// Harness tool span attributed to the LLM request that consumed it. Mirrors
105/// `tool_end` / `tool_error` fields from `dynamo.agent.trace.v1`.
106#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
107pub struct AgenticToolEvent {
108    pub tool_call_id: String,
109    pub tool_class: String,
110    pub started_at_unix_ms: u64,
111    pub ended_at_unix_ms: u64,
112    pub duration_ms: f64,
113    pub status: String,
114    #[serde(default, skip_serializing_if = "Option::is_none")]
115    pub output_bytes: Option<u64>,
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub output_tokens: Option<u64>,
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub error_type: Option<String>,
120}
121
122/// Maps sequence-aware block hashes to compact, stable `u64` ids.
123///
124/// The mapper is intentionally stateful and reusable across requests/turns: a
125/// block of tokens that appears at the same prefix position in two different
126/// requests will be assigned the same id. Equality of leading `hash_ids`
127/// between rows therefore signals shared prompt prefixes for replay purposes.
128///
129/// `hash_ids` here are workload identity labels, not literal Dynamo runtime
130/// KV-cache hashes. Producers should not try to reconcile them with a
131/// production cache.
132pub struct RollingHashIdMapper {
133    block_size: usize,
134    hash_to_id: FxHashMap<u64, u64>,
135    next_id: u64,
136}
137
138impl RollingHashIdMapper {
139    /// Create a new mapper for the given block size.
140    pub fn new(block_size: usize) -> Self {
141        Self {
142            block_size,
143            hash_to_id: FxHashMap::default(),
144            next_id: 0,
145        }
146    }
147
148    /// Block size that this mapper was constructed with.
149    pub fn block_size(&self) -> usize {
150        self.block_size
151    }
152
153    /// Hash a sequence of tokens into Mooncake `hash_ids`.
154    ///
155    /// Tokens are chunked by `block_size`; each block contributes one id. The
156    /// chained hash mixes the prior block's combined hash, so identical
157    /// prefixes across requests resolve to identical leading `hash_ids` once
158    /// the mapper has seen them.
159    pub fn hash_token_blocks(&mut self, tokens: &[u32]) -> Vec<u64> {
160        hash_token_blocks(self, tokens)
161    }
162
163    /// Map precomputed sequence-aware block hashes into compact Mooncake IDs.
164    ///
165    /// This is useful for producers that record stable block hashes in the
166    /// serving path and only compact them during offline trace conversion.
167    pub fn ids_for_sequence_hashes(&mut self, sequence_hashes: &[u64]) -> Vec<u64> {
168        ids_for_sequence_hashes(self, sequence_hashes)
169    }
170}
171
172/// Token-block hashing helper for the Mooncake replay schema.
173///
174/// Splits `tokens` into chunks of `mapper.block_size()`, computes a chained
175/// hash per block, and returns the compact ids assigned by `mapper`. Mirrors
176/// [`RollingHashIdMapper::hash_token_blocks`] as a free function so callers
177/// that already hold a mutable mapper reference can invoke it without
178/// re-borrowing.
179pub fn hash_token_blocks(mapper: &mut RollingHashIdMapper, tokens: &[u32]) -> Vec<u64> {
180    let block_size = mapper.block_size;
181    let mut hash_ids = Vec::with_capacity(tokens.len().div_ceil(block_size));
182    let mut parent_hash = 0_u64;
183    for block in tokens.chunks(block_size) {
184        let block_hash = compute_hash_v2(cast_slice(block), 0);
185        let combined_hash = compute_hash_v2(&block_hash.to_be_bytes(), parent_hash);
186        let id = *mapper.hash_to_id.entry(combined_hash).or_insert_with(|| {
187            let next_id = mapper.next_id;
188            mapper.next_id += 1;
189            next_id
190        });
191        hash_ids.push(id);
192        parent_hash = combined_hash;
193    }
194    hash_ids
195}
196
197/// Map stable sequence hashes to compact Mooncake IDs with a shared mapper.
198pub fn ids_for_sequence_hashes(
199    mapper: &mut RollingHashIdMapper,
200    sequence_hashes: &[u64],
201) -> Vec<u64> {
202    sequence_hashes
203        .iter()
204        .map(|sequence_hash| {
205            *mapper.hash_to_id.entry(*sequence_hash).or_insert_with(|| {
206                let next_id = mapper.next_id;
207                mapper.next_id += 1;
208                next_id
209            })
210        })
211        .collect()
212}
213
214/// Counters for what a [`MooncakeJsonlWriter`] has emitted.
215#[derive(Debug, Clone, Copy, Default)]
216pub struct WriterStats {
217    pub row_count: usize,
218    pub sidecar_count: usize,
219}
220
221/// JSONL writer for Mooncake rows plus an optional sidecar stream.
222///
223/// The sidecar stream is configured at construction time. Producers that do
224/// not emit sidecar metadata pass `None` for `sidecar_path` and never call
225/// [`Self::write_sidecar`]. When a sidecar path is configured, callers are
226/// responsible for choosing the path -- this writer does not enforce a naming
227/// convention.
228pub struct MooncakeJsonlWriter {
229    output: BufWriter<File>,
230    sidecar: Option<BufWriter<File>>,
231    stats: WriterStats,
232}
233
234impl MooncakeJsonlWriter {
235    /// Create a writer at `output_path`, optionally with a paired sidecar
236    /// JSONL file at `sidecar_path`. Parent directories are created as needed.
237    pub fn create(output_path: &Path, sidecar_path: Option<&Path>) -> Result<Self> {
238        if let Some(parent) = output_path.parent() {
239            std::fs::create_dir_all(parent)?;
240        }
241        let output = BufWriter::new(
242            File::create(output_path)
243                .with_context(|| format!("failed to create {}", output_path.display()))?,
244        );
245        let sidecar = if let Some(path) = sidecar_path {
246            if let Some(parent) = path.parent() {
247                std::fs::create_dir_all(parent)?;
248            }
249            Some(BufWriter::new(File::create(path).with_context(|| {
250                format!("failed to create {}", path.display())
251            })?))
252        } else {
253            None
254        };
255        Ok(Self {
256            output,
257            sidecar,
258            stats: WriterStats::default(),
259        })
260    }
261
262    /// Append one Mooncake row.
263    pub fn write_row(&mut self, row: &MooncakeRow) -> Result<()> {
264        serde_json::to_writer(&mut self.output, row)?;
265        self.output.write_all(b"\n")?;
266        self.stats.row_count += 1;
267        Ok(())
268    }
269
270    /// Append one agentic Mooncake row.
271    pub fn write_agentic_row(&mut self, row: &AgenticMooncakeRow) -> Result<()> {
272        serde_json::to_writer(&mut self.output, row)?;
273        self.output.write_all(b"\n")?;
274        self.stats.row_count += 1;
275        Ok(())
276    }
277
278    /// Append one sidecar entry. Errors if no sidecar was configured.
279    pub fn write_sidecar<S: Serialize>(&mut self, sidecar: &S) -> Result<()> {
280        let writer = self
281            .sidecar
282            .as_mut()
283            .ok_or_else(|| anyhow::anyhow!("sidecar was not configured for this writer"))?;
284        serde_json::to_writer(writer, sidecar)?;
285        let writer = self.sidecar.as_mut().unwrap();
286        writer.write_all(b"\n")?;
287        self.stats.sidecar_count += 1;
288        Ok(())
289    }
290
291    /// True if a sidecar stream is configured.
292    pub fn has_sidecar(&self) -> bool {
293        self.sidecar.is_some()
294    }
295
296    /// Snapshot of how many rows and sidecar entries have been written so far.
297    pub fn stats(&self) -> WriterStats {
298        self.stats
299    }
300
301    /// Flush both streams and return the final stats.
302    pub fn finish(mut self) -> Result<WriterStats> {
303        self.output.flush()?;
304        if let Some(sidecar) = self.sidecar.as_mut() {
305            sidecar.flush()?;
306        }
307        Ok(self.stats)
308    }
309}
310
311/// Create both files empty (touch-equivalent), preserving directory creation
312/// semantics for callers that want a "no rows produced" outcome to still emit
313/// well-formed (empty) JSONL files.
314pub fn write_empty_files(output_path: &Path, sidecar_path: Option<&Path>) -> Result<()> {
315    if let Some(parent) = output_path.parent() {
316        std::fs::create_dir_all(parent)?;
317    }
318    File::create(output_path)
319        .with_context(|| format!("failed to create {}", output_path.display()))?;
320    if let Some(path) = sidecar_path {
321        if let Some(parent) = path.parent() {
322            std::fs::create_dir_all(parent)?;
323        }
324        File::create(path).with_context(|| format!("failed to create {}", path.display()))?;
325    }
326    Ok(())
327}
328
329/// Sentinel used by callers that want to bail when neither block_size nor
330/// worker count is allowed to be zero. Producers may also enforce this on
331/// their own configuration types.
332pub fn require_positive(name: &str, value: usize) -> Result<()> {
333    if value == 0 {
334        bail!("{name} must be greater than 0");
335    }
336    Ok(())
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342    use serde_json::{Value, json};
343    use tempfile::TempDir;
344
345    #[test]
346    fn shared_prefix_yields_shared_leading_hash_ids() {
347        let mut mapper = RollingHashIdMapper::new(2);
348        let prefix = vec![1u32, 2, 3, 4];
349        let extended = vec![1u32, 2, 3, 4, 5, 6];
350
351        let prefix_ids = mapper.hash_token_blocks(&prefix);
352        let extended_ids = mapper.hash_token_blocks(&extended);
353
354        assert_eq!(prefix_ids.len(), 2);
355        assert_eq!(extended_ids.len(), 3);
356        assert_eq!(extended_ids[..2], prefix_ids[..]);
357    }
358
359    #[test]
360    fn mapper_state_is_reused_across_requests() {
361        let mut mapper = RollingHashIdMapper::new(4);
362        let request_a = vec![10u32, 20, 30, 40, 50, 60, 70, 80];
363        let request_b = vec![10u32, 20, 30, 40, 50, 60, 70, 80];
364        let request_c = vec![10u32, 20, 30, 40, 99, 99, 99, 99];
365
366        let ids_a = mapper.hash_token_blocks(&request_a);
367        let ids_b = mapper.hash_token_blocks(&request_b);
368        let ids_c = mapper.hash_token_blocks(&request_c);
369
370        assert_eq!(ids_a, ids_b);
371        assert_eq!(ids_c[0], ids_a[0], "shared first block should keep its id");
372        assert_ne!(
373            ids_c[1], ids_a[1],
374            "diverging tail block must get a fresh id"
375        );
376    }
377
378    #[test]
379    fn free_function_and_method_agree() {
380        let mut mapper_a = RollingHashIdMapper::new(2);
381        let mut mapper_b = RollingHashIdMapper::new(2);
382        let tokens = vec![7u32, 8, 9, 10, 11];
383
384        let via_method = mapper_a.hash_token_blocks(&tokens);
385        let via_function = hash_token_blocks(&mut mapper_b, &tokens);
386
387        assert_eq!(via_method, via_function);
388    }
389
390    #[test]
391    fn empty_token_input_yields_empty_hash_ids() {
392        let mut mapper = RollingHashIdMapper::new(4);
393        assert!(mapper.hash_token_blocks(&[]).is_empty());
394    }
395
396    #[test]
397    fn precomputed_sequence_hashes_map_to_stable_ids() {
398        let mut mapper = RollingHashIdMapper::new(64);
399
400        let first = mapper.ids_for_sequence_hashes(&[101, 202, 303]);
401        let second = mapper.ids_for_sequence_hashes(&[101, 202, 404]);
402
403        assert_eq!(first[..2], second[..2]);
404        assert_ne!(first[2], second[2]);
405    }
406
407    #[test]
408    fn row_omits_timestamp_and_delay_when_absent() {
409        let row = MooncakeRow {
410            session_id: Some("s".to_string()),
411            input_length: Some(4),
412            output_length: Some(1),
413            hash_ids: Some(vec![0, 1]),
414            timestamp: None,
415            delay: None,
416        };
417        let rendered: Value = serde_json::to_value(&row).unwrap();
418        assert!(rendered.get("timestamp").is_none());
419        assert!(rendered.get("delay").is_none());
420        assert_eq!(rendered["hash_ids"], json!([0, 1]));
421    }
422
423    #[test]
424    fn row_serializes_optional_fields_when_set() {
425        let with_timestamp = MooncakeRow {
426            session_id: Some("s".to_string()),
427            input_length: Some(4),
428            output_length: Some(1),
429            hash_ids: Some(vec![]),
430            timestamp: Some(0.0),
431            delay: None,
432        };
433        let with_delay = MooncakeRow {
434            session_id: Some("s".to_string()),
435            input_length: Some(4),
436            output_length: Some(1),
437            hash_ids: Some(vec![]),
438            timestamp: None,
439            delay: Some(123.0),
440        };
441        let v_ts: Value = serde_json::to_value(&with_timestamp).unwrap();
442        let v_dl: Value = serde_json::to_value(&with_delay).unwrap();
443        assert_eq!(v_ts["timestamp"], json!(0.0));
444        assert!(v_ts.get("delay").is_none());
445        assert_eq!(v_dl["delay"], json!(123.0));
446        assert!(v_dl.get("timestamp").is_none());
447    }
448
449    #[test]
450    fn row_deserializes_canonical_field_names() {
451        let raw = r#"{"session_id":"s","input_length":4,"output_length":1,"hash_ids":[0,1],"timestamp":12.5,"delay":3.0}"#;
452        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
453        assert_eq!(row.session_id.as_deref(), Some("s"));
454        assert_eq!(row.input_length, Some(4));
455        assert_eq!(row.output_length, Some(1));
456        assert_eq!(row.hash_ids, Some(vec![0, 1]));
457        assert_eq!(row.timestamp, Some(12.5));
458        assert_eq!(row.delay, Some(3.0));
459    }
460
461    #[test]
462    fn row_deserializes_upstream_mooncake_aliases() {
463        let raw = r#"{"input_tokens":4,"output_tokens":1,"hash_ids":[0,1],"created_time":12.5,"delay_ms":3.0}"#;
464        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
465        assert_eq!(row.input_length, Some(4));
466        assert_eq!(row.output_length, Some(1));
467        assert_eq!(row.timestamp, Some(12.5));
468        assert_eq!(row.delay, Some(3.0));
469    }
470
471    #[test]
472    fn row_deserializes_with_missing_optional_fields() {
473        let raw = r#"{"output_length":2}"#;
474        let row: MooncakeRow = serde_json::from_str(raw).unwrap();
475        assert_eq!(row.session_id, None);
476        assert_eq!(row.input_length, None);
477        assert_eq!(row.output_length, Some(2));
478        assert_eq!(row.hash_ids, None);
479        assert_eq!(row.timestamp, None);
480        assert_eq!(row.delay, None);
481    }
482
483    #[test]
484    fn agentic_row_defaults_workflow_fields() {
485        let raw = r#"{"request_id":"r1","input_length":4,"output_length":1,"hash_ids":[0,1],"timestamp":10.0}"#;
486        let row: AgenticMooncakeRow = serde_json::from_str(raw).unwrap();
487
488        assert_eq!(row.request_id, "r1");
489        assert!(row.wait_for.is_empty());
490        assert!(row.branches.is_empty());
491        assert_eq!(row.prefix_reset, None);
492        assert_eq!(row.dependency_delay_ms(), 0.0);
493    }
494
495    #[test]
496    fn agentic_row_delay_includes_tool_wait() {
497        let row = AgenticMooncakeRow {
498            request_id: "r2".to_string(),
499            session_id: Some("trajectory-a".to_string()),
500            input_length: Some(4),
501            output_length: Some(1),
502            hash_ids: Some(vec![0, 1]),
503            timestamp: Some(20.0),
504            delay: Some(3.0),
505            wait_for: vec!["r1".to_string()],
506            branches: vec!["r3".to_string()],
507            prefix_reset: Some(false),
508            tool_wait_ms: Some(7.0),
509            tool_events: Vec::new(),
510        };
511
512        assert_eq!(row.dependency_delay_ms(), 10.0);
513        let rendered: Value = serde_json::to_value(&row).unwrap();
514        assert_eq!(rendered["request_id"], json!("r2"));
515        assert_eq!(rendered["wait_for"], json!(["r1"]));
516        assert_eq!(rendered["branches"], json!(["r3"]));
517        assert_eq!(rendered["tool_wait_ms"], json!(7.0));
518        assert!(rendered.get("tool_events").is_none());
519    }
520
521    #[test]
522    fn agentic_row_round_trips_tool_events() {
523        let row = AgenticMooncakeRow {
524            request_id: "r1".to_string(),
525            session_id: Some("trajectory-a".to_string()),
526            input_length: Some(4),
527            output_length: Some(1),
528            hash_ids: Some(vec![0, 1]),
529            timestamp: Some(0.0),
530            delay: Some(0.0),
531            wait_for: Vec::new(),
532            branches: Vec::new(),
533            prefix_reset: Some(true),
534            tool_wait_ms: Some(8.0),
535            tool_events: vec![AgenticToolEvent {
536                tool_call_id: "call-1".to_string(),
537                tool_class: "web_search".to_string(),
538                started_at_unix_ms: 1_000,
539                ended_at_unix_ms: 1_008,
540                duration_ms: 8.0,
541                status: "succeeded".to_string(),
542                output_bytes: Some(512),
543                output_tokens: None,
544                error_type: None,
545            }],
546        };
547
548        let rendered = serde_json::to_string(&row).unwrap();
549        let decoded: AgenticMooncakeRow = serde_json::from_str(&rendered).unwrap();
550        assert_eq!(decoded.tool_events.len(), 1);
551        assert_eq!(decoded.tool_events[0].tool_class, "web_search");
552        assert_eq!(decoded.tool_events[0].output_bytes, Some(512));
553    }
554
555    #[test]
556    fn writer_writes_rows_and_sidecar_jsonl() {
557        let temp = TempDir::new().unwrap();
558        let output = temp.path().join("trace.jsonl");
559        let sidecar = temp.path().join("trace.sidecar.jsonl");
560
561        let mut writer = MooncakeJsonlWriter::create(&output, Some(&sidecar)).unwrap();
562        writer
563            .write_row(&MooncakeRow {
564                session_id: Some("s".to_string()),
565                input_length: Some(2),
566                output_length: Some(1),
567                hash_ids: Some(vec![0]),
568                timestamp: Some(0.0),
569                delay: None,
570            })
571            .unwrap();
572        writer.write_sidecar(&json!({"k": "v"})).unwrap();
573        let stats = writer.finish().unwrap();
574
575        assert_eq!(stats.row_count, 1);
576        assert_eq!(stats.sidecar_count, 1);
577
578        let row_lines: Vec<Value> = std::fs::read_to_string(&output)
579            .unwrap()
580            .lines()
581            .map(|line| serde_json::from_str(line).unwrap())
582            .collect();
583        let sidecar_lines: Vec<Value> = std::fs::read_to_string(&sidecar)
584            .unwrap()
585            .lines()
586            .map(|line| serde_json::from_str(line).unwrap())
587            .collect();
588        assert_eq!(row_lines.len(), 1);
589        assert_eq!(sidecar_lines, vec![json!({"k": "v"})]);
590        assert_eq!(row_lines[0]["session_id"], json!("s"));
591        assert!(row_lines[0].get("delay").is_none());
592    }
593
594    #[test]
595    fn writer_writes_agentic_rows() {
596        let temp = TempDir::new().unwrap();
597        let output = temp.path().join("agentic.jsonl");
598        let mut writer = MooncakeJsonlWriter::create(&output, None).unwrap();
599        writer
600            .write_agentic_row(&AgenticMooncakeRow {
601                request_id: "r1".to_string(),
602                session_id: None,
603                input_length: Some(2),
604                output_length: Some(1),
605                hash_ids: Some(vec![0]),
606                timestamp: Some(0.0),
607                delay: None,
608                wait_for: Vec::new(),
609                branches: Vec::new(),
610                prefix_reset: Some(true),
611                tool_wait_ms: None,
612                tool_events: Vec::new(),
613            })
614            .unwrap();
615        let stats = writer.finish().unwrap();
616
617        assert_eq!(stats.row_count, 1);
618        let row_lines: Vec<Value> = std::fs::read_to_string(&output)
619            .unwrap()
620            .lines()
621            .map(|line| serde_json::from_str(line).unwrap())
622            .collect();
623        assert_eq!(row_lines[0]["request_id"], json!("r1"));
624        assert_eq!(row_lines[0]["prefix_reset"], json!(true));
625    }
626
627    #[test]
628    fn writer_without_sidecar_rejects_sidecar_writes() {
629        let temp = TempDir::new().unwrap();
630        let output = temp.path().join("trace.jsonl");
631        let mut writer = MooncakeJsonlWriter::create(&output, None).unwrap();
632        assert!(!writer.has_sidecar());
633        let err = writer.write_sidecar(&json!({})).unwrap_err();
634        assert!(err.to_string().contains("sidecar was not configured"));
635    }
636}
dynamo_data_gen/mooncake.rs

dynamo_data_gen/
mooncake.rs
