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