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