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

switchyard_libsy/algorithms/util/
affinity.rs
