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