1use 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
16pub 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
26pub 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
49pub 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
65pub 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
75pub 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#[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 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 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 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 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 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 pub fn routable_slots(&self) -> impl Iterator<Item = RosterSlot> + '_ {
255 (0..self.active.len()).filter(|slot| self.routable(*slot))
256 }
257
258 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 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 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 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 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 if queued.is_none() {
436 return RelayDecision::Complete;
437 }
438 self.rounds = 0;
439 }
440 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 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 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 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 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 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 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 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 assert!(matches!(
1110 relay.begin("", 0),
1111 RelayDecision::Dispatch { slot: 0, .. }
1112 ));
1113 relay.finish(0, false, false);
1114 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 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}