notedeck

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

focus_queue.rs (21822B)


      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     pub fn as_str(&self) -> &'static str {
     31         match self {
     32             Self::NeedsInput => "needs_input",
     33             Self::Error => "error",
     34             Self::Done => "done",
     35         }
     36     }
     37 
     38     pub fn from_indicator_str(s: &str) -> Option<Self> {
     39         match s {
     40             "needs_input" => Some(Self::NeedsInput),
     41             "error" => Some(Self::Error),
     42             "done" => Some(Self::Done),
     43             _ => None,
     44         }
     45     }
     46 }
     47 
     48 pub struct FocusQueueUpdate {
     49     pub new_needs_input: bool,
     50     pub changed: bool,
     51 }
     52 
     53 /// Auto-steal focus state machine.
     54 ///
     55 /// - `Disabled`: auto-steal is off, user controls focus manually.
     56 /// - `Idle`: auto-steal is on but no pending work.
     57 /// - `Pending`: auto-steal is on and a focus-queue transition was
     58 ///   detected that hasn't been acted on yet (retries across frames
     59 ///   if temporarily suppressed, e.g. user is typing).
     60 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     61 pub enum AutoStealState {
     62     Disabled,
     63     Idle,
     64     Pending,
     65 }
     66 
     67 impl AutoStealState {
     68     pub fn is_enabled(self) -> bool {
     69         matches!(self, Self::Idle | Self::Pending)
     70     }
     71 }
     72 
     73 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     74 pub struct QueueEntry {
     75     pub session_id: SessionId,
     76     pub priority: FocusPriority,
     77 }
     78 
     79 pub struct FocusQueue {
     80     entries: Vec<QueueEntry>, // kept sorted: NeedsInput -> Error -> Done
     81     cursor: Option<usize>,    // index into entries
     82     previous_indicators: HashMap<SessionId, Option<FocusPriority>>,
     83 }
     84 
     85 impl Default for FocusQueue {
     86     fn default() -> Self {
     87         Self::new()
     88     }
     89 }
     90 
     91 impl FocusQueue {
     92     pub fn new() -> Self {
     93         Self {
     94             entries: Vec::new(),
     95             cursor: None,
     96             previous_indicators: HashMap::new(),
     97         }
     98     }
     99 
    100     pub fn len(&self) -> usize {
    101         self.entries.len()
    102     }
    103 
    104     pub fn is_empty(&self) -> bool {
    105         self.entries.is_empty()
    106     }
    107 
    108     fn sort_key(p: FocusPriority) -> i32 {
    109         // want NeedsInput first, then Error, then Done
    110         -(p as i32)
    111     }
    112 
    113     fn find(&self, session_id: SessionId) -> Option<usize> {
    114         self.entries.iter().position(|e| e.session_id == session_id)
    115     }
    116 
    117     fn normalize_cursor_after_remove(&mut self, removed_idx: usize) {
    118         match self.cursor {
    119             None => {}
    120             Some(_cur) if self.entries.is_empty() => self.cursor = None,
    121             Some(cur) if removed_idx < cur => self.cursor = Some(cur - 1),
    122             Some(cur) if removed_idx == cur => {
    123                 // keep cursor pointing at a valid item (same index if possible, else last)
    124                 let new_cur = cur.min(self.entries.len().saturating_sub(1));
    125                 self.cursor = Some(new_cur);
    126             }
    127             Some(_) => {}
    128         }
    129     }
    130 
    131     /// Insert entry in priority order (stable within same priority).
    132     fn insert_sorted(&mut self, entry: QueueEntry) {
    133         let key = Self::sort_key(entry.priority);
    134         let pos = self
    135             .entries
    136             .iter()
    137             .position(|e| Self::sort_key(e.priority) > key)
    138             .unwrap_or(self.entries.len());
    139         self.entries.insert(pos, entry);
    140 
    141         // initialize cursor if this is the first item
    142         if self.cursor.is_none() && self.entries.len() == 1 {
    143             self.cursor = Some(0);
    144         } else if let Some(cur) = self.cursor {
    145             // if we inserted before the cursor, shift cursor right
    146             if pos <= cur {
    147                 self.cursor = Some(cur + 1);
    148             }
    149         }
    150     }
    151 
    152     pub fn enqueue(&mut self, session_id: SessionId, priority: FocusPriority) {
    153         if let Some(i) = self.find(session_id) {
    154             if self.entries[i].priority == priority {
    155                 return;
    156             }
    157             // remove old entry, then reinsert at correct spot
    158             self.entries.remove(i);
    159             self.normalize_cursor_after_remove(i);
    160         }
    161         self.insert_sorted(QueueEntry {
    162             session_id,
    163             priority,
    164         });
    165     }
    166 
    167     pub fn dequeue(&mut self, session_id: SessionId) {
    168         if let Some(i) = self.find(session_id) {
    169             self.entries.remove(i);
    170             self.normalize_cursor_after_remove(i);
    171         }
    172     }
    173 
    174     /// Remove a session from the queue only if its priority is Done.
    175     pub fn dequeue_done(&mut self, session_id: SessionId) {
    176         if let Some(i) = self.find(session_id) {
    177             if self.entries[i].priority == FocusPriority::Done {
    178                 self.entries.remove(i);
    179                 self.normalize_cursor_after_remove(i);
    180             }
    181         }
    182     }
    183 
    184     pub fn next(&mut self) -> Option<SessionId> {
    185         if self.entries.is_empty() {
    186             self.cursor = None;
    187             return None;
    188         }
    189         let cur = self.cursor.unwrap_or(0);
    190         let current_priority = self.entries[cur].priority;
    191 
    192         // Find all entries with the same priority
    193         let same_priority_indices: Vec<usize> = self
    194             .entries
    195             .iter()
    196             .enumerate()
    197             .filter(|(_, e)| e.priority == current_priority)
    198             .map(|(i, _)| i)
    199             .collect();
    200 
    201         // Find current position within same-priority items
    202         let pos_in_group = same_priority_indices
    203             .iter()
    204             .position(|&i| i == cur)
    205             .unwrap_or(0);
    206 
    207         // If at last item in group, try to go up to highest priority level
    208         if pos_in_group == same_priority_indices.len() - 1 {
    209             // Check if there's a higher priority group (entries are sorted highest first)
    210             let first_same_priority = same_priority_indices[0];
    211             if first_same_priority > 0 {
    212                 // Go to the very first item (highest priority)
    213                 self.cursor = Some(0);
    214                 return Some(self.entries[0].session_id);
    215             }
    216             // No higher priority group, wrap to first item in current group
    217             let next = same_priority_indices[0];
    218             self.cursor = Some(next);
    219             Some(self.entries[next].session_id)
    220         } else {
    221             // Move forward within the same priority group
    222             let next = same_priority_indices[pos_in_group + 1];
    223             self.cursor = Some(next);
    224             Some(self.entries[next].session_id)
    225         }
    226     }
    227 
    228     pub fn prev(&mut self) -> Option<SessionId> {
    229         if self.entries.is_empty() {
    230             self.cursor = None;
    231             return None;
    232         }
    233         let cur = self.cursor.unwrap_or(0);
    234         let current_priority = self.entries[cur].priority;
    235 
    236         // Find all entries with the same priority
    237         let same_priority_indices: Vec<usize> = self
    238             .entries
    239             .iter()
    240             .enumerate()
    241             .filter(|(_, e)| e.priority == current_priority)
    242             .map(|(i, _)| i)
    243             .collect();
    244 
    245         // Find current position within same-priority items
    246         let pos_in_group = same_priority_indices
    247             .iter()
    248             .position(|&i| i == cur)
    249             .unwrap_or(0);
    250 
    251         // If at first item in group, try to drop to next lower priority level
    252         if pos_in_group == 0 {
    253             // Find first item with lower priority (higher index since sorted by priority desc)
    254             let last_same_priority = *same_priority_indices.last().unwrap();
    255             if last_same_priority + 1 < self.entries.len() {
    256                 // There's a lower priority group, go to first item of it
    257                 let next_idx = last_same_priority + 1;
    258                 self.cursor = Some(next_idx);
    259                 return Some(self.entries[next_idx].session_id);
    260             }
    261             // No lower priority group, wrap to last item in current group
    262             let prev = *same_priority_indices.last().unwrap();
    263             self.cursor = Some(prev);
    264             Some(self.entries[prev].session_id)
    265         } else {
    266             // Move backward within the same priority group
    267             let prev = same_priority_indices[pos_in_group - 1];
    268             self.cursor = Some(prev);
    269             Some(self.entries[prev].session_id)
    270         }
    271     }
    272 
    273     pub fn current(&self) -> Option<QueueEntry> {
    274         let i = self.cursor?;
    275         self.entries.get(i).copied()
    276     }
    277 
    278     pub fn current_position(&self) -> Option<usize> {
    279         Some(self.cursor? + 1) // 1-indexed
    280     }
    281 
    282     /// Get the raw cursor index (0-indexed)
    283     pub fn cursor_index(&self) -> Option<usize> {
    284         self.cursor
    285     }
    286 
    287     /// Set the cursor to a specific index, clamping to valid range
    288     pub fn set_cursor(&mut self, index: usize) {
    289         if self.entries.is_empty() {
    290             self.cursor = None;
    291         } else {
    292             self.cursor = Some(index.min(self.entries.len() - 1));
    293         }
    294     }
    295 
    296     /// Find the first entry with NeedsInput priority and return its index
    297     pub fn first_needs_input_index(&self) -> Option<usize> {
    298         self.entries
    299             .iter()
    300             .position(|e| e.priority == FocusPriority::NeedsInput)
    301     }
    302 
    303     /// Check if there are any NeedsInput items in the queue
    304     pub fn has_needs_input(&self) -> bool {
    305         self.entries
    306             .iter()
    307             .any(|e| e.priority == FocusPriority::NeedsInput)
    308     }
    309 
    310     /// Find the first entry with Done priority and return its index
    311     pub fn first_done_index(&self) -> Option<usize> {
    312         self.entries
    313             .iter()
    314             .position(|e| e.priority == FocusPriority::Done)
    315     }
    316 
    317     /// Check if there are any Done items in the queue
    318     pub fn has_done(&self) -> bool {
    319         self.entries
    320             .iter()
    321             .any(|e| e.priority == FocusPriority::Done)
    322     }
    323 
    324     pub fn ui_info(&self) -> Option<(usize, usize, FocusPriority)> {
    325         let entry = self.current()?;
    326         Some((self.current_position()?, self.len(), entry.priority))
    327     }
    328 
    329     /// Update focus queue based on session indicator fields.
    330     pub fn update_from_indicators(
    331         &mut self,
    332         sessions: impl Iterator<Item = (SessionId, Option<FocusPriority>)>,
    333     ) -> FocusQueueUpdate {
    334         let mut new_needs_input = false;
    335         let mut changed = false;
    336         for (session_id, indicator) in sessions {
    337             let prev = self.previous_indicators.get(&session_id).copied();
    338             if prev != Some(indicator) {
    339                 changed = true;
    340                 if let Some(priority) = indicator {
    341                     self.enqueue(session_id, priority);
    342                     if priority == FocusPriority::NeedsInput {
    343                         new_needs_input = true;
    344                     }
    345                 } else {
    346                     self.dequeue(session_id);
    347                 }
    348             }
    349             self.previous_indicators.insert(session_id, indicator);
    350         }
    351         FocusQueueUpdate {
    352             new_needs_input,
    353             changed,
    354         }
    355     }
    356 
    357     pub fn get_session_priority(&self, session_id: SessionId) -> Option<FocusPriority> {
    358         self.entries
    359             .iter()
    360             .find(|e| e.session_id == session_id)
    361             .map(|e| e.priority)
    362     }
    363 
    364     pub fn remove_session(&mut self, session_id: SessionId) {
    365         self.dequeue(session_id);
    366         self.previous_indicators.remove(&session_id);
    367     }
    368 }
    369 
    370 #[cfg(test)]
    371 mod tests {
    372     use super::*;
    373 
    374     fn session(id: u32) -> SessionId {
    375         id
    376     }
    377 
    378     #[test]
    379     fn test_empty_queue() {
    380         let mut queue = FocusQueue::new();
    381         assert!(queue.is_empty());
    382         assert_eq!(queue.next(), None);
    383         assert_eq!(queue.prev(), None);
    384         assert_eq!(queue.current(), None);
    385     }
    386 
    387     #[test]
    388     fn test_priority_ordering() {
    389         // Items should be sorted: NeedsInput -> Error -> Done
    390         let mut queue = FocusQueue::new();
    391 
    392         queue.enqueue(session(1), FocusPriority::Done);
    393         queue.enqueue(session(2), FocusPriority::NeedsInput);
    394         queue.enqueue(session(3), FocusPriority::Error);
    395 
    396         // Verify internal ordering: NeedsInput(2), Error(3), Done(1)
    397         assert_eq!(queue.entries[0].session_id, session(2));
    398         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    399         assert_eq!(queue.entries[1].session_id, session(3));
    400         assert_eq!(queue.entries[1].priority, FocusPriority::Error);
    401         assert_eq!(queue.entries[2].session_id, session(1));
    402         assert_eq!(queue.entries[2].priority, FocusPriority::Done);
    403 
    404         // Navigate to front (highest priority)
    405         queue.set_cursor(0);
    406         assert_eq!(queue.current().unwrap().session_id, session(2));
    407         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    408 
    409         // prev from NeedsInput should drop down to Error
    410         queue.prev();
    411         assert_eq!(queue.current().unwrap().priority, FocusPriority::Error);
    412 
    413         // prev from Error should drop down to Done
    414         queue.prev();
    415         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    416     }
    417 
    418     #[test]
    419     fn test_next_cycles_within_priority_then_wraps() {
    420         let mut queue = FocusQueue::new();
    421 
    422         // Add two NeedsInput items and one Done
    423         queue.enqueue(session(1), FocusPriority::NeedsInput);
    424         queue.enqueue(session(2), FocusPriority::NeedsInput);
    425         queue.enqueue(session(3), FocusPriority::Done);
    426 
    427         // Cursor starts at session 1 (first NeedsInput)
    428         queue.set_cursor(0);
    429         assert_eq!(queue.current().unwrap().session_id, session(1));
    430         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    431 
    432         // next should cycle to session 2 (also NeedsInput)
    433         let result = queue.next();
    434         assert_eq!(result, Some(session(2)));
    435         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    436 
    437         // next again should wrap back to session 1 (stays in NeedsInput, doesn't drop to Done)
    438         let result = queue.next();
    439         assert_eq!(result, Some(session(1)));
    440         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    441     }
    442 
    443     #[test]
    444     fn test_prev_drops_to_lower_priority() {
    445         let mut queue = FocusQueue::new();
    446 
    447         // Add two NeedsInput items and one Done
    448         queue.enqueue(session(1), FocusPriority::NeedsInput);
    449         queue.enqueue(session(2), FocusPriority::NeedsInput);
    450         queue.enqueue(session(3), FocusPriority::Done);
    451 
    452         // Start at first NeedsInput
    453         queue.set_cursor(0);
    454         assert_eq!(queue.current().unwrap().session_id, session(1));
    455 
    456         // prev from first in group should drop to Done (lower priority)
    457         let result = queue.prev();
    458         assert_eq!(result, Some(session(3)));
    459         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    460     }
    461 
    462     #[test]
    463     fn test_prev_cycles_within_priority_before_dropping() {
    464         let mut queue = FocusQueue::new();
    465 
    466         // Add two NeedsInput items and one Done
    467         queue.enqueue(session(1), FocusPriority::NeedsInput);
    468         queue.enqueue(session(2), FocusPriority::NeedsInput);
    469         queue.enqueue(session(3), FocusPriority::Done);
    470 
    471         // Start at second NeedsInput
    472         queue.set_cursor(1);
    473         assert_eq!(queue.current().unwrap().session_id, session(2));
    474 
    475         // prev should go to session 1 (earlier in same priority group)
    476         let result = queue.prev();
    477         assert_eq!(result, Some(session(1)));
    478         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    479 
    480         // prev again should now drop to Done
    481         let result = queue.prev();
    482         assert_eq!(result, Some(session(3)));
    483         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    484     }
    485 
    486     #[test]
    487     fn test_next_from_done_goes_to_needs_input() {
    488         let mut queue = FocusQueue::new();
    489 
    490         queue.enqueue(session(1), FocusPriority::Done);
    491         queue.enqueue(session(2), FocusPriority::NeedsInput);
    492         queue.enqueue(session(3), FocusPriority::Error);
    493 
    494         // Navigate to Done (only one item with this priority)
    495         queue.set_cursor(2); // Done is at index 2
    496         assert_eq!(queue.current().unwrap().session_id, session(1));
    497         assert_eq!(queue.current().unwrap().priority, FocusPriority::Done);
    498 
    499         // next from Done should go up to NeedsInput (highest priority)
    500         let result = queue.next();
    501         assert_eq!(result, Some(session(2)));
    502         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    503     }
    504 
    505     #[test]
    506     fn test_prev_from_lowest_wraps_within_group() {
    507         let mut queue = FocusQueue::new();
    508 
    509         // Only Done items, no lower priority to drop to
    510         queue.enqueue(session(1), FocusPriority::Done);
    511         queue.enqueue(session(2), FocusPriority::Done);
    512 
    513         queue.set_cursor(0);
    514         assert_eq!(queue.current().unwrap().session_id, session(1));
    515 
    516         // prev from first Done should wrap to last Done (no lower priority exists)
    517         let result = queue.prev();
    518         assert_eq!(result, Some(session(2)));
    519     }
    520 
    521     #[test]
    522     fn test_cursor_adjustment_on_higher_priority_insert() {
    523         let mut queue = FocusQueue::new();
    524 
    525         // Start with a Done item
    526         queue.enqueue(session(1), FocusPriority::Done);
    527         assert_eq!(queue.current().unwrap().session_id, session(1));
    528 
    529         // Insert a higher priority item - cursor should shift to keep pointing at same item
    530         queue.enqueue(session(2), FocusPriority::NeedsInput);
    531 
    532         // Cursor should still point to session 1 (now at index 1)
    533         assert_eq!(queue.current().unwrap().session_id, session(1));
    534 
    535         // next should go up to the new higher priority item
    536         queue.next();
    537         assert_eq!(queue.current().unwrap().session_id, session(2));
    538         assert_eq!(queue.current().unwrap().priority, FocusPriority::NeedsInput);
    539     }
    540 
    541     #[test]
    542     fn test_priority_upgrade() {
    543         let mut queue = FocusQueue::new();
    544 
    545         queue.enqueue(session(1), FocusPriority::Done);
    546         queue.enqueue(session(2), FocusPriority::Done);
    547 
    548         // Session 2 should be after session 1 (same priority, insertion order)
    549         assert_eq!(queue.entries[0].session_id, session(1));
    550         assert_eq!(queue.entries[1].session_id, session(2));
    551 
    552         // Upgrade session 2 to NeedsInput
    553         queue.enqueue(session(2), FocusPriority::NeedsInput);
    554 
    555         // Session 2 should now be first
    556         assert_eq!(queue.entries[0].session_id, session(2));
    557         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    558     }
    559 
    560     #[test]
    561     fn test_dequeue_adjusts_cursor() {
    562         let mut queue = FocusQueue::new();
    563 
    564         queue.enqueue(session(1), FocusPriority::NeedsInput);
    565         queue.enqueue(session(2), FocusPriority::Error);
    566         queue.enqueue(session(3), FocusPriority::Done);
    567 
    568         // Move cursor to session 2 (Error) using set_cursor
    569         queue.set_cursor(1);
    570         assert_eq!(queue.current().unwrap().session_id, session(2));
    571 
    572         // Remove session 1 (before cursor)
    573         queue.dequeue(session(1));
    574 
    575         // Cursor should adjust and still point to session 2
    576         assert_eq!(queue.current().unwrap().session_id, session(2));
    577     }
    578 
    579     #[test]
    580     fn test_single_item_navigation() {
    581         let mut queue = FocusQueue::new();
    582 
    583         queue.enqueue(session(1), FocusPriority::NeedsInput);
    584 
    585         // With single item, next and prev should both return that item
    586         assert_eq!(queue.next(), Some(session(1)));
    587         assert_eq!(queue.prev(), Some(session(1)));
    588         assert_eq!(queue.current().unwrap().session_id, session(1));
    589     }
    590 
    591     #[test]
    592     fn test_update_from_indicators() {
    593         let mut queue = FocusQueue::new();
    594 
    595         // Initial indicators - order matters for cursor position
    596         let indicators = vec![
    597             (session(1), Some(FocusPriority::Done)),
    598             (session(2), Some(FocusPriority::NeedsInput)),
    599             (session(3), None), // No indicator = no dot
    600         ];
    601         let update = queue.update_from_indicators(indicators.into_iter());
    602 
    603         assert!(update.changed);
    604         assert!(update.new_needs_input);
    605         assert_eq!(queue.len(), 2);
    606         // Verify NeedsInput is first in priority order
    607         assert_eq!(queue.entries[0].session_id, session(2));
    608         assert_eq!(queue.entries[0].priority, FocusPriority::NeedsInput);
    609 
    610         // No-op: same indicators again → no change
    611         let indicators = vec![
    612             (session(1), Some(FocusPriority::Done)),
    613             (session(2), Some(FocusPriority::NeedsInput)),
    614         ];
    615         let update = queue.update_from_indicators(indicators.into_iter());
    616         assert!(!update.changed);
    617         assert!(!update.new_needs_input);
    618 
    619         // Update: session 2 indicator cleared (should be removed from queue)
    620         let indicators = vec![(session(2), None)];
    621         let update = queue.update_from_indicators(indicators.into_iter());
    622 
    623         assert!(update.changed);
    624         assert!(!update.new_needs_input);
    625         assert_eq!(queue.len(), 1);
    626         assert_eq!(queue.current().unwrap().session_id, session(1));
    627     }
    628 }