orch/engine/router/
weights.rs

1//! Per-agent routing weights and rate-limit recovery.
2//!
3//! `RateLimitState` tracks a single agent's current routing weight and
4//! recovers it over time after rate-limit events.
5//!
6//! `AgentWeights` aggregates per-agent states and provides weighted
7//! probabilistic selection via `weighted_select`.
8
9use std::collections::HashMap;
10use std::time::{Duration, Instant};
11
12use super::selection::{simple_hash_fraction_for, simple_hash_index_for};
13
14/// Default weight for agents with full capacity.
15pub(super) const DEFAULT_WEIGHT: f64 = 1.0;
16
17/// Minimum weight — an agent never drops below this (still gets occasional tasks).
18pub(super) const MIN_WEIGHT: f64 = 0.05;
19
20/// How much to reduce weight on each rate limit hit (multiplicative decay).
21pub(super) const RATE_LIMIT_DECAY: f64 = 0.3;
22
23/// Weight to use for an agent that is absent from a configured `router.weights`
24/// map. Rather than the constant `DEFAULT_WEIGHT` (1.0) — which would silently
25/// outrank every agent the operator *did* configure — an omitted agent falls
26/// back to the minimum of the explicitly configured weights. An empty map
27/// means no weights are configured at all, so `DEFAULT_WEIGHT` applies to
28/// everyone equally.
29pub(super) fn default_missing_weight(configured_weights: &HashMap<String, f64>) -> f64 {
30    configured_weights
31        .values()
32        .copied()
33        .fold(None, |min, w| Some(min.map_or(w, |m: f64| m.min(w))))
34        .unwrap_or(DEFAULT_WEIGHT)
35}
36
37/// Generate jitter duration for recovery delay based on agent name.
38/// Uses a simple hash to create deterministic but varied jitter per agent.
39fn generate_recovery_jitter(agent: &str) -> Duration {
40    // Simple hash: sum of char values modulo a large prime
41    let hash: u64 = agent
42        .bytes()
43        .fold(0u64, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u64));
44    // Map hash to 0..RECOVERY_JITTER_MAX range
45    let max_secs = RECOVERY_JITTER_MAX.as_secs();
46    let jitter_secs = hash % max_secs;
47    Duration::from_secs(jitter_secs)
48}
49
50/// Duration after which a rate-limited agent starts recovering weight.
51pub(super) const RECOVERY_DELAY: Duration = Duration::from_secs(60);
52
53/// Maximum jitter added to recovery delay to stagger agent recoveries.
54/// This prevents thundering herd where all agents recover simultaneously.
55pub(super) const RECOVERY_JITTER_MAX: Duration = Duration::from_secs(30);
56
57/// Per-tick weight recovery amount (additive, applied each routing call).
58pub(super) const RECOVERY_RATE: f64 = 0.1;
59
60/// Rate limit state for a single agent.
61#[derive(Debug, Clone)]
62pub struct RateLimitState {
63    /// Current routing weight (0.0..=1.0). Higher = more tasks.
64    pub weight: f64,
65    /// When the last rate limit error was recorded.
66    pub last_limited_at: Option<Instant>,
67    /// Jitter added to recovery delay to stagger recoveries (avoid thundering herd).
68    pub recovery_jitter: Duration,
69    /// How many consecutive rate limit hits.
70    pub consecutive_hits: u32,
71    /// The configured base weight this agent should recover toward.
72    pub base_weight: f64,
73}
74
75impl Default for RateLimitState {
76    fn default() -> Self {
77        Self {
78            weight: DEFAULT_WEIGHT,
79            last_limited_at: None,
80            recovery_jitter: Duration::ZERO,
81            consecutive_hits: 0,
82            base_weight: DEFAULT_WEIGHT,
83        }
84    }
85}
86
87impl RateLimitState {
88    /// Record a rate limit event — decay the weight.
89    /// Note: Jitter is set at the AgentWeights level using the agent name.
90    pub fn record_rate_limit(&mut self) {
91        self.consecutive_hits += 1;
92        self.weight = (self.weight * RATE_LIMIT_DECAY).max(MIN_WEIGHT);
93        self.last_limited_at = Some(Instant::now());
94    }
95
96    /// Record a successful completion — bump weight back toward base_weight.
97    pub fn record_success(&mut self) {
98        self.consecutive_hits = 0;
99        self.weight = (self.weight + RECOVERY_RATE).min(self.base_weight);
100    }
101
102    /// Tick recovery: if enough time has passed since the last limit (plus jitter), gradually restore.
103    pub fn maybe_recover(&mut self) {
104        if let Some(last) = self.last_limited_at {
105            let recovery_threshold = RECOVERY_DELAY + self.recovery_jitter;
106            if last.elapsed() >= recovery_threshold {
107                self.weight = (self.weight + RECOVERY_RATE).min(self.base_weight);
108                if self.weight >= self.base_weight {
109                    self.last_limited_at = None;
110                    self.recovery_jitter = Duration::ZERO;
111                    self.consecutive_hits = 0;
112                }
113            }
114        }
115    }
116
117    /// Is this agent currently rate-limited (weight below its base weight)?
118    pub fn is_limited(&self) -> bool {
119        self.weight < self.base_weight
120    }
121}
122
123/// Tracks per-agent weights for weighted round-robin routing.
124#[derive(Debug, Clone, Default)]
125pub struct AgentWeights {
126    pub states: HashMap<String, RateLimitState>,
127    /// Configured base weights per agent (from config.yml).
128    /// Used as the recovery ceiling instead of DEFAULT_WEIGHT.
129    pub base_weights: HashMap<String, f64>,
130}
131
132impl AgentWeights {
133    /// Ensure all available agents have an entry, using configured base weights.
134    pub fn ensure_agents(&mut self, agents: &[String]) {
135        for agent in agents {
136            self.states.entry(agent.clone()).or_default();
137        }
138    }
139
140    /// Ensure all available agents have an entry, applying configured base weights.
141    ///
142    /// Agents with an explicit weight in `config.yml` use that as their base.
143    /// Agents without a configured weight default to the minimum of the
144    /// explicitly configured weights (see `default_missing_weight`), so an
145    /// omission never outranks an agent the operator actually configured.
146    pub fn ensure_agents_with_weights(
147        &mut self,
148        agents: &[String],
149        configured_weights: &std::collections::HashMap<String, f64>,
150    ) {
151        for agent in agents {
152            let base = configured_weights
153                .get(agent)
154                .copied()
155                .unwrap_or_else(|| default_missing_weight(configured_weights));
156            let state = self
157                .states
158                .entry(agent.clone())
159                .or_insert_with(|| RateLimitState {
160                    weight: base,
161                    base_weight: base,
162                    ..Default::default()
163                });
164            state.base_weight = base;
165            // Update base weight if agent was already initialized (e.g. after config reload)
166            // but only if the agent isn't currently rate-limited (don't override decay)
167            if !state.is_limited() {
168                state.weight = base;
169            }
170        }
171    }
172
173    /// Record a rate limit event for an agent.
174    pub fn record_rate_limit(&mut self, agent: &str) {
175        let state = self.states.entry(agent.to_string()).or_default();
176        state.record_rate_limit();
177        // Generate jitter to stagger recovery times across agents
178        state.recovery_jitter = generate_recovery_jitter(agent);
179        tracing::info!(
180            agent,
181            weight = state.weight,
182            hits = state.consecutive_hits,
183            jitter_secs = state.recovery_jitter.as_secs(),
184            "agent weight reduced (rate limit)"
185        );
186    }
187
188    /// Record a successful task completion for an agent.
189    pub fn record_success(&mut self, agent: &str) {
190        self.states
191            .entry(agent.to_string())
192            .or_default()
193            .record_success();
194    }
195
196    /// Tick recovery for all agents.
197    pub fn tick_recovery(&mut self) {
198        for (agent, state) in &mut self.states {
199            let was_limited = state.is_limited();
200            state.maybe_recover();
201            if was_limited && !state.is_limited() {
202                tracing::info!(agent, "agent weight fully recovered");
203            }
204        }
205    }
206
207    /// Select an agent by weighted probability from the given list.
208    ///
209    /// Uses a simple weighted random selection: each agent's probability is
210    /// proportional to its weight. If all weights are zero (shouldn't happen
211    /// due to MIN_WEIGHT), falls back to uniform selection.
212    pub fn weighted_select(&self, agents: &[String], task_id: &str) -> Option<String> {
213        if agents.is_empty() {
214            return None;
215        }
216
217        let weights: Vec<f64> = agents
218            .iter()
219            .map(|a| {
220                self.states.get(a).map(|s| s.weight).unwrap_or_else(|| {
221                    self.base_weights
222                        .get(a)
223                        .copied()
224                        .unwrap_or_else(|| default_missing_weight(&self.base_weights))
225                })
226            })
227            .collect();
228
229        let total: f64 = weights.iter().sum();
230        if total <= 0.0 {
231            // Safety fallback: uniform random
232            let idx = simple_hash_index_for(agents.len(), task_id);
233            return Some(agents[idx].clone());
234        }
235
236        // Deterministic-ish selection using a hash of the current time
237        // to avoid requiring rand crate. Good enough for load distribution.
238        let pick = simple_hash_fraction_for(task_id) * total;
239        let mut cumulative = 0.0;
240        for (i, w) in weights.iter().enumerate() {
241            cumulative += w;
242            if pick < cumulative {
243                return Some(agents[i].clone());
244            }
245        }
246
247        // Rounding edge case — return last agent (agents is guaranteed non-empty by caller)
248        debug_assert!(!agents.is_empty(), "agents must not be empty");
249        Some(agents[agents.len() - 1].clone())
250    }
251
252    /// Get the current weight for an agent.
253    pub fn get_weight(&self, agent: &str) -> f64 {
254        self.states.get(agent).map(|s| s.weight).unwrap_or_else(|| {
255            self.base_weights
256                .get(agent)
257                .copied()
258                .unwrap_or_else(|| default_missing_weight(&self.base_weights))
259        })
260    }
261
262    /// Get a snapshot of all agent weights (for logging/debugging).
263    pub fn snapshot(&self) -> Vec<(String, f64, u32)> {
264        let mut snap: Vec<_> = self
265            .states
266            .iter()
267            .map(|(a, s)| (a.clone(), s.weight, s.consecutive_hits))
268            .collect();
269        snap.sort_by(|a, b| a.0.cmp(&b.0));
270        snap
271    }
272}
orch/engine/router/weights.rs

orch/engine/router/
weights.rs
