atman_runtime/event_log/
reader.rs

1use std::io::BufRead;
2use std::path::Path;
3
4use crate::context_plan::{ContextCallPurpose, ContextCallScope};
5use crate::event::{Event, EventEnvelope};
6use crate::session::{ContextSnapshot, ContextUsageBucket, SessionOpenError};
7
8#[cfg(test)]
9thread_local! {
10    static PARSE_ATTEMPTS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
11}
12
13#[cfg(test)]
14pub(crate) fn reset_parse_attempts() {
15    PARSE_ATTEMPTS.with(|attempts| attempts.set(0));
16}
17
18#[cfg(test)]
19pub(crate) fn parse_attempts() -> u64 {
20    PARSE_ATTEMPTS.with(std::cell::Cell::get)
21}
22
23#[cfg(test)]
24fn record_parse_attempt() {
25    PARSE_ATTEMPTS.with(|attempts| attempts.set(attempts.get().saturating_add(1)));
26}
27
28#[derive(Debug, Clone)]
29pub(crate) struct ReplayRecord {
30    pub envelope: EventEnvelope,
31    pub persisted_ts: Option<chrono::DateTime<chrono::Utc>>,
32}
33
34pub(crate) fn scan_replay_records<R: BufRead>(mut reader: R) -> std::io::Result<Vec<ReplayRecord>> {
35    let mut records = Vec::new();
36    let mut line = String::new();
37    loop {
38        line.clear();
39        if reader.read_line(&mut line)? == 0 {
40            break;
41        }
42        let text = line.trim();
43        if text.is_empty() {
44            continue;
45        }
46        #[cfg(test)]
47        record_parse_attempt();
48        let Ok(value) = serde_json::from_str::<serde_json::Value>(text) else {
49            continue;
50        };
51        let persisted_ts = value
52            .get("ts")
53            .and_then(serde_json::Value::as_str)
54            .and_then(|text| chrono::DateTime::parse_from_rfc3339(text).ok())
55            .map(|ts| ts.with_timezone(&chrono::Utc));
56        let Ok(envelope) = EventEnvelope::from_json_value(value) else {
57            continue;
58        };
59        records.push(ReplayRecord {
60            envelope,
61            persisted_ts,
62        });
63    }
64    Ok(records)
65}
66
67pub(crate) fn read_replay_records(path: &Path) -> Result<Vec<ReplayRecord>, SessionOpenError> {
68    let file = match std::fs::File::open(path) {
69        Ok(file) => file,
70        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
71        Err(source) => {
72            return Err(SessionOpenError::Replay {
73                path: path.to_path_buf(),
74                source,
75            });
76        }
77    };
78    scan_replay_records(std::io::BufReader::new(file)).map_err(|source| SessionOpenError::Replay {
79        path: path.to_path_buf(),
80        source,
81    })
82}
83
84pub fn read_event_envelopes(path: &Path) -> Result<Vec<EventEnvelope>, SessionOpenError> {
85    Ok(read_replay_records(path)?
86        .into_iter()
87        .map(|record| record.envelope)
88        .collect())
89}
90
91pub fn parse_json_lines(text: &str) -> Vec<serde_json::Value> {
92    text.lines()
93        .filter_map(|line| {
94            let text = line.trim();
95            if text.is_empty() {
96                None
97            } else {
98                #[cfg(test)]
99                record_parse_attempt();
100                serde_json::from_str::<serde_json::Value>(text).ok()
101            }
102        })
103        .collect()
104}
105
106pub fn find_last_seq(path: &Path) -> Result<Option<u64>, SessionOpenError> {
107    Ok(read_replay_records(path)?
108        .last()
109        .map(|record| record.envelope.seq))
110}
111
112pub(crate) fn context_snapshot_from_records(records: &[ReplayRecord]) -> ContextSnapshot {
113    let mut snapshot = ContextSnapshot::default();
114    for record in records {
115        apply_context_record(&mut snapshot, &record.envelope.event);
116    }
117    snapshot
118}
119
120fn apply_context_record(snapshot: &mut ContextSnapshot, event: &Event) {
121    let Event::LlmCall {
122        model,
123        provider,
124        context_call_purpose,
125        context_call_identity,
126        usage,
127        ttft_ms,
128        tokens_per_second,
129        run_id,
130        ..
131    } = event
132    else {
133        return;
134    };
135    if context_call_purpose.is_none() && context_call_identity.is_none() && run_id.is_none() {
136        return;
137    }
138    let purpose = context_call_purpose.unwrap_or_default();
139    let scope = context_call_identity.as_ref().map_or_else(
140        || {
141            if run_id.is_none() {
142                ContextCallScope::Detached
143            } else {
144                ContextCallScope::Root
145            }
146        },
147        |identity| identity.scope,
148    );
149    let total_input = usage
150        .input
151        .saturating_add(usage.cached_input)
152        .saturating_add(usage.cache_write);
153    snapshot.tokens_in = snapshot.tokens_in.saturating_add(total_input);
154    snapshot.tokens_out = snapshot.tokens_out.saturating_add(usage.output);
155    snapshot.cache_read = snapshot.cache_read.saturating_add(usage.cached_input);
156    snapshot.cache_write = snapshot.cache_write.saturating_add(usage.cache_write);
157
158    let bucket_idx = snapshot
159        .usage_buckets
160        .iter()
161        .position(|bucket| {
162            bucket.provider == *provider
163                && bucket.model == *model
164                && bucket.call_purpose == purpose
165                && bucket.call_scope == scope
166        })
167        .unwrap_or_else(|| {
168            snapshot.usage_buckets.push(ContextUsageBucket {
169                provider: provider.clone(),
170                model: model.clone(),
171                call_purpose: purpose,
172                call_scope: scope,
173                ..Default::default()
174            });
175            snapshot.usage_buckets.len() - 1
176        });
177    let bucket = &mut snapshot.usage_buckets[bucket_idx];
178    bucket.calls = bucket.calls.saturating_add(1);
179    bucket.tokens_in = bucket.tokens_in.saturating_add(total_input);
180    bucket.tokens_out = bucket.tokens_out.saturating_add(usage.output);
181    bucket.cache_read = bucket.cache_read.saturating_add(usage.cached_input);
182    bucket.cache_write = bucket.cache_write.saturating_add(usage.cache_write);
183
184    if purpose == ContextCallPurpose::General && scope == ContextCallScope::Root {
185        snapshot.provider.clone_from(provider);
186        snapshot.model.clone_from(model);
187        snapshot.last_ttft_ms = ttft_ms.unwrap_or(0);
188        snapshot.last_tokens_per_sec = tokens_per_second.unwrap_or(0.0);
189    }
190}
191
192pub fn replay_context_snapshot_from(path: &Path) -> ContextSnapshot {
193    read_replay_records(path)
194        .map(|records| context_snapshot_from_records(&records))
195        .unwrap_or_default()
196}
atman_runtime/event_log/reader.rs

atman_runtime/event_log/
reader.rs
