oino_core/session/
queue.rs

1//! Single-item queue edits (`_oino/queue_edit`, Oino addition): remove, edit, reorder and
2//! send-now for one queued steering / follow-up message.
3//!
4//! The session keeps the queued texts (`queue_update`) and the agent keeps the queued
5//! messages; custom messages queued by extensions have no text entry. Each text is matched to
6//! its agent message (first unclaimed user message with the same text, in order), and both
7//! lists are edited together.
8
9use oino_ai::{AgentMessage, ImageContent, UserContent, UserContentField, UserMessage};
10use serde::{Deserialize, Serialize};
11
12use super::prompt::{InputSource, PromptOptions};
13use super::{AgentSession, SessionOpError, lock};
14
15/// Which queue.
16#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "camelCase")]
18pub enum QueueKind {
19    /// Steering (delivered after the current tool calls).
20    Steering,
21    /// Follow-up (delivered when the agent would stop).
22    FollowUp,
23}
24
25/// One queue edit; `index` is the position in the `queue_update` list of `queue`.
26#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(tag = "op", rename_all = "camelCase")]
28pub enum QueueEdit {
29    /// Drops the item.
30    Remove {
31        /// Queue.
32        queue: QueueKind,
33        /// Item.
34        index: usize,
35    },
36    /// Replaces the item's text (images are kept; prompt templates expand).
37    Edit {
38        /// Queue.
39        queue: QueueKind,
40        /// Item.
41        index: usize,
42        /// New text.
43        text: String,
44    },
45    /// Moves the item to position `to` of the same queue.
46    Move {
47        /// Queue.
48        queue: QueueKind,
49        /// Item.
50        from: usize,
51        /// Target position.
52        to: usize,
53    },
54    /// Delivers the item now: while running it becomes the first steering message; when
55    /// idle it is removed and prompted.
56    SendNow {
57        /// Queue.
58        queue: QueueKind,
59        /// Item.
60        index: usize,
61    },
62}
63
64fn user_text(message: &AgentMessage) -> Option<String> {
65    match message {
66        AgentMessage::User(user) => Some(user.content.text("")),
67        _ => None,
68    }
69}
70
71fn user_images(message: &AgentMessage) -> Vec<ImageContent> {
72    match message {
73        AgentMessage::User(UserMessage {
74            content: UserContentField::Blocks(blocks),
75            ..
76        }) => blocks
77            .iter()
78            .filter_map(|block| match block {
79                UserContent::Image(image) => Some(image.clone()),
80                UserContent::Text(_) => None,
81            })
82            .collect(),
83        _ => Vec::new(),
84    }
85}
86
87/// Agent queue index of every text (`None` when the message was already drained).
88fn match_texts(texts: &[String], messages: &[AgentMessage]) -> Vec<Option<usize>> {
89    let mut claimed = vec![false; messages.len()];
90    texts
91        .iter()
92        .map(|text| {
93            let found = messages.iter().enumerate().position(|(i, message)| {
94                !claimed[i] && user_text(message).as_deref() == Some(text.as_str())
95            });
96            if let Some(i) = found {
97                claimed[i] = true;
98            }
99            found
100        })
101        .collect()
102}
103
104fn out_of_range(index: usize, len: usize) -> SessionOpError {
105    SessionOpError::Failed(format!("Queue index {index} out of range ({len} queued)"))
106}
107
108impl AgentSession {
109    /// Applies one queue edit and publishes `queue_update`.
110    ///
111    /// # Errors
112    ///
113    /// [`SessionOpError::Failed`] for an index out of range; prompt errors for send-now while
114    /// idle.
115    pub async fn edit_queue(&self, edit: QueueEdit) -> Result<(), SessionOpError> {
116        if let QueueEdit::SendNow { queue, index } = edit {
117            return self.send_queued_now(queue, index).await;
118        }
119        let edit = match edit {
120            QueueEdit::Edit { queue, index, text } => QueueEdit::Edit {
121                queue,
122                index,
123                text: self.expand_queued_text(&text),
124            },
125            other => other,
126        };
127        self.apply_queue_edit(&edit)?;
128        self.emit_queue_update();
129        Ok(())
130    }
131
132    fn expand_queued_text(&self, text: &str) -> String {
133        let resources = self
134            .inner
135            .resources
136            .read()
137            .unwrap_or_else(std::sync::PoisonError::into_inner);
138        oino_resources::templates::expand_prompt_template(text, resources.prompts())
139            .unwrap_or_else(|| text.to_owned())
140    }
141
142    fn apply_queue_edit(&self, edit: &QueueEdit) -> Result<(), SessionOpError> {
143        let queue = match edit {
144            QueueEdit::Remove { queue, .. }
145            | QueueEdit::Edit { queue, .. }
146            | QueueEdit::Move { queue, .. }
147            | QueueEdit::SendNow { queue, .. } => *queue,
148        };
149        let steering = queue == QueueKind::Steering;
150        let mut book = lock(&self.inner.book);
151        let texts = if steering {
152            &mut book.steering
153        } else {
154            &mut book.follow_up
155        };
156        self.inner.agent.update_queue(steering, |messages| {
157            let matches = match_texts(texts, messages);
158            match edit {
159                QueueEdit::Remove { index, .. } | QueueEdit::SendNow { index, .. } => {
160                    if *index >= texts.len() {
161                        return Err(out_of_range(*index, texts.len()));
162                    }
163                    texts.remove(*index);
164                    if let Some(at) = matches[*index] {
165                        messages.remove(at);
166                    }
167                }
168                QueueEdit::Edit { index, text, .. } => {
169                    if *index >= texts.len() {
170                        return Err(out_of_range(*index, texts.len()));
171                    }
172                    texts[*index].clone_from(text);
173                    if let Some(at) = matches[*index] {
174                        let images = user_images(&messages[at]);
175                        let mut content = vec![UserContent::text(text.clone())];
176                        content.extend(images.into_iter().map(UserContent::Image));
177                        if let AgentMessage::User(user) = &mut messages[at] {
178                            user.content = UserContentField::Blocks(content);
179                        }
180                    }
181                }
182                QueueEdit::Move { from, to, .. } => {
183                    let len = texts.len();
184                    if *from >= len || *to >= len {
185                        return Err(out_of_range((*from).max(*to), len));
186                    }
187                    let item = texts.remove(*from);
188                    texts.insert(*to, item);
189                    let mut slots: Vec<usize> = matches.iter().flatten().copied().collect();
190                    if slots.len() == len {
191                        let mut ordered: Vec<AgentMessage> = matches
192                            .iter()
193                            .flatten()
194                            .map(|&i| messages[i].clone())
195                            .collect();
196                        let moved = ordered.remove(*from);
197                        ordered.insert(*to, moved);
198                        slots.sort_unstable();
199                        for (slot, message) in slots.into_iter().zip(ordered) {
200                            messages[slot] = message;
201                        }
202                    }
203                }
204            }
205            Ok(())
206        })
207    }
208
209    async fn send_queued_now(&self, queue: QueueKind, index: usize) -> Result<(), SessionOpError> {
210        let steering = queue == QueueKind::Steering;
211        let taken = {
212            let book = lock(&self.inner.book);
213            let texts = if steering {
214                &book.steering
215            } else {
216                &book.follow_up
217            };
218            let text = texts
219                .get(index)
220                .cloned()
221                .ok_or_else(|| out_of_range(index, texts.len()))?;
222            let message = self.inner.agent.update_queue(steering, |messages| {
223                match_texts(texts, messages)[index].map(|at| messages[at].clone())
224            });
225            (text, message)
226        };
227        let (text, message) = taken;
228        self.apply_queue_edit(&QueueEdit::Remove { queue, index })?;
229        if self.is_streaming() {
230            let message = message.unwrap_or_else(|| {
231                AgentMessage::User(UserMessage {
232                    content: UserContentField::Blocks(vec![UserContent::text(text.clone())]),
233                    timestamp: oino_agent::now_ms(),
234                })
235            });
236            lock(&self.inner.book).steering.insert(0, text);
237            self.inner
238                .agent
239                .update_queue(true, |messages| messages.insert(0, message));
240            self.emit_queue_update();
241            return Ok(());
242        }
243        self.emit_queue_update();
244        let images = message.as_ref().map(user_images).unwrap_or_default();
245        self.prompt(
246            text,
247            PromptOptions {
248                images,
249                source: InputSource::Rpc,
250                ..PromptOptions::default()
251            },
252        )
253        .await
254        .map(drop)
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261
262    fn user(text: &str) -> AgentMessage {
263        AgentMessage::User(UserMessage {
264            content: UserContentField::Blocks(vec![UserContent::text(text.to_owned())]),
265            timestamp: 0,
266        })
267    }
268
269    #[test]
270    fn texts_match_user_messages_in_order() {
271        let messages = vec![user("a"), user("b"), user("a")];
272        let texts = vec!["a".to_owned(), "a".to_owned(), "c".to_owned()];
273        assert_eq!(match_texts(&texts, &messages), vec![Some(0), Some(2), None]);
274    }
275
276    #[test]
277    fn edits_round_trip_as_json() {
278        let edit: QueueEdit = serde_json::from_value(
279            serde_json::json!({"op": "move", "queue": "followUp", "from": 1, "to": 0}),
280        )
281        .unwrap();
282        assert_eq!(
283            edit,
284            QueueEdit::Move {
285                queue: QueueKind::FollowUp,
286                from: 1,
287                to: 0
288            }
289        );
290    }
291}
oino_core/session/queue.rs

oino_core/session/
queue.rs
