Skip to main content

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