mesh_client/network/
affinity.rs

1//! Prefix affinity and sticky routing helpers for inference target selection.
2
3use crate::inference::election;
4use iroh::EndpointId;
5use serde::Serialize;
6use serde_json::Value;
7use std::collections::{HashMap, VecDeque};
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11const AFFINITY_TTL: Duration = Duration::from_secs(20 * 60);
12const AFFINITY_MAX_ENTRIES: usize = 4096;
13
14#[derive(Clone, Debug, Default, Serialize)]
15pub struct AffinityStatsSnapshot {
16    pub prefix_enabled: bool,
17    pub sticky_enabled: bool,
18    pub prefix_entries: usize,
19    pub prefix_lookups: u64,
20    pub prefix_hits: u64,
21    pub prefix_misses: u64,
22    pub prefix_stale: u64,
23    pub prefix_routes: u64,
24    pub sticky_routes: u64,
25    pub session_routes: u64,
26    pub learned: u64,
27    pub evicted: u64,
28}
29
30fn prefix_only_enabled() -> bool {
31    std::env::var("MESH_LLM_PREFIX_ONLY")
32        .map(|value| value == "1" || value.eq_ignore_ascii_case("true"))
33        .unwrap_or(false)
34}
35
36#[derive(Clone, Copy, Debug)]
37struct AffinityConfig {
38    prefix_enabled: bool,
39    sticky_enabled: bool,
40}
41
42impl AffinityConfig {
43    fn from_env() -> Self {
44        Self {
45            prefix_enabled: std::env::var_os("MESH_LLM_DISABLE_PREFIX_AFFINITY").is_none(),
46            sticky_enabled: std::env::var_os("MESH_LLM_DISABLE_STICKY_ROUTING").is_none(),
47        }
48    }
49}
50
51#[derive(Clone, Debug, Hash, PartialEq, Eq)]
52struct AffinityKey {
53    model: String,
54    prefix_hash: u64,
55}
56
57#[derive(Clone, Debug)]
58struct AffinityEntry {
59    target: election::InferenceTarget,
60    last_used: Instant,
61}
62
63#[derive(Default)]
64struct AffinityState {
65    entries: HashMap<AffinityKey, AffinityEntry>,
66    lru: VecDeque<AffinityKey>,
67    stats: AffinityStatsSnapshot,
68}
69
70#[derive(Clone)]
71pub struct AffinityRouter {
72    inner: Arc<Mutex<AffinityState>>,
73    config: Arc<AffinityConfig>,
74}
75
76impl AffinityRouter {
77    pub fn new() -> Self {
78        Self {
79            inner: Arc::new(Mutex::new(AffinityState::default())),
80            config: Arc::new(AffinityConfig::from_env()),
81        }
82    }
83
84    #[cfg(test)]
85    fn with_config(prefix_enabled: bool, sticky_enabled: bool) -> Self {
86        Self {
87            inner: Arc::new(Mutex::new(AffinityState::default())),
88            config: Arc::new(AffinityConfig {
89                prefix_enabled,
90                sticky_enabled,
91            }),
92        }
93    }
94
95    pub fn stats_snapshot(&self) -> AffinityStatsSnapshot {
96        let mut state = self.inner.lock().unwrap();
97        state.prune_expired();
98        let mut stats = state.stats.clone();
99        stats.prefix_entries = state.entries.len();
100        stats.prefix_enabled = self.config.prefix_enabled;
101        stats.sticky_enabled = self.config.sticky_enabled;
102        stats
103    }
104
105    pub fn sticky_enabled(&self) -> bool {
106        self.config.sticky_enabled
107    }
108
109    pub fn record_sticky_route(&self) {
110        let mut state = self.inner.lock().unwrap();
111        state.stats.sticky_routes += 1;
112    }
113
114    pub fn record_session_route(&self) {
115        let mut state = self.inner.lock().unwrap();
116        state.stats.session_routes += 1;
117    }
118
119    pub fn lookup_target(
120        &self,
121        model: &str,
122        prefix_hash: u64,
123        candidates: &[election::InferenceTarget],
124    ) -> Option<election::InferenceTarget> {
125        if !self.config.prefix_enabled {
126            return None;
127        }
128        let key = AffinityKey {
129            model: model.to_string(),
130            prefix_hash,
131        };
132        let mut state = self.inner.lock().unwrap();
133        state.prune_expired();
134        state.stats.prefix_lookups += 1;
135        let entry = match state.entries.get(&key).cloned() {
136            Some(entry) => entry,
137            None => {
138                state.stats.prefix_misses += 1;
139                return None;
140            }
141        };
142        if !candidates.contains(&entry.target) {
143            state.remove_key(&key);
144            state.stats.prefix_stale += 1;
145            state.stats.prefix_misses += 1;
146            return None;
147        }
148        state.touch_key(&key);
149        if let Some(existing) = state.entries.get_mut(&key) {
150            existing.last_used = Instant::now();
151        }
152        state.stats.prefix_hits += 1;
153        state.stats.prefix_routes += 1;
154        Some(entry.target)
155    }
156
157    pub fn learn_target(&self, model: &str, prefix_hash: u64, target: &election::InferenceTarget) {
158        if !self.config.prefix_enabled || matches!(target, election::InferenceTarget::None) {
159            return;
160        }
161
162        let key = AffinityKey {
163            model: model.to_string(),
164            prefix_hash,
165        };
166        let now = Instant::now();
167        let mut state = self.inner.lock().unwrap();
168        state.prune_expired();
169        state.entries.insert(
170            key.clone(),
171            AffinityEntry {
172                target: target.clone(),
173                last_used: now,
174            },
175        );
176        state.touch_key(&key);
177        state.stats.learned += 1;
178        while state.entries.len() > AFFINITY_MAX_ENTRIES {
179            if let Some(oldest) = state.lru.pop_front() {
180                if state.entries.remove(&oldest).is_some() {
181                    state.stats.evicted += 1;
182                }
183            } else {
184                break;
185            }
186        }
187    }
188
189    pub fn forget_target(&self, model: &str, prefix_hash: u64, target: &election::InferenceTarget) {
190        if !self.config.prefix_enabled {
191            return;
192        }
193        let key = AffinityKey {
194            model: model.to_string(),
195            prefix_hash,
196        };
197        let mut state = self.inner.lock().unwrap();
198        if state
199            .entries
200            .get(&key)
201            .map(|entry| &entry.target == target)
202            .unwrap_or(false)
203        {
204            state.remove_key(&key);
205            state.stats.prefix_stale += 1;
206        }
207    }
208}
209
210impl Default for AffinityRouter {
211    fn default() -> Self {
212        Self::new()
213    }
214}
215
216impl AffinityState {
217    fn prune_expired(&mut self) {
218        let now = Instant::now();
219
220        while let Some(front_key) = self.lru.front().cloned() {
221            match self.entries.get(&front_key) {
222                Some(entry) => {
223                    if now.duration_since(entry.last_used) > AFFINITY_TTL {
224                        self.lru.pop_front();
225                        if self.entries.remove(&front_key).is_some() {
226                            self.stats.prefix_stale += 1;
227                        }
228                    } else {
229                        break;
230                    }
231                }
232                None => {
233                    self.lru.pop_front();
234                }
235            }
236        }
237    }
238
239    fn touch_key(&mut self, key: &AffinityKey) {
240        if let Some(pos) = self.lru.iter().position(|existing| existing == key) {
241            self.lru.remove(pos);
242        }
243        self.lru.push_back(key.clone());
244    }
245
246    fn remove_key(&mut self, key: &AffinityKey) {
247        self.entries.remove(key);
248        if let Some(pos) = self.lru.iter().position(|existing| existing == key) {
249            self.lru.remove(pos);
250        }
251    }
252}
253
254#[derive(Clone, Debug, Default)]
255struct RoutingKeys {
256    session_hash: Option<u64>,
257    prefix_hash: Option<u64>,
258    sticky_hash: Option<u64>,
259}
260
261pub struct TargetSelection {
262    pub target: election::InferenceTarget,
263    pub learn_prefix_hash: Option<u64>,
264    pub cached_target: Option<election::InferenceTarget>,
265}
266
267pub struct PreparedTargets {
268    pub ordered: Vec<election::InferenceTarget>,
269    pub learn_prefix_hash: Option<u64>,
270    pub cached_target: Option<election::InferenceTarget>,
271}
272
273pub(crate) fn extract_session_hint_from_body(body: &Value) -> Option<String> {
274    top_level_string(body, "user").or_else(|| top_level_string(body, "session_id"))
275}
276
277fn top_level_string(body: &Value, key: &str) -> Option<String> {
278    body.get(key)
279        .and_then(|value| value.as_str())
280        .map(str::to_string)
281}
282
283fn message_text(msg: &Value) -> Option<String> {
284    if let Some(s) = msg.get("content").and_then(|c| c.as_str()) {
285        return Some(s.to_string());
286    }
287    if let Some(blocks) = msg.get("content").and_then(|c| c.as_array()) {
288        let mut out = String::new();
289        for block in blocks {
290            if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
291                out.push_str(text);
292                out.push('\n');
293            }
294        }
295        if !out.is_empty() {
296            return Some(out);
297        }
298    }
299    None
300}
301
302fn hash_bytes(bytes: &[u8]) -> u64 {
303    bytes.iter().fold(0xcbf29ce484222325u64, |acc, &b| {
304        (acc ^ b as u64).wrapping_mul(0x100000001b3)
305    })
306}
307
308fn hash_combine(a: u64, b: u64) -> u64 {
309    a.wrapping_mul(31).wrapping_add(b)
310}
311
312fn hash_tagged_text(mut acc: u64, tag: &str, text: &str) -> u64 {
313    acc = hash_combine(acc, hash_bytes(tag.as_bytes()));
314    hash_combine(acc, hash_bytes(text.as_bytes()))
315}
316
317fn hash_json_value(mut acc: u64, value: &Value) -> u64 {
318    match value {
319        Value::Null => hash_combine(acc, hash_bytes(b"null")),
320        Value::Bool(boolean) => {
321            acc = hash_combine(acc, hash_bytes(b"bool"));
322            hash_combine(acc, hash_bytes(boolean.to_string().as_bytes()))
323        }
324        Value::Number(number) => {
325            acc = hash_combine(acc, hash_bytes(b"number"));
326            hash_combine(acc, hash_bytes(number.to_string().as_bytes()))
327        }
328        Value::String(text) => {
329            acc = hash_combine(acc, hash_bytes(b"string"));
330            hash_combine(acc, hash_bytes(text.as_bytes()))
331        }
332        Value::Array(items) => {
333            acc = hash_combine(acc, hash_bytes(b"array"));
334            acc = hash_combine(acc, items.len() as u64);
335            for item in items {
336                acc = hash_json_value(acc, item);
337            }
338            acc
339        }
340        Value::Object(map) => {
341            acc = hash_combine(acc, hash_bytes(b"object"));
342            let mut keys: Vec<_> = map.keys().collect();
343            keys.sort_unstable();
344            for key in keys {
345                acc = hash_combine(acc, hash_bytes(key.as_bytes()));
346                acc = hash_json_value(acc, &map[key]);
347            }
348            acc
349        }
350    }
351}
352
353fn hash_tagged_json(mut acc: u64, tag: &str, value: &Value) -> u64 {
354    acc = hash_combine(acc, hash_bytes(tag.as_bytes()));
355    hash_json_value(acc, value)
356}
357
358fn scaffold_prefix_hash_from_body(body: &Value) -> Option<u64> {
359    let mut hash = 0u64;
360    let mut found = false;
361
362    for key in [
363        "tools",
364        "functions",
365        "response_format",
366        "tool_choice",
367        "parallel_tool_calls",
368    ] {
369        if let Some(value) = body.get(key) {
370            hash = hash_tagged_json(hash, key, value);
371            found = true;
372        }
373    }
374
375    if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
376        for msg in messages {
377            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
378            match role {
379                "system" | "developer" => {
380                    if let Some(text) = message_text(msg) {
381                        hash = hash_tagged_text(hash, role, &text);
382                        found = true;
383                    }
384                }
385                "user" => break,
386                _ => {}
387            }
388        }
389    }
390
391    found.then_some(hash)
392}
393
394fn first_user_hash_from_body(body: &Value) -> Option<u64> {
395    if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
396        for msg in messages {
397            if msg.get("role").and_then(|r| r.as_str()) == Some("user") {
398                return message_text(msg).map(|text| hash_tagged_text(0, "user", &text));
399            }
400        }
401    }
402    body.get("prompt")
403        .and_then(|value| value.as_str())
404        .map(|prompt| hash_tagged_text(0, "prompt", prompt))
405}
406
407fn routing_keys(parsed_body: Option<&Value>) -> RoutingKeys {
408    let Some(body) = parsed_body else {
409        return RoutingKeys::default();
410    };
411
412    let session_hash = extract_session_hint_from_body(body).map(|hint| hash_bytes(hint.as_bytes()));
413    let prefix_hash = scaffold_prefix_hash_from_body(body);
414    let sticky_hash = session_hash.or_else(|| {
415        let mut hash = 0u64;
416        let mut found = false;
417        if let Some(prefix_hash) = prefix_hash {
418            hash = hash_combine(hash, prefix_hash);
419            found = true;
420        }
421        if let Some(user_hash) = first_user_hash_from_body(body) {
422            hash = hash_combine(hash, user_hash);
423            found = true;
424        }
425        found.then_some(hash)
426    });
427
428    RoutingKeys {
429        session_hash,
430        prefix_hash,
431        sticky_hash,
432    }
433}
434
435fn rotate_targets_by_hash(targets: &mut [election::InferenceTarget], key: u64) {
436    if !targets.is_empty() {
437        let idx = key as usize % targets.len();
438        targets.rotate_left(idx);
439    }
440}
441
442fn move_target_first(
443    targets: &mut [election::InferenceTarget],
444    target: &election::InferenceTarget,
445) -> bool {
446    if let Some(pos) = targets.iter().position(|candidate| candidate == target) {
447        targets[..=pos].rotate_right(1);
448        true
449    } else {
450        false
451    }
452}
453
454/// Select an inference target for a model request from a caller-supplied candidate
455/// list instead of pulling it from `targets`. This avoids cloning the entire
456/// `ModelTargets` when the caller has already reordered the candidates (e.g. by
457/// context capacity).
458pub fn select_model_target_from_candidates(
459    targets: &election::ModelTargets,
460    candidates: &[election::InferenceTarget],
461    model: &str,
462    parsed_body: Option<&Value>,
463    affinity: &AffinityRouter,
464) -> TargetSelection {
465    let routing = routing_keys(parsed_body);
466
467    if let Some(session_hash) = routing.session_hash.filter(|_| affinity.sticky_enabled()) {
468        affinity.record_session_route();
469        return TargetSelection {
470            target: election::ModelTargets::pick_sticky_from(candidates, session_hash),
471            learn_prefix_hash: None,
472            cached_target: None,
473        };
474    }
475
476    if let Some(prefix_hash) = routing.prefix_hash {
477        if let Some(target) = affinity.lookup_target(model, prefix_hash, candidates) {
478            return TargetSelection {
479                target: target.clone(),
480                learn_prefix_hash: Some(prefix_hash),
481                cached_target: Some(target),
482            };
483        }
484
485        if prefix_only_enabled() {
486            return TargetSelection {
487                target: election::ModelTargets::pick_sticky_from(candidates, prefix_hash),
488                learn_prefix_hash: Some(prefix_hash),
489                cached_target: None,
490            };
491        }
492
493        if let Some(sticky_hash) = routing.sticky_hash.filter(|_| affinity.sticky_enabled()) {
494            affinity.record_sticky_route();
495            return TargetSelection {
496                target: election::ModelTargets::pick_sticky_from(candidates, sticky_hash),
497                learn_prefix_hash: Some(prefix_hash),
498                cached_target: None,
499            };
500        }
501
502        return TargetSelection {
503            target: targets.pick_from(candidates),
504            learn_prefix_hash: Some(prefix_hash),
505            cached_target: None,
506        };
507    }
508
509    if let Some(sticky_hash) = routing.sticky_hash.filter(|_| affinity.sticky_enabled()) {
510        affinity.record_sticky_route();
511        return TargetSelection {
512            target: election::ModelTargets::pick_sticky_from(candidates, sticky_hash),
513            learn_prefix_hash: None,
514            cached_target: None,
515        };
516    }
517
518    TargetSelection {
519        target: targets.pick_from(candidates),
520        learn_prefix_hash: None,
521        cached_target: None,
522    }
523}
524
525pub fn prepare_remote_targets_for_request(
526    model: &str,
527    hosts: &[EndpointId],
528    parsed_body: Option<&Value>,
529    affinity: &AffinityRouter,
530) -> PreparedTargets {
531    let routing = routing_keys(parsed_body);
532    let mut ordered: Vec<election::InferenceTarget> = hosts
533        .iter()
534        .copied()
535        .map(election::InferenceTarget::Remote)
536        .collect();
537    let mut cached_target = None;
538    let mut learn_prefix_hash = None;
539
540    if let Some(session_hash) = routing.session_hash.filter(|_| affinity.sticky_enabled()) {
541        affinity.record_session_route();
542        rotate_targets_by_hash(&mut ordered, session_hash);
543        return PreparedTargets {
544            ordered,
545            learn_prefix_hash: None,
546            cached_target: None,
547        };
548    }
549
550    if let Some(prefix_hash) = routing.prefix_hash {
551        learn_prefix_hash = Some(prefix_hash);
552        if let Some(target) = affinity.lookup_target(model, prefix_hash, &ordered) {
553            move_target_first(&mut ordered, &target);
554            cached_target = Some(target);
555        } else if prefix_only_enabled() {
556            rotate_targets_by_hash(&mut ordered, prefix_hash);
557        } else if let Some(sticky_hash) = routing.sticky_hash.filter(|_| affinity.sticky_enabled())
558        {
559            affinity.record_sticky_route();
560            rotate_targets_by_hash(&mut ordered, sticky_hash);
561        }
562    } else if let Some(sticky_hash) = routing.sticky_hash.filter(|_| affinity.sticky_enabled()) {
563        affinity.record_sticky_route();
564        rotate_targets_by_hash(&mut ordered, sticky_hash);
565    }
566
567    PreparedTargets {
568        ordered,
569        learn_prefix_hash,
570        cached_target,
571    }
572}
573
574#[cfg(test)]
575mod tests {
576    use super::*;
577    use iroh::SecretKey;
578
579    fn make_id(seed: u8) -> EndpointId {
580        let mut bytes = [0u8; 32];
581        bytes[0] = seed;
582        SecretKey::from_bytes(&bytes).public()
583    }
584
585    fn parse_body(body: &str) -> Value {
586        serde_json::from_str(body).unwrap()
587    }
588
589    #[test]
590    fn test_extract_session_hint_from_body_user_preferred() {
591        let body = parse_body(r#"{"user":"bob","session_id":"sess-1"}"#);
592        assert_eq!(
593            extract_session_hint_from_body(&body),
594            Some("bob".to_string())
595        );
596    }
597
598    #[test]
599    fn test_routing_keys_prefix_shared_across_first_user_changes() {
600        let req_a = parse_body(
601            r#"{"tools":[{"type":"function","function":{"name":"run"}}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"fix bug A"}]}"#,
602        );
603        let req_b = parse_body(
604            r#"{"tools":[{"type":"function","function":{"name":"run"}}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"fix bug B"}]}"#,
605        );
606
607        let keys_a = routing_keys(Some(&req_a));
608        let keys_b = routing_keys(Some(&req_b));
609
610        assert_eq!(keys_a.prefix_hash, keys_b.prefix_hash);
611        assert_ne!(keys_a.sticky_hash, keys_b.sticky_hash);
612    }
613
614    #[test]
615    fn test_routing_keys_prefix_ignores_object_key_order() {
616        let req_a = parse_body(
617            r#"{"tools":[{"type":"function","function":{"name":"run","description":"Run a command","parameters":{"type":"object","properties":{"path":{"type":"string"},"mode":{"type":"string"}},"required":["path","mode"]}}}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"fix bug A"}]}"#,
618        );
619        let req_b = parse_body(
620            r#"{"tools":[{"function":{"parameters":{"required":["path","mode"],"properties":{"mode":{"type":"string"},"path":{"type":"string"}},"type":"object"},"description":"Run a command","name":"run"},"type":"function"}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"fix bug B"}]}"#,
621        );
622
623        let keys_a = routing_keys(Some(&req_a));
624        let keys_b = routing_keys(Some(&req_b));
625
626        assert_eq!(keys_a.prefix_hash, keys_b.prefix_hash);
627        assert_ne!(keys_a.sticky_hash, keys_b.sticky_hash);
628    }
629
630    #[test]
631    fn test_select_model_target_uses_cached_prefix_target() {
632        let id_a = make_id(1);
633        let id_b = make_id(2);
634        let mut targets = election::ModelTargets::default();
635        targets.targets.insert(
636            "qwen".to_string(),
637            vec![
638                election::InferenceTarget::Remote(id_a),
639                election::InferenceTarget::Remote(id_b),
640            ],
641        );
642
643        let affinity = AffinityRouter::with_config(true, true);
644        let req_a = parse_body(
645            r#"{"tools":[{"type":"function","function":{"name":"run"}}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"task A"}]}"#,
646        );
647        let req_b = parse_body(
648            r#"{"tools":[{"type":"function","function":{"name":"run"}}],"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"task B"}]}"#,
649        );
650
651        let candidates = targets.candidates("qwen");
652        let first = select_model_target_from_candidates(
653            &targets,
654            &candidates,
655            "qwen",
656            Some(&req_a),
657            &affinity,
658        );
659        let prefix_hash = first.learn_prefix_hash.unwrap();
660        affinity.learn_target("qwen", prefix_hash, &first.target);
661
662        let second = select_model_target_from_candidates(
663            &targets,
664            &candidates,
665            "qwen",
666            Some(&req_b),
667            &affinity,
668        );
669        assert_eq!(Some(second.target.clone()), second.cached_target);
670        assert_eq!(first.target, second.target);
671    }
672
673    #[test]
674    fn test_prepare_remote_targets_prefers_cached_host() {
675        let id_a = make_id(1);
676        let id_b = make_id(2);
677        let hosts = vec![id_a, id_b];
678        let affinity = AffinityRouter::with_config(true, true);
679        let req = parse_body(
680            r#"{"messages":[{"role":"system","content":"You are an agent."},{"role":"user","content":"task A"}]}"#,
681        );
682
683        let prefix_hash = routing_keys(Some(&req)).prefix_hash.unwrap();
684        affinity.learn_target(
685            "qwen",
686            prefix_hash,
687            &election::InferenceTarget::Remote(id_b),
688        );
689
690        let prepared = prepare_remote_targets_for_request("qwen", &hosts, Some(&req), &affinity);
691        assert_eq!(
692            prepared.ordered.first(),
693            Some(&election::InferenceTarget::Remote(id_b))
694        );
695        assert_eq!(
696            prepared.cached_target,
697            Some(election::InferenceTarget::Remote(id_b))
698        );
699    }
700}
mesh_client/network/affinity.rs

mesh_client/network/
affinity.rs
