oino_core/session/
tree.rs

1//! Tree navigation with branch summaries (`navigateTree`, `getUserMessagesForForking`).
2
3use oino_ai::{AgentMessage, Usage};
4use oino_session::{BranchSummaryDraft, EntryId, EntryKind};
5use serde::{Deserialize, Serialize};
6use serde_json::json;
7use tokio_util::sync::CancellationToken;
8
9use super::{AgentSession, SessionOpError, lock};
10use crate::compaction::{
11    BranchSummaryOptions, collect_entries_for_branch_summary, generate_branch_summary,
12};
13use crate::events::SummarizationSource;
14
15/// Options of [`AgentSession::navigate_tree`].
16#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "camelCase")]
18pub struct NavigateTreeOptions {
19    /// Summarise the abandoned branch.
20    #[serde(default)]
21    pub summarize: bool,
22    /// Custom summariser instructions.
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub custom_instructions: Option<String>,
25    /// Instructions replace the default prompt.
26    #[serde(default)]
27    pub replace_instructions: bool,
28    /// Label for the summary (or the target without summary).
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub label: Option<String>,
31}
32
33/// `navigateTree` result.
34#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
35#[serde(rename_all = "camelCase")]
36pub struct NavigateTreeResult {
37    /// Text for the editor when the target was a user message.
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub editor_text: Option<String>,
40    /// Cancelled by an extension or the user.
41    pub cancelled: bool,
42    /// Summarisation was aborted.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub aborted: Option<bool>,
45    /// The created branch summary.
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub summary_entry: Option<oino_session::SessionEntry>,
48}
49
50/// `{entryId, text}` of a user message.
51#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(rename_all = "camelCase")]
53pub struct ForkMessage {
54    /// Entry id.
55    pub entry_id: String,
56    /// Message text.
57    pub text: String,
58}
59
60impl AgentSession {
61    /// User messages of the whole session.
62    #[must_use]
63    pub fn user_messages_for_forking(&self) -> Vec<ForkMessage> {
64        lock(&self.inner.manager)
65            .entries()
66            .iter()
67            .filter_map(|entry| match entry.message() {
68                Some(AgentMessage::User(user)) => {
69                    let text = user.content.text("");
70                    (!text.is_empty()).then(|| ForkMessage {
71                        entry_id: entry.id.as_str().to_owned(),
72                        text,
73                    })
74                }
75                _ => None,
76            })
77            .collect()
78    }
79
80    /// Navigates the entry tree.
81    ///
82    /// # Errors
83    ///
84    /// Fixed texts for streaming, compaction, unknown targets and summariser failures.
85    #[allow(clippy::too_many_lines, reason = "one-to-one port of navigateTree")]
86    pub async fn navigate_tree(
87        &self,
88        target_id: &EntryId,
89        options: NavigateTreeOptions,
90    ) -> Result<NavigateTreeResult, SessionOpError> {
91        if self.is_streaming() {
92            return Err(SessionOpError::Failed(
93                "Wait for the current response to finish before navigating the session tree."
94                    .into(),
95            ));
96        }
97        if self.is_compacting() {
98            return Err(SessionOpError::Failed(
99                "Wait for the current compaction or tree navigation to finish before navigating the session tree.".into(),
100            ));
101        }
102        let old_leaf = lock(&self.inner.manager).leaf_id().cloned();
103        if old_leaf.as_ref() == Some(target_id) {
104            return Ok(NavigateTreeResult::default());
105        }
106        if options.summarize && self.model().is_none() {
107            return Err(SessionOpError::Failed(
108                "No model available for summarization".into(),
109            ));
110        }
111        let (target, collected) = {
112            let manager = lock(&self.inner.manager);
113            let target = manager.entry(target_id).cloned().ok_or_else(|| {
114                SessionOpError::Failed(format!("Entry {} not found", target_id.as_str()))
115            })?;
116            let index = manager.index();
117            (
118                target,
119                collect_entries_for_branch_summary(index, old_leaf.as_ref(), target_id),
120            )
121        };
122        let mut custom_instructions = options.custom_instructions.clone();
123        let mut replace_instructions = options.replace_instructions;
124        let mut label = options.label.clone();
125        let cancel = CancellationToken::new();
126        lock(&self.inner.cancels).branch_summary = Some(cancel.clone());
127        self.mark_busy();
128        let result = async {
129            let mut extension_summary: Option<(String, Option<serde_json::Value>, Option<Usage>)> =
130                None;
131            let mut from_extension = false;
132            let runner = self.runner();
133            if runner.has_handlers("session_before_tree") {
134                let mut preparation = json!({
135                    "targetId": target_id.as_str(),
136                    "oldLeafId": old_leaf.as_ref().map(EntryId::as_str),
137                    "commonAncestorId": collected.common_ancestor_id.as_ref().map(EntryId::as_str),
138                    "entriesToSummarize": collected.entries,
139                    "userWantsSummary": options.summarize,
140                });
141                if let Some(instructions) = &custom_instructions {
142                    preparation["customInstructions"] = json!(instructions);
143                }
144                if options.replace_instructions {
145                    preparation["replaceInstructions"] = json!(true);
146                }
147                if let Some(label) = &label {
148                    preparation["label"] = json!(label);
149                }
150                if let Some(result) = runner
151                    .emit(json!({"type": "session_before_tree", "preparation": preparation}))
152                    .await
153                {
154                    if result.get("cancel") == Some(&serde_json::Value::Bool(true)) {
155                        return Ok(NavigateTreeResult {
156                            cancelled: true,
157                            ..NavigateTreeResult::default()
158                        });
159                    }
160                    if let (Some(summary), true) = (
161                        result.get("summary").filter(|s| s.is_object()),
162                        options.summarize,
163                    ) {
164                        extension_summary = Some((
165                            summary["summary"].as_str().unwrap_or_default().to_owned(),
166                            summary.get("details").cloned(),
167                            summary
168                                .get("usage")
169                                .and_then(|u| serde_json::from_value(u.clone()).ok()),
170                        ));
171                        from_extension = true;
172                    }
173                    if let Some(value) = result
174                        .get("customInstructions")
175                        .and_then(serde_json::Value::as_str)
176                    {
177                        custom_instructions = Some(value.to_owned());
178                    }
179                    if let Some(value) = result
180                        .get("replaceInstructions")
181                        .and_then(serde_json::Value::as_bool)
182                    {
183                        replace_instructions = value;
184                    }
185                    if let Some(value) = result.get("label").and_then(serde_json::Value::as_str) {
186                        label = Some(value.to_owned());
187                    }
188                }
189            }
190            let mut summary: Option<(String, Option<serde_json::Value>, Option<Usage>)> = None;
191            if options.summarize && !collected.entries.is_empty() && extension_summary.is_none() {
192                let model = self.model().ok_or_else(|| {
193                    SessionOpError::Failed("No model available for summarization".into())
194                })?;
195                let summarizer = self
196                    .summarizer(&model, cancel.clone(), SummarizationSource::BranchSummary)
197                    .await
198                    .map_err(SessionOpError::Failed)?;
199                let reserve = lock(&self.inner.settings)
200                    .settings()
201                    .branch_summary_settings()
202                    .reserve_tokens;
203                let result = generate_branch_summary(
204                    &collected.entries,
205                    &summarizer,
206                    &BranchSummaryOptions {
207                        custom_instructions: custom_instructions.clone(),
208                        replace_instructions,
209                        reserve_tokens: reserve,
210                    },
211                )
212                .await;
213                if result.aborted == Some(true) {
214                    return Ok(NavigateTreeResult {
215                        cancelled: true,
216                        aborted: Some(true),
217                        ..NavigateTreeResult::default()
218                    });
219                }
220                if let Some(error) = result.error {
221                    return Err(SessionOpError::Failed(error));
222                }
223                summary = Some((
224                    result.summary.unwrap_or_default(),
225                    Some(json!({
226                        "readFiles": result.read_files.unwrap_or_default(),
227                        "modifiedFiles": result.modified_files.unwrap_or_default(),
228                    })),
229                    result.usage,
230                ));
231            } else if let Some(extension) = extension_summary {
232                summary = Some(extension);
233            }
234            let (new_leaf, editor_text) = match &target.kind {
235                EntryKind::Message {
236                    message: AgentMessage::User(user),
237                } => (target.parent_id.clone(), Some(user.content.text(""))),
238                EntryKind::CustomMessage { content, .. } => {
239                    (target.parent_id.clone(), Some(content.text("")))
240                }
241                _ => (Some(target_id.clone()), None),
242            };
243            let mut summary_entry = None;
244            let summarized = summary
245                .as_ref()
246                .is_some_and(|(text, _, _)| !text.is_empty());
247            {
248                let mut manager = lock(&self.inner.manager);
249                if let (Some((text, details, usage)), true) = (summary, summarized) {
250                    let from_id = old_leaf.clone().unwrap_or_else(|| EntryId::new("root"));
251                    let id = manager
252                        .branch_with_summary(
253                            new_leaf.as_ref(),
254                            BranchSummaryDraft {
255                                from_id,
256                                summary: text,
257                                details,
258                                from_hook: Some(from_extension),
259                                usage,
260                            },
261                        )
262                        .map_err(|e| SessionOpError::Failed(e.to_string()))?;
263                    if let Some(label) = &label {
264                        manager
265                            .append_label_change(&id, Some(label))
266                            .map_err(|e| SessionOpError::Failed(e.to_string()))?;
267                    }
268                    summary_entry = manager.entry(&id).cloned();
269                } else {
270                    match &new_leaf {
271                        None => manager.reset_leaf(),
272                        Some(leaf) => manager
273                            .set_leaf(leaf)
274                            .map_err(|e| SessionOpError::Failed(e.to_string()))?,
275                    }
276                    if let Some(label) = &label {
277                        manager
278                            .append_label_change(target_id, Some(label))
279                            .map_err(|e| SessionOpError::Failed(e.to_string()))?;
280                    }
281                }
282            }
283            self.refresh_context();
284            self.restore_tools_from_transcript(&[]);
285            let new_leaf_id = lock(&self.inner.manager).leaf_id().cloned();
286            let runner = self.runner();
287            if runner.has_handlers("session_tree") {
288                let mut payload = json!({
289                    "type": "session_tree",
290                    "newLeafId": new_leaf_id.as_ref().map(EntryId::as_str),
291                    "oldLeafId": old_leaf.as_ref().map(EntryId::as_str),
292                });
293                if let Some(entry) = &summary_entry {
294                    payload["summaryEntry"] = json!(entry);
295                    payload["fromExtension"] = json!(from_extension);
296                }
297                runner.emit(payload).await;
298            }
299            Ok(NavigateTreeResult {
300                editor_text,
301                cancelled: false,
302                aborted: None,
303                summary_entry,
304            })
305        }
306        .await;
307        lock(&self.inner.cancels).branch_summary = None;
308        self.set_idle_if_idle();
309        result
310    }
311}
oino_core/session/tree.rs

oino_core/session/
tree.rs
