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

switchyard_libsy/algorithms/util/
affinity.rs
