notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

focus_queue.rs (19033B)


      1 use crate::agent_status::AgentStatus;
      2 use crate::session::SessionId;
      3 use std::collections::HashMap;
      4 
      5 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
      6 pub enum FocusPriority {
      7     Done = 0,
      8     Error = 1,
      9     NeedsInput = 2,
     10 }
     11 
     12 impl FocusPriority {
     13     pub fn from_status(status: AgentStatus) -> Option<Self> {
     14         match status {
     15             AgentStatus::NeedsInput => Some(Self::NeedsInput),
     16             AgentStatus::Error => Some(Self::Error),
     17             AgentStatus::Done => Some(Self::Done),
     18             AgentStatus::Idle | AgentStatus::Working => None,
     19         }
     20     }
     21 
     22     pub fn color(&self) -> egui::Color32 {
     23         match self {
     24             Self::NeedsInput => egui::Color32::from_rgb(255, 200, 0),
     25             Self::Error => egui::Color32::from_rgb(220, 60, 60),
     26             Self::Done => egui::Color32::from_rgb(70, 130, 220),
     27         }
     28     }
     29 }
     30 
     31 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     32 pub struct QueueEntry {
     33     pub session_id: SessionId,
     34     pub priority: FocusPriority,
     35 }
     36 
     37 pub struct FocusQueue {
     38     entries: Vec<QueueEntry>, // kept sorted: NeedsInput -> Error -> Done
     39     cursor: Option<usize>,    // index into entries
     40     previous_statuses: HashMap<SessionId, AgentStatus>,
     41 }
     42 
     43 impl Default for FocusQueue {
     44     fn default() -> Self {
     45         Self::new()
     46     }
     47 }
     48 
     49 impl FocusQueue {
     50     pub fn new() -> Self {
     51         Self {
     52             entries: Vec::new(),
     53             cursor: None,
     54             previous_statuses: HashMap::new(),
     55         }
     56     }
     57 
     58     pub fn len(&self) -> usize {
     59         self.entries.len()
     60     }
     61 
     62     pub fn is_empty(&self) -> bool {
     63         self.entries.is_empty()
     64     }
     65 
     66     fn sort_key(p: FocusPriority) -> i32 {
     67         // want NeedsInput first, then Error, then Done
     68         -(p as i32)
     69     }
     70 
     71     fn find(&self, session_id: SessionId) -> Option<usize> {
     72         self.entries.iter().position(|e| e.session_id == session_id)
     73     }
     74 
     75     fn normalize_cursor_after_remove(&mut self, removed_idx: usize) {
     76         match self.cursor {
     77             None => {}
     78             Some(_cur) if self.entries.is_empty() => self.cursor = None,
     79             Some(cur) if removed_idx < cur => self.cursor = Some(cur - 1),
     80             Some(cur) if removed_idx == cur => {
     81                 // keep cursor pointing at a valid item (same index if possible, else last)
     82                 let new_cur = cur.min(self.entries.len().saturating_sub(1));
     83                 self.cursor = Some(new_cur);
     84             }
     85             Some(_) => {}
     86         }
     87     }
     88 
     89     /// Insert entry in priority order (stable within same priority).
     90     fn insert_sorted(&mut self, entry: QueueEntry) {
     91         let key = Self::sort_key(entry.priority);
     92         let pos = self
     93             .entries
     94             .iter()
     95             .position(|e| Self::sort_key(e.priority) > key)
     96             .unwrap_or(self.entries.len());
     97         self.entries.insert(pos, entry);
     98 
     99         // initialize cursor if this is the first item
    100         if self.cursor.is_none() && self.entries.len() == 1 {
    101             self.cursor = Some(0);
    102         } else if let Some(cur) = self.cursor {
    103             // if we inserted before the cursor, shift cursor right
    104             if pos <= cur {
    105                 self.cursor = Some(cur + 1);
    106             }
    107         }
    108     }
    109 
    110     pub fn enqueue(&mut self, session_id: SessionId, priority: FocusPriority) {
    111         if let Some(i) = self.find(session_id) {
    112             if self.entries[i].priority == priority {
    113                 return;
    114             }
    115             // remove old entry, then reinsert at correct spot
    116             self.entries.remove(i);
    117             self.normalize_cursor_after_remove(i);
    118         }
    119         self.insert_sorted(QueueEntry {
    120             session_id,
    121             priority,
    122         });
    123     }
    124 
    125     pub fn dequeue(&mut self, session_id: SessionId) {
    126         if let Some(i) = self.find(session_id) {
    127             self.entries.remove(i);
    128             self.normalize_cursor_after_remove(i);
    129         }
    130     }
    131 
    132     pub fn next(&mut self) -> Option<SessionId> {
    133         if self.entries.is_empty() {
    134             self.cursor = None;
    135             return None;
    136         }
    137         let cur = self.cursor.unwrap_or(0);
    138         let current_priority = self.entries[cur].priority;
    139 
    140         // Find all entries with the same priority
    141         let same_priority_indices: Vec<usize> = self
    142             .entries
    143             .iter()
    144             .enumerate()
    145             .filter(|(_, e)| e.priority == current_priority)
    146             .map(|(i, _)| i)
    147             .collect();
    148 
    149         // Find current position within same-priority items
    150         let pos_in_group = same_priority_indices
    151             .iter()
    152             .position(|&i| i == cur)
    153             .unwrap_or(0);
    154 
    155         // If at last item in group, try to go up to highest priority level
    156         if pos_in_group == same_priority_indices.len() - 1 {
    157             // Check if there's a higher priority group (entries are sorted highest first)
    158             let first_same_priority = same_priority_indices[0];
    159             if first_same_priority > 0 {
    160                 // Go to the very first item (highest priority)
    161                 self.cursor = Some(0);
    162                 return Some(self.entries[0].session_id);
    163             }
    164             // No higher priority group, wrap to first item in current group
    165             let next = same_priority_indices[0];
    166             self.cursor = Some(next);
    167             Some(self.entries[next].session_id)
    168         } else {
    169             // Move forward within the same priority group
    170             let next = same_priority_indices[pos_in_group + 1];
    171             self.cursor = Some(next);
    172             Some(self.entries[next].session_id)
    173         }
    174     }
    175 
    176     pub fn prev(&mut self) -> Option<SessionId> {
    177         if self.entries.is_empty() {
    178             self.cursor = None;
    179             return None;
    180         }
    181         let cur = self.cursor.unwrap_or(0);
    182         let current_priority = self.entries[cur].priority;
    183 
    184         // Find all entries with the same priority
    185         let same_priority_indices: Vec<usize> = self
    186             .entries
    187             .iter()
    188             .enumerate()
    189             .filter(|(_, e)| e.priority == current_priority)
    190             .map(|(i, _)| i)
    191             .collect();
    192 
    193         // Find current position within same-priority items
    194         let pos_in_group = same_priority_indices
    195             .iter()
    196             .position(|&i| i == cur)
    197             .unwrap_or(0);
    198 
    199         // If at first item in group, try to drop to next lower priority level
    200         if pos_in_group == 0 {
    201             // Find first item with lower priority (higher index since sorted by priority desc)
    202             let last_same_priority = *same_priority_indices.last().unwrap();
    203             if last_same_priority + 1 < self.entries.len() {
    204                 // There's a lower priority group, go to first item of it
    205                 let next_idx = last_same_priority + 1;
    206                 self.cursor = Some(next_idx);
    207                 return Some(self.entries[next_idx].session_id);
    208             }
    209             // No lower priority group, wrap to last item in current group
    210             let prev = *same_priority_indices.last().unwrap();
    211             self.cursor = Some(prev);
    212             Some(self.entries[prev].session_id)
    213         } else {
    214             // Move backward within the same priority group
    215             let prev = same_priority_indices[pos_in_group - 1];
    216             self.cursor = Some(prev);
    217             Some(self.entries[prev].session_id)
    218         }
    219     }
    220 
    221     pub fn current(&self) -> Option<QueueEntry> {
    222         let i = self.cursor?;
    223         self.entries.get(i).copied()
    224     }
    225 
    226     pub fn current_position(&self) -> Option<usize> {
    227         Some(self.cursor? + 1) // 1-indexed
    228     }
    229 
    230     /// Get the raw cursor index (0-indexed)
    231     pub fn cursor_index(&self) -> Option<usize> {
    232         self.cursor
    233     }
    234 
    235     /// Set the cursor to a specific index, clamping to valid range
    236     pub fn set_cursor(&mut self, index: usize) {
    237         if self.entries.is_empty() {
    238             self.cursor = None;
    239         } else {
    240             self.cursor = Some(index.min(self.entries.len() - 1));
    241         }
    242     }
    243 
    244     /// Find the first entry with NeedsInput priority and return its index
    245     pub fn first_needs_input_index(&self) -> Option<usize> {
    246         self.entries
    247             .iter()
    248             .position(|e| e.priority == FocusPriority::NeedsInput)
    249     }
    250 
    251     /// Check if there are any NeedsInput items in the queue
    252     pub fn has_needs_input(&self) -> bool {
    253         self.entries
    254             .iter()
    255             .any(|e| e.priority == FocusPriority::NeedsInput)
    256     }
    257 
    258     pub fn ui_info(&self) -> Option<(usize, usize, FocusPriority)> {
    259         let entry = self.current()?;
    260         Some((self.current_position()?, self.len(), entry.priority))
    261     }
    262 
    263     pub fn update_from_statuses(
    264         &mut self,
    265         sessions: impl Iterator<Item = (SessionId, AgentStatus)>,
    266     ) {
    267         for (session_id, status) in sessions {
    268             let prev = self.previous_statuses.get(&session_id).copied();
    269             if prev != Some(status) {
    270                 if let Some(priority) = FocusPriority::from_status(status) {
    271                     self.enqueue(session_id, priority);
    272                 } else {
    273                     self.dequeue(session_id);
    274                 }
    275             }
    276             self.previous_statuses.insert(session_id, status);
    277         }
    278     }
    279 
    280     pub fn get_session_priority(&self, session_id: SessionId) -> Option<FocusPriority> {
    281         self.entries
    282             .iter()
    283             .find(|e| e.session_id == session_id)
    284             .map(|e| e.priority)
    285     }
    286 
    287     pub fn remove_session(&mut self, session_id: SessionId) {
    288         self.dequeue(session_id);
    289         self.previous_statuses.remove(&session_id);
    290     }
    291 }
    292 
    293 #[cfg(test)]
    294 mod tests {
    295     use super::*;
    296 
    297     fn session(id: u32) -> SessionId {
    298         id
    299     }
    300 
    301     #[test]
    302     fn test_empty_queue() {
    303         let mut queue = FocusQueue::new();
    304         assert!(queue.is_empty());
    305         assert_eq!(queue.next(), None);
    306         assert_eq!(queue.prev(), None);
    307         assert_eq!(queue.current(), None);
    308     }
    309 
    310     #[test]
    311     fn test_priority_ordering() {
    312         // Items should be sorted: NeedsInput -> Error -> Done
    313         let mut queue = FocusQueue::new();
    314 
    315         queue.enqueue(session(1), FocusPriority::Done);
    316         queue.enqueue(session(2), FocusPriority::NeedsInput);
    317         queue.enqueue(session(3), FocusPriority::Error);
    318 
    319         // Verify internal ordering: NeedsInput(2), Error(3), Done(1)
    320         assert_eq!(queue.entries[0].session_id, session(2));
    321         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    322         assert_eq!(queue.entries[1].session_id, session(3));
    323         assert_eq!(queue.entries[1].priority, FocusPriority::Error);
    324         assert_eq!(queue.entries[2].session_id, session(1));
    325         assert_eq!(queue.entries[2].priority, FocusPriority::Done);
    326 
    327         // Navigate to front (highest priority)
    328         queue.set_cursor(0);
    329         assert_eq!(queue.current().unwrap().session_id, session(2));
    330         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    331 
    332         // prev from NeedsInput should drop down to Error
    333         queue.prev();
    334         assert_eq!(queue.current().unwrap().priority, FocusPriority::Error);
    335 
    336         // prev from Error should drop down to Done
    337         queue.prev();
    338         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    339     }
    340 
    341     #[test]
    342     fn test_next_cycles_within_priority_then_wraps() {
    343         let mut queue = FocusQueue::new();
    344 
    345         // Add two NeedsInput items and one Done
    346         queue.enqueue(session(1), FocusPriority::NeedsInput);
    347         queue.enqueue(session(2), FocusPriority::NeedsInput);
    348         queue.enqueue(session(3), FocusPriority::Done);
    349 
    350         // Cursor starts at session 1 (first NeedsInput)
    351         queue.set_cursor(0);
    352         assert_eq!(queue.current().unwrap().session_id, session(1));
    353         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    354 
    355         // next should cycle to session 2 (also NeedsInput)
    356         let result = queue.next();
    357         assert_eq!(result, Some(session(2)));
    358         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    359 
    360         // next again should wrap back to session 1 (stays in NeedsInput, doesn't drop to Done)
    361         let result = queue.next();
    362         assert_eq!(result, Some(session(1)));
    363         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    364     }
    365 
    366     #[test]
    367     fn test_prev_drops_to_lower_priority() {
    368         let mut queue = FocusQueue::new();
    369 
    370         // Add two NeedsInput items and one Done
    371         queue.enqueue(session(1), FocusPriority::NeedsInput);
    372         queue.enqueue(session(2), FocusPriority::NeedsInput);
    373         queue.enqueue(session(3), FocusPriority::Done);
    374 
    375         // Start at first NeedsInput
    376         queue.set_cursor(0);
    377         assert_eq!(queue.current().unwrap().session_id, session(1));
    378 
    379         // prev from first in group should drop to Done (lower priority)
    380         let result = queue.prev();
    381         assert_eq!(result, Some(session(3)));
    382         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    383     }
    384 
    385     #[test]
    386     fn test_prev_cycles_within_priority_before_dropping() {
    387         let mut queue = FocusQueue::new();
    388 
    389         // Add two NeedsInput items and one Done
    390         queue.enqueue(session(1), FocusPriority::NeedsInput);
    391         queue.enqueue(session(2), FocusPriority::NeedsInput);
    392         queue.enqueue(session(3), FocusPriority::Done);
    393 
    394         // Start at second NeedsInput
    395         queue.set_cursor(1);
    396         assert_eq!(queue.current().unwrap().session_id, session(2));
    397 
    398         // prev should go to session 1 (earlier in same priority group)
    399         let result = queue.prev();
    400         assert_eq!(result, Some(session(1)));
    401         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    402 
    403         // prev again should now drop to Done
    404         let result = queue.prev();
    405         assert_eq!(result, Some(session(3)));
    406         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    407     }
    408 
    409     #[test]
    410     fn test_next_from_done_goes_to_needs_input() {
    411         let mut queue = FocusQueue::new();
    412 
    413         queue.enqueue(session(1), FocusPriority::Done);
    414         queue.enqueue(session(2), FocusPriority::NeedsInput);
    415         queue.enqueue(session(3), FocusPriority::Error);
    416 
    417         // Navigate to Done (only one item with this priority)
    418         queue.set_cursor(2); // Done is at index 2
    419         assert_eq!(queue.current().unwrap().session_id, session(1));
    420         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    421 
    422         // next from Done should go up to NeedsInput (highest priority)
    423         let result = queue.next();
    424         assert_eq!(result, Some(session(2)));
    425         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    426     }
    427 
    428     #[test]
    429     fn test_prev_from_lowest_wraps_within_group() {
    430         let mut queue = FocusQueue::new();
    431 
    432         // Only Done items, no lower priority to drop to
    433         queue.enqueue(session(1), FocusPriority::Done);
    434         queue.enqueue(session(2), FocusPriority::Done);
    435 
    436         queue.set_cursor(0);
    437         assert_eq!(queue.current().unwrap().session_id, session(1));
    438 
    439         // prev from first Done should wrap to last Done (no lower priority exists)
    440         let result = queue.prev();
    441         assert_eq!(result, Some(session(2)));
    442     }
    443 
    444     #[test]
    445     fn test_cursor_adjustment_on_higher_priority_insert() {
    446         let mut queue = FocusQueue::new();
    447 
    448         // Start with a Done item
    449         queue.enqueue(session(1), FocusPriority::Done);
    450         assert_eq!(queue.current().unwrap().session_id, session(1));
    451 
    452         // Insert a higher priority item - cursor should shift to keep pointing at same item
    453         queue.enqueue(session(2), FocusPriority::NeedsInput);
    454 
    455         // Cursor should still point to session 1 (now at index 1)
    456         assert_eq!(queue.current().unwrap().session_id, session(1));
    457 
    458         // next should go up to the new higher priority item
    459         queue.next();
    460         assert_eq!(queue.current().unwrap().session_id, session(2));
    461         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    462     }
    463 
    464     #[test]
    465     fn test_priority_upgrade() {
    466         let mut queue = FocusQueue::new();
    467 
    468         queue.enqueue(session(1), FocusPriority::Done);
    469         queue.enqueue(session(2), FocusPriority::Done);
    470 
    471         // Session 2 should be after session 1 (same priority, insertion order)
    472         assert_eq!(queue.entries[0].session_id, session(1));
    473         assert_eq!(queue.entries[1].session_id, session(2));
    474 
    475         // Upgrade session 2 to NeedsInput
    476         queue.enqueue(session(2), FocusPriority::NeedsInput);
    477 
    478         // Session 2 should now be first
    479         assert_eq!(queue.entries[0].session_id, session(2));
    480         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    481     }
    482 
    483     #[test]
    484     fn test_dequeue_adjusts_cursor() {
    485         let mut queue = FocusQueue::new();
    486 
    487         queue.enqueue(session(1), FocusPriority::NeedsInput);
    488         queue.enqueue(session(2), FocusPriority::Error);
    489         queue.enqueue(session(3), FocusPriority::Done);
    490 
    491         // Move cursor to session 2 (Error) using set_cursor
    492         queue.set_cursor(1);
    493         assert_eq!(queue.current().unwrap().session_id, session(2));
    494 
    495         // Remove session 1 (before cursor)
    496         queue.dequeue(session(1));
    497 
    498         // Cursor should adjust and still point to session 2
    499         assert_eq!(queue.current().unwrap().session_id, session(2));
    500     }
    501 
    502     #[test]
    503     fn test_single_item_navigation() {
    504         let mut queue = FocusQueue::new();
    505 
    506         queue.enqueue(session(1), FocusPriority::NeedsInput);
    507 
    508         // With single item, next and prev should both return that item
    509         assert_eq!(queue.next(), Some(session(1)));
    510         assert_eq!(queue.prev(), Some(session(1)));
    511         assert_eq!(queue.current().unwrap().session_id, session(1));
    512     }
    513 
    514     #[test]
    515     fn test_update_from_statuses() {
    516         let mut queue = FocusQueue::new();
    517 
    518         // Initial statuses - order matters for cursor position
    519         // First item added gets cursor, subsequent inserts shift it
    520         let statuses = vec![
    521             (session(1), AgentStatus::Done),
    522             (session(2), AgentStatus::NeedsInput),
    523             (session(3), AgentStatus::Working), // Should not be added (Idle/Working excluded)
    524         ];
    525         queue.update_from_statuses(statuses.into_iter());
    526 
    527         assert_eq!(queue.len(), 2);
    528         // Verify NeedsInput is first in priority order
    529         assert_eq!(queue.entries[0].session_id, session(2));
    530         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    531 
    532         // Update: session 2 becomes Idle (should be removed from queue)
    533         let statuses = vec![(session(2), AgentStatus::Idle)];
    534         queue.update_from_statuses(statuses.into_iter());
    535 
    536         assert_eq!(queue.len(), 1);
    537         assert_eq!(queue.current().unwrap().session_id, session(1));
    538     }
    539 }