1use std::collections::{HashMap, HashSet, hash_map::DefaultHasher};
21use std::hash::{Hash, Hasher};
22use std::sync::atomic::{AtomicBool, Ordering};
23
24use async_trait::async_trait;
25use parking_lot::Mutex;
26use switchyard_protocol::{ModelId, Request, Role};
27
28use crate::core::algorithm::{Driver, RoutingIdentity};
29use crate::core::classifier::{Classification, Classifier, Score};
30use crate::core::processor::{Event, Processor};
31
32const MAX_ASSIGNMENTS: usize = 4096;
35
36#[derive(Default)]
47pub struct AffinityRouter {
48 latch_only: Option<HashSet<ModelId>>,
50 subagents_only: bool,
52 message_hash_fallback: bool,
54 assignments: Mutex<HashMap<RoutingIdentity, ModelId>>,
59 unkeyed_warning_emitted: AtomicBool,
61}
62
63impl AffinityRouter {
64 pub fn new() -> Self {
66 Self::default()
67 }
68
69 pub fn for_subagents() -> Self {
74 Self {
75 subagents_only: true,
76 ..Self::default()
77 }
78 }
79
80 pub fn with_message_hash_fallback(mut self) -> Self {
82 self.message_hash_fallback = true;
83 self
84 }
85
86 pub fn with_latch_only(mut self, models: impl IntoIterator<Item = impl Into<ModelId>>) -> Self {
89 self.latch_only = Some(models.into_iter().map(Into::into).collect());
90 self
91 }
92
93 fn should_latch(&self, model: &str) -> bool {
95 self.latch_only
96 .as_ref()
97 .is_none_or(|set| set.contains(model))
98 }
99
100 fn affinity_key(&self, request: &Request) -> Option<RoutingIdentity> {
102 let metadata = request.metadata.as_ref();
103 let is_subagent = metadata.is_some_and(|metadata| metadata.is_subagent);
104 let is_subagent_work = metadata.is_some_and(|metadata| metadata.is_subagent_work());
105 if self.subagents_only && !is_subagent_work {
107 return None;
108 }
109
110 let key = match (RoutingIdentity::from_request(request), is_subagent) {
111 (Some(identity), _) => Some(identity),
112 (None, true) => None,
115 (None, false) if self.message_hash_fallback => {
116 first_user_message_hash(request).map(|hash| {
117 tracing::debug!(affinity_key = %hash, "affinity using message hash fallback");
118 RoutingIdentity::Session(hash)
119 })
120 }
121 (None, false) => None,
122 };
123 if key.is_none() && self.should_warn_unkeyed() {
126 tracing::warn!(
127 target: "libsy",
128 is_subagent,
129 message_hash_fallback = self.message_hash_fallback,
130 "affinity is enabled but this request carries no usable identity, so no \
131 affinity is applied; root requests need a session id or message-hash \
132 fallback with usable first-user text, and child requests need both session \
133 and agent ids"
134 );
135 }
136 key
137 }
138
139 fn should_warn_unkeyed(&self) -> bool {
141 !self.unkeyed_warning_emitted.swap(true, Ordering::Relaxed)
142 }
143}
144
145#[async_trait]
146impl<S> Processor<S> for AffinityRouter
147where
148 S: Send + 'static,
149{
150 async fn process(&self, _state: &mut S, event: Event<'_>) -> crate::Result<()> {
151 if let Event::Decision {
152 request,
153 selected_model_id,
154 } = event
155 && let Some(key) = self.affinity_key(request)
156 {
157 let mut assignments = self.assignments.lock();
158 if self.should_latch(selected_model_id) && !assignments.contains_key(&key) {
159 evict_if_full(&mut assignments);
160 assignments.insert(key, selected_model_id.clone());
161 }
162 }
163 Ok(())
164 }
165}
166
167fn first_user_message_hash(request: &Request) -> Option<String> {
171 let message = request
172 .llm_request
173 .messages
174 .iter()
175 .find(|message| message.role == Role::User)?;
176 let mut hasher = DefaultHasher::new();
177 message.text_content("")?.hash(&mut hasher);
178 Some(format!("{:016x}", hasher.finish()))
179}
180
181#[async_trait]
182impl<S> Classifier<S> for AffinityRouter
183where
184 S: Send + 'static,
185{
186 async fn score(
187 &self,
188 _state: &mut S,
189 request: &mut Request,
190 _driver: Option<&Driver>,
191 ) -> crate::Result<(Classification, Option<switchyard_protocol::Response>)> {
192 let Some(key) = self.affinity_key(request) else {
193 return Ok((Classification::Scores(Vec::new()), None));
194 };
195 let assigned = self.assignments.lock().get(&key).cloned();
196 Ok((
197 Classification::Scores(match assigned {
198 Some(target) => vec![Score {
199 confidence: 1.0,
200 target,
201 }],
202 None => Vec::new(),
203 }),
204 None,
205 ))
206 }
207}
208
209fn evict_if_full(assignments: &mut HashMap<RoutingIdentity, ModelId>) {
211 if assignments.len() >= MAX_ASSIGNMENTS
212 && let Some(evicted) = assignments.keys().next().cloned()
213 {
214 assignments.remove(&evicted);
215 }
216}
217
218#[cfg(test)]
219mod tests {
220 use super::*;
221
222 use std::sync::Arc;
223
224 use switchyard_protocol::{ContentBlock, LlmRequest, Message, Metadata, text_request};
225
226 type BoxErr = Box<dyn std::error::Error + Send + Sync>;
228
229 fn fixed_model(target: &str) -> ModelId {
230 ModelId::from(target)
231 }
232
233 fn request(metadata: Metadata) -> Request {
234 Request {
235 llm_request: text_request(Some("auto".to_string()), "hi"),
236 raw_request: None,
237 metadata: Some(metadata),
238 }
239 }
240
241 fn task_request(
242 metadata: Option<Metadata>,
243 first_user: &str,
244 follow_up: Option<&str>,
245 ) -> Request {
246 let mut messages = vec![
247 Message::text(Role::System, "follow repository instructions"),
248 Message::text(Role::User, first_user),
249 ];
250 if let Some(follow_up) = follow_up {
251 messages.push(Message::text(Role::Assistant, "I will inspect the code."));
252 messages.push(Message::text(Role::User, follow_up));
253 }
254 Request {
255 llm_request: LlmRequest {
256 model: Some("auto".to_string()),
257 messages,
258 ..LlmRequest::default()
259 },
260 raw_request: None,
261 metadata,
262 }
263 }
264
265 fn session(session_id: &str, agent_id: &str) -> Metadata {
266 Metadata {
267 session_id: Some(session_id.to_string()),
268 agent_id: Some(agent_id.to_string()),
269 ..Metadata::default()
270 }
271 }
272
273 fn subagent(agent_id: &str, task_id: &str) -> Metadata {
274 Metadata {
275 session_id: Some("session-1".to_string()),
276 agent_id: Some(agent_id.to_string()),
277 task_id: Some(task_id.to_string()),
278 is_subagent: true,
279 is_delegated_work: true,
280 ..Metadata::default()
281 }
282 }
283
284 async fn retain(
286 router: &AffinityRouter,
287 state: &mut (),
288 request: &mut Request,
289 model: &'static str,
290 ) -> Result<(), BoxErr> {
291 let selected_model_id = ModelId::from(model);
292 router
293 .process(
294 state,
295 Event::Decision {
296 request,
297 selected_model_id: &selected_model_id,
298 },
299 )
300 .await?;
301 Ok(())
302 }
303
304 async fn scores(
306 classifier: &dyn Classifier,
307 state: &mut (),
308 request: &mut Request,
309 ) -> Result<Vec<Score>, BoxErr> {
310 match classifier.score(state, request, None).await?.0 {
311 Classification::Scores(scores) => Ok(scores),
312 Classification::Ambiguous(_) => Err("affinity never returns ambiguous scores".into()),
313 }
314 }
315
316 #[tokio::test]
317 async fn session_retains_first_model_across_requests() -> Result<(), BoxErr> {
318 let router = AffinityRouter::new();
319 let mut state = ();
320
321 let mut first = request(session("session-1", "agent-a"));
322 retain(&router, &mut state, &mut first, "model-a").await?;
323
324 let mut second = request(session("session-1", "agent-b"));
326 let scores = scores(&router, &mut state, &mut second).await?;
327 assert_eq!(scores.len(), 1);
328 assert_eq!(scores[0].confidence, 1.0);
329 assert_eq!(scores[0].target, "model-a");
330 Ok(())
331 }
332
333 #[tokio::test]
334 async fn subagent_only_retains_children_without_latching_root_traffic() -> Result<(), BoxErr> {
335 let router = AffinityRouter::for_subagents();
336 let mut state = ();
337
338 let mut root = request(session("session-1", "root-agent"));
339 retain(&router, &mut state, &mut root, "model-a").await?;
340 assert!(scores(&router, &mut state, &mut root).await?.is_empty());
341
342 let mut first_child_turn = request(subagent("child-1", "task-1"));
343 retain(&router, &mut state, &mut first_child_turn, "model-b").await?;
344 let mut later_child_turn = request(subagent("child-1", "task-2"));
345 let scores = scores(&router, &mut state, &mut later_child_turn).await?;
346 assert_eq!(
347 scores.first().map(|score| score.target.as_str()),
348 Some("model-b")
349 );
350 Ok(())
351 }
352
353 #[tokio::test]
354 async fn first_decision_wins() -> Result<(), BoxErr> {
355 let router = AffinityRouter::new();
356 let mut state = ();
357
358 let mut req = request(session("session-1", "agent-a"));
359 retain(&router, &mut state, &mut req, "model-a").await?;
360 retain(&router, &mut state, &mut req, "model-b").await?;
362
363 let scores = scores(&router, &mut state, &mut req).await?;
364 assert_eq!(
365 scores.first().map(|score| score.target.as_str()),
366 Some("model-a")
367 );
368 Ok(())
369 }
370
371 #[tokio::test]
372 async fn subagent_is_keyed_by_agent_not_task() -> Result<(), BoxErr> {
373 let router = AffinityRouter::new();
374 let mut state = ();
375
376 let mut first = request(subagent("child-1", "task-1"));
377 retain(&router, &mut state, &mut first, "model-a").await?;
378
379 let mut second = request(subagent("child-1", "task-2"));
381 let scores = scores(&router, &mut state, &mut second).await?;
382 assert_eq!(
383 scores.first().map(|score| score.target.as_str()),
384 Some("model-a")
385 );
386 Ok(())
387 }
388
389 #[tokio::test]
390 async fn distinct_subagents_are_assigned_independently() -> Result<(), BoxErr> {
391 let router = AffinityRouter::new();
392 let mut state = ();
393
394 retain(
396 &router,
397 &mut state,
398 &mut request(subagent("child-1", "task-1")),
399 "model-a",
400 )
401 .await?;
402
403 let mut sibling = request(subagent("child-2", "task-1"));
405 assert!(scores(&router, &mut state, &mut sibling).await?.is_empty());
406 Ok(())
407 }
408
409 #[tokio::test]
410 async fn subagent_does_not_inherit_session_assignment() -> Result<(), BoxErr> {
411 let router = AffinityRouter::new();
412 let mut state = ();
413
414 retain(
416 &router,
417 &mut state,
418 &mut request(session("session-1", "root-1")),
419 "model-a",
420 )
421 .await?;
422
423 let mut child = request(subagent("child-1", "task-1"));
425 assert!(scores(&router, &mut state, &mut child).await?.is_empty());
426 Ok(())
427 }
428
429 #[tokio::test]
430 async fn classifier_abstains_without_a_session() -> Result<(), BoxErr> {
431 let router = AffinityRouter::new();
432 let mut state = ();
433
434 let mut req = request(Metadata::default());
436 assert!(scores(&router, &mut state, &mut req).await?.is_empty());
437 Ok(())
438 }
439
440 #[tokio::test]
441 async fn message_hash_fallback_uses_the_first_user_message() -> Result<(), BoxErr> {
442 let router = AffinityRouter::new().with_message_hash_fallback();
443 let mut state = ();
444
445 let mut first = task_request(
446 None,
447 "Add a unit test for this function.",
448 Some("Now run the test suite."),
449 );
450 retain(&router, &mut state, &mut first, "weak").await?;
451
452 let mut follow_up = task_request(
453 None,
454 "Add a unit test for this function.",
455 Some("Now file a pull request."),
456 );
457 assert_eq!(
458 scores(&router, &mut state, &mut follow_up)
459 .await?
460 .first()
461 .map(|score| score.target.as_str()),
462 Some("weak")
463 );
464
465 let mut other_task = task_request(
466 None,
467 "Reimplement this binary from two input/output pairs.",
468 Some("Now run the test suite."),
469 );
470 assert!(
471 scores(&router, &mut state, &mut other_task)
472 .await?
473 .is_empty()
474 );
475 Ok(())
476 }
477
478 #[tokio::test]
479 async fn subagents_only_root_traffic_does_not_warn() -> Result<(), BoxErr> {
480 let router = AffinityRouter::for_subagents();
482 let mut state = ();
483
484 let mut root = request(session("session-1", "agent-1"));
485 assert!(scores(&router, &mut state, &mut root).await?.is_empty());
486 assert!(
487 router.should_warn_unkeyed(),
488 "an intentional abstention should leave the warning unconsumed"
489 );
490 Ok(())
491 }
492
493 #[test]
494 fn user_message_hash_ignores_non_text_provider_payloads() {
495 let request = |user_message| Request {
496 llm_request: LlmRequest {
497 messages: vec![user_message],
498 ..LlmRequest::default()
499 },
500 raw_request: None,
501 metadata: None,
502 };
503 let text_only = request(Message::text(Role::User, "Implement the parser."));
504 let text_with_reasoning = request(Message {
505 role: Role::User,
506 content: vec![
507 ContentBlock::Text {
508 text: "Implement the parser.".to_string(),
509 },
510 ContentBlock::Reasoning {
511 text: "Internal provider reasoning.".to_string(),
512 signature: Some("provider-signature".to_string()),
513 details: Vec::new(),
514 },
515 ],
516 });
517
518 assert_eq!(
519 first_user_message_hash(&text_only),
520 first_user_message_hash(&text_with_reasoning)
521 );
522 }
523
524 #[tokio::test]
525 async fn metadata_session_takes_precedence_over_message_hash() -> Result<(), BoxErr> {
526 let router = AffinityRouter::new().with_message_hash_fallback();
527 let mut state = ();
528
529 let mut first = task_request(
530 Some(session("session-1", "agent-a")),
531 "Implement the parser.",
532 None,
533 );
534 retain(&router, &mut state, &mut first, "strong").await?;
535
536 let mut other_session = task_request(
537 Some(session("session-2", "agent-a")),
538 "Implement the parser.",
539 None,
540 );
541 assert!(
542 scores(&router, &mut state, &mut other_session)
543 .await?
544 .is_empty()
545 );
546 Ok(())
547 }
548
549 #[tokio::test]
550 async fn subagent_without_a_session_abstains_and_warns() -> Result<(), BoxErr> {
551 for session_id in [None, Some(String::new())] {
552 let router = AffinityRouter::new().with_message_hash_fallback();
553 let mut state = ();
554 let mut subagent = task_request(
555 Some(Metadata {
556 session_id,
557 agent_id: Some("agent-1".to_string()),
558 is_subagent: true,
559 ..Metadata::default()
560 }),
561 "Implement the parser.",
562 None,
563 );
564
565 retain(&router, &mut state, &mut subagent, "model-a").await?;
566 assert!(scores(&router, &mut state, &mut subagent).await?.is_empty());
567 assert!(
568 !router.should_warn_unkeyed(),
569 "an unidentifiable subagent should consume the warning"
570 );
571 }
572 Ok(())
573 }
574
575 #[tokio::test]
576 async fn one_router_serves_both_roles() -> Result<(), BoxErr> {
577 let router = Arc::new(AffinityRouter::new());
580 let processor: Arc<dyn Processor> = router.clone();
581 let classifier: Arc<dyn Classifier> = router;
582 let mut state = ();
583
584 let mut first = request(session("session-1", "agent-a"));
585 processor
586 .process(
587 &mut state,
588 Event::Decision {
589 request: &mut first,
590 selected_model_id: &fixed_model("model-a"),
591 },
592 )
593 .await?;
594
595 let mut second = request(session("session-1", "agent-b"));
596 let scores = scores(classifier.as_ref(), &mut state, &mut second).await?;
597 assert_eq!(
598 scores.first().map(|score| score.target.as_str()),
599 Some("model-a")
600 );
601 Ok(())
602 }
603
604 #[tokio::test]
605 async fn decision_without_an_affinity_identity_is_ignored() -> Result<(), BoxErr> {
606 let router = AffinityRouter::new();
607 let mut state = ();
608 let mut unkeyed = request(Metadata::default());
609
610 router
611 .process(
612 &mut state,
613 Event::Decision {
614 request: &mut unkeyed,
615 selected_model_id: &fixed_model("model-a"),
616 },
617 )
618 .await?;
619
620 let mut req = request(session("session-1", "agent-a"));
621 assert!(scores(&router, &mut state, &mut req).await?.is_empty());
622 Ok(())
623 }
624
625 #[tokio::test]
626 async fn decisions_retain_their_originating_request_identity() -> Result<(), BoxErr> {
627 let router = AffinityRouter::new();
628 let mut state = ();
629 let mut first = request(session("session-1", "agent-a"));
630 let mut second = request(session("session-2", "agent-b"));
631
632 router
635 .process(
636 &mut state,
637 Event::Decision {
638 request: &mut second,
639 selected_model_id: &fixed_model("model-b"),
640 },
641 )
642 .await?;
643 router
644 .process(
645 &mut state,
646 Event::Decision {
647 request: &mut first,
648 selected_model_id: &fixed_model("model-a"),
649 },
650 )
651 .await?;
652
653 let first_scores = scores(&router, &mut state, &mut first).await?;
654 let second_scores = scores(&router, &mut state, &mut second).await?;
655 assert_eq!(
656 first_scores.first().map(|score| score.target.as_str()),
657 Some("model-a")
658 );
659 assert_eq!(
660 second_scores.first().map(|score| score.target.as_str()),
661 Some("model-b")
662 );
663 Ok(())
664 }
665
666 #[tokio::test]
667 async fn distinct_sessions_are_assigned_independently() -> Result<(), BoxErr> {
668 let router = AffinityRouter::new();
669 let mut state = ();
670
671 retain(
672 &router,
673 &mut state,
674 &mut request(session("session-1", "agent-a")),
675 "model-a",
676 )
677 .await?;
678 retain(
679 &router,
680 &mut state,
681 &mut request(session("session-2", "agent-a")),
682 "model-b",
683 )
684 .await?;
685
686 let first = scores(
687 &router,
688 &mut state,
689 &mut request(session("session-1", "other")),
690 )
691 .await?;
692 let second = scores(
693 &router,
694 &mut state,
695 &mut request(session("session-2", "other")),
696 )
697 .await?;
698 assert_eq!(
699 first.first().map(|score| score.target.as_str()),
700 Some("model-a")
701 );
702 assert_eq!(
703 second.first().map(|score| score.target.as_str()),
704 Some("model-b")
705 );
706 Ok(())
707 }
708
709 #[tokio::test]
710 async fn subagent_without_an_agent_id_abstains_and_warns() -> Result<(), BoxErr> {
711 let router = AffinityRouter::new();
712 let mut state = ();
713
714 let metadata = Metadata {
717 session_id: Some("session-1".to_string()),
718 is_subagent: true,
719 ..Metadata::default()
720 };
721 let mut req = request(metadata);
722 retain(&router, &mut state, &mut req, "model-a").await?;
723 assert!(scores(&router, &mut state, &mut req).await?.is_empty());
724 assert!(
725 !router.should_warn_unkeyed(),
726 "an unidentifiable subagent should consume the warning"
727 );
728 Ok(())
729 }
730
731 #[tokio::test]
732 async fn assignments_are_bounded_by_the_cap() -> Result<(), BoxErr> {
733 let router = AffinityRouter::new();
734 let mut state = ();
735
736 for index in 0..=MAX_ASSIGNMENTS {
738 let session_id = format!("session-{index}");
739 retain(
740 &router,
741 &mut state,
742 &mut request(session(&session_id, "agent-a")),
743 "model-a",
744 )
745 .await?;
746 }
747
748 let len = router.assignments.lock().len();
749 assert_eq!(len, MAX_ASSIGNMENTS);
750 Ok(())
751 }
752
753 #[tokio::test]
754 async fn latch_only_retains_matching_models() -> Result<(), BoxErr> {
755 let router = AffinityRouter::new().with_latch_only(["strong"]);
756 let mut state = ();
757 let mut req = request(session("session-1", "agent-a"));
758
759 retain(&router, &mut state, &mut req, "weak").await?;
761 assert!(scores(&router, &mut state, &mut req).await?.is_empty());
762
763 retain(&router, &mut state, &mut req, "strong").await?;
765 assert_eq!(
766 scores(&router, &mut state, &mut req)
767 .await?
768 .first()
769 .map(|s| s.target.as_str()),
770 Some("strong")
771 );
772 Ok(())
773 }
774}