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