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 strip_stop_token(response: &str) -> (String, bool) {
16 let trimmed = response.trim_end();
17 let requested = trimmed.ends_with(STOP_TOKEN);
18 let visible = if requested {
19 trimmed[..trimmed.len() - STOP_TOKEN.len()]
20 .trim_end()
21 .to_owned()
22 } else {
23 response.to_owned()
24 };
25 (visible, requested)
26}
27
28#[derive(Clone, Copy, Debug, Eq, PartialEq)]
29pub enum QueuedKind {
30 Steering,
31 Direct,
32}
33
34#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
42pub enum CollaborationStrategy {
43 #[default]
44 Roster,
45 Manual,
46 Pair,
47}
48
49#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct QueuedPrompt {
51 pub slot: RosterSlot,
52 pub prompt: String,
53 pub kind: QueuedKind,
54}
55
56#[derive(Clone, Debug, Eq, PartialEq)]
57pub enum RelayDecision {
58 Dispatch {
59 slot: RosterSlot,
60 prompt: String,
61 direct: bool,
62 can_stop: bool,
63 },
64 Paused,
65 Collapsed,
66 Complete,
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct Relay {
71 active: Vec<bool>,
72 max_rounds: usize,
73 rounds: usize,
74 stopped: bool,
75 paused: bool,
76 last_active: RosterSlot,
77 next: Option<RosterSlot>,
78 steering: VecDeque<QueuedPrompt>,
79 direct: VecDeque<QueuedPrompt>,
80 previous_slot: Option<RosterSlot>,
81 context: CollaborationContext,
82 strategy: CollaborationStrategy,
83 pair_partner: Option<RosterSlot>,
84}
85
86impl Relay {
87 pub fn new(roster_size: usize, max_rounds: usize) -> Self {
88 assert!(roster_size >= 1);
89 assert!(max_rounds >= 1);
90 Self {
91 active: vec![true; roster_size],
92 max_rounds,
93 rounds: 0,
94 stopped: false,
95 paused: false,
96 last_active: 0,
97 next: None,
98 steering: VecDeque::new(),
99 direct: VecDeque::new(),
100 previous_slot: None,
101 context: CollaborationContext::new(roster_size),
102 strategy: CollaborationStrategy::Roster,
103 pair_partner: None,
104 }
105 }
106
107 pub fn strategy(&self) -> CollaborationStrategy {
108 self.strategy
109 }
110
111 pub fn set_strategy(&mut self, strategy: CollaborationStrategy) {
115 if self.strategy != strategy {
116 self.strategy = strategy;
117 self.pair_partner = None;
118 }
119 }
120
121 pub fn active_slots(&self) -> impl Iterator<Item = RosterSlot> + '_ {
122 self.active
123 .iter()
124 .enumerate()
125 .filter_map(|(slot, active)| active.then_some(slot))
126 }
127
128 pub fn pause(&mut self) {
129 self.paused = true;
130 }
131
132 pub fn resume(&mut self) {
133 self.paused = false;
134 }
135
136 pub fn tombstone(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
137 let active = self.active.get_mut(slot).ok_or("slot out of range")?;
138 *active = false;
139 Ok(())
140 }
141
142 pub fn reactivate(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
143 let active = self.active.get_mut(slot).ok_or("slot out of range")?;
144 *active = true;
145 self.context.rewind(slot);
146 Ok(())
147 }
148
149 pub fn drop_agent(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
150 if slot == 0 {
151 return Err("owner cannot be dropped");
152 }
153 self.tombstone(slot)?;
154 self.direct.retain(|queued| queued.slot != slot);
155 self.steering.retain(|queued| queued.slot != slot);
156 Ok(())
157 }
158
159 pub fn promote_owner(&mut self, slot: RosterSlot) -> Result<(), &'static str> {
163 if slot == 0 {
164 return Ok(());
165 }
166 if !self.active.get(slot).copied().ok_or("slot out of range")? {
167 return Err("replacement owner is not active");
168 }
169
170 fn swap_targets(queue: &mut VecDeque<QueuedPrompt>, first: usize, second: usize) {
174 for queued in queue {
175 if queued.slot == first {
176 queued.slot = second;
177 } else if queued.slot == second {
178 queued.slot = first;
179 }
180 }
181 }
182 swap_targets(&mut self.direct, 0, slot);
183 swap_targets(&mut self.steering, 0, slot);
184 if self.last_active == 0 {
185 self.last_active = slot;
186 } else if self.last_active == slot {
187 self.last_active = 0;
188 }
189 if self.next == Some(0) {
190 self.next = Some(slot);
191 } else if self.next == Some(slot) {
192 self.next = Some(0);
193 }
194 if self.previous_slot == Some(0) {
195 self.previous_slot = Some(slot);
196 } else if self.previous_slot == Some(slot) {
197 self.previous_slot = Some(0);
198 }
199 if self.pair_partner == Some(slot) {
200 self.pair_partner = Some(0);
201 } else if self.pair_partner == Some(0) {
202 self.pair_partner = Some(slot);
203 }
204
205 self.context.swap_agents(0, slot);
206 self.context.rewind(0);
207 self.active[0] = true;
208 self.active[slot] = false;
209 Ok(())
210 }
211
212 pub fn swap_agents(
215 &mut self,
216 first: RosterSlot,
217 second: RosterSlot,
218 ) -> Result<(), &'static str> {
219 if first == second {
220 return Ok(());
221 }
222 if first >= self.active.len() || second >= self.active.len() {
223 return Err("roster slot out of range");
224 }
225 if !self.active[first] || !self.active[second] {
226 return Err("both roster slots must be active");
227 }
228 fn swap_targets(queue: &mut VecDeque<QueuedPrompt>, first: usize, second: usize) {
229 for queued in queue {
230 if queued.slot == first {
231 queued.slot = second;
232 } else if queued.slot == second {
233 queued.slot = first;
234 }
235 }
236 }
237 swap_targets(&mut self.direct, first, second);
238 swap_targets(&mut self.steering, first, second);
239 if self.last_active == first {
240 self.last_active = second;
241 } else if self.last_active == second {
242 self.last_active = first;
243 }
244 fn swap_option(cursor: &mut Option<usize>, first: usize, second: usize) {
245 if *cursor == Some(first) {
246 *cursor = Some(second);
247 } else if *cursor == Some(second) {
248 *cursor = Some(first);
249 }
250 }
251 swap_option(&mut self.next, first, second);
252 swap_option(&mut self.previous_slot, first, second);
253 swap_option(&mut self.pair_partner, first, second);
254 self.active.swap(first, second);
255 self.context.swap_agents(first, second);
256 Ok(())
257 }
258
259 pub fn enqueue_human(
260 &mut self,
261 prompt: impl Into<String>,
262 selected: Option<RosterSlot>,
263 ) -> bool {
264 let prompt = prompt.into();
265 if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
266 return false;
267 }
268 let slot = selected.unwrap_or(self.last_active);
269 if !self.active.get(slot).copied().unwrap_or(false) {
270 return false;
271 }
272 self.steering.push_back(QueuedPrompt {
273 slot,
274 prompt,
275 kind: QueuedKind::Steering,
276 });
277 true
278 }
279
280 pub fn enqueue_direct(
281 &mut self,
282 slot: RosterSlot,
283 prompt: impl Into<String>,
284 ) -> Result<bool, &'static str> {
285 let prompt = prompt.into();
286 if !self.active.get(slot).copied().unwrap_or(false) {
287 return Err("direct target is not active");
288 }
289 if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
290 return Ok(false);
291 }
292 self.direct.push_back(QueuedPrompt {
293 slot,
294 prompt,
295 kind: QueuedKind::Direct,
296 });
297 Ok(true)
298 }
299
300 pub fn queued_count(&self) -> usize {
301 self.direct.len() + self.steering.len()
302 }
303
304 pub fn set_shared_task(&mut self, task: impl Into<String>) {
305 self.context.set_shared_task(task);
306 }
307
308 pub fn shared_task(&self) -> Option<&str> {
309 self.context.shared_task()
310 }
311
312 pub fn record_public(&mut self, speaker: impl Into<String>, text: impl Into<String>) {
313 self.context.record(speaker, text, &self.active);
314 }
315
316 pub fn mark_context_seen(&mut self, slot: RosterSlot) {
317 self.context.mark_seen(slot);
318 }
319
320 pub fn unseen_context(&mut self, slot: RosterSlot) -> String {
321 self.context.unseen(slot)
322 }
323
324 pub fn add_agent(&mut self) {
325 self.active.push(true);
326 self.context.add_agent();
327 }
328
329 pub fn begin(&mut self, initial_prompt: impl Into<String>, first: RosterSlot) -> RelayDecision {
331 if self.paused {
332 return RelayDecision::Paused;
333 }
334 if self.active_slots().next().is_none() {
335 return RelayDecision::Collapsed;
336 }
337 let queued = Self::pop_active(&self.active, &mut self.direct)
338 .or_else(|| Self::pop_active(&self.active, &mut self.steering));
339 if self.stopped {
343 if queued.is_none() {
344 return RelayDecision::Complete;
345 }
346 self.stopped = false;
347 self.rounds = 0;
348 }
349 if self.rounds >= self.max_rounds {
350 if queued.is_none() {
353 return RelayDecision::Complete;
354 }
355 self.rounds = 0;
356 }
357 if self.strategy == CollaborationStrategy::Manual
362 && queued.is_none()
363 && self.previous_slot.is_some()
364 {
365 return RelayDecision::Complete;
366 }
367 let (slot, prompt, direct, human_prompt) = match queued {
368 Some(queued) => (
369 queued.slot,
370 queued.prompt,
371 queued.kind == QueuedKind::Direct,
372 queued.kind == QueuedKind::Steering,
373 ),
374 None => {
375 let slot = self.next_automatic_slot(first);
376 (slot, initial_prompt.into(), false, false)
377 }
378 };
379 let can_stop =
385 !direct && !human_prompt && self.previous_slot.is_some_and(|previous| previous != slot);
386 self.last_active = slot;
387 self.rounds += 1;
388 RelayDecision::Dispatch {
389 slot,
390 prompt,
391 direct,
392 can_stop,
393 }
394 }
395
396 pub fn finish(&mut self, slot: RosterSlot, direct: bool, accepted_stop: bool) {
399 self.next = Some(self.next_active(slot));
400 if !direct {
401 self.previous_slot = Some(slot);
402 }
403 if accepted_stop && self.direct.is_empty() && self.steering.is_empty() {
404 self.stopped = true;
405 }
406 }
407
408 fn pop_active(active: &[bool], queue: &mut VecDeque<QueuedPrompt>) -> Option<QueuedPrompt> {
409 let position = queue
410 .iter()
411 .position(|queued| active.get(queued.slot).copied().unwrap_or(false))?;
412 queue.remove(position)
413 }
414
415 fn first_active_from(&self, start: RosterSlot) -> RosterSlot {
416 (0..self.active.len())
417 .map(|offset| (start + offset) % self.active.len())
418 .find(|slot| self.active[*slot])
419 .expect("callers require an active roster")
420 }
421
422 fn next_active(&self, slot: RosterSlot) -> RosterSlot {
423 (1..=self.active.len())
424 .map(|offset| (slot + offset) % self.active.len())
425 .find(|candidate| self.active[*candidate])
426 .expect("callers require an active roster")
427 }
428
429 fn next_automatic_slot(&mut self, first: RosterSlot) -> RosterSlot {
430 match self.strategy {
431 CollaborationStrategy::Roster | CollaborationStrategy::Manual => self
432 .next
433 .filter(|slot| self.active[*slot])
434 .unwrap_or_else(|| self.first_active_from(first)),
435 CollaborationStrategy::Pair => {
436 let partner = if let Some(partner) = self.pair_partner {
440 partner
441 } else {
442 let partner = if first != 0 && self.active.get(first).copied().unwrap_or(false)
443 {
444 first
445 } else {
446 self.first_active_from(1)
447 };
448 self.pair_partner = Some(partner);
449 partner
450 };
451 match self.previous_slot {
452 Some(previous) if previous == 0 && self.active[partner] => partner,
453 Some(previous) if previous == partner && self.active[0] => 0,
454 Some(previous) if previous == 0 || previous == partner => {
455 self.first_active_from(0)
456 }
457 _ => self.first_active_from(first),
458 }
459 }
460 }
461 }
462}
463
464#[cfg(test)]
465mod tests {
466 use super::{CollaborationStrategy, Relay, RelayDecision, STOP_TOKEN, strip_stop_token};
467
468 #[test]
469 fn relay_moves_around_the_ring_without_self_review() {
470 let mut relay = Relay::new(3, 10);
471 let first = relay.begin("task", 0);
472 assert!(matches!(
473 first,
474 RelayDecision::Dispatch {
475 slot: 0,
476 can_stop: false,
477 ..
478 }
479 ));
480 relay.finish(0, false, false);
481 let second = relay.begin("response", 0);
482 assert!(matches!(
483 second,
484 RelayDecision::Dispatch {
485 slot: 1,
486 can_stop: true,
487 ..
488 }
489 ));
490 }
491
492 #[test]
493 fn explicit_human_target_beats_ring_order() {
494 let mut relay = Relay::new(3, 10);
495 relay.begin("task", 0);
496 assert!(relay.enqueue_human("correction", Some(2)));
497 relay.finish(0, false, false);
498 assert!(matches!(
499 relay.begin("response", 0),
500 RelayDecision::Dispatch { slot: 2, prompt, direct: false, can_stop: false } if prompt == "correction"
501 ));
502 }
503
504 #[test]
505 fn human_prompt_cannot_stop_even_after_a_previous_relay_batch() {
506 let mut relay = Relay::new(2, 10);
507 assert!(matches!(
508 relay.begin("first task", 0),
509 RelayDecision::Dispatch {
510 slot: 0,
511 can_stop: false,
512 ..
513 }
514 ));
515 relay.finish(0, false, false);
516 assert!(matches!(
517 relay.begin("review", 0),
518 RelayDecision::Dispatch {
519 slot: 1,
520 can_stop: true,
521 ..
522 }
523 ));
524 relay.finish(1, false, false);
525
526 assert!(relay.enqueue_human("new task", Some(0)));
530 assert!(matches!(
531 relay.begin("", 0),
532 RelayDecision::Dispatch {
533 slot: 0,
534 can_stop: false,
535 ..
536 }
537 ));
538 }
539
540 #[test]
541 fn direct_work_has_priority_and_owner_is_not_droppable() {
542 let mut relay = Relay::new(3, 10);
543 relay.enqueue_human("ordinary", Some(1));
544 assert_eq!(relay.enqueue_direct(2, "private"), Ok(true));
545 assert!(matches!(
546 relay.begin("task", 0),
547 RelayDecision::Dispatch {
548 slot: 2,
549 direct: true,
550 ..
551 }
552 ));
553 assert_eq!(relay.drop_agent(0), Err("owner cannot be dropped"));
554 }
555
556 #[test]
557 fn relay_context_tracks_public_updates_per_slot() {
558 let mut relay = Relay::new(2, 10);
559 relay.set_shared_task("refactor");
560 relay.record_public("Agent 0", "first answer");
561 relay.mark_context_seen(0);
562 assert_eq!(relay.unseen_context(0), "");
563 assert_eq!(relay.unseen_context(1), "Agent 0:\nfirst answer");
564 assert_eq!(relay.shared_task(), Some("refactor"));
565 relay.add_agent();
566 assert_eq!(relay.active_slots().count(), 3);
567 }
568
569 #[test]
570 fn stop_token_is_stripped_only_from_the_response_suffix() {
571 let (visible, requested) = strip_stop_token(&format!("looks good\n{STOP_TOKEN}"));
572 assert_eq!(visible, "looks good");
573 assert!(requested);
574 let (visible, requested) = strip_stop_token("ordinary response");
575 assert_eq!(visible, "ordinary response");
576 assert!(!requested);
577 }
578
579 #[test]
580 fn accepted_stop_ends_the_batch_but_a_new_prompt_can_start_one() {
581 let mut relay = Relay::new(2, 10);
582 assert!(matches!(
583 relay.begin("task", 0),
584 RelayDecision::Dispatch { slot: 0, .. }
585 ));
586 relay.finish(0, false, false);
587 assert!(matches!(
588 relay.begin("review", 0),
589 RelayDecision::Dispatch {
590 slot: 1,
591 can_stop: true,
592 ..
593 }
594 ));
595 relay.finish(1, false, true);
596 assert_eq!(relay.begin("", 0), RelayDecision::Complete);
597
598 assert!(relay.enqueue_human("new task", Some(0)));
599 assert!(matches!(
600 relay.begin("", 0),
601 RelayDecision::Dispatch { slot: 0, prompt, .. } if prompt == "new task"
602 ));
603 }
604
605 #[test]
606 fn owner_promotion_follows_targets_and_replays_context_to_replacement() {
607 let mut relay = Relay::new(3, 10);
608 relay.set_shared_task("shared");
609 relay.record_public("owner", "context");
610 relay.mark_context_seen(0);
611 assert_eq!(relay.unseen_context(0), "");
612 assert!(relay.enqueue_direct(0, "old owner work").expect("queue"));
613 assert!(relay.enqueue_direct(1, "replacement work").expect("queue"));
614 assert!(matches!(
615 relay.begin("task", 0),
616 RelayDecision::Dispatch { slot: 0, .. }
617 ));
618 relay.finish(0, false, false);
619
620 relay.promote_owner(1).expect("promote active peer");
621 assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![0, 2]);
622 assert!(matches!(
625 relay.begin("", 0),
626 RelayDecision::Dispatch { slot: 0, direct: true, prompt, .. }
627 if prompt == "replacement work"
628 ));
629 assert_eq!(relay.unseen_context(0), "owner:\ncontext");
630 }
631
632 #[test]
633 fn live_slot_swap_follows_queued_targets_and_runtime_cursors() {
634 let mut relay = Relay::new(3, 10);
635 relay.record_public("Agent 0", "owner work");
636 relay.mark_context_seen(0);
637 assert!(matches!(
638 relay.begin("task", 0),
639 RelayDecision::Dispatch { slot: 0, .. }
640 ));
641 relay.finish(0, false, false);
642 relay.enqueue_human("to owner", Some(0));
643 relay.enqueue_direct(2, "to third").expect("queue direct");
644
645 relay.swap_agents(0, 2).expect("swap live slots");
646 assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![0, 1, 2]);
647 assert!(matches!(
648 relay.begin("", 0),
649 RelayDecision::Dispatch { slot: 0, direct: true, prompt, .. }
650 if prompt == "to third"
651 ));
652 relay.finish(0, true, false);
653 assert!(matches!(
654 relay.begin("", 0),
655 RelayDecision::Dispatch { slot: 2, direct: false, prompt, .. }
656 if prompt == "to owner"
657 ));
658 assert_eq!(relay.unseen_context(0), "Agent 0:\nowner work");
659 }
660
661 #[test]
662 fn stop_token_is_not_allowed_on_the_first_response() {
663 let mut relay = Relay::new(2, 10);
664 assert!(matches!(
665 relay.begin("task", 0),
666 RelayDecision::Dispatch {
667 slot: 0,
668 can_stop: false,
669 ..
670 }
671 ));
672 relay.finish(0, false, false);
676 assert!(matches!(
677 relay.begin("", 0),
678 RelayDecision::Dispatch { slot: 1, .. }
679 ));
680 }
681
682 #[test]
683 fn a_healthy_peer_continues_after_the_other_slot_is_tombstoned() {
684 let mut relay = Relay::new(2, 10);
685 relay.tombstone(0).expect("owner failure");
686 assert!(matches!(
687 relay.begin("continue", 0),
688 RelayDecision::Dispatch {
689 slot: 1,
690 can_stop: false,
691 ..
692 }
693 ));
694 }
695
696 #[test]
697 fn manual_strategy_requires_an_explicit_follow_up_prompt() {
698 let mut relay = Relay::new(3, 10);
699 relay.set_strategy(CollaborationStrategy::Manual);
700 assert!(matches!(
701 relay.begin("task", 0),
702 RelayDecision::Dispatch { slot: 0, .. }
703 ));
704 relay.finish(0, false, false);
705 assert_eq!(
706 relay.begin("would auto advance", 0),
707 RelayDecision::Complete
708 );
709 assert!(relay.enqueue_human("review", Some(2)));
710 assert!(
711 matches!(relay.begin("", 0), RelayDecision::Dispatch { slot: 2, prompt, .. } if prompt == "review")
712 );
713 }
714
715 #[test]
716 fn pair_strategy_alternates_owner_and_selected_reviewer() {
717 let mut relay = Relay::new(4, 10);
718 relay.set_strategy(CollaborationStrategy::Pair);
719 assert!(matches!(
720 relay.begin("task", 2),
721 RelayDecision::Dispatch { slot: 2, .. }
722 ));
723 relay.finish(2, false, false);
724 assert!(matches!(
725 relay.begin("review", 2),
726 RelayDecision::Dispatch { slot: 0, .. }
727 ));
728 relay.finish(0, false, false);
729 assert!(matches!(
730 relay.begin("next", 2),
731 RelayDecision::Dispatch { slot: 2, .. }
732 ));
733 }
734}