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

switchyard_libsy/algorithms/util/
affinity.rs
