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, RoutingIdentity};
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/// Retains a model per request identity and forces it on later matching requests.
36///
37/// Register the same instance as both a processor and a classifier; the two roles share
38/// the retained assignments through this instance's interior storage. Decision events carry
39/// their originating request, so concurrent turns cannot bind one request's identity to
40/// another request's selected model.
41///
42/// [`with_latch_only`](Self::with_latch_only) narrows *which* models are retained — a
43/// decision for any other model routes normally but is not latched (the escalation latch:
44/// retain only the strong tier, never the weak one).
45#[derive(Default)]
46pub struct AffinityRouter {
47    /// When set, only these models are retained; a decision for any other model is not latched.
48    latch_only: Option<HashSet<String>>,
49    /// Whether root-session requests should abstain instead of being retained.
50    subagents_only: bool,
51    /// In absence of headers, use the message hash based fallback key to do task based routing
52    message_hash_fallback: bool,
53    /// Retained assignments, shared across this router's processor and classifier roles.
54    ///
55    /// Held on the instance so the two roles share one process-local map through a
56    /// single registered [`Arc`](std::sync::Arc); bounded by [`MAX_ASSIGNMENTS`].
57    assignments: Mutex<HashMap<RoutingIdentity, String>>,
58}
59
60impl AffinityRouter {
61    /// Creates a router that latches every decision.
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    /// Creates a router that retains assignments only for explicitly identified child agents.
67    ///
68    /// Root-agent requests always abstain, so a later classifier selects them on every turn.
69    pub fn for_subagents() -> Self {
70        Self {
71            subagents_only: true,
72            ..Self::default()
73        }
74    }
75
76    /// Uses the first user-message text as a fallback key when metadata has no session.
77    pub fn with_message_hash_fallback(mut self) -> Self {
78        self.message_hash_fallback = true;
79        self
80    }
81
82    /// Restricts latching to `models`; a decision for any other model routes but is not
83    /// retained.
84    pub fn with_latch_only(mut self, models: impl IntoIterator<Item = impl Into<String>>) -> Self {
85        self.latch_only = Some(models.into_iter().map(Into::into).collect());
86        self
87    }
88
89    /// Whether a decision for `model` should be retained.
90    fn should_latch(&self, model: &str) -> bool {
91        self.latch_only
92            .as_ref()
93            .is_none_or(|set| set.contains(model))
94    }
95
96    /// Derives the stable identity this router should retain for `request`.
97    fn affinity_key(&self, request: &Request) -> Option<RoutingIdentity> {
98        if let Some(identity) = RoutingIdentity::from_request(request) {
99            return match identity {
100                RoutingIdentity::Session(_) if self.subagents_only => None,
101                identity => Some(identity),
102            };
103        }
104
105        // If headers are not present and we are not a subagent, use the message hash based fallback key to do task based routing
106        let is_subagent = request
107            .metadata
108            .as_ref()
109            .is_some_and(|metadata| metadata.is_subagent);
110        (!self.subagents_only && !is_subagent && self.message_hash_fallback)
111            .then(|| {
112                first_user_message_hash(request).map(|hash| {
113                    tracing::debug!(affinity_key = %hash, "affinity using message hash fallback");
114                    RoutingIdentity::Session(hash)
115                })
116            })
117            .flatten()
118    }
119}
120
121#[async_trait]
122impl<S> Processor<S> for AffinityRouter
123where
124    S: Send + 'static,
125{
126    async fn process(&self, _state: &mut S, event: Event<'_>) -> crate::Result<()> {
127        if let Event::Decision { request, decision } = event
128            && let Some(key) = self.affinity_key(request)
129        {
130            let model = decision.selected_model();
131            let mut assignments = self.assignments.lock();
132            if self.should_latch(model) && !assignments.contains_key(&key) {
133                evict_if_full(&mut assignments);
134                assignments.insert(key, model.to_string());
135            }
136        }
137        Ok(())
138    }
139}
140
141/// Hashes the first user message so later turns retain the initial task's affinity.
142/// For benchmarking purpose with harnesses, task instructions are added as a user prompt to the request so we hash the initial user message.
143/// TODO: Have not considered multi-modal payloads yet. That needs to be handled separately.
144fn first_user_message_hash(request: &Request) -> Option<String> {
145    let message = request
146        .llm_request
147        .messages
148        .iter()
149        .find(|message| message.role == Role::User)?;
150    let mut hasher = DefaultHasher::new();
151    message.text_content("")?.hash(&mut hasher);
152    Some(format!("{:016x}", hasher.finish()))
153}
154
155#[async_trait]
156impl<S> Classifier<S> for AffinityRouter
157where
158    S: Send + 'static,
159{
160    fn target_unavailable(&self, request: &Request, target: &str) {
161        let Some(key) = self.affinity_key(request) else {
162            return;
163        };
164        let mut assignments = self.assignments.lock();
165        if assignments
166            .get(&key)
167            .is_some_and(|assigned| assigned == target)
168        {
169            assignments.remove(&key);
170        }
171    }
172
173    async fn score(
174        &self,
175        _state: &mut S,
176        request: &mut Request,
177        _driver: Option<&Driver>,
178    ) -> crate::Result<(Classification, Option<switchyard_protocol::Response>)> {
179        let Some(key) = self.affinity_key(request) else {
180            return Ok((Classification::Scores(Vec::new()), None));
181        };
182        let assigned = self.assignments.lock().get(&key).cloned();
183        Ok((
184            Classification::Scores(match assigned {
185                Some(target) => vec![Score {
186                    confidence: 1.0,
187                    target,
188                }],
189                None => Vec::new(),
190            }),
191            None,
192        ))
193    }
194}
195
196/// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`].
197fn evict_if_full(assignments: &mut HashMap<RoutingIdentity, String>) {
198    if assignments.len() >= MAX_ASSIGNMENTS
199        && let Some(evicted) = assignments.keys().next().cloned()
200    {
201        assignments.remove(&evicted);
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    use std::sync::Arc;
210
211    use switchyard_protocol::{
212        ContentBlock, Decision, LlmRequest, Message, Metadata, text_request,
213    };
214
215    /// Boxed, thread-safe error type keeping the test helpers ergonomic.
216    type BoxErr = Box<dyn std::error::Error + Send + Sync>;
217
218    /// A decision that reports a fixed selected model.
219    struct FixedDecision(&'static str);
220
221    impl Decision for FixedDecision {
222        fn selected_model(&self) -> &str {
223            self.0
224        }
225
226        fn reasoning(&self) -> Option<&str> {
227            None
228        }
229
230        fn as_any(&self) -> &dyn std::any::Any {
231            self
232        }
233    }
234
235    fn request(metadata: Metadata) -> Request {
236        Request {
237            llm_request: text_request(Some("auto".to_string()), "hi"),
238            raw_request: None,
239            metadata: Some(metadata),
240        }
241    }
242
243    fn task_request(
244        metadata: Option<Metadata>,
245        first_user: &str,
246        follow_up: Option<&str>,
247    ) -> Request {
248        let mut messages = vec![
249            Message::text(Role::System, "follow repository instructions"),
250            Message::text(Role::User, first_user),
251        ];
252        if let Some(follow_up) = follow_up {
253            messages.push(Message::text(Role::Assistant, "I will inspect the code."));
254            messages.push(Message::text(Role::User, follow_up));
255        }
256        Request {
257            llm_request: LlmRequest {
258                model: Some("auto".to_string()),
259                messages,
260                ..LlmRequest::default()
261            },
262            raw_request: None,
263            metadata,
264        }
265    }
266
267    fn session(session_id: &str, agent_id: &str) -> Metadata {
268        Metadata {
269            session_id: Some(session_id.to_string()),
270            agent_id: Some(agent_id.to_string()),
271            ..Metadata::default()
272        }
273    }
274
275    fn subagent(agent_id: &str, task_id: &str) -> Metadata {
276        Metadata {
277            session_id: Some("session-1".to_string()),
278            agent_id: Some(agent_id.to_string()),
279            task_id: Some(task_id.to_string()),
280            is_subagent: true,
281            ..Metadata::default()
282        }
283    }
284
285    /// Folds a request and its decision through the router, retaining `model`.
286    async fn retain(
287        router: &AffinityRouter,
288        state: &mut (),
289        request: &mut Request,
290        model: &'static str,
291    ) -> Result<(), BoxErr> {
292        router
293            .process(
294                state,
295                Event::Decision {
296                    request,
297                    decision: &FixedDecision(model),
298                },
299            )
300            .await?;
301        Ok(())
302    }
303
304    /// Scores through the definitive classification variant used by affinity.
305    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        // A different agent in the same session is scored onto the retained model.
325        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        // A later decision for the same identity must not overwrite the first.
361        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        // Same child, different task: still scored onto the retained model.
380        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        // One child in the session is pinned...
395        retain(
396            &router,
397            &mut state,
398            &mut request(subagent("child-1", "task-1")),
399            "model-a",
400        )
401        .await?;
402
403        // ...a sibling child in the same session has no assignment of its own yet.
404        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        // The session root is pinned, but a sub-agent is keyed separately...
415        retain(
416            &router,
417            &mut state,
418            &mut request(session("session-1", "root-1")),
419            "model-a",
420        )
421        .await?;
422
423        // ...so the sub-agent abstains until it is assigned in its own right.
424        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        // No session id at all: nothing to key on.
435        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    #[test]
479    fn user_message_hash_ignores_non_text_provider_payloads() {
480        let request = |user_message| Request {
481            llm_request: LlmRequest {
482                messages: vec![user_message],
483                ..LlmRequest::default()
484            },
485            raw_request: None,
486            metadata: None,
487        };
488        let text_only = request(Message::text(Role::User, "Implement the parser."));
489        let text_with_reasoning = request(Message {
490            role: Role::User,
491            content: vec![
492                ContentBlock::Text {
493                    text: "Implement the parser.".to_string(),
494                },
495                ContentBlock::Reasoning {
496                    text: "Internal provider reasoning.".to_string(),
497                    signature: Some("provider-signature".to_string()),
498                },
499            ],
500        });
501
502        assert_eq!(
503            first_user_message_hash(&text_only),
504            first_user_message_hash(&text_with_reasoning)
505        );
506    }
507
508    #[tokio::test]
509    async fn metadata_session_takes_precedence_over_message_hash() -> Result<(), BoxErr> {
510        let router = AffinityRouter::new().with_message_hash_fallback();
511        let mut state = ();
512
513        let mut first = task_request(
514            Some(session("session-1", "agent-a")),
515            "Implement the parser.",
516            None,
517        );
518        retain(&router, &mut state, &mut first, "strong").await?;
519
520        let mut other_session = task_request(
521            Some(session("session-2", "agent-a")),
522            "Implement the parser.",
523            None,
524        );
525        assert!(
526            scores(&router, &mut state, &mut other_session)
527                .await?
528                .is_empty()
529        );
530        Ok(())
531    }
532
533    #[tokio::test]
534    async fn message_hash_fallback_abstains_for_subagents() -> Result<(), BoxErr> {
535        let router = AffinityRouter::new().with_message_hash_fallback();
536        let mut state = ();
537        let mut subagent = task_request(
538            Some(Metadata {
539                is_subagent: true,
540                ..Metadata::default()
541            }),
542            "Implement the parser.",
543            None,
544        );
545
546        assert!(scores(&router, &mut state, &mut subagent).await?.is_empty());
547        Ok(())
548    }
549
550    #[tokio::test]
551    async fn one_router_serves_both_roles() -> Result<(), BoxErr> {
552        // The same instance is registered under both SDK roles; a decision folded in via
553        // the processor handle is read back via the classifier handle.
554        let router = Arc::new(AffinityRouter::new());
555        let processor: Arc<dyn Processor> = router.clone();
556        let classifier: Arc<dyn Classifier> = router;
557        let mut state = ();
558
559        let mut first = request(session("session-1", "agent-a"));
560        processor
561            .process(
562                &mut state,
563                Event::Decision {
564                    request: &mut first,
565                    decision: &FixedDecision("model-a"),
566                },
567            )
568            .await?;
569
570        let mut second = request(session("session-1", "agent-b"));
571        let scores = scores(classifier.as_ref(), &mut state, &mut second).await?;
572        assert_eq!(
573            scores.first().map(|score| score.target.as_str()),
574            Some("model-a")
575        );
576        Ok(())
577    }
578
579    #[tokio::test]
580    async fn decision_without_an_affinity_identity_is_ignored() -> Result<(), BoxErr> {
581        let router = AffinityRouter::new();
582        let mut state = ();
583        let mut unkeyed = request(Metadata::default());
584
585        router
586            .process(
587                &mut state,
588                Event::Decision {
589                    request: &mut unkeyed,
590                    decision: &FixedDecision("model-a"),
591                },
592            )
593            .await?;
594
595        let mut req = request(session("session-1", "agent-a"));
596        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
597        Ok(())
598    }
599
600    #[tokio::test]
601    async fn decisions_retain_their_originating_request_identity() -> Result<(), BoxErr> {
602        let router = AffinityRouter::new();
603        let mut state = ();
604        let mut first = request(session("session-1", "agent-a"));
605        let mut second = request(session("session-2", "agent-b"));
606
607        // Replay decisions in the opposite order. Each decision carries its request,
608        // so the assignments cannot cross.
609        router
610            .process(
611                &mut state,
612                Event::Decision {
613                    request: &mut second,
614                    decision: &FixedDecision("model-b"),
615                },
616            )
617            .await?;
618        router
619            .process(
620                &mut state,
621                Event::Decision {
622                    request: &mut first,
623                    decision: &FixedDecision("model-a"),
624                },
625            )
626            .await?;
627
628        let first_scores = scores(&router, &mut state, &mut first).await?;
629        let second_scores = scores(&router, &mut state, &mut second).await?;
630        assert_eq!(
631            first_scores.first().map(|score| score.target.as_str()),
632            Some("model-a")
633        );
634        assert_eq!(
635            second_scores.first().map(|score| score.target.as_str()),
636            Some("model-b")
637        );
638        Ok(())
639    }
640
641    #[tokio::test]
642    async fn distinct_sessions_are_assigned_independently() -> Result<(), BoxErr> {
643        let router = AffinityRouter::new();
644        let mut state = ();
645
646        retain(
647            &router,
648            &mut state,
649            &mut request(session("session-1", "agent-a")),
650            "model-a",
651        )
652        .await?;
653        retain(
654            &router,
655            &mut state,
656            &mut request(session("session-2", "agent-a")),
657            "model-b",
658        )
659        .await?;
660
661        let first = scores(
662            &router,
663            &mut state,
664            &mut request(session("session-1", "other")),
665        )
666        .await?;
667        let second = scores(
668            &router,
669            &mut state,
670            &mut request(session("session-2", "other")),
671        )
672        .await?;
673        assert_eq!(
674            first.first().map(|score| score.target.as_str()),
675            Some("model-a")
676        );
677        assert_eq!(
678            second.first().map(|score| score.target.as_str()),
679            Some("model-b")
680        );
681        Ok(())
682    }
683
684    #[tokio::test]
685    async fn subagent_without_an_agent_id_is_not_keyed() -> Result<(), BoxErr> {
686        let router = AffinityRouter::new();
687        let mut state = ();
688
689        // The sub-agent flag is set but no agent id is present, so no key can be formed;
690        // the request is neither retained nor scored.
691        let metadata = Metadata {
692            session_id: Some("session-1".to_string()),
693            is_subagent: true,
694            ..Metadata::default()
695        };
696        let mut req = request(metadata);
697        retain(&router, &mut state, &mut req, "model-a").await?;
698        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
699        Ok(())
700    }
701
702    #[tokio::test]
703    async fn assignments_are_bounded_by_the_cap() -> Result<(), BoxErr> {
704        let router = AffinityRouter::new();
705        let mut state = ();
706
707        // One distinct session past the cap forces exactly one eviction.
708        for index in 0..=MAX_ASSIGNMENTS {
709            let session_id = format!("session-{index}");
710            retain(
711                &router,
712                &mut state,
713                &mut request(session(&session_id, "agent-a")),
714                "model-a",
715            )
716            .await?;
717        }
718
719        let len = router.assignments.lock().len();
720        assert_eq!(len, MAX_ASSIGNMENTS);
721        Ok(())
722    }
723
724    #[tokio::test]
725    async fn latch_only_retains_matching_models() -> Result<(), BoxErr> {
726        let router = AffinityRouter::new().with_latch_only(["strong"]);
727        let mut state = ();
728        let mut req = request(session("session-1", "agent-a"));
729
730        // A "weak" decision is not retained — a later turn is not latched.
731        retain(&router, &mut state, &mut req, "weak").await?;
732        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
733
734        // A "strong" decision is retained — later turns latch onto it.
735        retain(&router, &mut state, &mut req, "strong").await?;
736        assert_eq!(
737            scores(&router, &mut state, &mut req)
738                .await?
739                .first()
740                .map(|s| s.target.as_str()),
741            Some("strong")
742        );
743        Ok(())
744    }
745}
switchyard_libsy/algorithms/util/affinity.rs

switchyard_libsy/algorithms/util/
affinity.rs
