oino_session/
context.rs

1//! Context building and projection (`buildContextEntries`, `buildSessionProjection`,
2//! `buildSessionContext`). Pure functions over an [`EntryIndex`].
3//!
4//! Leaf convention: `leaf = None` is the `null` leaf (empty path, e.g. after `resetLeaf`);
5//! an id that is not in the index falls back to the last entry. Use
6//! [`EntryIndex::last_id`] for the "default leaf".
7
8use std::collections::{HashMap, HashSet};
9
10use oino_ai::transcript::SystemState;
11use oino_ai::{
12    AgentMessage, AssistantContent, SystemMessage, TextContent, UserContent, UserContentField,
13};
14use serde::{Deserialize, Serialize};
15
16use crate::entry::{ContextReplacement, EditableContent, EntryId, EntryKind, SessionEntry};
17use crate::messages::{branch_summary_message, compaction_summary_message, custom_message};
18
19/// Entries in file order plus an id index (latest entry wins on duplicate ids).
20#[derive(Clone, Debug, Default, PartialEq)]
21pub struct EntryIndex {
22    entries: Vec<SessionEntry>,
23    by_id: HashMap<EntryId, usize>,
24}
25
26impl EntryIndex {
27    /// Indexes `entries` (file order).
28    #[must_use]
29    pub fn new(entries: Vec<SessionEntry>) -> Self {
30        let by_id = entries
31            .iter()
32            .enumerate()
33            .map(|(index, entry)| (entry.id.clone(), index))
34            .collect();
35        Self { entries, by_id }
36    }
37
38    /// Appends an entry.
39    pub fn push(&mut self, entry: SessionEntry) {
40        self.by_id.insert(entry.id.clone(), self.entries.len());
41        self.entries.push(entry);
42    }
43
44    /// Entry by id.
45    #[must_use]
46    pub fn get(&self, id: &EntryId) -> Option<&SessionEntry> {
47        self.by_id.get(id).map(|&index| &self.entries[index])
48    }
49
50    /// Whether `id` is present.
51    #[must_use]
52    pub fn contains(&self, id: &EntryId) -> bool {
53        self.by_id.contains_key(id)
54    }
55
56    /// Whether the id text is present.
57    #[must_use]
58    pub fn contains_str(&self, id: &str) -> bool {
59        self.by_id.contains_key(&EntryId::new(id))
60    }
61
62    /// All entries in file order.
63    #[must_use]
64    pub fn entries(&self) -> &[SessionEntry] {
65        &self.entries
66    }
67
68    /// Consumes the index.
69    #[must_use]
70    pub fn into_entries(self) -> Vec<SessionEntry> {
71        self.entries
72    }
73
74    /// Number of distinct ids.
75    #[must_use]
76    pub fn id_count(&self) -> usize {
77        self.by_id.len()
78    }
79
80    /// Number of entries.
81    #[must_use]
82    pub fn len(&self) -> usize {
83        self.entries.len()
84    }
85
86    /// Whether there are no entries.
87    #[must_use]
88    pub fn is_empty(&self) -> bool {
89        self.entries.is_empty()
90    }
91
92    /// Id of the last entry in file order (the default leaf).
93    #[must_use]
94    pub fn last_id(&self) -> Option<&EntryId> {
95        self.entries.last().map(|entry| &entry.id)
96    }
97}
98
99impl FromIterator<SessionEntry> for EntryIndex {
100    fn from_iter<I: IntoIterator<Item = SessionEntry>>(iter: I) -> Self {
101        Self::new(iter.into_iter().collect())
102    }
103}
104
105/// `{provider, modelId}` selection.
106#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
107#[serde(rename_all = "camelCase")]
108pub struct ModelRef {
109    /// Provider id.
110    pub provider: String,
111    /// Model id.
112    pub model_id: String,
113}
114
115/// One selected entry and the messages it contributes.
116#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
117#[serde(rename_all = "camelCase")]
118pub struct ProjectedEntry {
119    /// Selected raw entry.
120    pub source_entry: SessionEntry,
121    /// Model-visible messages after context edits.
122    pub messages: Vec<AgentMessage>,
123}
124
125/// The session projection record.
126#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
127#[serde(rename_all = "camelCase")]
128pub struct SessionProjection {
129    /// Selected entries with provenance.
130    pub entries: Vec<ProjectedEntry>,
131    /// Flattened messages.
132    pub messages: Vec<AgentMessage>,
133    /// Thinking level string (default `"off"`).
134    pub thinking_level: String,
135    /// Current model.
136    pub model: Option<ModelRef>,
137}
138
139/// The session context record.
140#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
141#[serde(rename_all = "camelCase")]
142pub struct SessionContext {
143    /// Messages for the agent.
144    pub messages: Vec<AgentMessage>,
145    /// Thinking level string (default `"off"`).
146    pub thinking_level: String,
147    /// Current model.
148    pub model: Option<ModelRef>,
149}
150
151/// Latest explicit selections on a path (input to virtual-model aware resolution).
152#[derive(Clone, Debug, Default, PartialEq, Eq)]
153pub struct BranchSelection {
154    /// Latest `model_change`.
155    pub model: Option<ModelRef>,
156    /// Id of that `model_change` entry.
157    pub model_entry_id: Option<EntryId>,
158    /// Latest `thinking_level_change`.
159    pub thinking_level: Option<String>,
160}
161
162/// Root → leaf path. Stops at a missing parent (orphans become roots)
163/// and at cycles.
164#[must_use]
165pub fn path_to<'a>(entries: &'a EntryIndex, leaf: Option<&EntryId>) -> Vec<&'a SessionEntry> {
166    let Some(leaf) = leaf else {
167        return Vec::new();
168    };
169    let Some(start) = entries.get(leaf).or_else(|| entries.entries.last()) else {
170        return Vec::new();
171    };
172    walk_up(entries, start)
173}
174
175pub(crate) fn walk_up<'a>(
176    entries: &'a EntryIndex,
177    start: &'a SessionEntry,
178) -> Vec<&'a SessionEntry> {
179    let mut path = Vec::new();
180    let mut seen = HashSet::new();
181    let mut current = Some(start);
182    while let Some(entry) = current {
183        if !seen.insert(&entry.id) {
184            break;
185        }
186        path.push(entry);
187        current = entry
188            .parent_id
189            .as_ref()
190            .and_then(|parent| entries.get(parent));
191    }
192    path.reverse();
193    path
194}
195
196/// Latest `compaction` entry on `path`.
197#[must_use]
198pub fn latest_compaction<'a>(path: &[&'a SessionEntry]) -> Option<&'a SessionEntry> {
199    path.iter()
200        .rev()
201        .find(|entry| matches!(entry.kind, EntryKind::Compaction(_)))
202        .copied()
203}
204
205/// Latest explicit `model_change` and `thinking_level_change` on `path`.
206#[must_use]
207pub fn branch_selection(path: &[&SessionEntry]) -> BranchSelection {
208    let mut selection = BranchSelection::default();
209    for entry in path {
210        match &entry.kind {
211            EntryKind::ModelChange { provider, model_id } => {
212                selection.model = Some(ModelRef {
213                    provider: provider.clone(),
214                    model_id: model_id.clone(),
215                });
216                selection.model_entry_id = Some(entry.id.clone());
217            }
218            EntryKind::ThinkingLevelChange { thinking_level } => {
219                selection.thinking_level = Some(thinking_level.clone());
220            }
221            _ => {}
222        }
223    }
224    selection
225}
226
227/// Thinking level (default `"off"`) and model from the full
228/// path; assistant messages also update the model (the physical model that answered).
229fn context_settings(path: &[&SessionEntry]) -> (String, Option<ModelRef>) {
230    let mut thinking_level = "off".to_owned();
231    let mut model = None;
232    for entry in path {
233        match &entry.kind {
234            EntryKind::ThinkingLevelChange {
235                thinking_level: level,
236            } => {
237                thinking_level.clone_from(level);
238            }
239            EntryKind::ModelChange { provider, model_id } => {
240                model = Some(ModelRef {
241                    provider: provider.clone(),
242                    model_id: model_id.clone(),
243                });
244            }
245            EntryKind::Message {
246                message: AgentMessage::Assistant(assistant),
247            } => {
248                model = Some(ModelRef {
249                    provider: assistant.provider.clone(),
250                    model_id: assistant.model.clone(),
251                });
252            }
253            _ => {}
254        }
255    }
256    (thinking_level, model)
257}
258
259/// Honours the latest compaction on the path.
260#[must_use]
261pub fn build_context_entries(entries: &EntryIndex, leaf: Option<&EntryId>) -> Vec<SessionEntry> {
262    select_context(&path_to(entries, leaf))
263        .into_iter()
264        .cloned()
265        .collect()
266}
267
268fn select_context<'a>(path: &[&'a SessionEntry]) -> Vec<&'a SessionEntry> {
269    let Some(compaction_index) = path
270        .iter()
271        .rposition(|entry| matches!(entry.kind, EntryKind::Compaction(_)))
272    else {
273        return path.to_vec();
274    };
275    let compaction = path[compaction_index];
276    let EntryKind::Compaction(data) = &compaction.kind else {
277        return path.to_vec();
278    };
279    let mut selected = vec![compaction];
280    let mut found_first_kept = false;
281    for entry in &path[..compaction_index] {
282        if entry.id == data.first_kept_entry_id {
283            found_first_kept = true;
284        }
285        if found_first_kept && !entry.is_system_message() {
286            selected.push(entry);
287        }
288    }
289    selected.extend_from_slice(&path[compaction_index + 1..]);
290    selected
291}
292
293/// Converts entries to context messages. `compaction_system` is the checkpoint used when a
294/// compaction entry has no `systemMessage` (projection passes `None`).
295#[must_use]
296pub fn entry_to_context_messages(
297    entry: &SessionEntry,
298    compaction_system: Option<&SystemMessage>,
299) -> Vec<AgentMessage> {
300    let timestamp = entry.timestamp.to_ms().unwrap_or(0);
301    match &entry.kind {
302        EntryKind::Message { message } => vec![message.clone()],
303        EntryKind::CustomMessage {
304            custom_type,
305            content,
306            display,
307            details,
308        } => vec![custom_message(
309            custom_type,
310            content.clone(),
311            *display,
312            details.clone(),
313            timestamp,
314        )],
315        EntryKind::BranchSummary(summary) if !summary.summary.is_empty() => {
316            vec![branch_summary_message(
317                &summary.summary,
318                Some(summary.from_id.as_str()),
319                timestamp,
320            )]
321        }
322        EntryKind::Compaction(compaction) => {
323            let system = compaction.system_message.as_ref().or(compaction_system);
324            system
325                .cloned()
326                .map(AgentMessage::System)
327                .into_iter()
328                .chain(std::iter::once(compaction_summary_message(
329                    &compaction.summary,
330                    compaction.tokens_before,
331                    timestamp,
332                )))
333                .collect()
334        }
335        _ => Vec::new(),
336    }
337}
338
339fn user_blocks(content: &EditableContent) -> UserContentField {
340    match content {
341        EditableContent::Text(text) => UserContentField::Text(text.clone()),
342        EditableContent::User(blocks) => UserContentField::Blocks(blocks.clone()),
343        EditableContent::Assistant(blocks) => UserContentField::Blocks(
344            blocks
345                .iter()
346                .filter_map(|block| match block {
347                    AssistantContent::Text(text) => Some(UserContent::Text(text.clone())),
348                    _ => None,
349                })
350                .collect(),
351        ),
352    }
353}
354
355fn result_blocks(content: &EditableContent) -> Vec<UserContent> {
356    match user_blocks(content) {
357        UserContentField::Text(text) => vec![UserContent::Text(TextContent::new(text))],
358        UserContentField::Blocks(blocks) => blocks,
359    }
360}
361
362fn assistant_blocks(content: &EditableContent) -> Vec<AssistantContent> {
363    match content {
364        EditableContent::Text(text) => vec![AssistantContent::Text(TextContent::new(text.clone()))],
365        EditableContent::Assistant(blocks) => blocks.clone(),
366        EditableContent::User(blocks) => blocks
367            .iter()
368            .filter_map(|block| match block {
369                UserContent::Text(text) => Some(AssistantContent::Text(text.clone())),
370                UserContent::Image(_) => None,
371            })
372            .collect(),
373    }
374}
375
376/// Applies a replacement to one message, keeping role and metadata.
377/// Blocks that the target role cannot hold (images on assistants, thinking on users) are
378/// dropped.
379fn replace_content(message: AgentMessage, replacement: &ContextReplacement) -> AgentMessage {
380    let content = &replacement.content;
381    match message {
382        AgentMessage::User(mut user) => {
383            user.content = user_blocks(content);
384            AgentMessage::User(user)
385        }
386        AgentMessage::Custom(mut custom) => {
387            custom.content = user_blocks(content);
388            AgentMessage::Custom(custom)
389        }
390        AgentMessage::Assistant(mut assistant) => {
391            assistant.content = assistant_blocks(content);
392            AgentMessage::Assistant(assistant)
393        }
394        AgentMessage::ToolResult(mut result) => {
395            result.content = result_blocks(content);
396            AgentMessage::ToolResult(result)
397        }
398        other => other,
399    }
400}
401
402fn project_entry(
403    entry: &SessionEntry,
404    edit: Option<&Option<ContextReplacement>>,
405) -> Vec<AgentMessage> {
406    let messages = entry_to_context_messages(entry, None);
407    match edit {
408        None => messages,
409        Some(None) => Vec::new(),
410        Some(Some(replacement)) => messages
411            .into_iter()
412            .map(|message| replace_content(message, replacement))
413            .collect(),
414    }
415}
416
417/// Context selection, latest context edit per target (among the
418/// selected entries), entry → messages.
419#[must_use]
420pub fn build_session_projection(entries: &EntryIndex, leaf: Option<&EntryId>) -> SessionProjection {
421    let path = path_to(entries, leaf);
422    let (thinking_level, model) = context_settings(&path);
423    let selected = select_context(&path);
424    let mut edits: HashMap<&EntryId, &Option<ContextReplacement>> = HashMap::new();
425    for entry in &selected {
426        if let EntryKind::ContextEdit {
427            target_id,
428            replacement,
429        } = &entry.kind
430        {
431            edits.insert(target_id, replacement);
432        }
433    }
434    let projected: Vec<ProjectedEntry> = selected
435        .iter()
436        .enumerate()
437        .map(|(index, entry)| ProjectedEntry {
438            source_entry: (*entry).clone(),
439            messages: if index > 0 && matches!(entry.kind, EntryKind::Compaction(_)) {
440                Vec::new()
441            } else {
442                project_entry(entry, edits.get(&entry.id).copied())
443            },
444        })
445        .collect();
446    let messages = projected
447        .iter()
448        .flat_map(|entry| entry.messages.iter().cloned())
449        .collect();
450    SessionProjection {
451        entries: projected,
452        messages,
453        thinking_level,
454        model,
455    }
456}
457
458/// B u i l d   S e s s i o n   C o n t e x t.
459#[must_use]
460pub fn build_session_context(entries: &EntryIndex, leaf: Option<&EntryId>) -> SessionContext {
461    let SessionProjection {
462        messages,
463        thinking_level,
464        model,
465        ..
466    } = build_session_projection(entries, leaf);
467    SessionContext {
468        messages,
469        thinking_level,
470        model,
471    }
472}
473
474/// Folded system checkpoint of a projection.
475#[must_use]
476pub fn current_system_message(messages: &[AgentMessage]) -> Option<SystemMessage> {
477    SystemState::fold(messages.iter().filter_map(AgentMessage::as_system)).into_message()
478}
479
480#[cfg(test)]
481mod tests {
482    use super::*;
483    use serde_json::json;
484
485    fn index(lines: &serde_json::Value) -> EntryIndex {
486        serde_json::from_value::<Vec<SessionEntry>>(lines.clone())
487            .unwrap()
488            .into_iter()
489            .collect()
490    }
491
492    fn user(id: &str, parent: Option<&str>, text: &str) -> serde_json::Value {
493        json!({"type":"message","id":id,"parentId":parent,"timestamp":"2024-01-01T00:00:00.000Z",
494            "message":{"role":"user","content":text,"timestamp":1}})
495    }
496
497    fn texts(context: &SessionContext) -> Vec<String> {
498        context
499            .messages
500            .iter()
501            .map(|m| match m {
502                AgentMessage::User(u) => u.content.text(""),
503                other => other.role().to_owned(),
504            })
505            .collect()
506    }
507
508    #[test]
509    fn compaction_keeps_range_and_drops_system() {
510        let entries = index(&json!([
511            {"type":"message","id":"s","parentId":null,"timestamp":"2024-01-01T00:00:00.000Z",
512             "message":{"role":"system","content":"sys","timestamp":0}},
513            user("a", Some("s"), "a"),
514            user("b", Some("a"), "b"),
515            {"type":"message","id":"s2","parentId":"b","timestamp":"2024-01-01T00:00:00.000Z",
516             "message":{"role":"system","content":"patch","timestamp":0}},
517            {"type":"compaction","id":"c","parentId":"s2","timestamp":"2024-01-01T00:00:01.000Z",
518             "summary":"sum","firstKeptEntryId":"b","tokensBefore":5,
519             "systemMessage":{"role":"system","content":"sys\n\npatch","timestamp":1000}},
520            user("d", Some("c"), "d"),
521        ]));
522        let context = build_session_context(&entries, entries.last_id());
523        assert_eq!(texts(&context), ["system", "compactionSummary", "b", "d"]);
524        let selected: Vec<_> = build_context_entries(&entries, Some(&"d".into()))
525            .into_iter()
526            .map(|e| e.id.to_string())
527            .collect();
528        assert_eq!(selected, ["c", "b", "d"]);
529        assert_eq!(build_session_context(&entries, None).messages, []);
530    }
531
532    #[test]
533    fn context_edits_are_branch_relative() {
534        let entries = index(&json!([
535            user("a", None, "a"),
536            {"type":"context_edit","id":"e","parentId":"a","timestamp":"2024-01-01T00:00:00.000Z","targetId":"a",
537             "replacement":{"content":"edited"}},
538            {"type":"context_edit","id":"f","parentId":"e","timestamp":"2024-01-01T00:00:00.000Z","targetId":"a","replacement":null},
539        ]));
540        assert_eq!(
541            texts(&build_session_context(&entries, Some(&"e".into()))),
542            ["edited"]
543        );
544        assert_eq!(
545            build_session_context(&entries, Some(&"f".into())).messages,
546            []
547        );
548        assert_eq!(
549            texts(&build_session_context(&entries, Some(&"a".into()))),
550            ["a"]
551        );
552    }
553
554    #[test]
555    fn settings_come_from_full_path() {
556        let entries = index(&json!([
557            {"type":"model_change","id":"m","parentId":null,"timestamp":"t","provider":"p","modelId":"x"},
558            {"type":"thinking_level_change","id":"t","parentId":"m","timestamp":"t","thinkingLevel":"high"},
559        ]));
560        let context = build_session_context(&entries, entries.last_id());
561        assert_eq!(context.thinking_level, "high");
562        assert_eq!(context.model.map(|m| m.model_id), Some("x".to_owned()));
563        let path = path_to(&entries, entries.last_id());
564        assert_eq!(branch_selection(&path).model_entry_id, Some("m".into()));
565    }
566
567    #[test]
568    fn cycles_terminate() {
569        let entries = index(&json!([
570            user("a", Some("b"), "a"),
571            user("b", Some("a"), "b")
572        ]));
573        assert_eq!(path_to(&entries, Some(&"b".into())).len(), 2);
574    }
575}
oino_session/context.rs

oino_session/
context.rs
