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