dynamo_bench/coding/
replay.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Usage-shaped replay tokens and request-trace rows shared by coding-agent exporters.
5
6use anyhow::{Result, bail};
7use dynamo_data_gen::{
8    AgenticDependencyRelation, AgenticDependencyTrigger, extend_sequence_hashes,
9};
10use rustc_hash::FxHashMap;
11use serde::Serialize;
12use serde_json::Value;
13
14pub const REQUEST_TRACE_SCHEMA: &str = "dynamo.request.trace.v1";
15pub const HARNESS_EVENT_SOURCE: &str = "harness";
16
17/// FNV-1a seed of one synthetic token stream.
18pub fn synthetic_stream_seed(stream_id: &str) -> u32 {
19    let mut hash = 0x811c_9dc5_u32;
20    for byte in stream_id.bytes() {
21        hash = (hash ^ u32::from(byte)).wrapping_mul(0x0100_0193);
22    }
23    hash
24}
25
26/// Deterministic stand-in for a token whose identity the source does not record.
27pub fn synthetic_token(stream_seed: u32, turn_index: usize, position: usize) -> u32 {
28    let hash = (stream_seed ^ turn_index as u32).wrapping_mul(0x0100_0193);
29    (hash ^ position as u32).wrapping_mul(0x0100_0193)
30}
31
32/// Builds `input_length` replay tokens whose first `cached_length` tokens are a cached prefix.
33///
34/// The prefix is copied from `prefix_source` as far as it reaches; the rest of the cached prefix
35/// and the uncached suffix are synthetic tokens of `turn_index - 1` and `turn_index`. Returns the
36/// tokens and how many were copied.
37pub fn usage_shaped_tokens(
38    stream_seed: u32,
39    turn_index: usize,
40    input_length: usize,
41    cached_length: usize,
42    prefix_source: Option<&[u32]>,
43) -> (Vec<u32>, usize) {
44    let cached_length = cached_length.min(input_length);
45    let mut tokens = Vec::with_capacity(input_length);
46    if let Some(prefix_source) = prefix_source {
47        tokens.extend_from_slice(&prefix_source[..cached_length.min(prefix_source.len())]);
48    }
49    let shared_length = tokens.len();
50    tokens.extend(
51        (shared_length..cached_length)
52            .map(|position| synthetic_token(stream_seed, turn_index.saturating_sub(1), position)),
53    );
54    let start = tokens.len();
55    tokens.extend(
56        (start..input_length).map(|position| synthetic_token(stream_seed, turn_index, position)),
57    );
58    (tokens, shared_length)
59}
60
61/// Key of the shared prefix stream for sessions of one harness, model, and working directory.
62pub fn prefix_pool_key(harness: &str, model: &str, cwd: &str) -> String {
63    format!("{harness}\u{1f}{model}\u{1f}{cwd}")
64}
65
66/// Replay tokens of one request and their sequence hashes, including a trailing partial block.
67#[derive(Debug, Default, Clone)]
68pub struct ReplayBase {
69    pub tokens: Vec<u32>,
70    pub hashes: Vec<u64>,
71}
72
73impl ReplayBase {
74    /// Hashes `tokens`, whose first `shared_tokens` tokens were copied from `source`.
75    ///
76    /// Sequence hashes chain from the first block, so every full block inside the copied prefix
77    /// keeps the source's hash and only the remaining blocks are hashed.
78    pub fn derive(
79        source: Option<&ReplayBase>,
80        shared_tokens: usize,
81        tokens: Vec<u32>,
82        block_size: usize,
83    ) -> Result<Self> {
84        let reused = source.map_or(&[][..], |source| {
85            let blocks = (shared_tokens / block_size).min(source.tokens.len() / block_size);
86            &source.hashes[..blocks]
87        });
88        // Reused hashes are trusted rather than recomputed, so check the copied boundary block.
89        let boundary = reused.len().saturating_sub(1) * block_size..reused.len() * block_size;
90        if let Some(source) = source
91            && source.tokens.get(boundary.clone()) != tokens.get(boundary)
92        {
93            bail!(
94                "replay prefix diverges from the {} reused hash blocks",
95                reused.len()
96            );
97        }
98        debug_assert_eq!(
99            source.map(|source| &source.tokens[..reused.len() * block_size]),
100            source.map(|_| &tokens[..reused.len() * block_size])
101        );
102        let hashes = extend_sequence_hashes(reused, &tokens, block_size)?;
103        Ok(Self { tokens, hashes })
104    }
105}
106
107/// Synthetic prefixes shared by sessions of one harness, model, and working directory.
108///
109/// Usage reports how much of a session's first request was already cached, but not which earlier
110/// request wrote it. The pool attributes that prefix to context such sessions share, such as the
111/// system prompt and tool definitions, so replay can reuse it across sessions.
112#[derive(Debug, Default)]
113pub struct PrefixPool {
114    streams: FxHashMap<String, ReplayBase>,
115}
116
117impl PrefixPool {
118    /// Returns the shared stream for `key`, extended to at least `length` tokens.
119    pub fn prefix(&mut self, key: &str, length: usize, block_size: usize) -> Result<&ReplayBase> {
120        if !self.streams.contains_key(key) {
121            self.streams.insert(key.to_string(), ReplayBase::default());
122        }
123        let stream = self.streams.get_mut(key).expect("stream was inserted");
124        if stream.tokens.len() < length {
125            let seed = synthetic_stream_seed(key);
126            let start = stream.tokens.len();
127            stream
128                .tokens
129                .extend((start..length).map(|position| synthetic_token(seed, 0, position)));
130            stream.hashes.truncate(start / block_size);
131            stream.hashes = extend_sequence_hashes(&stream.hashes, &stream.tokens, block_size)?;
132        }
133        Ok(stream)
134    }
135}
136
137/// Exporter-only edge from a request to an earlier request it could not start before.
138///
139/// Live request traces cannot know these edges; an exporter that reads complete sessions can.
140/// Lowering adds them to the per-session sequence edges it always derives.
141#[derive(Debug, Clone, PartialEq, Serialize)]
142pub struct ReplayDependency {
143    pub request_id: String,
144    pub relation: AgenticDependencyRelation,
145    pub trigger: AgenticDependencyTrigger,
146}
147
148#[derive(Serialize)]
149pub struct TraceLine<E> {
150    pub timestamp: u64,
151    pub event: E,
152}
153
154#[derive(Serialize)]
155pub struct AgentContext<'a> {
156    pub session_id: &'a str,
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub parent_session_id: Option<&'a str>,
159}
160
161#[derive(Serialize)]
162pub struct RequestEndEvent<'a> {
163    pub schema: &'static str,
164    pub event_type: &'static str,
165    pub event_time_unix_ms: u64,
166    pub event_source: &'static str,
167    pub agent_context: &'a AgentContext<'a>,
168    pub request: RequestFields<'a>,
169}
170
171#[derive(Serialize)]
172pub struct RequestFields<'a> {
173    pub request_id: &'a str,
174    pub model: &'a str,
175    pub input_tokens: usize,
176    pub output_tokens: usize,
177    pub request_received_ms: u64,
178    pub total_time_ms: f64,
179    pub replay: ReplayFields<'a>,
180    #[serde(skip_serializing_if = "Option::is_none")]
181    pub cached_tokens: Option<usize>,
182    #[serde(skip_serializing_if = "Option::is_none")]
183    pub claude: Option<Value>,
184}
185
186#[derive(Serialize)]
187pub struct ReplayFields<'a> {
188    pub trace_block_size: usize,
189    pub input_length: usize,
190    pub input_sequence_hashes: &'a [u64],
191    #[serde(skip_serializing_if = "<[_]>::is_empty")]
192    pub dependencies: &'a [ReplayDependency],
193}
194
195#[derive(Serialize)]
196pub struct ToolEvent<'a, M> {
197    pub schema: &'static str,
198    pub event_type: &'static str,
199    pub event_time_unix_ms: u64,
200    pub event_source: &'static str,
201    pub agent_context: &'a AgentContext<'a>,
202    pub tool: ToolFields<'a, M>,
203}
204
205#[derive(Serialize)]
206pub struct ToolFields<'a, M> {
207    pub tool_call_id: &'a str,
208    pub tool_class: &'a str,
209    #[serde(skip_serializing_if = "Option::is_none")]
210    pub claude: Option<M>,
211    pub started_at_unix_ms: u64,
212    pub ended_at_unix_ms: u64,
213    pub duration_ms: f64,
214    pub status: &'static str,
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub output_bytes: Option<usize>,
217    pub error_type: Option<&'static str>,
218}
219
220#[cfg(test)]
221mod tests {
222    use super::{PrefixPool, ReplayBase};
223    use dynamo_data_gen::sequence_hashes_for_tokens;
224
225    #[test]
226    fn pooled_prefixes_are_stable_as_the_stream_grows() {
227        let mut pool = PrefixPool::default();
228        let short = pool.prefix("claude|model|/repo", 100, 16).unwrap().tokens[..100].to_vec();
229        let long = pool.prefix("claude|model|/repo", 300, 16).unwrap();
230        assert_eq!(long.tokens[..100], short);
231        assert_eq!(
232            long.hashes,
233            sequence_hashes_for_tokens(&long.tokens, 16).unwrap()
234        );
235        let other = pool.prefix("claude|model|/other", 100, 16).unwrap();
236        assert_ne!(other.tokens[..100], short);
237    }
238
239    #[test]
240    fn derived_hashes_match_full_hashing() {
241        let source = ReplayBase {
242            tokens: (0..70).collect(),
243            hashes: sequence_hashes_for_tokens(&(0..70).collect::<Vec<_>>(), 16).unwrap(),
244        };
245        let mut tokens = source.tokens[..50].to_vec();
246        tokens.extend(1_000..1_040);
247        let derived = ReplayBase::derive(Some(&source), 50, tokens.clone(), 16).unwrap();
248        assert_eq!(
249            derived.hashes,
250            sequence_hashes_for_tokens(&tokens, 16).unwrap()
251        );
252        // Claiming more shared tokens than were copied must not reuse a mismatched block hash.
253        assert!(ReplayBase::derive(Some(&source), 64, tokens, 16).is_err());
254    }
255}
dynamo_bench/coding/replay.rs

dynamo_bench/coding/
replay.rs
