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 chosen for a request's stable identity and
7//! forces that model 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 explicitly identified
18//! child agents, leaving root traffic to later classifiers on every turn.
19
20use std::collections::{HashMap, HashSet, hash_map::DefaultHasher};
21use std::hash::{Hash, Hasher};
22
23use async_trait::async_trait;
24use parking_lot::Mutex;
25use switchyard_protocol::{Request, Role};
26
27use crate::core::algorithm::{Driver, RoutingIdentity};
28use crate::core::classifier::{Classification, Classifier, Score};
29use crate::core::processor::{Event, Processor};
30
31/// Upper bound on retained assignments, keeping the process-local map from growing
32/// without limit; the oldest entry is evicted once the bound is reached.
33const MAX_ASSIGNMENTS: usize = 4096;
34
35/// Retains a model per request identity and forces it on later matching requests.
36///
37/// Register the same instance as both a processor and a classifier; the two roles share
38/// the retained assignments through this instance's interior storage. Decision events carry
39/// their originating request, so concurrent turns cannot bind one request's identity to
40/// another request's selected model.
41///
42/// [`with_latch_only`](Self::with_latch_only) narrows *which* models are retained — a
43/// decision for any other model routes normally but is not latched (the escalation latch:
44/// retain only the strong tier, never the weak one).
45#[derive(Default)]
46pub struct AffinityRouter {
47 /// When set, only these models are retained; a decision for any other model is not latched.
48 latch_only: Option<HashSet<String>>,
49 /// Whether root-session requests should abstain instead of being retained.
50 subagents_only: bool,
51 /// In absence of headers, use the message hash based fallback key to do task based routing
52 message_hash_fallback: bool,
53 /// Retained assignments, shared across this router's processor and classifier roles.
54 ///
55 /// Held on the instance so the two roles share one process-local map through a
56 /// single registered [`Arc`](std::sync::Arc); bounded by [`MAX_ASSIGNMENTS`].
57 assignments: Mutex<HashMap<RoutingIdentity, String>>,
58}
59
60impl AffinityRouter {
61 /// Creates a router that latches every decision.
62 pub fn new() -> Self {
63 Self::default()
64 }
65
66 /// Creates a router that retains assignments only for explicitly identified child agents.
67 ///
68 /// Root-agent requests always abstain, so a later classifier selects them on every turn.
69 pub fn for_subagents() -> Self {
70 Self {
71 subagents_only: true,
72 ..Self::default()
73 }
74 }
75
76 /// Uses the first user-message text as a fallback key when metadata has no session.
77 pub fn with_message_hash_fallback(mut self) -> Self {
78 self.message_hash_fallback = true;
79 self
80 }
81
82 /// Restricts latching to `models`; a decision for any other model routes but is not
83 /// retained.
84 pub fn with_latch_only(mut self, models: impl IntoIterator<Item = impl Into<String>>) -> Self {
85 self.latch_only = Some(models.into_iter().map(Into::into).collect());
86 self
87 }
88
89 /// Whether a decision for `model` should be retained.
90 fn should_latch(&self, model: &str) -> bool {
91 self.latch_only
92 .as_ref()
93 .is_none_or(|set| set.contains(model))
94 }
95
96 /// Derives the stable identity this router should retain for `request`.
97 fn affinity_key(&self, request: &Request) -> Option<RoutingIdentity> {
98 if let Some(identity) = RoutingIdentity::from_request(request) {
99 return match identity {
100 RoutingIdentity::Session(_) if self.subagents_only => None,
101 identity => Some(identity),
102 };
103 }
104
105 // If headers are not present and we are not a subagent, use the message hash based fallback key to do task based routing
106 let is_subagent = request
107 .metadata
108 .as_ref()
109 .is_some_and(|metadata| metadata.is_subagent);
110 (!self.subagents_only && !is_subagent && self.message_hash_fallback)
111 .then(|| {
112 first_user_message_hash(request).map(|hash| {
113 tracing::debug!(affinity_key = %hash, "affinity using message hash fallback");
114 RoutingIdentity::Session(hash)
115 })
116 })
117 .flatten()
118 }
119}
120
121#[async_trait]
122impl<S> Processor<S> for AffinityRouter
123where
124 S: Send + 'static,
125{
126 async fn process(&self, _state: &mut S, event: Event<'_>) -> crate::Result<()> {
127 if let Event::Decision { request, decision } = event
128 && let Some(key) = self.affinity_key(request)
129 {
130 let model = decision.selected_model();
131 let mut assignments = self.assignments.lock();
132 if self.should_latch(model) && !assignments.contains_key(&key) {
133 evict_if_full(&mut assignments);
134 assignments.insert(key, model.to_string());
135 }
136 }
137 Ok(())
138 }
139}
140
141/// Hashes the first user message so later turns retain the initial task's affinity.
142/// For benchmarking purpose with harnesses, task instructions are added as a user prompt to the request so we hash the initial user message.
143/// TODO: Have not considered multi-modal payloads yet. That needs to be handled separately.
144fn first_user_message_hash(request: &Request) -> Option<String> {
145 let message = request
146 .llm_request
147 .messages
148 .iter()
149 .find(|message| message.role == Role::User)?;
150 let mut hasher = DefaultHasher::new();
151 message.text_content("")?.hash(&mut hasher);
152 Some(format!("{:016x}", hasher.finish()))
153}
154
155#[async_trait]
156impl<S> Classifier<S> for AffinityRouter
157where
158 S: Send + 'static,
159{
160 fn target_unavailable(&self, request: &Request, target: &str) {
161 let Some(key) = self.affinity_key(request) else {
162 return;
163 };
164 let mut assignments = self.assignments.lock();
165 if assignments
166 .get(&key)
167 .is_some_and(|assigned| assigned == target)
168 {
169 assignments.remove(&key);
170 }
171 }
172
173 async fn score(
174 &self,
175 _state: &mut S,
176 request: &mut Request,
177 _driver: Option<&Driver>,
178 ) -> crate::Result<(Classification, Option<switchyard_protocol::Response>)> {
179 let Some(key) = self.affinity_key(request) else {
180 return Ok((Classification::Scores(Vec::new()), None));
181 };
182 let assigned = self.assignments.lock().get(&key).cloned();
183 Ok((
184 Classification::Scores(match assigned {
185 Some(target) => vec![Score {
186 confidence: 1.0,
187 target,
188 }],
189 None => Vec::new(),
190 }),
191 None,
192 ))
193 }
194}
195
196/// Evicts one arbitrary assignment when the map has reached [`MAX_ASSIGNMENTS`].
197fn evict_if_full(assignments: &mut HashMap<RoutingIdentity, String>) {
198 if assignments.len() >= MAX_ASSIGNMENTS
199 && let Some(evicted) = assignments.keys().next().cloned()
200 {
201 assignments.remove(&evicted);
202 }
203}
204
205#[cfg(test)]
206#[path = "affinity_tests.rs"]
207mod tests;