1use std::collections::HashSet;
2
3use crate::session_model::{ConversationRecord, ProtocolEvent, SessionHistoryRecord};
4use crate::{Message, MessageOrigin, MessageRole, MessageSequence, Part};
5
6#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
7pub struct ChronologicalProjection {
8 entries: Vec<ChronologicalEntry>,
9}
10
11#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
12pub struct ChronologicalEntry {
13 pub index: usize,
14 pub payload: ChronologicalPayload,
15}
16
17#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
18#[serde(tag = "kind", rename_all = "snake_case")]
19#[allow(clippy::large_enum_variant)]
21pub enum ChronologicalPayload {
22 Message(Message),
23 ProtocolEvent(crate::session_model::ProtocolEvent),
24}
25
26#[derive(Clone, Copy, Debug)]
27pub struct BorrowedChronologicalEntry<'a> {
28 pub index: usize,
29 pub payload: BorrowedChronologicalPayload<'a>,
30}
31
32#[derive(Clone, Copy, Debug)]
33pub enum BorrowedChronologicalPayload<'a> {
34 Message(BorrowedChronologicalMessage<'a>),
35 ProtocolEvent(&'a ProtocolEvent),
36}
37
38#[derive(Clone, Copy, Debug)]
39pub struct BorrowedChronologicalMessage<'a> {
40 pub id: &'a str,
41 pub role: MessageRole,
42 pub parts: &'a [Part],
43 pub origin: Option<&'a MessageOrigin>,
44}
45
46impl<'a> BorrowedChronologicalMessage<'a> {
47 fn from_message(message: &'a Message) -> Self {
48 Self {
49 id: &message.id,
50 role: message.role,
51 parts: message.parts.as_slice(),
52 origin: message.origin.as_ref(),
53 }
54 }
55
56 fn from_record(record: &'a ConversationRecord) -> Self {
57 Self {
58 id: &record.id,
59 role: record.role,
60 parts: record.parts.as_slice(),
61 origin: record.origin.as_ref(),
62 }
63 }
64
65 pub fn is_transient(&self) -> bool {
66 matches!(
67 self.origin,
68 Some(MessageOrigin::Plugin {
69 transient: true,
70 ..
71 })
72 )
73 }
74
75 fn to_owned(self) -> Message {
76 Message {
77 id: self.id.to_string(),
78 role: self.role,
79 parts: std::sync::Arc::new(self.parts.to_vec()),
80 origin: self.origin.cloned(),
81 }
82 }
83}
84
85impl ChronologicalProjection {
86 pub(crate) fn from_read_model(read_model: &crate::session_graph::SessionReadModel) -> Self {
87 Self::from_active_read(
88 read_model.active_events.as_slice(),
89 read_model.messages.as_slice(),
90 )
91 }
92
93 pub fn from_turn_view(events: &[SessionHistoryRecord], messages: &MessageSequence) -> Self {
94 Self::from_active_read(events, messages.as_slice())
95 }
96
97 fn from_active_read(active_events: &[SessionHistoryRecord], messages: &[Message]) -> Self {
98 let mut projection = Self::default();
99 projection
100 .entries
101 .reserve(active_events.len().saturating_add(messages.len()));
102 visit_active_read(active_events, messages, |entry| {
103 projection.push(match entry.payload {
104 BorrowedChronologicalPayload::Message(message) => {
105 ChronologicalPayload::Message(message.to_owned())
106 }
107 BorrowedChronologicalPayload::ProtocolEvent(event) => {
108 ChronologicalPayload::ProtocolEvent(event.clone())
109 }
110 });
111 });
112 projection
113 }
114
115 fn push(&mut self, payload: ChronologicalPayload) {
116 let index = self.entries.len();
117 self.entries.push(ChronologicalEntry { index, payload });
118 }
119
120 pub fn entries(&self) -> &[ChronologicalEntry] {
121 self.entries.as_slice()
122 }
123
124 pub fn into_entries(self) -> Vec<ChronologicalEntry> {
125 self.entries
126 }
127}
128
129pub fn visit_turn_view<'a>(
130 events: &'a [SessionHistoryRecord],
131 messages: &'a MessageSequence,
132 visit: impl FnMut(BorrowedChronologicalEntry<'a>),
133) {
134 visit_active_read(events, messages.as_slice(), visit);
135}
136
137fn visit_active_read<'a>(
138 active_events: &'a [SessionHistoryRecord],
139 messages: &'a [Message],
140 mut visit: impl FnMut(BorrowedChronologicalEntry<'a>),
141) {
142 if active_events.is_empty() {
143 visit_transcript(messages, visit);
144 return;
145 }
146
147 let mut index = 0;
148 let mut seen_messages = HashSet::new();
149
150 for event in active_events {
151 match event {
152 SessionHistoryRecord::Conversation(record) => {
153 let message = BorrowedChronologicalMessage::from_record(record);
154 if !message.is_transient() && seen_messages.insert(message.id.to_string()) {
155 visit(BorrowedChronologicalEntry {
156 index,
157 payload: BorrowedChronologicalPayload::Message(message),
158 });
159 index += 1;
160 }
161 }
162 SessionHistoryRecord::Protocol(event) => {
163 visit(BorrowedChronologicalEntry {
164 index,
165 payload: BorrowedChronologicalPayload::ProtocolEvent(event),
166 });
167 index += 1;
168 }
169 }
170 }
171
172 seen_messages.reserve(messages.len());
173 for message in messages {
174 let message = BorrowedChronologicalMessage::from_message(message);
175 if !message.is_transient() && seen_messages.insert(message.id.to_string()) {
176 visit(BorrowedChronologicalEntry {
177 index,
178 payload: BorrowedChronologicalPayload::Message(message),
179 });
180 index += 1;
181 }
182 }
183}
184
185fn visit_transcript<'a>(
186 messages: &'a [Message],
187 mut visit: impl FnMut(BorrowedChronologicalEntry<'a>),
188) {
189 for (index, message) in messages
190 .iter()
191 .filter(|message| !message.is_transient())
192 .enumerate()
193 {
194 visit(BorrowedChronologicalEntry {
195 index,
196 payload: BorrowedChronologicalPayload::Message(
197 BorrowedChronologicalMessage::from_message(message),
198 ),
199 });
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 use super::*;
206 use crate::session_model::ConversationRecord;
207 use crate::{PartKind, PruneState, shared_parts};
208
209 fn text_message(id: &str, role: MessageRole, text: &str) -> Message {
210 Message {
211 id: id.to_string(),
212 role,
213 parts: shared_parts(vec![Part {
214 id: format!("{id}.p0"),
215 kind: PartKind::Text,
216 content: text.to_string(),
217 attachment: None,
218 tool_call_id: None,
219 tool_name: None,
220 tool_replay: None,
221 prune_state: PruneState::Intact,
222 reasoning_meta: None,
223 response_meta: None,
224 }]),
225 origin: None,
226 }
227 }
228
229 fn tool_result_message(id: &str, call_id: &str) -> Message {
230 let mut message = text_message(id, MessageRole::User, "tool result");
231 std::sync::Arc::make_mut(&mut message.parts)[0].tool_call_id = Some(call_id.to_string());
232 message
233 }
234
235 fn transient_message(id: &str) -> Message {
236 let mut message = text_message(id, MessageRole::System, "transient");
237 message.origin = Some(MessageOrigin::Plugin {
238 plugin_id: "test".to_string(),
239 transient: true,
240 });
241 message
242 }
243
244 fn protocol_event(value: &str) -> ProtocolEvent {
245 ProtocolEvent {
246 plugin_id: "test_protocol".to_string(),
247 payload: serde_json::json!({ "value": value }),
248 }
249 }
250
251 fn owned_summary(projection: &ChronologicalProjection) -> Vec<String> {
252 projection
253 .entries()
254 .iter()
255 .map(|entry| match &entry.payload {
256 ChronologicalPayload::Message(message) => {
257 format!("{}:message:{}", entry.index, message.id)
258 }
259 ChronologicalPayload::ProtocolEvent(event) => {
260 format!("{}:protocol:{}", entry.index, event.payload)
261 }
262 })
263 .collect()
264 }
265
266 fn borrowed_summary(
267 events: &[SessionHistoryRecord],
268 messages: &MessageSequence,
269 ) -> Vec<String> {
270 let mut summary = Vec::new();
271 visit_turn_view(events, messages, |entry| {
272 summary.push(match entry.payload {
273 BorrowedChronologicalPayload::Message(message) => {
274 format!("{}:message:{}", entry.index, message.id)
275 }
276 BorrowedChronologicalPayload::ProtocolEvent(event) => {
277 format!("{}:protocol:{}", entry.index, event.payload)
278 }
279 });
280 });
281 summary
282 }
283
284 #[test]
285 fn borrowed_turn_view_matches_owned_active_event_projection() {
286 let m1 = text_message("m1", MessageRole::User, "first");
287 let m2 = text_message("m2", MessageRole::Assistant, "second");
288 let events = vec![
289 SessionHistoryRecord::Conversation(ConversationRecord::from_message(m1.clone())),
290 SessionHistoryRecord::Protocol(protocol_event("step")),
291 SessionHistoryRecord::Conversation(ConversationRecord::from_message(m1.clone())),
292 ];
293 let messages = MessageSequence::from_owned(vec![m1, m2]);
294 let projection = ChronologicalProjection::from_turn_view(&events, &messages);
295
296 assert_eq!(
297 borrowed_summary(&events, &messages),
298 owned_summary(&projection)
299 );
300 }
301
302 #[test]
303 fn borrowed_turn_view_matches_owned_transcript_fallback_projection() {
304 let messages = MessageSequence::from_owned(vec![
305 tool_result_message("m1", "call-1"),
306 transient_message("transient"),
307 text_message("m2", MessageRole::Assistant, "second"),
308 ]);
309 let events = Vec::new();
310 let projection = ChronologicalProjection::from_turn_view(&events, &messages);
311
312 assert_eq!(
313 borrowed_summary(&events, &messages),
314 owned_summary(&projection)
315 );
316 }
317}