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    use crate::text::text_request;
216
217    use std::sync::Arc;
218
219    use switchyard_protocol::{ContentBlock, Decision, LlmRequest, Message, Metadata};
220
221    /// Boxed, thread-safe error type keeping the test helpers ergonomic.
222    type BoxErr = Box<dyn std::error::Error + Send + Sync>;
223
224    /// A decision that reports a fixed selected model.
225    struct FixedDecision(&'static str);
226
227    impl Decision for FixedDecision {
228        fn selected_model(&self) -> &str {
229            self.0
230        }
231
232        fn reasoning(&self) -> Option<&str> {
233            None
234        }
235
236        fn as_any(&self) -> &dyn std::any::Any {
237            self
238        }
239    }
240
241    fn request(metadata: Metadata) -> Request {
242        Request {
243            llm_request: text_request(Some("auto".to_string()), "hi"),
244            raw_request: None,
245            metadata: Some(metadata),
246        }
247    }
248
249    fn task_request(
250        metadata: Option<Metadata>,
251        first_user: &str,
252        follow_up: Option<&str>,
253    ) -> Request {
254        let mut messages = vec![
255            Message::text(Role::System, "follow repository instructions"),
256            Message::text(Role::User, first_user),
257        ];
258        if let Some(follow_up) = follow_up {
259            messages.push(Message::text(Role::Assistant, "I will inspect the code."));
260            messages.push(Message::text(Role::User, follow_up));
261        }
262        Request {
263            llm_request: LlmRequest {
264                model: Some("auto".to_string()),
265                messages,
266                ..LlmRequest::default()
267            },
268            raw_request: None,
269            metadata,
270        }
271    }
272
273    fn session(session_id: &str, agent_id: &str) -> Metadata {
274        Metadata {
275            session_id: Some(session_id.to_string()),
276            agent_id: Some(agent_id.to_string()),
277            ..Metadata::default()
278        }
279    }
280
281    fn subagent(agent_id: &str, task_id: &str) -> Metadata {
282        Metadata {
283            session_id: Some("session-1".to_string()),
284            agent_id: Some(agent_id.to_string()),
285            task_id: Some(task_id.to_string()),
286            is_subagent: true,
287            ..Metadata::default()
288        }
289    }
290
291    /// Folds a request and its decision through the router, retaining `model`.
292    async fn retain(
293        router: &AffinityRouter,
294        state: &mut (),
295        request: &mut Request,
296        model: &'static str,
297    ) -> Result<(), BoxErr> {
298        router
299            .process(
300                state,
301                Event::Decision {
302                    request,
303                    decision: &FixedDecision(model),
304                },
305            )
306            .await?;
307        Ok(())
308    }
309
310    /// Scores through the definitive classification variant used by affinity.
311    async fn scores(
312        classifier: &dyn Classifier,
313        state: &mut (),
314        request: &mut Request,
315    ) -> Result<Vec<Score>, BoxErr> {
316        match classifier.score(state, request, None).await?.0 {
317            Classification::Scores(scores) => Ok(scores),
318            Classification::Ambiguous(_) => Err("affinity never returns ambiguous scores".into()),
319        }
320    }
321
322    #[tokio::test]
323    async fn session_retains_first_model_across_requests() -> Result<(), BoxErr> {
324        let router = AffinityRouter::new();
325        let mut state = ();
326
327        let mut first = request(session("session-1", "agent-a"));
328        retain(&router, &mut state, &mut first, "model-a").await?;
329
330        // A different agent in the same session is scored onto the retained model.
331        let mut second = request(session("session-1", "agent-b"));
332        let scores = scores(&router, &mut state, &mut second).await?;
333        assert_eq!(scores.len(), 1);
334        assert_eq!(scores[0].confidence, 1.0);
335        assert_eq!(scores[0].target, "model-a");
336        Ok(())
337    }
338
339    #[tokio::test]
340    async fn subagent_only_retains_children_without_latching_root_traffic() -> Result<(), BoxErr> {
341        let router = AffinityRouter::for_subagents();
342        let mut state = ();
343
344        let mut root = request(session("session-1", "root-agent"));
345        retain(&router, &mut state, &mut root, "model-a").await?;
346        assert!(scores(&router, &mut state, &mut root).await?.is_empty());
347
348        let mut first_child_turn = request(subagent("child-1", "task-1"));
349        retain(&router, &mut state, &mut first_child_turn, "model-b").await?;
350        let mut later_child_turn = request(subagent("child-1", "task-2"));
351        let scores = scores(&router, &mut state, &mut later_child_turn).await?;
352        assert_eq!(
353            scores.first().map(|score| score.target.as_str()),
354            Some("model-b")
355        );
356        Ok(())
357    }
358
359    #[tokio::test]
360    async fn first_decision_wins() -> Result<(), BoxErr> {
361        let router = AffinityRouter::new();
362        let mut state = ();
363
364        let mut req = request(session("session-1", "agent-a"));
365        retain(&router, &mut state, &mut req, "model-a").await?;
366        // A later decision for the same identity must not overwrite the first.
367        retain(&router, &mut state, &mut req, "model-b").await?;
368
369        let scores = scores(&router, &mut state, &mut req).await?;
370        assert_eq!(
371            scores.first().map(|score| score.target.as_str()),
372            Some("model-a")
373        );
374        Ok(())
375    }
376
377    #[tokio::test]
378    async fn subagent_is_keyed_by_agent_not_task() -> Result<(), BoxErr> {
379        let router = AffinityRouter::new();
380        let mut state = ();
381
382        let mut first = request(subagent("child-1", "task-1"));
383        retain(&router, &mut state, &mut first, "model-a").await?;
384
385        // Same child, different task: still scored onto the retained model.
386        let mut second = request(subagent("child-1", "task-2"));
387        let scores = scores(&router, &mut state, &mut second).await?;
388        assert_eq!(
389            scores.first().map(|score| score.target.as_str()),
390            Some("model-a")
391        );
392        Ok(())
393    }
394
395    #[tokio::test]
396    async fn distinct_subagents_are_assigned_independently() -> Result<(), BoxErr> {
397        let router = AffinityRouter::new();
398        let mut state = ();
399
400        // One child in the session is pinned...
401        retain(
402            &router,
403            &mut state,
404            &mut request(subagent("child-1", "task-1")),
405            "model-a",
406        )
407        .await?;
408
409        // ...a sibling child in the same session has no assignment of its own yet.
410        let mut sibling = request(subagent("child-2", "task-1"));
411        assert!(scores(&router, &mut state, &mut sibling).await?.is_empty());
412        Ok(())
413    }
414
415    #[tokio::test]
416    async fn subagent_does_not_inherit_session_assignment() -> Result<(), BoxErr> {
417        let router = AffinityRouter::new();
418        let mut state = ();
419
420        // The session root is pinned, but a sub-agent is keyed separately...
421        retain(
422            &router,
423            &mut state,
424            &mut request(session("session-1", "root-1")),
425            "model-a",
426        )
427        .await?;
428
429        // ...so the sub-agent abstains until it is assigned in its own right.
430        let mut child = request(subagent("child-1", "task-1"));
431        assert!(scores(&router, &mut state, &mut child).await?.is_empty());
432        Ok(())
433    }
434
435    #[tokio::test]
436    async fn classifier_abstains_without_a_session() -> Result<(), BoxErr> {
437        let router = AffinityRouter::new();
438        let mut state = ();
439
440        // No session id at all: nothing to key on.
441        let mut req = request(Metadata::default());
442        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
443        Ok(())
444    }
445
446    #[tokio::test]
447    async fn message_hash_fallback_uses_the_first_user_message() -> Result<(), BoxErr> {
448        let router = AffinityRouter::new().with_message_hash_fallback();
449        let mut state = ();
450
451        let mut first = task_request(
452            None,
453            "Add a unit test for this function.",
454            Some("Now run the test suite."),
455        );
456        retain(&router, &mut state, &mut first, "weak").await?;
457
458        let mut follow_up = task_request(
459            None,
460            "Add a unit test for this function.",
461            Some("Now file a pull request."),
462        );
463        assert_eq!(
464            scores(&router, &mut state, &mut follow_up)
465                .await?
466                .first()
467                .map(|score| score.target.as_str()),
468            Some("weak")
469        );
470
471        let mut other_task = task_request(
472            None,
473            "Reimplement this binary from two input/output pairs.",
474            Some("Now run the test suite."),
475        );
476        assert!(
477            scores(&router, &mut state, &mut other_task)
478                .await?
479                .is_empty()
480        );
481        Ok(())
482    }
483
484    #[test]
485    fn user_message_hash_ignores_non_text_provider_payloads() {
486        let request = |user_message| Request {
487            llm_request: LlmRequest {
488                messages: vec![user_message],
489                ..LlmRequest::default()
490            },
491            raw_request: None,
492            metadata: None,
493        };
494        let text_only = request(Message::text(Role::User, "Implement the parser."));
495        let text_with_reasoning = request(Message {
496            role: Role::User,
497            content: vec![
498                ContentBlock::Text {
499                    text: "Implement the parser.".to_string(),
500                },
501                ContentBlock::Reasoning {
502                    text: "Internal provider reasoning.".to_string(),
503                    signature: Some("provider-signature".to_string()),
504                },
505            ],
506        });
507
508        assert_eq!(
509            first_user_message_hash(&text_only),
510            first_user_message_hash(&text_with_reasoning)
511        );
512    }
513
514    #[tokio::test]
515    async fn metadata_session_takes_precedence_over_message_hash() -> Result<(), BoxErr> {
516        let router = AffinityRouter::new().with_message_hash_fallback();
517        let mut state = ();
518
519        let mut first = task_request(
520            Some(session("session-1", "agent-a")),
521            "Implement the parser.",
522            None,
523        );
524        retain(&router, &mut state, &mut first, "strong").await?;
525
526        let mut other_session = task_request(
527            Some(session("session-2", "agent-a")),
528            "Implement the parser.",
529            None,
530        );
531        assert!(
532            scores(&router, &mut state, &mut other_session)
533                .await?
534                .is_empty()
535        );
536        Ok(())
537    }
538
539    #[tokio::test]
540    async fn message_hash_fallback_abstains_for_subagents() -> Result<(), BoxErr> {
541        let router = AffinityRouter::new().with_message_hash_fallback();
542        let mut state = ();
543        let mut subagent = task_request(
544            Some(Metadata {
545                is_subagent: true,
546                ..Metadata::default()
547            }),
548            "Implement the parser.",
549            None,
550        );
551
552        assert!(scores(&router, &mut state, &mut subagent).await?.is_empty());
553        Ok(())
554    }
555
556    #[tokio::test]
557    async fn one_router_serves_both_roles() -> Result<(), BoxErr> {
558        // The same instance is registered under both SDK roles; a decision folded in via
559        // the processor handle is read back via the classifier handle.
560        let router = Arc::new(AffinityRouter::new());
561        let processor: Arc<dyn Processor> = router.clone();
562        let classifier: Arc<dyn Classifier> = router;
563        let mut state = ();
564
565        let mut first = request(session("session-1", "agent-a"));
566        processor
567            .process(
568                &mut state,
569                Event::Decision {
570                    request: &mut first,
571                    decision: &FixedDecision("model-a"),
572                },
573            )
574            .await?;
575
576        let mut second = request(session("session-1", "agent-b"));
577        let scores = scores(classifier.as_ref(), &mut state, &mut second).await?;
578        assert_eq!(
579            scores.first().map(|score| score.target.as_str()),
580            Some("model-a")
581        );
582        Ok(())
583    }
584
585    #[tokio::test]
586    async fn decision_without_an_affinity_identity_is_ignored() -> Result<(), BoxErr> {
587        let router = AffinityRouter::new();
588        let mut state = ();
589        let mut unkeyed = request(Metadata::default());
590
591        router
592            .process(
593                &mut state,
594                Event::Decision {
595                    request: &mut unkeyed,
596                    decision: &FixedDecision("model-a"),
597                },
598            )
599            .await?;
600
601        let mut req = request(session("session-1", "agent-a"));
602        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
603        Ok(())
604    }
605
606    #[tokio::test]
607    async fn decisions_retain_their_originating_request_identity() -> Result<(), BoxErr> {
608        let router = AffinityRouter::new();
609        let mut state = ();
610        let mut first = request(session("session-1", "agent-a"));
611        let mut second = request(session("session-2", "agent-b"));
612
613        // Replay decisions in the opposite order. Each decision carries its request,
614        // so the assignments cannot cross.
615        router
616            .process(
617                &mut state,
618                Event::Decision {
619                    request: &mut second,
620                    decision: &FixedDecision("model-b"),
621                },
622            )
623            .await?;
624        router
625            .process(
626                &mut state,
627                Event::Decision {
628                    request: &mut first,
629                    decision: &FixedDecision("model-a"),
630                },
631            )
632            .await?;
633
634        let first_scores = scores(&router, &mut state, &mut first).await?;
635        let second_scores = scores(&router, &mut state, &mut second).await?;
636        assert_eq!(
637            first_scores.first().map(|score| score.target.as_str()),
638            Some("model-a")
639        );
640        assert_eq!(
641            second_scores.first().map(|score| score.target.as_str()),
642            Some("model-b")
643        );
644        Ok(())
645    }
646
647    #[tokio::test]
648    async fn distinct_sessions_are_assigned_independently() -> Result<(), BoxErr> {
649        let router = AffinityRouter::new();
650        let mut state = ();
651
652        retain(
653            &router,
654            &mut state,
655            &mut request(session("session-1", "agent-a")),
656            "model-a",
657        )
658        .await?;
659        retain(
660            &router,
661            &mut state,
662            &mut request(session("session-2", "agent-a")),
663            "model-b",
664        )
665        .await?;
666
667        let first = scores(
668            &router,
669            &mut state,
670            &mut request(session("session-1", "other")),
671        )
672        .await?;
673        let second = scores(
674            &router,
675            &mut state,
676            &mut request(session("session-2", "other")),
677        )
678        .await?;
679        assert_eq!(
680            first.first().map(|score| score.target.as_str()),
681            Some("model-a")
682        );
683        assert_eq!(
684            second.first().map(|score| score.target.as_str()),
685            Some("model-b")
686        );
687        Ok(())
688    }
689
690    #[tokio::test]
691    async fn subagent_without_an_agent_id_is_not_keyed() -> Result<(), BoxErr> {
692        let router = AffinityRouter::new();
693        let mut state = ();
694
695        // The sub-agent flag is set but no agent id is present, so no key can be formed;
696        // the request is neither retained nor scored.
697        let metadata = Metadata {
698            session_id: Some("session-1".to_string()),
699            is_subagent: true,
700            ..Metadata::default()
701        };
702        let mut req = request(metadata);
703        retain(&router, &mut state, &mut req, "model-a").await?;
704        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
705        Ok(())
706    }
707
708    #[tokio::test]
709    async fn assignments_are_bounded_by_the_cap() -> Result<(), BoxErr> {
710        let router = AffinityRouter::new();
711        let mut state = ();
712
713        // One distinct session past the cap forces exactly one eviction.
714        for index in 0..=MAX_ASSIGNMENTS {
715            let session_id = format!("session-{index}");
716            retain(
717                &router,
718                &mut state,
719                &mut request(session(&session_id, "agent-a")),
720                "model-a",
721            )
722            .await?;
723        }
724
725        let len = router.assignments.lock().len();
726        assert_eq!(len, MAX_ASSIGNMENTS);
727        Ok(())
728    }
729
730    #[tokio::test]
731    async fn latch_only_retains_matching_models() -> Result<(), BoxErr> {
732        let router = AffinityRouter::new().with_latch_only(["strong"]);
733        let mut state = ();
734        let mut req = request(session("session-1", "agent-a"));
735
736        // A "weak" decision is not retained — a later turn is not latched.
737        retain(&router, &mut state, &mut req, "weak").await?;
738        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
739
740        // A "strong" decision is retained — later turns latch onto it.
741        retain(&router, &mut state, &mut req, "strong").await?;
742        assert_eq!(
743            scores(&router, &mut state, &mut req)
744                .await?
745                .first()
746                .map(|s| s.target.as_str()),
747            Some("strong")
748        );
749        Ok(())
750    }
751}
switchyard_libsy/algorithms/util/affinity.rs

switchyard_libsy/algorithms/util/
affinity.rs
