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