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

switchyard_libsy/algorithms/util/
affinity.rs
