switchyard_libsy/algorithms/util/
affinity.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Model affinity as a single SDK component.
5//!
6//! [`AffinityRouter`] retains the first model chosen for a request's stable identity and
7//! forces that model on later requests sharing the identity. It is one object that plays
8//! both SDK roles, so registering it as a processor and a classifier cannot drift apart:
9//!
10//! - As a [`Processor`] it *writes* the assignment: [`Event::Decision`] carries the request
11//!   whose chosen model is retained, with the first assignment for an identity winning.
12//! - As a [`Classifier`] it *reads* the assignment: it scores the retained model with
13//!   confidence `1.0`, or returns no scores to abstain.
14//!
15//! Identity is derived from correlation metadata: by default, a request is keyed by its
16//! session, and a sub-agent is keyed more finely by `session + agent` — a subset of its
17//! session. [`AffinityRouter::for_subagents`] narrows affinity to explicitly identified
18//! child agents, leaving root traffic to later classifiers on every turn.
19
20use std::collections::{HashMap, HashSet, hash_map::DefaultHasher};
21use std::hash::{Hash, Hasher};
22
23use async_trait::async_trait;
24use parking_lot::Mutex;
25use switchyard_protocol::{Request, Role};
26
27use crate::core::algorithm::Driver;
28use crate::core::classifier::{Classification, Classifier, Score};
29use crate::core::processor::{Event, Processor};
30
31/// Upper bound on retained assignments, keeping the process-local map from growing
32/// without limit; the oldest entry is evicted once the bound is reached.
33const MAX_ASSIGNMENTS: usize = 4096;
34
35/// The stable identity a model assignment is retained against.
36///
37/// A sub-agent request is keyed by `session + agent` and a root request by session alone,
38/// so a sub-agent's assignment is scoped within — but distinct from — its session's.
39#[derive(Clone, Hash, PartialEq, Eq)]
40enum AffinityKey {
41    /// One model per session, for root-agent traffic.
42    Session(String),
43    /// One model per identified child agent within a session.
44    Subagent { session: String, agent: String },
45}
46
47/// Retains a model per request identity and forces it on later matching requests.
48///
49/// Register the same instance as both a processor and a classifier; the two roles share
50/// the retained assignments through this instance's interior storage. Decision events carry
51/// their originating request, so concurrent turns cannot bind one request's identity to
52/// another request's selected model.
53///
54/// [`with_latch_only`](Self::with_latch_only) narrows *which* models are retained — a
55/// decision for any other model routes normally but is not latched (the escalation latch:
56/// retain only the strong tier, never the weak one).
57#[derive(Default)]
58pub struct AffinityRouter {
59    /// When set, only these models are retained; a decision for any other model is not latched.
60    latch_only: Option<HashSet<String>>,
61    /// Whether root-session requests should abstain instead of being retained.
62    subagents_only: bool,
63    /// In absence of headers, use the message hash based fallback key to do task based routing
64    message_hash_fallback: bool,
65    /// Retained assignments, shared across this router's processor and classifier roles.
66    ///
67    /// Held on the instance so the two roles share one process-local map through a
68    /// single registered [`Arc`](std::sync::Arc); bounded by [`MAX_ASSIGNMENTS`].
69    assignments: Mutex<HashMap<AffinityKey, String>>,
70}
71
72impl AffinityRouter {
73    /// Creates a router that latches every decision.
74    pub fn new() -> Self {
75        Self::default()
76    }
77
78    /// Creates a router that retains assignments only for explicitly identified child agents.
79    ///
80    /// Root-agent requests always abstain, so a later classifier selects them on every turn.
81    pub fn for_subagents() -> Self {
82        Self {
83            subagents_only: true,
84            ..Self::default()
85        }
86    }
87
88    /// Uses the first user-message text as a fallback key when metadata has no session.
89    pub fn with_message_hash_fallback(mut self) -> Self {
90        self.message_hash_fallback = true;
91        self
92    }
93
94    /// Restricts latching to `models`; a decision for any other model routes but is not
95    /// retained.
96    pub fn with_latch_only(mut self, models: impl IntoIterator<Item = impl Into<String>>) -> Self {
97        self.latch_only = Some(models.into_iter().map(Into::into).collect());
98        self
99    }
100
101    /// Whether a decision for `model` should be retained.
102    fn should_latch(&self, model: &str) -> bool {
103        self.latch_only
104            .as_ref()
105            .is_none_or(|set| set.contains(model))
106    }
107
108    /// Derives the stable identity this router should retain for `request`.
109    fn affinity_key(&self, request: &Request) -> Option<AffinityKey> {
110        if let Some(metadata) = request.metadata.as_ref()
111            && let Some(session) = metadata.session_id.clone()
112        {
113            return if metadata.is_subagent {
114                metadata
115                    .agent_id
116                    .clone()
117                    .map(|agent| AffinityKey::Subagent { session, agent })
118            } else if self.subagents_only {
119                None
120            } else {
121                Some(AffinityKey::Session(session))
122            };
123        }
124
125        // If headers are not present and we are not a subagent, use the message hash based fallback key to do task based routing
126        let is_subagent = request
127            .metadata
128            .as_ref()
129            .is_some_and(|metadata| metadata.is_subagent);
130        (!self.subagents_only && !is_subagent && self.message_hash_fallback)
131            .then(|| {
132                first_user_message_hash(request).map(|hash| {
133                    tracing::debug!(affinity_key = %hash, "affinity using message hash fallback");
134                    AffinityKey::Session(hash)
135                })
136            })
137            .flatten()
138    }
139}
140
141#[async_trait]
142impl<S> Processor<S> for AffinityRouter
143where
144    S: Send + 'static,
145{
146    async fn process(&self, _state: &mut S, event: Event<'_>) -> crate::Result<()> {
147        if let Event::Decision { request, decision } = event
148            && let Some(key) = self.affinity_key(request)
149        {
150            let model = decision.selected_model();
151            let mut assignments = self.assignments.lock();
152            if self.should_latch(model) && !assignments.contains_key(&key) {
153                evict_if_full(&mut assignments);
154                assignments.insert(key, model.to_string());
155            }
156        }
157        Ok(())
158    }
159}
160
161/// Hashes the first user message so later turns retain the initial task's affinity.
162/// For benchmarking purpose with harnesses, task instructions are added as a user prompt to the request so we hash the initial user message.
163/// TODO: Have not considered multi-modal payloads yet. That needs to be handled separately.
164fn first_user_message_hash(request: &Request) -> Option<String> {
165    let message = request
166        .llm_request
167        .messages
168        .iter()
169        .find(|message| message.role == Role::User)?;
170    let mut hasher = DefaultHasher::new();
171    message.text_content("")?.hash(&mut hasher);
172    Some(format!("{:016x}", hasher.finish()))
173}
174
175#[async_trait]
176impl<S> Classifier<S> for AffinityRouter
177where
178    S: Send + 'static,
179{
180    async fn score(
181        &self,
182        _state: &mut S,
183        request: &mut Request,
184        _driver: Option<&Driver>,
185    ) -> crate::Result<(Classification, Option<switchyard_protocol::Response>)> {
186        let Some(key) = self.affinity_key(request) else {
187            return Ok((Classification::Scores(Vec::new()), None));
188        };
189        let assigned = self.assignments.lock().get(&key).cloned();
190        Ok((
191            Classification::Scores(match assigned {
192                Some(target) => vec![Score {
193                    confidence: 1.0,
194                    target,
195                }],
196                None => Vec::new(),
197            }),
198            None,
199        ))
200    }
201}
202
203/// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`].
204fn evict_if_full(assignments: &mut HashMap<AffinityKey, String>) {
205    if assignments.len() >= MAX_ASSIGNMENTS
206        && let Some(evicted) = assignments.keys().next().cloned()
207    {
208        assignments.remove(&evicted);
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    use std::sync::Arc;
217
218    use switchyard_protocol::{
219        ContentBlock, Decision, LlmRequest, Message, Metadata, text_request,
220    };
221
222    /// Boxed, thread-safe error type keeping the test helpers ergonomic.
223    type BoxErr = Box<dyn std::error::Error + Send + Sync>;
224
225    /// A decision that reports a fixed selected model.
226    struct FixedDecision(&'static str);
227
228    impl Decision for FixedDecision {
229        fn selected_model(&self) -> &str {
230            self.0
231        }
232
233        fn reasoning(&self) -> Option<&str> {
234            None
235        }
236
237        fn as_any(&self) -> &dyn std::any::Any {
238            self
239        }
240    }
241
242    fn request(metadata: Metadata) -> Request {
243        Request {
244            llm_request: text_request(Some("auto".to_string()), "hi"),
245            raw_request: None,
246            metadata: Some(metadata),
247        }
248    }
249
250    fn task_request(
251        metadata: Option<Metadata>,
252        first_user: &str,
253        follow_up: Option<&str>,
254    ) -> Request {
255        let mut messages = vec![
256            Message::text(Role::System, "follow repository instructions"),
257            Message::text(Role::User, first_user),
258        ];
259        if let Some(follow_up) = follow_up {
260            messages.push(Message::text(Role::Assistant, "I will inspect the code."));
261            messages.push(Message::text(Role::User, follow_up));
262        }
263        Request {
264            llm_request: LlmRequest {
265                model: Some("auto".to_string()),
266                messages,
267                ..LlmRequest::default()
268            },
269            raw_request: None,
270            metadata,
271        }
272    }
273
274    fn session(session_id: &str, agent_id: &str) -> Metadata {
275        Metadata {
276            session_id: Some(session_id.to_string()),
277            agent_id: Some(agent_id.to_string()),
278            ..Metadata::default()
279        }
280    }
281
282    fn subagent(agent_id: &str, task_id: &str) -> Metadata {
283        Metadata {
284            session_id: Some("session-1".to_string()),
285            agent_id: Some(agent_id.to_string()),
286            task_id: Some(task_id.to_string()),
287            is_subagent: true,
288            ..Metadata::default()
289        }
290    }
291
292    /// Folds a request and its decision through the router, retaining `model`.
293    async fn retain(
294        router: &AffinityRouter,
295        state: &mut (),
296        request: &mut Request,
297        model: &'static str,
298    ) -> Result<(), BoxErr> {
299        router
300            .process(
301                state,
302                Event::Decision {
303                    request,
304                    decision: &FixedDecision(model),
305                },
306            )
307            .await?;
308        Ok(())
309    }
310
311    /// Scores through the definitive classification variant used by affinity.
312    async fn scores(
313        classifier: &dyn Classifier,
314        state: &mut (),
315        request: &mut Request,
316    ) -> Result<Vec<Score>, BoxErr> {
317        match classifier.score(state, request, None).await?.0 {
318            Classification::Scores(scores) => Ok(scores),
319            Classification::Ambiguous(_) => Err("affinity never returns ambiguous scores".into()),
320        }
321    }
322
323    #[tokio::test]
324    async fn session_retains_first_model_across_requests() -> Result<(), BoxErr> {
325        let router = AffinityRouter::new();
326        let mut state = ();
327
328        let mut first = request(session("session-1", "agent-a"));
329        retain(&router, &mut state, &mut first, "model-a").await?;
330
331        // A different agent in the same session is scored onto the retained model.
332        let mut second = request(session("session-1", "agent-b"));
333        let scores = scores(&router, &mut state, &mut second).await?;
334        assert_eq!(scores.len(), 1);
335        assert_eq!(scores[0].confidence, 1.0);
336        assert_eq!(scores[0].target, "model-a");
337        Ok(())
338    }
339
340    #[tokio::test]
341    async fn subagent_only_retains_children_without_latching_root_traffic() -> Result<(), BoxErr> {
342        let router = AffinityRouter::for_subagents();
343        let mut state = ();
344
345        let mut root = request(session("session-1", "root-agent"));
346        retain(&router, &mut state, &mut root, "model-a").await?;
347        assert!(scores(&router, &mut state, &mut root).await?.is_empty());
348
349        let mut first_child_turn = request(subagent("child-1", "task-1"));
350        retain(&router, &mut state, &mut first_child_turn, "model-b").await?;
351        let mut later_child_turn = request(subagent("child-1", "task-2"));
352        let scores = scores(&router, &mut state, &mut later_child_turn).await?;
353        assert_eq!(
354            scores.first().map(|score| score.target.as_str()),
355            Some("model-b")
356        );
357        Ok(())
358    }
359
360    #[tokio::test]
361    async fn first_decision_wins() -> Result<(), BoxErr> {
362        let router = AffinityRouter::new();
363        let mut state = ();
364
365        let mut req = request(session("session-1", "agent-a"));
366        retain(&router, &mut state, &mut req, "model-a").await?;
367        // A later decision for the same identity must not overwrite the first.
368        retain(&router, &mut state, &mut req, "model-b").await?;
369
370        let scores = scores(&router, &mut state, &mut req).await?;
371        assert_eq!(
372            scores.first().map(|score| score.target.as_str()),
373            Some("model-a")
374        );
375        Ok(())
376    }
377
378    #[tokio::test]
379    async fn subagent_is_keyed_by_agent_not_task() -> Result<(), BoxErr> {
380        let router = AffinityRouter::new();
381        let mut state = ();
382
383        let mut first = request(subagent("child-1", "task-1"));
384        retain(&router, &mut state, &mut first, "model-a").await?;
385
386        // Same child, different task: still scored onto the retained model.
387        let mut second = request(subagent("child-1", "task-2"));
388        let scores = scores(&router, &mut state, &mut second).await?;
389        assert_eq!(
390            scores.first().map(|score| score.target.as_str()),
391            Some("model-a")
392        );
393        Ok(())
394    }
395
396    #[tokio::test]
397    async fn distinct_subagents_are_assigned_independently() -> Result<(), BoxErr> {
398        let router = AffinityRouter::new();
399        let mut state = ();
400
401        // One child in the session is pinned...
402        retain(
403            &router,
404            &mut state,
405            &mut request(subagent("child-1", "task-1")),
406            "model-a",
407        )
408        .await?;
409
410        // ...a sibling child in the same session has no assignment of its own yet.
411        let mut sibling = request(subagent("child-2", "task-1"));
412        assert!(scores(&router, &mut state, &mut sibling).await?.is_empty());
413        Ok(())
414    }
415
416    #[tokio::test]
417    async fn subagent_does_not_inherit_session_assignment() -> Result<(), BoxErr> {
418        let router = AffinityRouter::new();
419        let mut state = ();
420
421        // The session root is pinned, but a sub-agent is keyed separately...
422        retain(
423            &router,
424            &mut state,
425            &mut request(session("session-1", "root-1")),
426            "model-a",
427        )
428        .await?;
429
430        // ...so the sub-agent abstains until it is assigned in its own right.
431        let mut child = request(subagent("child-1", "task-1"));
432        assert!(scores(&router, &mut state, &mut child).await?.is_empty());
433        Ok(())
434    }
435
436    #[tokio::test]
437    async fn classifier_abstains_without_a_session() -> Result<(), BoxErr> {
438        let router = AffinityRouter::new();
439        let mut state = ();
440
441        // No session id at all: nothing to key on.
442        let mut req = request(Metadata::default());
443        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
444        Ok(())
445    }
446
447    #[tokio::test]
448    async fn message_hash_fallback_uses_the_first_user_message() -> Result<(), BoxErr> {
449        let router = AffinityRouter::new().with_message_hash_fallback();
450        let mut state = ();
451
452        let mut first = task_request(
453            None,
454            "Add a unit test for this function.",
455            Some("Now run the test suite."),
456        );
457        retain(&router, &mut state, &mut first, "weak").await?;
458
459        let mut follow_up = task_request(
460            None,
461            "Add a unit test for this function.",
462            Some("Now file a pull request."),
463        );
464        assert_eq!(
465            scores(&router, &mut state, &mut follow_up)
466                .await?
467                .first()
468                .map(|score| score.target.as_str()),
469            Some("weak")
470        );
471
472        let mut other_task = task_request(
473            None,
474            "Reimplement this binary from two input/output pairs.",
475            Some("Now run the test suite."),
476        );
477        assert!(
478            scores(&router, &mut state, &mut other_task)
479                .await?
480                .is_empty()
481        );
482        Ok(())
483    }
484
485    #[test]
486    fn user_message_hash_ignores_non_text_provider_payloads() {
487        let request = |user_message| Request {
488            llm_request: LlmRequest {
489                messages: vec![user_message],
490                ..LlmRequest::default()
491            },
492            raw_request: None,
493            metadata: None,
494        };
495        let text_only = request(Message::text(Role::User, "Implement the parser."));
496        let text_with_reasoning = request(Message {
497            role: Role::User,
498            content: vec![
499                ContentBlock::Text {
500                    text: "Implement the parser.".to_string(),
501                },
502                ContentBlock::Reasoning {
503                    text: "Internal provider reasoning.".to_string(),
504                    signature: Some("provider-signature".to_string()),
505                },
506            ],
507        });
508
509        assert_eq!(
510            first_user_message_hash(&text_only),
511            first_user_message_hash(&text_with_reasoning)
512        );
513    }
514
515    #[tokio::test]
516    async fn metadata_session_takes_precedence_over_message_hash() -> Result<(), BoxErr> {
517        let router = AffinityRouter::new().with_message_hash_fallback();
518        let mut state = ();
519
520        let mut first = task_request(
521            Some(session("session-1", "agent-a")),
522            "Implement the parser.",
523            None,
524        );
525        retain(&router, &mut state, &mut first, "strong").await?;
526
527        let mut other_session = task_request(
528            Some(session("session-2", "agent-a")),
529            "Implement the parser.",
530            None,
531        );
532        assert!(
533            scores(&router, &mut state, &mut other_session)
534                .await?
535                .is_empty()
536        );
537        Ok(())
538    }
539
540    #[tokio::test]
541    async fn message_hash_fallback_abstains_for_subagents() -> Result<(), BoxErr> {
542        let router = AffinityRouter::new().with_message_hash_fallback();
543        let mut state = ();
544        let mut subagent = task_request(
545            Some(Metadata {
546                is_subagent: true,
547                ..Metadata::default()
548            }),
549            "Implement the parser.",
550            None,
551        );
552
553        assert!(scores(&router, &mut state, &mut subagent).await?.is_empty());
554        Ok(())
555    }
556
557    #[tokio::test]
558    async fn one_router_serves_both_roles() -> Result<(), BoxErr> {
559        // The same instance is registered under both SDK roles; a decision folded in via
560        // the processor handle is read back via the classifier handle.
561        let router = Arc::new(AffinityRouter::new());
562        let processor: Arc<dyn Processor> = router.clone();
563        let classifier: Arc<dyn Classifier> = router;
564        let mut state = ();
565
566        let mut first = request(session("session-1", "agent-a"));
567        processor
568            .process(
569                &mut state,
570                Event::Decision {
571                    request: &mut first,
572                    decision: &FixedDecision("model-a"),
573                },
574            )
575            .await?;
576
577        let mut second = request(session("session-1", "agent-b"));
578        let scores = scores(classifier.as_ref(), &mut state, &mut second).await?;
579        assert_eq!(
580            scores.first().map(|score| score.target.as_str()),
581            Some("model-a")
582        );
583        Ok(())
584    }
585
586    #[tokio::test]
587    async fn decision_without_an_affinity_identity_is_ignored() -> Result<(), BoxErr> {
588        let router = AffinityRouter::new();
589        let mut state = ();
590        let mut unkeyed = request(Metadata::default());
591
592        router
593            .process(
594                &mut state,
595                Event::Decision {
596                    request: &mut unkeyed,
597                    decision: &FixedDecision("model-a"),
598                },
599            )
600            .await?;
601
602        let mut req = request(session("session-1", "agent-a"));
603        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
604        Ok(())
605    }
606
607    #[tokio::test]
608    async fn decisions_retain_their_originating_request_identity() -> Result<(), BoxErr> {
609        let router = AffinityRouter::new();
610        let mut state = ();
611        let mut first = request(session("session-1", "agent-a"));
612        let mut second = request(session("session-2", "agent-b"));
613
614        // Replay decisions in the opposite order. Each decision carries its request,
615        // so the assignments cannot cross.
616        router
617            .process(
618                &mut state,
619                Event::Decision {
620                    request: &mut second,
621                    decision: &FixedDecision("model-b"),
622                },
623            )
624            .await?;
625        router
626            .process(
627                &mut state,
628                Event::Decision {
629                    request: &mut first,
630                    decision: &FixedDecision("model-a"),
631                },
632            )
633            .await?;
634
635        let first_scores = scores(&router, &mut state, &mut first).await?;
636        let second_scores = scores(&router, &mut state, &mut second).await?;
637        assert_eq!(
638            first_scores.first().map(|score| score.target.as_str()),
639            Some("model-a")
640        );
641        assert_eq!(
642            second_scores.first().map(|score| score.target.as_str()),
643            Some("model-b")
644        );
645        Ok(())
646    }
647
648    #[tokio::test]
649    async fn distinct_sessions_are_assigned_independently() -> Result<(), BoxErr> {
650        let router = AffinityRouter::new();
651        let mut state = ();
652
653        retain(
654            &router,
655            &mut state,
656            &mut request(session("session-1", "agent-a")),
657            "model-a",
658        )
659        .await?;
660        retain(
661            &router,
662            &mut state,
663            &mut request(session("session-2", "agent-a")),
664            "model-b",
665        )
666        .await?;
667
668        let first = scores(
669            &router,
670            &mut state,
671            &mut request(session("session-1", "other")),
672        )
673        .await?;
674        let second = scores(
675            &router,
676            &mut state,
677            &mut request(session("session-2", "other")),
678        )
679        .await?;
680        assert_eq!(
681            first.first().map(|score| score.target.as_str()),
682            Some("model-a")
683        );
684        assert_eq!(
685            second.first().map(|score| score.target.as_str()),
686            Some("model-b")
687        );
688        Ok(())
689    }
690
691    #[tokio::test]
692    async fn subagent_without_an_agent_id_is_not_keyed() -> Result<(), BoxErr> {
693        let router = AffinityRouter::new();
694        let mut state = ();
695
696        // The sub-agent flag is set but no agent id is present, so no key can be formed;
697        // the request is neither retained nor scored.
698        let metadata = Metadata {
699            session_id: Some("session-1".to_string()),
700            is_subagent: true,
701            ..Metadata::default()
702        };
703        let mut req = request(metadata);
704        retain(&router, &mut state, &mut req, "model-a").await?;
705        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
706        Ok(())
707    }
708
709    #[tokio::test]
710    async fn assignments_are_bounded_by_the_cap() -> Result<(), BoxErr> {
711        let router = AffinityRouter::new();
712        let mut state = ();
713
714        // One distinct session past the cap forces exactly one eviction.
715        for index in 0..=MAX_ASSIGNMENTS {
716            let session_id = format!("session-{index}");
717            retain(
718                &router,
719                &mut state,
720                &mut request(session(&session_id, "agent-a")),
721                "model-a",
722            )
723            .await?;
724        }
725
726        let len = router.assignments.lock().len();
727        assert_eq!(len, MAX_ASSIGNMENTS);
728        Ok(())
729    }
730
731    #[tokio::test]
732    async fn latch_only_retains_matching_models() -> Result<(), BoxErr> {
733        let router = AffinityRouter::new().with_latch_only(["strong"]);
734        let mut state = ();
735        let mut req = request(session("session-1", "agent-a"));
736
737        // A "weak" decision is not retained — a later turn is not latched.
738        retain(&router, &mut state, &mut req, "weak").await?;
739        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
740
741        // A "strong" decision is retained — later turns latch onto it.
742        retain(&router, &mut state, &mut req, "strong").await?;
743        assert_eq!(
744            scores(&router, &mut state, &mut req)
745                .await?
746                .first()
747                .map(|s| s.target.as_str()),
748            Some("strong")
749        );
750        Ok(())
751    }
752}
switchyard_libsy/algorithms/util/affinity.rs

switchyard_libsy/algorithms/util/
affinity.rs
