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

switchyard_libsy/algorithms/util/
affinity.rs
