Skip to main content

atman_runtime/projection/
message_window.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use crate::event;
5use crate::event_log::reader::{parse_json_lines, read_event_envelopes};
6use crate::message::{Message, MessagePart};
7use crate::nodegraph;
8use crate::provider;
9use crate::session::SessionOpenError;
10use serde_json;
11
12#[derive(Debug, Clone)]
13pub enum TranscriptEntry {
14    Message {
15        message: Message,
16        flow_run_id: Option<String>,
17    },
18    CompactionSummary {
19        range_start: usize,
20        range_end: usize,
21        compacted_count: usize,
22        before_tokens: u64,
23        after_tokens: u64,
24        summary: String,
25        ts: Option<chrono::DateTime<chrono::Utc>>,
26    },
27    DiffPreview {
28        title: String,
29        old_content: Option<String>,
30        new_content: Option<String>,
31        unified_diff: Option<String>,
32    },
33    FlowGraph {
34        run_id: String,
35        flow_name: String,
36        graph: nodegraph::FlowGraph,
37        ts: Option<chrono::DateTime<chrono::Utc>>,
38    },
39    FlowStart {
40        run_id: String,
41        flow_name: String,
42        parent_run_id: Option<String>,
43        parent_node_id: Option<String>,
44        ts: Option<chrono::DateTime<chrono::Utc>>,
45    },
46    FlowNodeStart {
47        run_id: String,
48        node_id: String,
49        kind: nodegraph::NodeKind,
50        label: String,
51        parent_node_id: Option<String>,
52        ts: Option<chrono::DateTime<chrono::Utc>>,
53    },
54    FlowNodeEnd {
55        run_id: String,
56        node_id: String,
57        status: event::FlowNodeStatus,
58        output_preview: Option<String>,
59        ts: Option<chrono::DateTime<chrono::Utc>>,
60    },
61    ToolNode {
62        run_id: String,
63        parent_node_id: String,
64        tool_use_id: String,
65        tool_name: String,
66        args_preview: String,
67        ts: Option<chrono::DateTime<chrono::Utc>>,
68    },
69    FlowDone {
70        run_id: String,
71        ok: bool,
72        cancelled: bool,
73        ts: Option<chrono::DateTime<chrono::Utc>>,
74    },
75    LlmCall {
76        model: String,
77        usage: provider::TokenUsage,
78        wallclock_ms: u64,
79        ttft_ms: Option<u64>,
80        tokens_per_second: Option<f64>,
81        run_id: Option<event::FlowRunId>,
82        node_id: Option<String>,
83        ts: Option<chrono::DateTime<chrono::Utc>>,
84    },
85}
86
87pub fn replay_messages_from(path: &Path) -> Result<Vec<Message>, SessionOpenError> {
88    Ok(replay_messages_with_seq(path)?
89        .into_iter()
90        .map(|(_, msg)| msg)
91        .collect())
92}
93
94pub fn replay_messages_with_seq(path: &Path) -> Result<Vec<(u64, Message)>, SessionOpenError> {
95    let envelopes = read_event_envelopes(path)?;
96    Ok(envelopes.as_slice().to_messages_with_seq())
97}
98
99pub fn replay_all_messages_with_seq(path: &Path) -> Result<Vec<(u64, Message)>, SessionOpenError> {
100    let envelopes = read_event_envelopes(path)?;
101    Ok(envelopes
102        .iter()
103        .filter_map(|env| match &env.event {
104            crate::event::Event::UserMsg { message, .. }
105            | crate::event::Event::AssistantMsg { message, .. }
106            | crate::event::Event::ToolResultMsg { message, .. }
107            | crate::event::Event::SystemMsg { message, .. } => Some((env.seq, message.clone())),
108            _ => None,
109        })
110        .collect())
111}
112
113#[derive(Debug, Clone)]
114pub struct AttachmentPatch {
115    part_index: usize,
116    file_basename: String,
117    reason: String,
118}
119
120pub fn parse_ts(v: &serde_json::Value) -> Option<chrono::DateTime<chrono::Utc>> {
121    v.get("ts")?
122        .as_str()
123        .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
124        .map(|dt| dt.with_timezone(&chrono::Utc))
125}
126
127pub fn parse_context_compact_event(v: &serde_json::Value) -> Option<CompactReplayEvent> {
128    if v["type"].as_str() != Some("context_compact") {
129        return None;
130    }
131    Some(CompactReplayEvent {
132        range_start: v["compacted_range_start"].as_u64().unwrap_or(0) as usize,
133        range_end: v["compacted_range_end"].as_u64().unwrap_or(0) as usize,
134        replacement_msg_seq: v["replacement_msg_seq"].as_u64(),
135    })
136}
137
138#[derive(Debug, Clone)]
139pub struct CompactReplayEvent {
140    range_start: usize,
141    range_end: usize,
142    replacement_msg_seq: Option<u64>,
143}
144
145pub fn collect_attachment_patches(
146    values: &[serde_json::Value],
147) -> HashMap<u64, Vec<AttachmentPatch>> {
148    let mut map: HashMap<u64, Vec<AttachmentPatch>> = HashMap::new();
149    for v in values {
150        if v["type"].as_str() == Some("attachment_degraded") {
151            let Some(msg_seq) = v["message_seq"].as_u64() else {
152                continue;
153            };
154            let Some(part_index) = v["part_index"].as_u64() else {
155                continue;
156            };
157            let file_basename = v["file_basename"].as_str().unwrap_or("").to_string();
158            let reason = v["reason"].as_str().unwrap_or("degraded").to_string();
159            map.entry(msg_seq).or_default().push(AttachmentPatch {
160                part_index: part_index as usize,
161                file_basename,
162                reason,
163            });
164        }
165    }
166    map
167}
168
169pub fn apply_attachment_patches(msg: &mut Message, patches: &[AttachmentPatch]) {
170    for p in patches {
171        if let Some(part) = msg.parts.get_mut(p.part_index) {
172            *part = MessagePart::Text {
173                text: format!(
174                    "[attachment unavailable: {} — {}]",
175                    p.file_basename, p.reason
176                ),
177            };
178        }
179    }
180}
181
182pub fn replay_transcript_from(path: &Path) -> Result<Vec<TranscriptEntry>, SessionOpenError> {
183    let text = match std::fs::read_to_string(path) {
184        Ok(t) => t,
185        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
186        Err(e) => {
187            return Err(SessionOpenError::Replay {
188                path: path.to_path_buf(),
189                source: e,
190            });
191        }
192    };
193    let values = parse_json_lines(&text);
194    let patches = collect_attachment_patches(&values);
195    let mut out = Vec::new();
196    let mut msg_indices: Vec<usize> = Vec::new();
197    let mut msg_seqs: Vec<u64> = Vec::new();
198    for v in &values {
199        let ty = v["type"].as_str().unwrap_or("");
200        match ty {
201            "user_msg" | "assistant_msg" | "tool_result_msg" | "system_msg" => {
202                if let Some(m) = v.get("message")
203                    && let Ok(mut msg) = serde_json::from_value::<Message>(m.clone())
204                {
205                    let seq = v["seq"].as_u64().unwrap_or(0);
206                    if let Some(ps) = patches.get(&seq) {
207                        apply_attachment_patches(&mut msg, ps);
208                    }
209                    let flow_run_id = v["flow_run_id"].as_str().map(String::from);
210                    msg_indices.push(out.len());
211                    msg_seqs.push(seq);
212                    out.push(TranscriptEntry::Message {
213                        message: msg,
214                        flow_run_id,
215                    });
216                }
217            }
218            "context_compact" => {
219                let Some(event) = parse_context_compact_event(v) else {
220                    continue;
221                };
222                if event.range_start > event.range_end || event.range_end >= msg_indices.len() {
223                    continue;
224                }
225                let Some(replacement_seq) = event.replacement_msg_seq else {
226                    continue;
227                };
228                let Some(replacement_pos) = msg_seqs.iter().position(|seq| *seq == replacement_seq)
229                else {
230                    continue;
231                };
232                let replacement_out_idx = msg_indices[replacement_pos];
233                let replacement_entry = out.remove(replacement_out_idx);
234                let removed_out_start = msg_indices[event.range_start];
235                let removed_count = event.range_end - event.range_start + 1;
236                for _ in 0..removed_count {
237                    out.remove(removed_out_start);
238                }
239                msg_indices.drain(event.range_start..=event.range_end);
240                msg_seqs.drain(event.range_start..=event.range_end);
241                out.insert(removed_out_start, replacement_entry);
242                msg_indices.insert(event.range_start, removed_out_start);
243                msg_seqs.insert(event.range_start, replacement_seq);
244                for (i, ordinal_out_idx) in msg_indices.iter_mut().enumerate() {
245                    if i > event.range_start {
246                        *ordinal_out_idx =
247                            ordinal_out_idx.saturating_sub(removed_count.saturating_sub(1));
248                    }
249                }
250            }
251            "compaction_summary" => {
252                out.push(TranscriptEntry::CompactionSummary {
253                    range_start: v["range_start"].as_u64().unwrap_or(0) as usize,
254                    range_end: v["range_end"].as_u64().unwrap_or(0) as usize,
255                    compacted_count: v["compacted_count"].as_u64().unwrap_or(0) as usize,
256                    before_tokens: v["before_tokens"].as_u64().unwrap_or(0),
257                    after_tokens: v["after_tokens"].as_u64().unwrap_or(0),
258                    summary: v["summary"].as_str().unwrap_or("").to_string(),
259                    ts: parse_ts(v),
260                });
261            }
262            "diff_preview" => {
263                out.push(TranscriptEntry::DiffPreview {
264                    title: v["title"].as_str().unwrap_or("").to_string(),
265                    old_content: v["old_content"].as_str().map(String::from),
266                    new_content: v["new_content"].as_str().map(String::from),
267                    unified_diff: v["unified_diff"].as_str().map(String::from),
268                });
269            }
270            "flow_graph" => {
271                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
272                let flow_name = v
273                    .get("graph")
274                    .and_then(|g| g["flow_name"].as_str())
275                    .unwrap_or("")
276                    .to_string();
277                let ts = parse_ts(v);
278                if let Some(g) = v.get("graph")
279                    && let Ok(graph) = serde_json::from_value::<nodegraph::FlowGraph>(g.clone())
280                {
281                    out.push(TranscriptEntry::FlowGraph {
282                        run_id,
283                        flow_name,
284                        graph,
285                        ts,
286                    });
287                }
288            }
289            "flow_start" => {
290                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
291                let flow_name = v["flow_name"].as_str().unwrap_or("").to_string();
292                let parent_run_id = v["parent_run_id"].as_str().map(String::from);
293                let parent_node_id = v["parent_node_id"].as_str().map(String::from);
294                let ts = parse_ts(v);
295                out.push(TranscriptEntry::FlowStart {
296                    run_id,
297                    flow_name,
298                    parent_run_id,
299                    parent_node_id,
300                    ts,
301                });
302            }
303            "flow_node_start" => {
304                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
305                let node_id = v["node_id"].as_str().unwrap_or("").to_string();
306                let label = v["label"].as_str().unwrap_or(&node_id).to_string();
307                let parent_node_id = v["parent_node_id"].as_str().map(String::from);
308                let kind = v
309                    .get("kind")
310                    .and_then(|k| serde_json::from_value(k.clone()).ok())
311                    .unwrap_or(nodegraph::NodeKind::UserConfirm);
312                let ts = parse_ts(v);
313                out.push(TranscriptEntry::FlowNodeStart {
314                    run_id,
315                    node_id,
316                    kind,
317                    label,
318                    parent_node_id,
319                    ts,
320                });
321            }
322            "flow_node_end" => {
323                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
324                let node_id = v["node_id"].as_str().unwrap_or("").to_string();
325                let status: event::FlowNodeStatus = v
326                    .get("status")
327                    .and_then(|s| serde_json::from_value(s.clone()).ok())
328                    .unwrap_or(event::FlowNodeStatus::Ok);
329                let output_preview = v["output_preview"].as_str().map(String::from);
330                let ts = parse_ts(v);
331                out.push(TranscriptEntry::FlowNodeEnd {
332                    run_id,
333                    node_id,
334                    status,
335                    output_preview,
336                    ts,
337                });
338            }
339            "tool_node" => {
340                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
341                let parent_node_id = v["parent_node_id"].as_str().unwrap_or("").to_string();
342                let tool_use_id = v["tool_use_id"].as_str().unwrap_or("").to_string();
343                let tool_name = v["tool_name"].as_str().unwrap_or("").to_string();
344                let args_preview = v["args_preview"].as_str().unwrap_or("").to_string();
345                let ts = parse_ts(v);
346                out.push(TranscriptEntry::ToolNode {
347                    run_id,
348                    parent_node_id,
349                    tool_use_id,
350                    tool_name,
351                    args_preview,
352                    ts,
353                });
354            }
355            "flow_end" => {
356                let run_id = v["run_id"].as_str().unwrap_or("").to_string();
357                let ok = v["status"]["kind"].as_str() == Some("ok");
358                let cancelled = v["status"]["kind"].as_str() == Some("cancelled");
359                let ts = parse_ts(v);
360                out.push(TranscriptEntry::FlowDone {
361                    run_id,
362                    ok,
363                    cancelled,
364                    ts,
365                });
366            }
367            "llm_call" => {
368                let model = v["model"].as_str().unwrap_or("").to_string();
369                let usage: provider::TokenUsage = v
370                    .get("usage")
371                    .and_then(|u| serde_json::from_value(u.clone()).ok())
372                    .unwrap_or_default();
373                let wallclock_ms = v["wallclock_ms"].as_u64().unwrap_or(0);
374                let ttft_ms = v["ttft_ms"].as_u64();
375                let tokens_per_second = v["tokens_per_second"].as_f64();
376                let run_id = v["run_id"]
377                    .as_str()
378                    .and_then(|s| uuid::Uuid::parse_str(s).ok())
379                    .map(event::FlowRunId);
380                let node_id = v["node_id"].as_str().map(String::from);
381                let ts = parse_ts(v);
382                out.push(TranscriptEntry::LlmCall {
383                    model,
384                    usage,
385                    wallclock_ms,
386                    ttft_ms,
387                    tokens_per_second,
388                    run_id,
389                    node_id,
390                    ts,
391                });
392            }
393            _ => {}
394        }
395    }
396    Ok(out)
397}
398
399pub trait MessageProjection {
400    fn to_messages(&self) -> Vec<Message>;
401    fn to_messages_with_seq(&self) -> Vec<(u64, Message)>;
402}
403
404impl MessageProjection for [crate::event::EventEnvelope] {
405    fn to_messages(&self) -> Vec<Message> {
406        self.to_messages_with_seq()
407            .into_iter()
408            .map(|(_, msg)| msg)
409            .collect()
410    }
411
412    fn to_messages_with_seq(&self) -> Vec<(u64, Message)> {
413        let mut acc: Vec<(u64, Message)> = Vec::new();
414        for env in self {
415            apply_envelope_to_messages(env, &mut acc);
416        }
417        acc
418    }
419}
420
421pub(crate) fn apply_envelope_to_messages(
422    env: &crate::event::EventEnvelope,
423    acc: &mut Vec<(u64, Message)>,
424) {
425    match &env.event {
426        crate::event::Event::UserMsg { message, .. }
427        | crate::event::Event::AssistantMsg { message, .. }
428        | crate::event::Event::ToolResultMsg { message, .. }
429        | crate::event::Event::SystemMsg { message, .. } => {
430            acc.push((env.seq, message.clone()));
431        }
432        crate::event::Event::ContextCompact {
433            compacted_range_start,
434            compacted_range_end,
435            replacement_msg_seq,
436            summary_text,
437            after_tokens,
438            before_tokens,
439            ..
440        } => {
441            let range_start = *compacted_range_start as usize;
442            let range_end = *compacted_range_end as usize;
443            if range_start > range_end || range_end >= acc.len() {
444                return;
445            }
446            let Some(rep_seq) = replacement_msg_seq else {
447                return;
448            };
449            let Some(rep_idx) = acc.iter().position(|(s, _)| *s == *rep_seq) else {
450                return;
451            };
452            if *after_tokens >= *before_tokens {
453                return;
454            }
455            let replacement = acc.remove(rep_idx);
456            let removed_count = range_end - range_start + 1;
457            for _ in 0..removed_count {
458                acc.remove(range_start);
459            }
460            let insertion_idx = range_start.min(acc.len());
461            if let Some(summary) = summary_text {
462                acc.insert(
463                    insertion_idx,
464                    (
465                        *rep_seq,
466                        Message::system_compact_summary(
467                            crate::event::TurnId::now(),
468                            summary.clone(),
469                            range_start as u64,
470                            range_end as u64,
471                            removed_count,
472                        ),
473                    ),
474                );
475            } else {
476                acc.insert(insertion_idx, replacement);
477            }
478        }
479        crate::event::Event::AttachmentDegraded {
480            message_seq,
481            part_index,
482            file_basename,
483            reason,
484            ..
485        } => {
486            if let Some((_, msg)) = acc.iter_mut().find(|(s, _)| *s == *message_seq) {
487                if let Some(part) = msg.parts.get_mut(*part_index) {
488                    *part = MessagePart::Text {
489                        text: format!("[attachment unavailable: {} — {}]", file_basename, reason),
490                    };
491                }
492            }
493        }
494        _ => {}
495    }
496}