codeswarm_core/
relay.rs

1//! Deterministic sequential roster scheduling.
2//!
3//! This owns turn selection only. Prompt construction and adapter I/O remain
4//! outside the scheduler, making the relay safe to replay and test.
5
6use std::collections::VecDeque;
7
8use crate::RosterSlot;
9use crate::collaboration::CollaborationContext;
10
11pub const MAX_QUEUED_PROMPTS: usize = 100;
12pub const STOP_TOKEN: &str = "[CODESWARM:STOP]";
13pub const 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/// How a multi-agent session chooses its next non-direct recipient.
35///
36/// `Roster` is the normal sequential ring. `Pair` keeps the first two active
37/// agents in a tight review loop, which is useful when a
38/// larger saved roster is available but the user wants focused two-agent
39/// collaboration. `Manual` never advances on its own after the first turn;
40/// every subsequent turn must be explicitly targeted or queued by the user.
41#[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    /// Change routing for future turns. This does not discard queued work or
112    /// alter the public context journal; it only affects the next automatic
113    /// recipient. Pair selection is re-derived from the next `first` value.
114    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 !self.active.get(slot).copied().ok_or("slot out of range")? {
151            return Ok(());
152        }
153        if self.active_slots().count() == 1 {
154            return Err("last active agent cannot be dropped");
155        }
156        self.tombstone(slot)?;
157        self.direct.retain(|queued| queued.slot != slot);
158        self.steering.retain(|queued| queued.slot != slot);
159        Ok(())
160    }
161
162    /// Exchange two live roster slots while preserving adapter identity in
163    /// queued work, routing cursors, and per-agent context watermarks.
164    pub fn swap_agents(
165        &mut self,
166        first: RosterSlot,
167        second: RosterSlot,
168    ) -> Result<(), &'static str> {
169        if first == second {
170            return Ok(());
171        }
172        if first >= self.active.len() || second >= self.active.len() {
173            return Err("roster slot out of range");
174        }
175        if !self.active[first] || !self.active[second] {
176            return Err("both roster slots must be active");
177        }
178        fn swap_targets(queue: &mut VecDeque<QueuedPrompt>, first: usize, second: usize) {
179            for queued in queue {
180                if queued.slot == first {
181                    queued.slot = second;
182                } else if queued.slot == second {
183                    queued.slot = first;
184                }
185            }
186        }
187        swap_targets(&mut self.direct, first, second);
188        swap_targets(&mut self.steering, first, second);
189        if self.last_active == first {
190            self.last_active = second;
191        } else if self.last_active == second {
192            self.last_active = first;
193        }
194        fn swap_option(cursor: &mut Option<usize>, first: usize, second: usize) {
195            if *cursor == Some(first) {
196                *cursor = Some(second);
197            } else if *cursor == Some(second) {
198                *cursor = Some(first);
199            }
200        }
201        swap_option(&mut self.next, first, second);
202        swap_option(&mut self.previous_slot, first, second);
203        swap_option(&mut self.pair_partner, first, second);
204        self.active.swap(first, second);
205        self.context.swap_agents(first, second);
206        Ok(())
207    }
208
209    pub fn enqueue_human(
210        &mut self,
211        prompt: impl Into<String>,
212        selected: Option<RosterSlot>,
213    ) -> bool {
214        let prompt = prompt.into();
215        if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
216            return false;
217        }
218        let slot = selected.unwrap_or(self.last_active);
219        if !self.active.get(slot).copied().unwrap_or(false) {
220            return false;
221        }
222        self.steering.push_back(QueuedPrompt {
223            slot,
224            prompt,
225            kind: QueuedKind::Steering,
226        });
227        true
228    }
229
230    pub fn enqueue_direct(
231        &mut self,
232        slot: RosterSlot,
233        prompt: impl Into<String>,
234    ) -> Result<bool, &'static str> {
235        let prompt = prompt.into();
236        if !self.active.get(slot).copied().unwrap_or(false) {
237            return Err("direct target is not active");
238        }
239        if prompt.trim().is_empty() || self.queued_count() >= MAX_QUEUED_PROMPTS {
240            return Ok(false);
241        }
242        self.direct.push_back(QueuedPrompt {
243            slot,
244            prompt,
245            kind: QueuedKind::Direct,
246        });
247        Ok(true)
248    }
249
250    pub fn queued_count(&self) -> usize {
251        self.direct.len() + self.steering.len()
252    }
253
254    pub fn set_shared_task(&mut self, task: impl Into<String>) {
255        self.context.set_shared_task(task);
256    }
257
258    pub fn shared_task(&self) -> Option<&str> {
259        self.context.shared_task()
260    }
261
262    pub fn record_public(&mut self, speaker: impl Into<String>, text: impl Into<String>) {
263        self.context.record(speaker, text, &self.active);
264    }
265
266    pub fn mark_context_seen(&mut self, slot: RosterSlot) {
267        self.context.mark_seen(slot);
268    }
269
270    pub fn unseen_context(&mut self, slot: RosterSlot) -> String {
271        self.context.unseen(slot)
272    }
273
274    pub fn add_agent(&mut self) {
275        self.active.push(true);
276        self.context.add_agent();
277    }
278
279    /// Select the next causal turn. Direct work always precedes steering work.
280    pub fn begin(&mut self, initial_prompt: impl Into<String>, first: RosterSlot) -> RelayDecision {
281        if self.paused {
282            return RelayDecision::Paused;
283        }
284        if self.active_slots().next().is_none() {
285            return RelayDecision::Collapsed;
286        }
287        let queued = Self::pop_active(&self.active, &mut self.direct)
288            .or_else(|| Self::pop_active(&self.active, &mut self.steering));
289        // A reviewer stop ends only the current automatic batch. A later
290        // queued/user prompt starts a fresh batch without rebuilding the
291        // relay, while an unprompted handoff remains complete.
292        if self.stopped {
293            if queued.is_none() {
294                return RelayDecision::Complete;
295            }
296            self.stopped = false;
297            self.rounds = 0;
298        }
299        if self.rounds >= self.max_rounds {
300            // A queued human/direct prompt is a new batch; do not strand it
301            // behind the safety limit reached by the previous batch.
302            if queued.is_none() {
303                return RelayDecision::Complete;
304            }
305            self.rounds = 0;
306        }
307        // Manual mode is deliberately input-driven. A queued prompt (which
308        // includes a newly submitted human prompt) is still dispatched, but
309        // an unprompted call after a completed turn must not silently hand the
310        // conversation to another agent.
311        if self.strategy == CollaborationStrategy::Manual
312            && queued.is_none()
313            && self.previous_slot.is_some()
314        {
315            return RelayDecision::Complete;
316        }
317        let (slot, prompt, direct, human_prompt) = match queued {
318            Some(queued) => (
319                queued.slot,
320                queued.prompt,
321                queued.kind == QueuedKind::Direct,
322                queued.kind == QueuedKind::Steering,
323            ),
324            None => {
325                let slot = self.next_automatic_slot(first);
326                (slot, initial_prompt.into(), false, false)
327            }
328        };
329        // A human steering prompt starts a fresh review batch. Even when it
330        // targets a different slot from the preceding relay turn, that first
331        // responder must not be allowed to terminate the batch with the safe
332        // word. Only an automatic handoff after another agent's response is
333        // eligible to review-stop.
334        let can_stop =
335            !direct && !human_prompt && self.previous_slot.is_some_and(|previous| previous != slot);
336        self.last_active = slot;
337        self.rounds += 1;
338        RelayDecision::Dispatch {
339            slot,
340            prompt,
341            direct,
342            can_stop,
343        }
344    }
345
346    /// Finalize a dispatched turn and choose the next ring position. Direct
347    /// turns never become shared relay context.
348    pub fn finish(&mut self, slot: RosterSlot, direct: bool, accepted_stop: bool) {
349        self.next = Some(self.next_active(slot));
350        if !direct {
351            self.previous_slot = Some(slot);
352        }
353        if accepted_stop && self.direct.is_empty() && self.steering.is_empty() {
354            self.stopped = true;
355        }
356    }
357
358    fn pop_active(active: &[bool], queue: &mut VecDeque<QueuedPrompt>) -> Option<QueuedPrompt> {
359        let position = queue
360            .iter()
361            .position(|queued| active.get(queued.slot).copied().unwrap_or(false))?;
362        queue.remove(position)
363    }
364
365    fn first_active_from(&self, start: RosterSlot) -> RosterSlot {
366        (0..self.active.len())
367            .map(|offset| (start + offset) % self.active.len())
368            .find(|slot| self.active[*slot])
369            .expect("callers require an active roster")
370    }
371
372    fn next_active(&self, slot: RosterSlot) -> RosterSlot {
373        (1..=self.active.len())
374            .map(|offset| (slot + offset) % self.active.len())
375            .find(|candidate| self.active[*candidate])
376            .expect("callers require an active roster")
377    }
378
379    fn next_automatic_slot(&mut self, first: RosterSlot) -> RosterSlot {
380        match self.strategy {
381            CollaborationStrategy::Roster | CollaborationStrategy::Manual => self
382                .next
383                .filter(|slot| self.active[*slot])
384                .unwrap_or_else(|| self.first_active_from(first)),
385            CollaborationStrategy::Pair => {
386                let primary = self.first_active_from(0);
387                let partner = if let Some(partner) = self.pair_partner {
388                    if self.active.get(partner).copied().unwrap_or(false) && partner != primary {
389                        partner
390                    } else {
391                        let partner = self.next_active(primary);
392                        self.pair_partner = Some(partner);
393                        partner
394                    }
395                } else {
396                    let partner =
397                        if first != primary && self.active.get(first).copied().unwrap_or(false) {
398                            first
399                        } else {
400                            self.next_active(primary)
401                        };
402                    self.pair_partner = Some(partner);
403                    partner
404                };
405                match self.previous_slot {
406                    Some(previous) if previous == primary && self.active[partner] => partner,
407                    Some(previous) if previous == partner && self.active[primary] => primary,
408                    Some(previous) if previous == primary || previous == partner => primary,
409                    _ => self.first_active_from(first),
410                }
411            }
412        }
413    }
414}
415
416#[cfg(test)]
417mod tests {
418    use super::{CollaborationStrategy, Relay, RelayDecision, STOP_TOKEN, strip_stop_token};
419
420    #[test]
421    fn relay_moves_around_the_ring_without_self_review() {
422        let mut relay = Relay::new(3, 10);
423        let first = relay.begin("task", 0);
424        assert!(matches!(
425            first,
426            RelayDecision::Dispatch {
427                slot: 0,
428                can_stop: false,
429                ..
430            }
431        ));
432        relay.finish(0, false, false);
433        let second = relay.begin("response", 0);
434        assert!(matches!(
435            second,
436            RelayDecision::Dispatch {
437                slot: 1,
438                can_stop: true,
439                ..
440            }
441        ));
442    }
443
444    #[test]
445    fn explicit_human_target_beats_ring_order() {
446        let mut relay = Relay::new(3, 10);
447        relay.begin("task", 0);
448        assert!(relay.enqueue_human("correction", Some(2)));
449        relay.finish(0, false, false);
450        assert!(matches!(
451            relay.begin("response", 0),
452            RelayDecision::Dispatch { slot: 2, prompt, direct: false, can_stop: false } if prompt == "correction"
453        ));
454    }
455
456    #[test]
457    fn human_prompt_cannot_stop_even_after_a_previous_relay_batch() {
458        let mut relay = Relay::new(2, 10);
459        assert!(matches!(
460            relay.begin("first task", 0),
461            RelayDecision::Dispatch {
462                slot: 0,
463                can_stop: false,
464                ..
465            }
466        ));
467        relay.finish(0, false, false);
468        assert!(matches!(
469            relay.begin("review", 0),
470            RelayDecision::Dispatch {
471                slot: 1,
472                can_stop: true,
473                ..
474            }
475        ));
476        relay.finish(1, false, false);
477
478        // The next user prompt targets the first agent, which differs from the
479        // previous reviewer. It is still the first response to a human turn,
480        // so it must not receive reviewer stop permission.
481        assert!(relay.enqueue_human("new task", Some(0)));
482        assert!(matches!(
483            relay.begin("", 0),
484            RelayDecision::Dispatch {
485                slot: 0,
486                can_stop: false,
487                ..
488            }
489        ));
490    }
491
492    #[test]
493    fn direct_work_has_priority_and_any_agent_except_the_last_can_be_dropped() {
494        let mut relay = Relay::new(3, 10);
495        relay.enqueue_human("ordinary", Some(1));
496        assert_eq!(relay.enqueue_direct(2, "private"), Ok(true));
497        assert!(matches!(
498            relay.begin("task", 0),
499            RelayDecision::Dispatch {
500                slot: 2,
501                direct: true,
502                ..
503            }
504        ));
505        assert_eq!(relay.drop_agent(0), Ok(()));
506        assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![1, 2]);
507        assert_eq!(relay.drop_agent(1), Ok(()));
508        assert_eq!(
509            relay.drop_agent(2),
510            Err("last active agent cannot be dropped")
511        );
512    }
513
514    #[test]
515    fn relay_context_tracks_public_updates_per_slot() {
516        let mut relay = Relay::new(2, 10);
517        relay.set_shared_task("refactor");
518        relay.record_public("Agent 0", "first answer");
519        relay.mark_context_seen(0);
520        assert_eq!(relay.unseen_context(0), "");
521        assert_eq!(relay.unseen_context(1), "Agent 0:\nfirst answer");
522        assert_eq!(relay.shared_task(), Some("refactor"));
523        relay.add_agent();
524        assert_eq!(relay.active_slots().count(), 3);
525    }
526
527    #[test]
528    fn stop_token_is_stripped_only_from_the_response_suffix() {
529        let (visible, requested) = strip_stop_token(&format!("looks good\n{STOP_TOKEN}"));
530        assert_eq!(visible, "looks good");
531        assert!(requested);
532        let (visible, requested) = strip_stop_token("ordinary response");
533        assert_eq!(visible, "ordinary response");
534        assert!(!requested);
535    }
536
537    #[test]
538    fn accepted_stop_ends_the_batch_but_a_new_prompt_can_start_one() {
539        let mut relay = Relay::new(2, 10);
540        assert!(matches!(
541            relay.begin("task", 0),
542            RelayDecision::Dispatch { slot: 0, .. }
543        ));
544        relay.finish(0, false, false);
545        assert!(matches!(
546            relay.begin("review", 0),
547            RelayDecision::Dispatch {
548                slot: 1,
549                can_stop: true,
550                ..
551            }
552        ));
553        relay.finish(1, false, true);
554        assert_eq!(relay.begin("", 0), RelayDecision::Complete);
555
556        assert!(relay.enqueue_human("new task", Some(0)));
557        assert!(matches!(
558            relay.begin("", 0),
559            RelayDecision::Dispatch { slot: 0, prompt, .. } if prompt == "new task"
560        ));
561    }
562
563    #[test]
564    fn live_slot_swap_follows_queued_targets_and_runtime_cursors() {
565        let mut relay = Relay::new(3, 10);
566        relay.record_public("Agent 0", "first work");
567        relay.mark_context_seen(0);
568        assert!(matches!(
569            relay.begin("task", 0),
570            RelayDecision::Dispatch { slot: 0, .. }
571        ));
572        relay.finish(0, false, false);
573        relay.enqueue_human("to first", Some(0));
574        relay.enqueue_direct(2, "to third").expect("queue direct");
575
576        relay.swap_agents(0, 2).expect("swap live slots");
577        assert_eq!(relay.active_slots().collect::<Vec<_>>(), vec![0, 1, 2]);
578        assert!(matches!(
579            relay.begin("", 0),
580            RelayDecision::Dispatch { slot: 0, direct: true, prompt, .. }
581                if prompt == "to third"
582        ));
583        relay.finish(0, true, false);
584        assert!(matches!(
585            relay.begin("", 0),
586            RelayDecision::Dispatch { slot: 2, direct: false, prompt, .. }
587                if prompt == "to first"
588        ));
589        assert_eq!(relay.unseen_context(0), "Agent 0:\nfirst work");
590    }
591
592    #[test]
593    fn stop_token_is_not_allowed_on_the_first_response() {
594        let mut relay = Relay::new(2, 10);
595        assert!(matches!(
596            relay.begin("task", 0),
597            RelayDecision::Dispatch {
598                slot: 0,
599                can_stop: false,
600                ..
601            }
602        ));
603        // RelayHost validates the token against `can_stop`; the first turn
604        // therefore finalizes as a normal response even if the agent tried
605        // to include the token.
606        relay.finish(0, false, false);
607        assert!(matches!(
608            relay.begin("", 0),
609            RelayDecision::Dispatch { slot: 1, .. }
610        ));
611    }
612
613    #[test]
614    fn a_healthy_peer_continues_after_the_other_slot_is_tombstoned() {
615        let mut relay = Relay::new(2, 10);
616        relay.tombstone(0).expect("first agent failure");
617        assert!(matches!(
618            relay.begin("continue", 0),
619            RelayDecision::Dispatch {
620                slot: 1,
621                can_stop: false,
622                ..
623            }
624        ));
625    }
626
627    #[test]
628    fn manual_strategy_requires_an_explicit_follow_up_prompt() {
629        let mut relay = Relay::new(3, 10);
630        relay.set_strategy(CollaborationStrategy::Manual);
631        assert!(matches!(
632            relay.begin("task", 0),
633            RelayDecision::Dispatch { slot: 0, .. }
634        ));
635        relay.finish(0, false, false);
636        assert_eq!(
637            relay.begin("would auto advance", 0),
638            RelayDecision::Complete
639        );
640        assert!(relay.enqueue_human("review", Some(2)));
641        assert!(
642            matches!(relay.begin("", 0), RelayDecision::Dispatch { slot: 2, prompt, .. } if prompt == "review")
643        );
644    }
645
646    #[test]
647    fn pair_strategy_alternates_the_first_two_active_agents() {
648        let mut relay = Relay::new(4, 10);
649        relay.set_strategy(CollaborationStrategy::Pair);
650        assert!(matches!(
651            relay.begin("task", 2),
652            RelayDecision::Dispatch { slot: 2, .. }
653        ));
654        relay.finish(2, false, false);
655        assert!(matches!(
656            relay.begin("review", 2),
657            RelayDecision::Dispatch { slot: 0, .. }
658        ));
659        relay.finish(0, false, false);
660        assert!(matches!(
661            relay.begin("next", 2),
662            RelayDecision::Dispatch { slot: 2, .. }
663        ));
664
665        relay.drop_agent(0).expect("remove first agent");
666        relay.finish(2, false, false);
667        assert!(matches!(
668            relay.begin("after removal", 2),
669            RelayDecision::Dispatch { slot: 1, .. }
670        ));
671    }
672}
codeswarm_core/relay.rs

codeswarm_core/
relay.rs
