Skip to main content

codeswarm_adapters/
relay.rs

1//! Deterministic sequential roster scheduling.
2//!
3//! This owns turn selection only. Prompt construction and adapter I/O remain
4//! outside the scheduler, making the relay safe to replay and test.
5
6use std::collections::VecDeque;
7
8use crate::RosterSlot;
9use crate::collaboration::CollaborationContext;
10
11pub const MAX_QUEUED_PROMPTS: usize = 100;
12pub const STOP_TOKEN: &str = "[CODESWARM:STOP]";
13pub const DEFAULT_STOP_ACKNOWLEDGMENT: &str = "👍";
14
15/// Recognize a provider usage-limit reply (for example an exhausted Codex
16/// plan). Kept deliberately narrow so ordinary conversation mentioning
17/// "limits" is never misclassified.
18pub fn is_usage_limit_response(text: &str) -> bool {
19    let haystack = text.to_lowercase();
20    let exhausted_usage = haystack.contains("usage limit")
21        && ["hit", "reached", "exceeded"]
22            .iter()
23            .any(|marker| haystack.contains(marker));
24    exhausted_usage
25        || haystack.contains("insufficient_quota")
26        || haystack.contains("quota exceeded")
27        || haystack.contains("insufficient credits")
28}
29
30pub fn strip_stop_token(response: &str) -> (String, bool) {
31    let trimmed = response.trim_end();
32    let requested = trimmed.ends_with(STOP_TOKEN);
33    let visible = if requested {
34        trimmed[..trimmed.len() - STOP_TOKEN.len()]
35            .trim_end()
36            .to_owned()
37    } else {
38        response.to_owned()
39    };
40    (visible, requested)
41}
42
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub enum QueuedKind {
45    Steering,
46    Direct,
47}
48
49/// How a multi-agent session chooses its next non-direct recipient.
50///
51/// `Roster` is the normal sequential ring. `Pair` keeps the first two active
52/// agents in a tight review loop, which is useful when a
53/// larger saved roster is available but the user wants focused two-agent
54/// collaboration. `Manual` never advances on its own after the first turn;
55/// every subsequent turn must be explicitly targeted or queued by the user.
56#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
57pub enum CollaborationStrategy {
58    #[default]
59    Roster,
60    Manual,
61    Pair,
62}
63
64#[derive(Clone, Debug, Eq, PartialEq)]
65pub struct QueuedPrompt {
66    pub slot: RosterSlot,
67    pub prompt: String,
68    pub kind: QueuedKind,
69}
70
71#[derive(Clone, Debug, Eq, PartialEq)]
72pub enum RelayDecision {
73    Dispatch {
74        slot: RosterSlot,
75        prompt: String,
76        direct: bool,
77        can_stop: bool,
78    },
79    Paused,
80    Collapsed,
81    Complete,
82}
83
84#[derive(Clone, Debug, Eq, PartialEq)]
85pub struct Relay {
86    active: Vec<bool>,
87    /// Slots whose provider plan is exhausted. Unlike a tombstone this is
88    /// expected to clear (recharge or reload), so queued prompts targeting a
89    /// limited slot are preserved instead of discarded.
90    limited: Vec<bool>,
91    max_rounds: usize,
92    rounds: usize,
93    stopped: bool,
94    paused: bool,
95    last_active: RosterSlot,
96    next: Option<RosterSlot>,
97    steering: VecDeque<QueuedPrompt>,
98    direct: VecDeque<QueuedPrompt>,
99    previous_slot: Option<RosterSlot>,
100    context: CollaborationContext,
101    strategy: CollaborationStrategy,
102    pair_partner: Option<RosterSlot>,
103}
104
105impl Relay {
106    pub fn new(roster_size: usize, max_rounds: usize) -> Self {
107        assert!(roster_size >= 1);
108        assert!(max_rounds >= 1);
109        Self {
110            active: vec![true; roster_size],
111            limited: vec![false; roster_size],
112            max_rounds,
113            rounds: 0,
114            stopped: false,
115            paused: false,
116            last_active: 0,
117            next: None,
118            steering: VecDeque::new(),
119            direct: VecDeque::new(),
120            previous_slot: None,
121            context: CollaborationContext::new(roster_size),
122            strategy: CollaborationStrategy::Roster,
123            pair_partner: None,
124        }
125    }
126
127    pub fn strategy(&self) -> CollaborationStrategy {
128        self.strategy
129    }
130
131    /// Change routing for future turns. This does not discard queued work or
132    /// alter the public context journal; it only affects the next automatic
133    /// recipient. Pair selection is re-derived from the next `first` value.
134    pub fn set_strategy(&mut self, strategy: CollaborationStrategy) {
135        if self.strategy != strategy {
136            self.strategy = strategy;
137            self.pair_partner = None;
138        }
139    }
140
141    pub fn active_slots(&self) -> impl Iterator<Item = RosterSlot> + '_ {
142        self.active
143            .iter()
144            .enumerate()
145            .filter_map(|(slot, active)| active.then_some(slot))
146    }
147
148    pub fn pause(&mut self) {
149        self.paused = true;
150    }
151
152    pub fn resume(&mut self) {
153        self.paused = false;
154    }
155
156    pub fn tombstone(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
157        let active = self.active.get_mut(slot).ok_or("slot out of range")?;
158        *active = false;
159        Ok(())
160    }
161
162    /// Flag a slot whose provider plan is exhausted. Routing skips it but its
163    /// queued prompts and roster identity are preserved for a later recharge.
164    pub fn mark_limited(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
165        let limited = self.limited.get_mut(slot).ok_or("slot out of range")?;
166        *limited = true;
167        Ok(())
168    }
169
170    /// Clear a usage-limit flag after a recharge or adapter reload.
171    pub fn clear_limited(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
172        let limited = self.limited.get_mut(slot).ok_or("slot out of range")?;
173        *limited = false;
174        Ok(())
175    }
176
177    pub fn is_limited(&self, slot: RosterSlot) -> bool {
178        self.limited.get(slot).copied().unwrap_or(false)
179    }
180
181    /// A slot is routable when it is active and not usage-limited.
182    fn routable(&self, slot: RosterSlot) -> bool {
183        self.active.get(slot).copied().unwrap_or(false)
184            && !self.limited.get(slot).copied().unwrap_or(false)
185    }
186
187    /// Slots that can currently receive a turn.
188    pub fn routable_slots(&self) -> impl Iterator<Item = RosterSlot> + '_ {
189        (0..self.active.len()).filter(|slot| self.routable(*slot))
190    }
191
192    /// Whether any slot other than `excluded` can still receive a turn.
193    fn any_routable_except(&self, excluded: RosterSlot) -> bool {
194        self.routable_slots().any(|slot| slot != excluded)
195    }
196
197    pub fn reactivate(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
198        let active = self.active.get_mut(slot).ok_or("slot out of range")?;
199        *active = true;
200        self.context.rewind(slot);
201        Ok(())
202    }
203
204    pub fn drop_agent(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
205        if !self.active.get(slot).copied().ok_or("slot out of range")? {
206            return Ok(());
207        }
208        if self.active_slots().count() == 1 {
209            return Err("last active agent cannot be dropped");
210        }
211        self.tombstone(slot)?;
212        self.direct.retain(|queued| queued.slot != slot);
213        self.steering.retain(|queued| queued.slot != slot);
214        Ok(())
215    }
216
217    /// Exchange two live roster slots while preserving adapter identity in
218    /// queued work, routing cursors, and per-agent context watermarks.
219    pub fn swap_agents(
220        &mut self,
221        first: RosterSlot,
222        second: RosterSlot,
223    ) -> Result<(), &'static str> {
224        if first == second {
225            return Ok(());
226        }
227        if first >= self.active.len() || second >= self.active.len() {
228            return Err("roster slot out of range");
229        }
230        if !self.active[first] || !self.active[second] {
231            return Err("both roster slots must be active");
232        }
233        fn swap_targets(queue: &mut VecDeque<QueuedPrompt>, first: usize, second: usize) {
234            for queued in queue {
235                if queued.slot == first {
236                    queued.slot = second;
237                } else if queued.slot == second {
238                    queued.slot = first;
239                }
240            }
241        }
242        swap_targets(&mut self.direct, first, second);
243        swap_targets(&mut self.steering, first, second);
244        if self.last_active == first {
245            self.last_active = second;
246        } else if self.last_active == second {
247            self.last_active = first;
248        }
249        fn swap_option(cursor: &mut Option<usize>, first: usize, second: usize) {
250            if *cursor == Some(first) {
251                *cursor = Some(second);
252            } else if *cursor == Some(second) {
253                *cursor = Some(first);
254            }
255        }
256        swap_option(&mut self.next, first, second);
257        swap_option(&mut self.previous_slot, first, second);
258        swap_option(&mut self.pair_partner, first, second);
259        self.active.swap(first, second);
260        self.limited.swap(first, second);
261        self.context.swap_agents(first, second);
262        Ok(())
263    }
264
265    pub fn enqueue_human(
266        &mut self,
267        prompt: impl Into<String>,
268        selected: Option<RosterSlot>,
269    ) -> bool {
270        let prompt = prompt.into();
271        if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
272            return false;
273        }
274        let slot = selected.unwrap_or(self.last_active);
275        if !self.active.get(slot).copied().unwrap_or(false) {
276            return false;
277        }
278        self.steering.push_back(QueuedPrompt {
279            slot,
280            prompt,
281            kind: QueuedKind::Steering,
282        });
283        true
284    }
285
286    pub fn enqueue_direct(
287        &mut self,
288        slot: RosterSlot,
289        prompt: impl Into<String>,
290    ) -> Result<bool, &'static str> {
291        let prompt = prompt.into();
292        if !self.active.get(slot).copied().unwrap_or(false) {
293            return Err("direct target is not active");
294        }
295        if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
296            return Ok(false);
297        }
298        self.direct.push_back(QueuedPrompt {
299            slot,
300            prompt,
301            kind: QueuedKind::Direct,
302        });
303        Ok(true)
304    }
305
306    pub fn queued_count(&self) -> usize {
307        self.direct.len() + self.steering.len()
308    }
309
310    pub fn set_shared_task(&mut self, task: impl Into<String>) {
311        self.context.set_shared_task(task);
312    }
313
314    pub fn shared_task(&self) -> Option<&str> {
315        self.context.shared_task()
316    }
317
318    pub fn record_public(&mut self, speaker: impl Into<String>, text: impl Into<String>) {
319        self.context.record(speaker, text, &self.active);
320    }
321
322    pub fn mark_context_seen(&mut self, slot: RosterSlot) {
323        self.context.mark_seen(slot);
324    }
325
326    pub fn unseen_context(&mut self, slot: RosterSlot) -> String {
327        self.context.unseen(slot)
328    }
329
330    pub fn add_agent(&mut self) {
331        self.active.push(true);
332        self.limited.push(false);
333        self.context.add_agent();
334    }
335
336    /// Select the next causal turn. Direct work always precedes steering work.
337    pub fn begin(&mut self, initial_prompt: impl Into<String>, first: RosterSlot) -> RelayDecision {
338        if self.paused {
339            return RelayDecision::Paused;
340        }
341        if self.active_slots().next().is_none() {
342            return RelayDecision::Collapsed;
343        }
344        // Every live slot may be usage-limited; never spin the batch against
345        // an exhausted plan. The queued work survives for the next begin().
346        if !self.any_routable_except(usize::MAX) {
347            self.stopped = true;
348            return RelayDecision::Paused;
349        }
350        let queued = Self::pop_routable(&self.active, &self.limited, &mut self.direct)
351            .or_else(|| Self::pop_routable(&self.active, &self.limited, &mut self.steering));
352        // A reviewer stop ends only the current automatic batch. A later
353        // queued/user prompt starts a fresh batch without rebuilding the
354        // relay, while an unprompted handoff remains complete.
355        if self.stopped {
356            if queued.is_none() {
357                return RelayDecision::Complete;
358            }
359            self.stopped = false;
360            self.rounds = 0;
361        }
362        if self.rounds >= self.max_rounds {
363            // A queued human/direct prompt is a new batch; do not strand it
364            // behind the safety limit reached by the previous batch.
365            if queued.is_none() {
366                return RelayDecision::Complete;
367            }
368            self.rounds = 0;
369        }
370        // Manual mode is deliberately input-driven. A queued prompt (which
371        // includes a newly submitted human prompt) is still dispatched, but
372        // an unprompted call after a completed turn must not silently hand the
373        // conversation to another agent.
374        if self.strategy == CollaborationStrategy::Manual
375            && queued.is_none()
376            && self.previous_slot.is_some()
377        {
378            return RelayDecision::Complete;
379        }
380        let (slot, prompt, direct, human_prompt) = match queued {
381            Some(queued) => (
382                queued.slot,
383                queued.prompt,
384                queued.kind == QueuedKind::Direct,
385                queued.kind == QueuedKind::Steering,
386            ),
387            None => {
388                let slot = self.next_automatic_slot(first);
389                (slot, initial_prompt.into(), false, false)
390            }
391        };
392        // A human steering prompt starts a fresh review batch. Even when it
393        // targets a different slot from the preceding relay turn, that first
394        // responder must not be allowed to terminate the batch with the safe
395        // word. Only an automatic handoff after another agent's response is
396        // eligible to review-stop.
397        let can_stop =
398            !direct && !human_prompt && self.previous_slot.is_some_and(|previous| previous != slot);
399        self.last_active = slot;
400        self.rounds += 1;
401        RelayDecision::Dispatch {
402            slot,
403            prompt,
404            direct,
405            can_stop,
406        }
407    }
408
409    /// Finalize a dispatched turn and choose the next ring position. Direct
410    /// turns never become shared relay context.
411    pub fn finish(&mut self, slot: RosterSlot, direct: bool, accepted_stop: bool) {
412        self.next = Some(self.next_active(slot));
413        if !direct {
414            self.previous_slot = Some(slot);
415        }
416        if accepted_stop && self.direct.is_empty() && self.steering.is_empty() {
417            self.stopped = true;
418        }
419    }
420
421    /// Pop the first queued prompt whose target is routable. Prompts aimed at
422    /// a limited slot stay queued until the slot recovers.
423    fn pop_routable(
424        active: &[bool],
425        limited: &[bool],
426        queue: &mut VecDeque<QueuedPrompt>,
427    ) -> Option<QueuedPrompt> {
428        let position = queue.iter().position(|queued| {
429            active.get(queued.slot).copied().unwrap_or(false)
430                && !limited.get(queued.slot).copied().unwrap_or(false)
431        })?;
432        queue.remove(position)
433    }
434
435    fn first_active_from(&self, start: RosterSlot) -> RosterSlot {
436        (0..self.active.len())
437            .map(|offset| (start + offset) % self.active.len())
438            .find(|slot| self.routable(*slot))
439            .expect("callers require a routable roster")
440    }
441
442    fn next_active(&self, slot: RosterSlot) -> RosterSlot {
443        (1..=self.active.len())
444            .map(|offset| (slot + offset) % self.active.len())
445            .find(|candidate| self.routable(*candidate))
446            .expect("callers require a routable roster")
447    }
448
449    fn next_automatic_slot(&mut self, first: RosterSlot) -> RosterSlot {
450        match self.strategy {
451            CollaborationStrategy::Roster | CollaborationStrategy::Manual => self
452                .next
453                .filter(|slot| self.routable(*slot))
454                .unwrap_or_else(|| self.first_active_from(first)),
455            CollaborationStrategy::Pair => {
456                let primary = self.first_active_from(0);
457                let partner = if let Some(partner) = self.pair_partner {
458                    if self.routable(partner) && partner != primary {
459                        partner
460                    } else {
461                        let partner = self.next_active(primary);
462                        self.pair_partner = Some(partner);
463                        partner
464                    }
465                } else {
466                    let partner = if first != primary && self.routable(first) {
467                        first
468                    } else {
469                        self.next_active(primary)
470                    };
471                    self.pair_partner = Some(partner);
472                    partner
473                };
474                match self.previous_slot {
475                    Some(previous) if previous == primary && self.routable(partner) => partner,
476                    Some(previous) if previous == partner && self.routable(primary) => primary,
477                    Some(previous) if previous == primary || previous == partner => primary,
478                    _ => self.first_active_from(first),
479                }
480            }
481        }
482    }
483}
484
485#[cfg(test)]
486mod tests {
487    use super::{CollaborationStrategy, Relay, RelayDecision, STOP_TOKEN, strip_stop_token};
488
489    #[test]
490    fn relay_moves_around_the_ring_without_self_review() {
491        let mut relay = Relay::new(3, 10);
492        let first = relay.begin("task", 0);
493        assert!(matches!(
494            first,
495            RelayDecision::Dispatch {
496                slot: 0,
497                can_stop: false,
498                ..
499            }
500        ));
501        relay.finish(0, false, false);
502        let second = relay.begin("response", 0);
503        assert!(matches!(
504            second,
505            RelayDecision::Dispatch {
506                slot: 1,
507                can_stop: true,
508                ..
509            }
510        ));
511    }
512
513    #[test]
514    fn explicit_human_target_beats_ring_order() {
515        let mut relay = Relay::new(3, 10);
516        relay.begin("task", 0);
517        assert!(relay.enqueue_human("correction", Some(2)));
518        relay.finish(0, false, false);
519        assert!(matches!(
520            relay.begin("response", 0),
521            RelayDecision::Dispatch { slot: 2, prompt, direct: false, can_stop: false } if prompt == "correction"
522        ));
523    }
524
525    #[test]
526    fn human_prompt_cannot_stop_even_after_a_previous_relay_batch() {
527        let mut relay = Relay::new(2, 10);
528        assert!(matches!(
529            relay.begin("first task", 0),
530            RelayDecision::Dispatch {
531                slot: 0,
532                can_stop: false,
533                ..
534            }
535        ));
536        relay.finish(0, false, false);
537        assert!(matches!(
538            relay.begin("review", 0),
539            RelayDecision::Dispatch {
540                slot: 1,
541                can_stop: true,
542                ..
543            }
544        ));
545        relay.finish(1, false, false);
546
547        // The next user prompt targets the first agent, which differs from the
548        // previous reviewer. It is still the first response to a human turn,
549        // so it must not receive reviewer stop permission.
550        assert!(relay.enqueue_human("new task", Some(0)));
551        assert!(matches!(
552            relay.begin("", 0),
553            RelayDecision::Dispatch {
554                slot: 0,
555                can_stop: false,
556                ..
557            }
558        ));
559    }
560
561    #[test]
562    fn direct_work_has_priority_and_any_agent_except_the_last_can_be_dropped() {
563        let mut relay = Relay::new(3, 10);
564        relay.enqueue_human("ordinary", Some(1));
565        assert_eq!(relay.enqueue_direct(2, "private"), Ok(true));
566        assert!(matches!(
567            relay.begin("task", 0),
568            RelayDecision::Dispatch {
569                slot: 2,
570                direct: true,
571                ..
572            }
573        ));
574        assert_eq!(relay.drop_agent(0), Ok(()));
575        assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![1, 2]);
576        assert_eq!(relay.drop_agent(1), Ok(()));
577        assert_eq!(
578            relay.drop_agent(2),
579            Err("last active agent cannot be dropped")
580        );
581    }
582
583    #[test]
584    fn relay_context_tracks_public_updates_per_slot() {
585        let mut relay = Relay::new(2, 10);
586        relay.set_shared_task("refactor");
587        relay.record_public("Agent 0", "first answer");
588        relay.mark_context_seen(0);
589        assert_eq!(relay.unseen_context(0), "");
590        assert_eq!(relay.unseen_context(1), "Agent 0:\nfirst answer");
591        assert_eq!(relay.shared_task(), Some("refactor"));
592        relay.add_agent();
593        assert_eq!(relay.active_slots().count(), 3);
594    }
595
596    #[test]
597    fn stop_token_is_stripped_only_from_the_response_suffix() {
598        let (visible, requested) = strip_stop_token(&format!("looks good\n{STOP_TOKEN}"));
599        assert_eq!(visible, "looks good");
600        assert!(requested);
601        let (visible, requested) = strip_stop_token("ordinary response");
602        assert_eq!(visible, "ordinary response");
603        assert!(!requested);
604    }
605
606    #[test]
607    fn accepted_stop_ends_the_batch_but_a_new_prompt_can_start_one() {
608        let mut relay = Relay::new(2, 10);
609        assert!(matches!(
610            relay.begin("task", 0),
611            RelayDecision::Dispatch { slot: 0, .. }
612        ));
613        relay.finish(0, false, false);
614        assert!(matches!(
615            relay.begin("review", 0),
616            RelayDecision::Dispatch {
617                slot: 1,
618                can_stop: true,
619                ..
620            }
621        ));
622        relay.finish(1, false, true);
623        assert_eq!(relay.begin("", 0), RelayDecision::Complete);
624
625        assert!(relay.enqueue_human("new task", Some(0)));
626        assert!(matches!(
627            relay.begin("", 0),
628            RelayDecision::Dispatch { slot: 0, prompt, .. } if prompt == "new task"
629        ));
630    }
631
632    #[test]
633    fn live_slot_swap_follows_queued_targets_and_runtime_cursors() {
634        let mut relay = Relay::new(3, 10);
635        relay.record_public("Agent 0", "first work");
636        relay.mark_context_seen(0);
637        assert!(matches!(
638            relay.begin("task", 0),
639            RelayDecision::Dispatch { slot: 0, .. }
640        ));
641        relay.finish(0, false, false);
642        relay.enqueue_human("to first", Some(0));
643        relay.enqueue_direct(2, "to third").expect("queue direct");
644
645        relay.swap_agents(0, 2).expect("swap live slots");
646        assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![0, 1, 2]);
647        assert!(matches!(
648            relay.begin("", 0),
649            RelayDecision::Dispatch { slot: 0, direct: true, prompt, .. }
650                if prompt == "to third"
651        ));
652        relay.finish(0, true, false);
653        assert!(matches!(
654            relay.begin("", 0),
655            RelayDecision::Dispatch { slot: 2, direct: false, prompt, .. }
656                if prompt == "to first"
657        ));
658        assert_eq!(relay.unseen_context(0), "Agent 0:\nfirst work");
659    }
660
661    #[test]
662    fn stop_token_is_not_allowed_on_the_first_response() {
663        let mut relay = Relay::new(2, 10);
664        assert!(matches!(
665            relay.begin("task", 0),
666            RelayDecision::Dispatch {
667                slot: 0,
668                can_stop: false,
669                ..
670            }
671        ));
672        // RelayHost validates the token against `can_stop`; the first turn
673        // therefore finalizes as a normal response even if the agent tried
674        // to include the token.
675        relay.finish(0, false, false);
676        assert!(matches!(
677            relay.begin("", 0),
678            RelayDecision::Dispatch { slot: 1, .. }
679        ));
680    }
681
682    #[test]
683    fn a_healthy_peer_continues_after_the_other_slot_is_tombstoned() {
684        let mut relay = Relay::new(2, 10);
685        relay.tombstone(0).expect("first agent failure");
686        assert!(matches!(
687            relay.begin("continue", 0),
688            RelayDecision::Dispatch {
689                slot: 1,
690                can_stop: false,
691                ..
692            }
693        ));
694    }
695
696    #[test]
697    fn manual_strategy_requires_an_explicit_follow_up_prompt() {
698        let mut relay = Relay::new(3, 10);
699        relay.set_strategy(CollaborationStrategy::Manual);
700        assert!(matches!(
701            relay.begin("task", 0),
702            RelayDecision::Dispatch { slot: 0, .. }
703        ));
704        relay.finish(0, false, false);
705        assert_eq!(
706            relay.begin("would auto advance", 0),
707            RelayDecision::Complete
708        );
709        assert!(relay.enqueue_human("review", Some(2)));
710        assert!(
711            matches!(relay.begin("", 0), RelayDecision::Dispatch { slot: 2, prompt, .. } if prompt == "review")
712        );
713    }
714
715    #[test]
716    fn pair_strategy_alternates_the_first_two_active_agents() {
717        let mut relay = Relay::new(4, 10);
718        relay.set_strategy(CollaborationStrategy::Pair);
719        assert!(matches!(
720            relay.begin("task", 2),
721            RelayDecision::Dispatch { slot: 2, .. }
722        ));
723        relay.finish(2, false, false);
724        assert!(matches!(
725            relay.begin("review", 2),
726            RelayDecision::Dispatch { slot: 0, .. }
727        ));
728        relay.finish(0, false, false);
729        assert!(matches!(
730            relay.begin("next", 2),
731            RelayDecision::Dispatch { slot: 2, .. }
732        ));
733
734        relay.drop_agent(0).expect("remove first agent");
735        relay.finish(2, false, false);
736        assert!(matches!(
737            relay.begin("after removal", 2),
738            RelayDecision::Dispatch { slot: 1, .. }
739        ));
740    }
741
742    #[test]
743    fn usage_limit_detection_matches_provider_copy() {
744        assert!(super::is_usage_limit_response(
745            "You've hit your usage limit. Visit chatgpt.com to purchase more credits."
746        ));
747        assert!(super::is_usage_limit_response("Error: insufficient_quota"));
748        assert!(super::is_usage_limit_response(
749            "Monthly quota exceeded for this plan"
750        ));
751        assert!(!super::is_usage_limit_response(
752            "The rate limit on the build job slowed things down."
753        ));
754        assert!(!super::is_usage_limit_response(
755            "I updated the usage-limit documentation and billing upgrade flow."
756        ));
757        assert!(!super::is_usage_limit_response("Ready to review the diff."));
758    }
759
760    #[test]
761    fn hot_added_and_swapped_slots_keep_limit_state_with_the_agent() {
762        let mut relay = Relay::new(2, 10);
763        relay.add_agent();
764        relay.mark_limited(2).expect("mark hot-added slot limited");
765        assert!(relay.is_limited(2));
766
767        relay.swap_agents(0, 2).expect("swap limited agent");
768        assert!(relay.is_limited(0));
769        assert!(!relay.is_limited(2));
770        assert!(matches!(
771            relay.begin("task", 0),
772            RelayDecision::Dispatch { slot: 1, .. }
773        ));
774    }
775
776    #[test]
777    fn a_limited_slot_is_routed_around_until_cleared() {
778        let mut relay = Relay::new(2, 10);
779        relay.mark_limited(0).expect("mark limited");
780        assert!(relay.is_limited(0));
781        assert!(matches!(
782            relay.begin("task", 0),
783            RelayDecision::Dispatch {
784                slot: 1,
785                can_stop: false,
786                ..
787            }
788        ));
789        relay.finish(1, false, false);
790        // The ring still skips the limited slot after a full loop.
791        assert!(matches!(
792            relay.begin("again", 1),
793            RelayDecision::Dispatch { slot: 1, .. }
794        ));
795        relay.clear_limited(0).expect("clear limited");
796        assert!(!relay.is_limited(0));
797        relay.finish(1, false, false);
798        assert!(matches!(
799            relay.begin("next", 1),
800            RelayDecision::Dispatch { slot: 0, .. }
801        ));
802    }
803
804    #[test]
805    fn prompts_targeting_a_limited_slot_wait_for_recovery() {
806        let mut relay = Relay::new(2, 10);
807        relay.mark_limited(1).expect("mark limited");
808        assert_eq!(relay.enqueue_direct(1, "private work"), Ok(true));
809        assert!(matches!(
810            relay.begin("task", 0),
811            RelayDecision::Dispatch { slot: 0, .. }
812        ));
813        relay.finish(0, false, false);
814        // The direct prompt is not dropped while slot 1 is limited...
815        assert!(matches!(
816            relay.begin("", 0),
817            RelayDecision::Dispatch { slot: 0, .. }
818        ));
819        relay.finish(0, false, false);
820        // ...and dispatches once the slot recovers.
821        relay.clear_limited(1).expect("clear limited");
822        assert!(matches!(
823            relay.begin("", 0),
824            RelayDecision::Dispatch { slot: 1, direct: true, prompt, .. } if prompt == "private work"
825        ));
826    }
827
828    #[test]
829    fn an_all_limited_roster_pauses_instead_of_spinning() {
830        let mut relay = Relay::new(2, 10);
831        relay.mark_limited(0).expect("mark limited");
832        relay.mark_limited(1).expect("mark limited");
833        assert_eq!(relay.begin("task", 0), RelayDecision::Paused);
834        // Queued work is preserved for the next begin().
835        assert!(relay.enqueue_human("queued", Some(1)));
836        relay.clear_limited(1).expect("recharge one agent");
837        assert!(matches!(
838            relay.begin("", 0),
839            RelayDecision::Dispatch { slot: 1, prompt, .. } if prompt == "queued"
840        ));
841    }
842}