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 serde::Deserialize;
27use switchyard_protocol::{Category, ContentBlock, Message, ModelId, Request, Role};
28
29use crate::core::algorithm::{Driver, RoutingIdentity};
30use crate::core::classifier::{Classification, Classifier, Score};
31use crate::core::processor::{Event, Processor};
32
33/// Upper bound on retained assignments, keeping the process-local map from growing
34/// without limit; the oldest entry is evicted once the bound is reached.
35const MAX_ASSIGNMENTS: usize = 4096;
36
37/// How often the classifier re-decides a session's target.
38#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)]
39#[serde(rename_all = "snake_case")]
40pub enum ClassifyTrigger {
41    /// Judge every request, tool continuations included.
42    #[default]
43    EveryRequest,
44    /// Judge each new user message, holding that target across the tool calls between.
45    UserTurn,
46    /// Judge once and reuse that target for the session.
47    NewSession,
48}
49
50/// Anthropic carries tool results as a `Role::User` message, so role alone cannot tell a
51/// human turn from a tool continuation.
52fn is_user_turn(message: &Message) -> bool {
53    message.role == Role::User
54        && !message
55            .content
56            .iter()
57            .all(|block| matches!(block, ContentBlock::ToolResult(_)))
58}
59
60pub(crate) fn has_new_user_turn(messages: &[Message]) -> bool {
61    messages.last().is_some_and(is_user_turn)
62}
63
64/// Retains a model per request identity and forces it on later matching requests.
65///
66/// Register the same instance as both a processor and a classifier; the two roles share
67/// the retained assignments through this instance's interior storage. Decision events carry
68/// their originating request, so concurrent turns cannot bind one request's identity to
69/// another request's selected model.
70///
71/// [`with_latch_only`](Self::with_latch_only) narrows *which* models are retained — a
72/// decision for any other model routes normally but is not latched (the escalation latch:
73/// retain only the strong tier, never the weak one).
74#[derive(Default)]
75pub struct AffinityRouter {
76    /// When set, only these models are retained; a decision for any other model is not latched.
77    latch_only: Option<HashSet<ModelId>>,
78    /// Whether requests other than delegated sub-agent work should abstain.
79    subagents_only: bool,
80    /// In absence of headers, use the message hash based fallback key to do task based routing
81    message_hash_fallback: bool,
82    /// Whether a new user message drops the assignment so the judge runs again.
83    release_on_user_turn: bool,
84    /// Retained assignments, shared across this router's processor and classifier roles.
85    ///
86    /// Held on the instance so the two roles share one process-local map through a
87    /// single registered [`Arc`](std::sync::Arc); bounded by [`MAX_ASSIGNMENTS`].
88    assignments: Mutex<HashMap<RoutingIdentity, ModelId>>,
89    /// Whether the "no identity to key on" warning has already been emitted.
90    unkeyed_warning_emitted: AtomicBool,
91}
92
93impl AffinityRouter {
94    /// Creates a router that latches every decision.
95    pub fn new() -> Self {
96        Self::default()
97    }
98
99    /// Creates a router that retains assignments only for delegated child-agent work.
100    ///
101    /// Root and harness-maintenance requests always abstain, so a later classifier selects
102    /// them on every turn.
103    pub fn for_subagents() -> Self {
104        Self {
105            subagents_only: true,
106            ..Self::default()
107        }
108    }
109
110    /// Uses the first user-message text as a fallback key when metadata has no session.
111    pub fn with_message_hash_fallback(mut self) -> Self {
112        self.message_hash_fallback = true;
113        self
114    }
115
116    /// Re-decides per user turn, holding the target across the tool calls between.
117    pub fn with_release_on_user_turn(mut self) -> Self {
118        self.release_on_user_turn = true;
119        self
120    }
121
122    /// Restricts latching to `models`; a decision for any other model routes but is not
123    /// retained.
124    pub fn with_latch_only(mut self, models: impl IntoIterator<Item = impl Into<ModelId>>) -> Self {
125        self.latch_only = Some(models.into_iter().map(Into::into).collect());
126        self
127    }
128
129    /// Whether a decision for `model` should be retained.
130    fn should_latch(&self, model: &str) -> bool {
131        self.latch_only
132            .as_ref()
133            .is_none_or(|set| set.contains(model))
134    }
135
136    /// Derives the stable identity this router should retain for `request`.
137    fn affinity_key(&self, request: &Request) -> Option<RoutingIdentity> {
138        let metadata = request.metadata.as_ref();
139        let is_subagent = metadata.is_some_and(|metadata| metadata.is_subagent);
140        let is_subagent_work = metadata.is_some_and(|metadata| metadata.is_subagent_work());
141        // This mode handles only delegated work; root and maintenance requests fall through.
142        if self.subagents_only && !is_subagent_work {
143            return None;
144        }
145
146        // Child requests require their explicit session + agent identity and never fall
147        // back to task text.
148        let key = retention_key(request, self.message_hash_fallback && !is_subagent);
149        // Affinity that never keys anything is silent otherwise: the route reports itself as
150        // configured while every turn is classified afresh. Say so once rather than per turn.
151        if key.is_none() && self.should_warn_unkeyed() {
152            tracing::warn!(
153                target: "libsy",
154                is_subagent,
155                message_hash_fallback = self.message_hash_fallback,
156                "affinity is enabled but this request carries no usable identity, so no \
157                 affinity is applied; root requests need a session id or message-hash \
158                 fallback with usable first-user text, and child requests need both session \
159                 and agent ids"
160            );
161        }
162        key
163    }
164
165    /// Reports whether this call owns the one-time warning for an unkeyable request.
166    fn should_warn_unkeyed(&self) -> bool {
167        !self.unkeyed_warning_emitted.swap(true, Ordering::Relaxed)
168    }
169}
170
171#[async_trait]
172impl<S> Processor<S> for AffinityRouter
173where
174    S: Send + 'static,
175{
176    async fn process(&self, _state: &mut S, event: Event<'_>) -> crate::Result<()> {
177        if let Event::Decision {
178            request,
179            selected_model_id,
180            ..
181        } = event
182            && let Some(key) = self.affinity_key(request)
183        {
184            let mut assignments = self.assignments.lock();
185            let writable = self.release_on_user_turn || !assignments.contains_key(&key);
186            if self.should_latch(selected_model_id) && writable {
187                evict_if_full(&mut assignments);
188                assignments.insert(key, selected_model_id.clone());
189            }
190        }
191        Ok(())
192    }
193}
194
195/// The key a request's retained value is stored under.
196///
197/// Prefers the request's own identity and falls back to hashing the first user
198/// message when the caller allows it, which is the only handle on a client that
199/// sends no session ID.
200pub(crate) fn retention_key(request: &Request, hash_fallback: bool) -> Option<RoutingIdentity> {
201    if let Some(identity) = RoutingIdentity::from_request(request) {
202        return Some(identity);
203    }
204    if !hash_fallback {
205        return None;
206    }
207    first_user_message_hash(request).map(|hash| {
208        tracing::debug!(retention_key = %hash, "retaining by message hash fallback");
209        RoutingIdentity::Session(hash)
210    })
211}
212
213/// Hashes the first user message so later turns retain the initial task's affinity.
214/// For benchmarking purpose with harnesses, task instructions are added as a user prompt to the request so we hash the initial user message.
215/// TODO: Have not considered multi-modal payloads yet. That needs to be handled separately.
216fn first_user_message_hash(request: &Request) -> Option<String> {
217    let message = request
218        .llm_request
219        .messages
220        .iter()
221        .find(|message| message.role == Role::User)?;
222    let mut hasher = DefaultHasher::new();
223    message.text_content("")?.hash(&mut hasher);
224    Some(format!("{:016x}", hasher.finish()))
225}
226
227#[async_trait]
228impl<S> Classifier<S> for AffinityRouter
229where
230    S: Send + 'static,
231{
232    async fn score(
233        &self,
234        _state: &mut S,
235        request: &mut Request,
236        driver: &Driver,
237    ) -> crate::Result<(Classification, Option<switchyard_protocol::Response>)> {
238        let Some(key) = self.affinity_key(request) else {
239            return Ok((Classification::Scores(Vec::new()), None));
240        };
241        // Abstaining rather than clearing keeps this read path free of writes, so the judge
242        // decides the turn whatever order concurrent latches land in.
243        if self.release_on_user_turn && has_new_user_turn(&request.llm_request.messages) {
244            return Ok((Classification::Scores(Vec::new()), None));
245        }
246        // An empty `any` group carries no information about which models are still
247        // available, so it must not be read as "every assignment is now stale".
248        let available = driver.models_for(&Category::Any);
249        let mut assignments = self.assignments.lock();
250        let assigned = assignments.get(&key).cloned();
251        let assigned = match assigned.as_ref() {
252            Some(target) if !available.is_empty() && !available.contains(target) => {
253                assignments.remove(&key);
254                None
255            }
256            assigned => assigned,
257        };
258        if assigned.is_some() {
259            driver.set_evidence(serde_json::json!({"source": "retained"}));
260        }
261        Ok((
262            Classification::Scores(match assigned {
263                Some(target) => vec![Score {
264                    confidence: 1.0,
265                    target: target.clone(),
266                    category: None,
267                }],
268                None => Vec::new(),
269            }),
270            None,
271        ))
272    }
273}
274
275/// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`].
276pub(crate) fn evict_if_full<V>(retained: &mut HashMap<RoutingIdentity, V>) {
277    if retained.len() >= MAX_ASSIGNMENTS
278        && let Some(evicted) = retained.keys().next().cloned()
279    {
280        retained.remove(&evicted);
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    use std::sync::Arc;
289
290    use switchyard_protocol::{LlmRequest, Metadata, ToolResult, text_request};
291
292    use crate::core::algorithm::RuntimeModels;
293
294    /// Boxed, thread-safe error type keeping the test helpers ergonomic.
295    type BoxErr = Box<dyn std::error::Error + Send + Sync>;
296
297    fn fixed_model(target: &str) -> ModelId {
298        ModelId::from(target)
299    }
300
301    fn driver() -> Driver {
302        Driver::new(
303            "test",
304            Arc::new(RuntimeModels::new(
305                [(
306                    Category::Any,
307                    ["model-a", "model-b", "weak", "strong"]
308                        .map(ModelId::from)
309                        .to_vec(),
310                )]
311                .into(),
312            )),
313        )
314        .0
315    }
316
317    fn request(metadata: Metadata) -> Request {
318        Request {
319            llm_request: text_request(Some("auto".to_string()), "hi"),
320            raw_request: None,
321            metadata: Some(metadata),
322        }
323    }
324
325    fn task_request(
326        metadata: Option<Metadata>,
327        first_user: &str,
328        follow_up: Option<&str>,
329    ) -> Request {
330        let mut messages = vec![
331            Message::text(Role::System, "follow repository instructions"),
332            Message::text(Role::User, first_user),
333        ];
334        if let Some(follow_up) = follow_up {
335            messages.push(Message::text(Role::Assistant, "I will inspect the code."));
336            messages.push(Message::text(Role::User, follow_up));
337        }
338        Request {
339            llm_request: LlmRequest {
340                model: Some("auto".to_string()),
341                messages,
342                ..LlmRequest::default()
343            },
344            raw_request: None,
345            metadata,
346        }
347    }
348
349    fn session(session_id: &str, agent_id: &str) -> Metadata {
350        Metadata {
351            session_id: Some(session_id.to_string()),
352            agent_id: Some(agent_id.to_string()),
353            ..Metadata::default()
354        }
355    }
356
357    fn subagent(agent_id: &str, task_id: &str) -> Metadata {
358        Metadata {
359            session_id: Some("session-1".to_string()),
360            agent_id: Some(agent_id.to_string()),
361            task_id: Some(task_id.to_string()),
362            is_subagent: true,
363            is_delegated_work: true,
364            ..Metadata::default()
365        }
366    }
367
368    /// Folds a request and its decision through the router, retaining `model`.
369    async fn retain(
370        router: &AffinityRouter,
371        state: &mut (),
372        request: &mut Request,
373        model: &'static str,
374    ) -> Result<(), BoxErr> {
375        let selected_model_id = ModelId::from(model);
376        router
377            .process(
378                state,
379                Event::Decision {
380                    request,
381                    selected_model_id: &selected_model_id,
382                    category: None,
383                    driver: &driver(),
384                },
385            )
386            .await?;
387        Ok(())
388    }
389
390    /// Scores through the definitive classification variant used by affinity.
391    async fn scores(
392        classifier: &dyn Classifier,
393        state: &mut (),
394        request: &mut Request,
395    ) -> Result<Vec<Score>, BoxErr> {
396        match classifier.score(state, request, &driver()).await?.0 {
397            Classification::Scores(scores) => Ok(scores),
398            Classification::Ambiguous(_) => Err("affinity never returns ambiguous scores".into()),
399        }
400    }
401
402    #[tokio::test]
403    async fn session_retains_first_model_across_requests() -> Result<(), BoxErr> {
404        let router = AffinityRouter::new();
405        let mut state = ();
406
407        let mut first = request(session("session-1", "agent-a"));
408        retain(&router, &mut state, &mut first, "model-a").await?;
409
410        // A different agent in the same session is scored onto the retained model.
411        let mut second = request(session("session-1", "agent-b"));
412        let scores = scores(&router, &mut state, &mut second).await?;
413        assert_eq!(scores.len(), 1);
414        assert_eq!(scores[0].confidence, 1.0);
415        assert_eq!(scores[0].target, "model-a");
416        Ok(())
417    }
418
419    #[tokio::test]
420    async fn subagent_only_retains_children_without_latching_root_traffic() -> Result<(), BoxErr> {
421        let router = AffinityRouter::for_subagents();
422        let mut state = ();
423
424        let mut root = request(session("session-1", "root-agent"));
425        retain(&router, &mut state, &mut root, "model-a").await?;
426        assert!(scores(&router, &mut state, &mut root).await?.is_empty());
427
428        let mut first_child_turn = request(subagent("child-1", "task-1"));
429        retain(&router, &mut state, &mut first_child_turn, "model-b").await?;
430        let mut later_child_turn = request(subagent("child-1", "task-2"));
431        let scores = scores(&router, &mut state, &mut later_child_turn).await?;
432        assert_eq!(
433            scores.first().map(|score| score.target.as_str()),
434            Some("model-b")
435        );
436        Ok(())
437    }
438
439    #[tokio::test]
440    async fn first_decision_wins() -> Result<(), BoxErr> {
441        let router = AffinityRouter::new();
442        let mut state = ();
443
444        let mut req = request(session("session-1", "agent-a"));
445        retain(&router, &mut state, &mut req, "model-a").await?;
446        // A later decision for the same identity must not overwrite the first.
447        retain(&router, &mut state, &mut req, "model-b").await?;
448
449        let scores = scores(&router, &mut state, &mut req).await?;
450        assert_eq!(
451            scores.first().map(|score| score.target.as_str()),
452            Some("model-a")
453        );
454        Ok(())
455    }
456
457    #[tokio::test]
458    async fn subagent_is_keyed_by_agent_not_task() -> Result<(), BoxErr> {
459        let router = AffinityRouter::new();
460        let mut state = ();
461
462        let mut first = request(subagent("child-1", "task-1"));
463        retain(&router, &mut state, &mut first, "model-a").await?;
464
465        // Same child, different task: still scored onto the retained model.
466        let mut second = request(subagent("child-1", "task-2"));
467        let scores = scores(&router, &mut state, &mut second).await?;
468        assert_eq!(
469            scores.first().map(|score| score.target.as_str()),
470            Some("model-a")
471        );
472        Ok(())
473    }
474
475    #[tokio::test]
476    async fn distinct_subagents_are_assigned_independently() -> Result<(), BoxErr> {
477        let router = AffinityRouter::new();
478        let mut state = ();
479
480        // One child in the session is pinned...
481        retain(
482            &router,
483            &mut state,
484            &mut request(subagent("child-1", "task-1")),
485            "model-a",
486        )
487        .await?;
488
489        // ...a sibling child in the same session has no assignment of its own yet.
490        let mut sibling = request(subagent("child-2", "task-1"));
491        assert!(scores(&router, &mut state, &mut sibling).await?.is_empty());
492        Ok(())
493    }
494
495    #[tokio::test]
496    async fn subagent_does_not_inherit_session_assignment() -> Result<(), BoxErr> {
497        let router = AffinityRouter::new();
498        let mut state = ();
499
500        // The session root is pinned, but a sub-agent is keyed separately...
501        retain(
502            &router,
503            &mut state,
504            &mut request(session("session-1", "root-1")),
505            "model-a",
506        )
507        .await?;
508
509        // ...so the sub-agent abstains until it is assigned in its own right.
510        let mut child = request(subagent("child-1", "task-1"));
511        assert!(scores(&router, &mut state, &mut child).await?.is_empty());
512        Ok(())
513    }
514
515    #[tokio::test]
516    async fn classifier_abstains_without_a_session() -> Result<(), BoxErr> {
517        let router = AffinityRouter::new();
518        let mut state = ();
519
520        // No session id at all: nothing to key on.
521        let mut req = request(Metadata::default());
522        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
523        Ok(())
524    }
525
526    #[tokio::test]
527    async fn message_hash_fallback_uses_the_first_user_message() -> Result<(), BoxErr> {
528        let router = AffinityRouter::new().with_message_hash_fallback();
529        let mut state = ();
530
531        let mut first = task_request(
532            None,
533            "Add a unit test for this function.",
534            Some("Now run the test suite."),
535        );
536        retain(&router, &mut state, &mut first, "weak").await?;
537
538        let mut follow_up = task_request(
539            None,
540            "Add a unit test for this function.",
541            Some("Now file a pull request."),
542        );
543        assert_eq!(
544            scores(&router, &mut state, &mut follow_up)
545                .await?
546                .first()
547                .map(|score| score.target.as_str()),
548            Some("weak")
549        );
550
551        let mut other_task = task_request(
552            None,
553            "Reimplement this binary from two input/output pairs.",
554            Some("Now run the test suite."),
555        );
556        assert!(
557            scores(&router, &mut state, &mut other_task)
558                .await?
559                .is_empty()
560        );
561        Ok(())
562    }
563
564    #[tokio::test]
565    async fn subagents_only_root_traffic_does_not_warn() -> Result<(), BoxErr> {
566        // Abstaining on root traffic is this mode's contract, so it must not warn.
567        let router = AffinityRouter::for_subagents();
568        let mut state = ();
569
570        let mut root = request(session("session-1", "agent-1"));
571        assert!(scores(&router, &mut state, &mut root).await?.is_empty());
572        assert!(
573            router.should_warn_unkeyed(),
574            "an intentional abstention should leave the warning unconsumed"
575        );
576        Ok(())
577    }
578
579    #[test]
580    fn user_message_hash_ignores_non_text_provider_payloads() {
581        let request = |user_message| Request {
582            llm_request: LlmRequest {
583                messages: vec![user_message],
584                ..LlmRequest::default()
585            },
586            raw_request: None,
587            metadata: None,
588        };
589        let text_only = request(Message::text(Role::User, "Implement the parser."));
590        let text_with_reasoning = request(Message {
591            role: Role::User,
592            content: vec![
593                ContentBlock::Text {
594                    text: "Implement the parser.".to_string(),
595                },
596                ContentBlock::Reasoning {
597                    text: "Internal provider reasoning.".to_string(),
598                    signature: Some("provider-signature".to_string()),
599                    details: Vec::new(),
600                },
601            ],
602        });
603
604        assert_eq!(
605            first_user_message_hash(&text_only),
606            first_user_message_hash(&text_with_reasoning)
607        );
608    }
609
610    #[tokio::test]
611    async fn metadata_session_takes_precedence_over_message_hash() -> Result<(), BoxErr> {
612        let router = AffinityRouter::new().with_message_hash_fallback();
613        let mut state = ();
614
615        let mut first = task_request(
616            Some(session("session-1", "agent-a")),
617            "Implement the parser.",
618            None,
619        );
620        retain(&router, &mut state, &mut first, "strong").await?;
621
622        let mut other_session = task_request(
623            Some(session("session-2", "agent-a")),
624            "Implement the parser.",
625            None,
626        );
627        assert!(
628            scores(&router, &mut state, &mut other_session)
629                .await?
630                .is_empty()
631        );
632        Ok(())
633    }
634
635    #[tokio::test]
636    async fn subagent_without_a_session_abstains_and_warns() -> Result<(), BoxErr> {
637        for session_id in [None, Some(String::new())] {
638            let router = AffinityRouter::new().with_message_hash_fallback();
639            let mut state = ();
640            let mut subagent = task_request(
641                Some(Metadata {
642                    session_id,
643                    agent_id: Some("agent-1".to_string()),
644                    is_subagent: true,
645                    ..Metadata::default()
646                }),
647                "Implement the parser.",
648                None,
649            );
650
651            retain(&router, &mut state, &mut subagent, "model-a").await?;
652            assert!(scores(&router, &mut state, &mut subagent).await?.is_empty());
653            assert!(
654                !router.should_warn_unkeyed(),
655                "an unidentifiable subagent should consume the warning"
656            );
657        }
658        Ok(())
659    }
660
661    #[tokio::test]
662    async fn one_router_serves_both_roles() -> Result<(), BoxErr> {
663        // The same instance is registered under both SDK roles; a decision folded in via
664        // the processor handle is read back via the classifier handle.
665        let router = Arc::new(AffinityRouter::new());
666        let processor: Arc<dyn Processor> = router.clone();
667        let classifier: Arc<dyn Classifier> = router;
668        let mut state = ();
669
670        let mut first = request(session("session-1", "agent-a"));
671        processor
672            .process(
673                &mut state,
674                Event::Decision {
675                    request: &mut first,
676                    selected_model_id: &fixed_model("model-a"),
677                    category: None,
678                    driver: &driver(),
679                },
680            )
681            .await?;
682
683        let mut second = request(session("session-1", "agent-b"));
684        let scores = scores(classifier.as_ref(), &mut state, &mut second).await?;
685        assert_eq!(
686            scores.first().map(|score| score.target.as_str()),
687            Some("model-a")
688        );
689        Ok(())
690    }
691
692    #[tokio::test]
693    async fn decision_without_an_affinity_identity_is_ignored() -> Result<(), BoxErr> {
694        let router = AffinityRouter::new();
695        let mut state = ();
696        let mut unkeyed = request(Metadata::default());
697
698        router
699            .process(
700                &mut state,
701                Event::Decision {
702                    request: &mut unkeyed,
703                    selected_model_id: &fixed_model("model-a"),
704                    category: None,
705                    driver: &driver(),
706                },
707            )
708            .await?;
709
710        let mut req = request(session("session-1", "agent-a"));
711        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
712        Ok(())
713    }
714
715    #[tokio::test]
716    async fn decisions_retain_their_originating_request_identity() -> Result<(), BoxErr> {
717        let router = AffinityRouter::new();
718        let mut state = ();
719        let mut first = request(session("session-1", "agent-a"));
720        let mut second = request(session("session-2", "agent-b"));
721
722        // Replay decisions in the opposite order. Each decision carries its request,
723        // so the assignments cannot cross.
724        router
725            .process(
726                &mut state,
727                Event::Decision {
728                    request: &mut second,
729                    selected_model_id: &fixed_model("model-b"),
730                    category: None,
731                    driver: &driver(),
732                },
733            )
734            .await?;
735        router
736            .process(
737                &mut state,
738                Event::Decision {
739                    request: &mut first,
740                    selected_model_id: &fixed_model("model-a"),
741                    category: None,
742                    driver: &driver(),
743                },
744            )
745            .await?;
746
747        let first_scores = scores(&router, &mut state, &mut first).await?;
748        let second_scores = scores(&router, &mut state, &mut second).await?;
749        assert_eq!(
750            first_scores.first().map(|score| score.target.as_str()),
751            Some("model-a")
752        );
753        assert_eq!(
754            second_scores.first().map(|score| score.target.as_str()),
755            Some("model-b")
756        );
757        Ok(())
758    }
759
760    #[tokio::test]
761    async fn distinct_sessions_are_assigned_independently() -> Result<(), BoxErr> {
762        let router = AffinityRouter::new();
763        let mut state = ();
764
765        retain(
766            &router,
767            &mut state,
768            &mut request(session("session-1", "agent-a")),
769            "model-a",
770        )
771        .await?;
772        retain(
773            &router,
774            &mut state,
775            &mut request(session("session-2", "agent-a")),
776            "model-b",
777        )
778        .await?;
779
780        let first = scores(
781            &router,
782            &mut state,
783            &mut request(session("session-1", "other")),
784        )
785        .await?;
786        let second = scores(
787            &router,
788            &mut state,
789            &mut request(session("session-2", "other")),
790        )
791        .await?;
792        assert_eq!(
793            first.first().map(|score| score.target.as_str()),
794            Some("model-a")
795        );
796        assert_eq!(
797            second.first().map(|score| score.target.as_str()),
798            Some("model-b")
799        );
800        Ok(())
801    }
802
803    #[tokio::test]
804    async fn subagent_without_an_agent_id_abstains_and_warns() -> Result<(), BoxErr> {
805        let router = AffinityRouter::new();
806        let mut state = ();
807
808        // The sub-agent flag is set but no agent id is present, so no key can be formed;
809        // the request is neither retained nor scored.
810        let metadata = Metadata {
811            session_id: Some("session-1".to_string()),
812            is_subagent: true,
813            ..Metadata::default()
814        };
815        let mut req = request(metadata);
816        retain(&router, &mut state, &mut req, "model-a").await?;
817        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
818        assert!(
819            !router.should_warn_unkeyed(),
820            "an unidentifiable subagent should consume the warning"
821        );
822        Ok(())
823    }
824
825    #[tokio::test]
826    async fn assignments_are_bounded_by_the_cap() -> Result<(), BoxErr> {
827        let router = AffinityRouter::new();
828        let mut state = ();
829
830        // One distinct session past the cap forces exactly one eviction.
831        for index in 0..=MAX_ASSIGNMENTS {
832            let session_id = format!("session-{index}");
833            retain(
834                &router,
835                &mut state,
836                &mut request(session(&session_id, "agent-a")),
837                "model-a",
838            )
839            .await?;
840        }
841
842        let len = router.assignments.lock().len();
843        assert_eq!(len, MAX_ASSIGNMENTS);
844        Ok(())
845    }
846
847    #[tokio::test]
848    async fn release_on_user_turn_drops_the_assignment_only_when_the_user_speaks()
849    -> Result<(), BoxErr> {
850        let router = AffinityRouter::new().with_release_on_user_turn();
851        let mut state = ();
852        let mut opening = task_request(Some(session("session-1", "agent-a")), "add caching", None);
853        retain(&router, &mut state, &mut opening, "weak").await?;
854
855        // A tool continuation holds the assignment, so no judge call.
856        let mut continued = opening.clone();
857        continued.llm_request.messages.push(Message {
858            role: Role::User,
859            content: vec![ContentBlock::ToolResult(ToolResult {
860                tool_call_id: "call-1".to_string(),
861                content: Vec::new(),
862                is_error: None,
863            })],
864        });
865        assert_eq!(
866            scores(&router, &mut state, &mut continued)
867                .await?
868                .first()
869                .map(|s| s.target.as_str()),
870            Some("weak")
871        );
872
873        // A new user message releases it, so the turn abstains and the judge runs.
874        let mut spoke = task_request(
875            Some(session("session-1", "agent-a")),
876            "add caching",
877            Some("no, shared across processes"),
878        );
879        assert!(scores(&router, &mut state, &mut spoke).await?.is_empty());
880        Ok(())
881    }
882
883    #[tokio::test]
884    async fn latch_only_retains_matching_models() -> Result<(), BoxErr> {
885        let router = AffinityRouter::new().with_latch_only(["strong"]);
886        let mut state = ();
887        let mut req = request(session("session-1", "agent-a"));
888
889        // A "weak" decision is not retained — a later turn is not latched.
890        retain(&router, &mut state, &mut req, "weak").await?;
891        assert!(scores(&router, &mut state, &mut req).await?.is_empty());
892
893        // A "strong" decision is retained — later turns latch onto it.
894        retain(&router, &mut state, &mut req, "strong").await?;
895        assert_eq!(
896            scores(&router, &mut state, &mut req)
897                .await?
898                .first()
899                .map(|s| s.target.as_str()),
900            Some("strong")
901        );
902        Ok(())
903    }
904
905    #[tokio::test]
906    async fn an_unprovisioned_any_group_does_not_evict_assignments() -> Result<(), BoxErr> {
907        let router = AffinityRouter::new();
908        let mut state = ();
909        let mut req = request(session("session-1", "agent-a"));
910        retain(&router, &mut state, &mut req, "model-a").await?;
911
912        let empty = Driver::new("test", Arc::new(RuntimeModels::default())).0;
913        let (classification, _) = router.score(&mut state, &mut req, &empty).await?;
914        let Classification::Scores(retained) = classification else {
915            return Err("affinity never returns ambiguous scores".into());
916        };
917        assert_eq!(retained.first().map(|s| s.target.as_str()), Some("model-a"));
918
919        // The assignment survived, so a later turn with the group present still latches.
920        assert_eq!(
921            scores(&router, &mut state, &mut req)
922                .await?
923                .first()
924                .map(|s| s.target.as_str()),
925            Some("model-a")
926        );
927        Ok(())
928    }
929}
switchyard_libsy/algorithms/util/affinity.rs

switchyard_libsy/algorithms/util/
affinity.rs
