tenant.openagents/omega

1use std::collections::VecDeque;
2
3use super::*;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub struct QueueEntryId(usize);
7
8pub struct QueueEntry {
9    pub id: QueueEntryId,
10    pub content: Vec<acp::ContentBlock>,
11    pub tracked_buffers: Vec<Entity<Buffer>>,
12    /// When true, this message interrupts the agent at the next turn boundary
13    /// instead of waiting for generation to fully complete. Only the front
14    /// entry's value matters, since messages are delivered in FIFO order.
15    pub steer: bool,
16    pub editor: Entity<MessageEditor>,
17    pub _subscription: Subscription,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21enum ProcessingState {
22    AutoProcess,
23    Paused,
24    // Sending a message out of turn cancelled the current generation; we must
25    // absorb the Stopped event from that cancellation before resuming
26    // auto-processing, otherwise the queue would double-send.
27    AbsorbingCancel,
28}
29
30/// Holds follow-up messages typed while the agent is generating, along with
31/// the state machine that decides when they're auto-sent.
32pub struct MessageQueue {
33    entries: VecDeque<QueueEntry>,
34    processing_state: ProcessingState,
35    can_fast_track: bool,
36    next_id: usize,
37}
38
39impl Default for MessageQueue {
40    fn default() -> Self {
41        Self {
42            entries: VecDeque::new(),
43            processing_state: ProcessingState::AutoProcess,
44            can_fast_track: false,
45            next_id: 0,
46        }
47    }
48}
49
50impl MessageQueue {
51    pub fn is_empty(&self) -> bool {
52        self.entries.is_empty()
53    }
54
55    pub fn len(&self) -> usize {
56        self.entries.len()
57    }
58
59    pub fn first(&self) -> Option<&QueueEntry> {
60        self.entries.front()
61    }
62
63    pub fn first_id(&self) -> Option<QueueEntryId> {
64        self.entries.front().map(|entry| entry.id)
65    }
66
67    pub fn last_id(&self) -> Option<QueueEntryId> {
68        self.entries.back().map(|entry| entry.id)
69    }
70
71    /// Whether the next message should interrupt the agent at the next turn
72    /// boundary. Drives the native thread's boundary flag.
73    pub fn front_wants_steer(&self) -> bool {
74        self.entries.front().is_some_and(|entry| entry.steer)
75    }
76
77    pub fn toggle_steer(&mut self, id: QueueEntryId) {
78        if let Some(entry) = self.entries.iter_mut().find(|entry| entry.id == id) {
79            entry.steer = !entry.steer;
80        }
81    }
82
83    pub fn iter(&self) -> impl Iterator<Item = &QueueEntry> {
84        self.entries.iter()
85    }
86
87    pub fn can_fast_track(&self) -> bool {
88        self.can_fast_track && !self.entries.is_empty()
89    }
90
91    pub fn entry_by_id(&self, id: QueueEntryId) -> Option<&QueueEntry> {
92        self.entries.iter().find(|entry| entry.id == id)
93    }
94
95    pub fn entry_by_id_mut(&mut self, id: QueueEntryId) -> Option<&mut QueueEntry> {
96        self.entries.iter_mut().find(|entry| entry.id == id)
97    }
98
99    /// Allocates a stable ID for a new entry. This is separate from `enqueue`
100    /// because the editor event subscription must capture the ID before the
101    /// `QueueEntry` (which owns that subscription) can be constructed.
102    pub fn next_id(&mut self) -> QueueEntryId {
103        let id = QueueEntryId(self.next_id);
104        self.next_id += 1;
105        id
106    }
107
108    /// Queuing a message is active engagement, so it also resumes
109    /// auto-processing if the queue was paused.
110    pub fn enqueue(&mut self, entry: QueueEntry) {
111        self.entries.push_back(entry);
112        self.processing_state = ProcessingState::AutoProcess;
113        self.can_fast_track = true;
114    }
115
116    pub fn remove(&mut self, id: QueueEntryId) -> Option<QueueEntry> {
117        let index = self.entries.iter().position(|entry| entry.id == id)?;
118        self.entries.remove(index)
119    }
120
121    pub fn clear(&mut self) {
122        self.entries.clear();
123        self.can_fast_track = false;
124    }
125
126    /// Pops the front entry if a fast-track send is allowed (the user just
127    /// queued a message and pressed Enter on an empty main editor).
128    ///
129    /// This works even when paused — pressing Enter is an explicit user
130    /// action, distinct from auto-processing. If a generation is in flight,
131    /// the dispatch will cancel it, so we must absorb that cancellation's
132    /// Stopped event to avoid double-sending the next entry.
133    pub fn try_fast_track(&mut self, is_generating: bool) -> Option<QueueEntry> {
134        if !self.can_fast_track {
135            return None;
136        }
137        self.can_fast_track = false;
138        let entry = self.entries.pop_front()?;
139        self.processing_state = if is_generating {
140            ProcessingState::AbsorbingCancel
141        } else {
142            ProcessingState::AutoProcess
143        };
144        Some(entry)
145    }
146
147    /// Handles a generation Stopped event, returning the entry to auto-send,
148    /// if any.
149    pub fn on_generation_stopped(&mut self, is_first_editor_focused: bool) -> Option<QueueEntry> {
150        match self.processing_state {
151            ProcessingState::AbsorbingCancel => {
152                // This Stopped event came from a cancellation we initiated
153                // ourselves (e.g. "Send Now"); swallow it and resume.
154                self.processing_state = ProcessingState::AutoProcess;
155                None
156            }
157            ProcessingState::Paused => None,
158            ProcessingState::AutoProcess => {
159                // Don't auto-send while the user is editing the next message.
160                if is_first_editor_focused {
161                    None
162                } else {
163                    self.entries.pop_front()
164                }
165            }
166        }
167    }
168
169    /// Removes an entry for an explicit "Send Now". If a generation is in
170    /// flight, the dispatch will cancel it, so we must absorb that
171    /// cancellation's Stopped event.
172    pub fn send_now(&mut self, id: QueueEntryId, is_generating: bool) -> Option<QueueEntry> {
173        let entry = self.remove(id)?;
174        if is_generating {
175            self.processing_state = ProcessingState::AbsorbingCancel;
176        }
177        Some(entry)
178    }
179
180    /// Called when the user stops generation; queued messages stay put until
181    /// the user re-engages.
182    pub fn pause(&mut self) {
183        self.processing_state = ProcessingState::Paused;
184    }
185
186    /// Called when the user sends a new message, re-enabling auto-processing.
187    /// This is what un-freezes the queue after a manual stop.
188    pub fn resume(&mut self) {
189        self.processing_state = ProcessingState::AutoProcess;
190    }
191}
192