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        status,
130        run_id,
131        ..
132    } = event
133    else {
134        return;
135    };
136    if !matches!(status, crate::event::LlmCallStatus::Ok) {
137        return;
138    }
139    if context_call_purpose.is_none() && context_call_identity.is_none() && run_id.is_none() {
140        return;
141    }
142    let purpose = context_call_purpose.unwrap_or_default();
143    let scope = context_call_identity.as_ref().map_or_else(
144        || {
145            if run_id.is_none() {
146                ContextCallScope::Detached
147            } else {
148                ContextCallScope::Root
149            }
150        },
151        |identity| identity.scope,
152    );
153    let total_input = usage
154        .input
155        .saturating_add(usage.cached_input)
156        .saturating_add(usage.cache_write);
157    snapshot.tokens_in = snapshot.tokens_in.saturating_add(total_input);
158    snapshot.tokens_out = snapshot.tokens_out.saturating_add(usage.output);
159    snapshot.cache_read = snapshot.cache_read.saturating_add(usage.cached_input);
160    snapshot.cache_write = snapshot.cache_write.saturating_add(usage.cache_write);
161
162    let bucket_idx = snapshot
163        .usage_buckets
164        .iter()
165        .position(|bucket| {
166            bucket.provider == *provider
167                && bucket.model == *model
168                && bucket.call_purpose == purpose
169                && bucket.call_scope == scope
170        })
171        .unwrap_or_else(|| {
172            snapshot.usage_buckets.push(ContextUsageBucket {
173                provider: provider.clone(),
174                model: model.clone(),
175                call_purpose: purpose,
176                call_scope: scope,
177                ..Default::default()
178            });
179            snapshot.usage_buckets.len() - 1
180        });
181    let bucket = &mut snapshot.usage_buckets[bucket_idx];
182    bucket.calls = bucket.calls.saturating_add(1);
183    bucket.tokens_in = bucket.tokens_in.saturating_add(total_input);
184    bucket.tokens_out = bucket.tokens_out.saturating_add(usage.output);
185    bucket.cache_read = bucket.cache_read.saturating_add(usage.cached_input);
186    bucket.cache_write = bucket.cache_write.saturating_add(usage.cache_write);
187
188    if purpose == ContextCallPurpose::General && scope == ContextCallScope::Root {
189        snapshot.provider.clone_from(provider);
190        snapshot.model.clone_from(model);
191        snapshot.last_ttft_ms = ttft_ms.unwrap_or(0);
192        snapshot.last_tokens_per_sec = tokens_per_second.unwrap_or(0.0);
193    }
194}
195
196pub fn replay_context_snapshot_from(path: &Path) -> ContextSnapshot {
197    read_replay_records(path)
198        .map(|records| context_snapshot_from_records(&records))
199        .unwrap_or_default()
200}
atman_runtime/event_log/reader.rs

atman_runtime/event_log/
reader.rs
