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