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

switchyard_libsy/algorithms/util/
affinity.rs
