hatchery_server/
routing.rs

1//! hq-core router: roles, presets, seats and a limit-aware fallback.
2//!
3//! Config (`<data_dir>/routing.json`, owner-edited, hot-reloaded on each call):
4//! ```json
5//! {
6//!   "seats": [{ "id": "box-claude", "node": "hatchery-box", "provider": "claude" }],
7//!   "roles": {
8//!     "dev":     { "preset": ["claude/opus:high", "codex/gpt-5:high"], "auto": true },
9//!     "advisor": { "preset": ["codex/gpt-5:medium"], "auto": false }
10//!   },
11//!   "limit_cooldown_secs": 3600,
12//!   "advisor": { "min_interval_secs": 300, "dedup_secs": 3600 }
13//! }
14//! ```
15//! A preset entry is `provider/model[:effort]`. A seat is one place a provider can run
16//! (a node + provider, optionally a named account). The router walks the role's preset in
17//! order and, for each entry, the seats of that provider; a seat with a recent `limit`
18//! signal is skipped until a `ready` signal, a parsed reset time, or the cooldown passes.
19//! Signals are hints (screen heuristics): they only reorder and never delete a seat, and
20//! when every seat is limited the router still answers with the soonest-free one.
21//! The router recommends; it picks on its own only in roles with `"auto": true`.
22
23use serde::{Deserialize, Serialize};
24use std::collections::BTreeMap;
25
26#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
27#[serde(deny_unknown_fields)]
28pub struct Seat {
29    pub id: String,
30    pub node: String,
31    pub provider: String,
32    #[serde(default)]
33    pub account: Option<String>,
34}
35
36#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
37#[serde(deny_unknown_fields)]
38pub struct Role {
39    pub preset: Vec<String>,
40    #[serde(default)]
41    pub auto: bool,
42}
43
44#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
45#[serde(deny_unknown_fields)]
46pub struct AdvisorPolicy {
47    #[serde(default = "d_interval")]
48    pub min_interval_secs: u64,
49    #[serde(default = "d_dedup")]
50    pub dedup_secs: u64,
51}
52fn d_interval() -> u64 { 300 }
53fn d_dedup() -> u64 { 3600 }
54impl Default for AdvisorPolicy {
55    fn default() -> Self { Self { min_interval_secs: d_interval(), dedup_secs: d_dedup() } }
56}
57
58#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
59#[serde(deny_unknown_fields)]
60pub struct RoutingConfig {
61    #[serde(default)]
62    pub seats: Vec<Seat>,
63    #[serde(default)]
64    pub roles: BTreeMap<String, Role>,
65    #[serde(default)]
66    pub limit_cooldown_secs: Option<u64>,
67    #[serde(default)]
68    pub advisor: AdvisorPolicy,
69}
70
71/// One model a node's provider CLI exposes (as reported by the node).
72#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
73pub struct InventoryModel {
74    pub id: String,
75    #[serde(default, skip_serializing_if = "Vec::is_empty")]
76    pub efforts: Vec<String>,
77    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
78    pub default: bool,
79}
80
81/// What a node reports for one provider.
82#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
83pub struct ProviderInventory {
84    pub provider: String,
85    #[serde(default)]
86    pub source: String,
87    #[serde(default)]
88    pub models: Vec<InventoryModel>,
89}
90
91/// Stored inventory row: one provider on one node.
92#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
93pub struct NodeInventory {
94    pub node: String,
95    pub provider: String,
96    pub source: String,
97    pub models: Vec<InventoryModel>,
98    pub updated_at_ms: u64,
99}
100
101fn ok_name(x: &str) -> bool {
102    !x.is_empty() && x.len() <= 128 && x.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.[]".contains(&b))
103}
104
105/// Validate a node's report (bounded, plain identifiers only).
106pub fn clean_report(providers: Vec<ProviderInventory>) -> Result<Vec<ProviderInventory>, &'static str> {
107    if providers.len() > 16 {
108        return Err("too many providers");
109    }
110    for p in &providers {
111        if !ok_name(&p.provider) || p.source.len() > 32 || p.models.len() > 256 {
112            return Err("bad provider entry");
113        }
114        if p.models.iter().any(|m| !ok_name(&m.id) || m.efforts.len() > 16 || m.efforts.iter().any(|e| !ok_name(e))) {
115            return Err("bad model entry");
116        }
117    }
118    Ok(providers)
119}
120
121impl RoutingConfig {
122    /// Seats are the node-reported providers; roles are the user's stored choices.
123    pub fn from_store(inventory: &[NodeInventory], roles: BTreeMap<String, Role>) -> Self {
124        let seats = inventory
125            .iter()
126            .map(|i| Seat { id: format!("{}/{}", i.node, i.provider), node: i.node.clone(), provider: i.provider.clone(), account: None })
127            .collect();
128        Self { seats, roles, limit_cooldown_secs: None, advisor: AdvisorPolicy::default() }
129    }
130    fn cooldown_ms(&self) -> u64 {
131        self.limit_cooldown_secs.unwrap_or(3600) * 1000
132    }
133}
134
135/// `provider/model[:effort]`
136#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
137pub struct ModelRef {
138    pub provider: String,
139    pub model: String,
140    pub effort: Option<String>,
141}
142
143pub fn parse_model_ref(s: &str) -> Option<ModelRef> {
144    let (provider, rest) = s.split_once('/')?;
145    let (model, effort) = match rest.split_once(':') {
146        Some((m, e)) => (m, Some(e.to_owned()).filter(|e| !e.is_empty())),
147        None => (rest, None),
148    };
149    let ok = |x: &str| !x.is_empty() && x.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b));
150    (ok(provider) && ok(model) && effort.as_deref().is_none_or(ok)).then(|| ModelRef { provider: provider.into(), model: model.into(), effort })
151}
152
153/// Latest limit knowledge per (node, provider), built from session signals.
154#[derive(Clone, Debug, Default, Serialize, PartialEq, Eq)]
155pub struct SeatState {
156    /// When the seat is expected to be free again (ms); `None` = free.
157    pub limited_until_ms: Option<u64>,
158    pub hint: Option<String>,
159}
160
161/// Fold signals (oldest first) into per-(node, provider) state.
162pub fn fold_signals<'a>(signals: impl IntoIterator<Item = &'a crate::tasks::SessionSignal>, cfg: &RoutingConfig) -> BTreeMap<(String, String), SeatState> {
163    let mut out: BTreeMap<(String, String), SeatState> = BTreeMap::new();
164    for s in signals {
165        let key = (s.node_id.clone(), s.provider.clone());
166        match s.kind.as_str() {
167            "limit" => {
168                let until = s.hint.as_deref().and_then(|h| parse_reset(h, s.at_ms)).unwrap_or(s.at_ms + cfg.cooldown_ms());
169                out.insert(key, SeatState { limited_until_ms: Some(until), hint: s.hint.clone() });
170            }
171            "ready" => {
172                out.insert(key, SeatState::default());
173            }
174            _ => {}
175        }
176    }
177    out
178}
179
180/// "try again in 2 hours", "in 45 minutes": relative reset, capped at 8 days.
181pub fn parse_reset(hint: &str, at_ms: u64) -> Option<u64> {
182    let lower = hint.to_lowercase();
183    let idx = lower.find(" in ")?;
184    let mut words = lower[idx + 4..].split_whitespace();
185    let n: f64 = words.next()?.parse().ok()?;
186    let unit = words.next()?;
187    let secs = if unit.starts_with("min") { 60.0 } else if unit.starts_with("hour") || unit == "h" || unit.starts_with("hr") { 3600.0 } else if unit.starts_with("day") { 86400.0 } else { return None };
188    let ms = (n * secs * 1000.0) as u64;
189    (ms > 0 && ms <= 8 * 86_400_000).then_some(at_ms + ms)
190}
191
192#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
193pub struct Candidate {
194    pub seat: Seat,
195    pub model: ModelRef,
196    pub available: bool,
197    pub limited_until_ms: Option<u64>,
198    pub hint: Option<String>,
199}
200
201#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
202pub struct Recommendation {
203    pub role: String,
204    /// First available candidate (or the soonest-free one when all are limited).
205    pub pick: Option<Candidate>,
206    /// Whether the router may act on `pick` without the operator (role allows auto).
207    pub auto: bool,
208    /// The preferred candidate when `pick` is a fallback from it.
209    pub fallback_from: Option<Candidate>,
210    pub candidates: Vec<Candidate>,
211}
212
213/// A role preset entry is valid only if some node reports that provider/model (and effort).
214pub fn offered(inventory: &[NodeInventory], m: &ModelRef) -> bool {
215    inventory.iter().any(|i| {
216        i.provider == m.provider
217            && i.models.iter().any(|x| x.id == m.model && m.effort.as_ref().is_none_or(|e| x.efforts.iter().any(|y| y == e)))
218    })
219}
220
221pub fn recommend(cfg: &RoutingConfig, state: &BTreeMap<(String, String), SeatState>, role: &str, now_ms: u64) -> Result<Recommendation, String> {
222    recommend_in(cfg, None, state, role, now_ms)
223}
224
225/// Like [`recommend`], but only seats whose node reports the model count (when `inventory` is given).
226pub fn recommend_in(cfg: &RoutingConfig, inventory: Option<&[NodeInventory]>, state: &BTreeMap<(String, String), SeatState>, role: &str, now_ms: u64) -> Result<Recommendation, String> {
227    let r = cfg.roles.get(role).ok_or("unknown role")?;
228    let mut candidates = Vec::new();
229    for entry in &r.preset {
230        let model = parse_model_ref(entry).ok_or_else(|| format!("bad preset entry: {entry}"))?;
231        let has = |seat: &Seat| {
232            inventory.is_none_or(|inv| {
233                inv.iter().any(|i| i.node == seat.node && i.provider == seat.provider && i.models.iter().any(|x| x.id == model.model))
234            })
235        };
236        for seat in cfg.seats.iter().filter(|s| s.provider == model.provider && has(s)) {
237            let st = state.get(&(seat.node.clone(), seat.provider.clone())).cloned().unwrap_or_default();
238            let limited = st.limited_until_ms.filter(|u| *u > now_ms);
239            candidates.push(Candidate { seat: seat.clone(), model: model.clone(), available: limited.is_none(), limited_until_ms: limited, hint: limited.and(st.hint) });
240        }
241    }
242    let pick = candidates
243        .iter()
244        .find(|c| c.available)
245        .or_else(|| candidates.iter().min_by_key(|c| c.limited_until_ms.unwrap_or(u64::MAX)))
246        .cloned();
247    let fallback_from = match (&pick, candidates.first()) {
248        (Some(p), Some(first)) if p != first => Some(first.clone()),
249        _ => None,
250    };
251    Ok(Recommendation { role: role.into(), pick, auto: r.auto, fallback_from, candidates })
252}
253
254/// Advisor anti-noise gate: one advice per `min_interval` per advisor, and the same topic
255/// not again within `dedup`. Memory only (a restart forgets; that is acceptable noise).
256#[derive(Default)]
257pub struct AdviceGate {
258    last: BTreeMap<String, u64>,
259    topics: BTreeMap<(String, String), u64>,
260}
261
262impl AdviceGate {
263    pub fn check(&mut self, policy: &AdvisorPolicy, advisor: &str, topic: &str, now_ms: u64) -> Result<(), &'static str> {
264        let topic = topic.trim().to_lowercase();
265        if topic.is_empty() {
266            return Err("topic required");
267        }
268        if self.last.get(advisor).is_some_and(|t| now_ms < t + policy.min_interval_secs * 1000) {
269            return Err("too soon");
270        }
271        let key = (advisor.to_owned(), topic);
272        if self.topics.get(&key).is_some_and(|t| now_ms < t + policy.dedup_secs * 1000) {
273            return Err("duplicate topic");
274        }
275        self.last.insert(advisor.to_owned(), now_ms);
276        self.topics.insert(key, now_ms);
277        let horizon = now_ms.saturating_sub(policy.dedup_secs * 1000);
278        self.topics.retain(|_, t| *t >= horizon);
279        Ok(())
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use crate::tasks::SessionSignal;
287
288    fn cfg() -> RoutingConfig {
289        serde_json::from_value(serde_json::json!({
290            "seats": [
291                { "id": "a-claude", "node": "a", "provider": "claude" },
292                { "id": "b-claude", "node": "b", "provider": "claude" },
293                { "id": "a-codex", "node": "a", "provider": "codex" }
294            ],
295            "roles": {
296                "dev": { "preset": ["claude/opus:high", "codex/gpt-5:high"], "auto": true },
297                "advisor": { "preset": ["codex/gpt-5:medium"] }
298            },
299            "limit_cooldown_secs": 600
300        }))
301        .unwrap()
302    }
303
304    fn sig(id: i64, node: &str, provider: &str, kind: &str, hint: Option<&str>, at: u64) -> SessionSignal {
305        SessionSignal { id, node_id: node.into(), workspace_id: "w".into(), instance_id: 1, provider: provider.into(), kind: kind.into(), hint: hint.map(Into::into), at_ms: at }
306    }
307
308    #[test]
309    fn parses_model_refs() {
310        assert_eq!(parse_model_ref("claude/opus:high").unwrap().effort.as_deref(), Some("high"));
311        assert_eq!(parse_model_ref("codex/gpt-5").unwrap().effort, None);
312        assert!(parse_model_ref("bad").is_none());
313        assert!(parse_model_ref("a/b c").is_none());
314    }
315
316    #[test]
317    fn picks_first_free_seat_and_falls_back_on_limit() {
318        let c = cfg();
319        let r = recommend(&c, &BTreeMap::new(), "dev", 0).unwrap();
320        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
321        assert!(r.auto && r.fallback_from.is_none());
322        let sigs = [sig(1, "a", "claude", "limit", Some("usage limit, try again in 2 hours"), 1000)];
323        let st = fold_signals(&sigs, &c);
324        let r = recommend(&c, &st, "dev", 5000).unwrap();
325        assert_eq!(r.pick.as_ref().unwrap().seat.id, "b-claude");
326        assert_eq!(r.fallback_from.as_ref().unwrap().seat.id, "a-claude");
327        assert_eq!(r.candidates[0].limited_until_ms, Some(1000 + 7_200_000));
328        // Return after reset.
329        let r = recommend(&c, &st, "dev", 1000 + 7_200_001).unwrap();
330        assert_eq!(r.pick.unwrap().seat.id, "a-claude");
331    }
332
333    #[test]
334    fn ready_clears_and_cooldown_applies_without_hint() {
335        let c = cfg();
336        let sigs = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 0)];
337        let st = fold_signals(&sigs, &c);
338        let r = recommend(&c, &st, "dev", 1).unwrap();
339        assert_eq!(r.pick.as_ref().unwrap().seat.provider, "codex");
340        // All limited: soonest-free is still offered.
341        let sigs2 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 10), sig(3, "a", "codex", "limit", None, 20)];
342        let r = recommend(&c, &fold_signals(&sigs2, &c), "dev", 30).unwrap();
343        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
344        assert!(!r.pick.unwrap().available);
345        // ready clears.
346        let sigs3 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "a", "claude", "ready", None, 5)];
347        assert!(recommend(&c, &fold_signals(&sigs3, &c), "dev", 6).unwrap().pick.unwrap().available);
348        // After cooldown (600 s).
349        assert!(recommend(&c, &st, "dev", 600_001).unwrap().candidates.iter().all(|x| x.available));
350    }
351
352    #[test]
353    fn non_auto_role_and_unknown_role() {
354        let c = cfg();
355        assert!(!recommend(&c, &BTreeMap::new(), "advisor", 0).unwrap().auto);
356        assert!(recommend(&c, &BTreeMap::new(), "nope", 0).is_err());
357    }
358
359    #[test]
360    fn advice_gate_limits_noise() {
361        let p = AdvisorPolicy { min_interval_secs: 60, dedup_secs: 600 };
362        let mut g = AdviceGate::default();
363        assert!(g.check(&p, "adv", "tests failing", 0).is_ok());
364        assert_eq!(g.check(&p, "adv", "other", 1000), Err("too soon"));
365        assert_eq!(g.check(&p, "adv", "Tests failing ", 61_000), Err("duplicate topic"));
366        assert!(g.check(&p, "adv", "other", 61_000).is_ok());
367        assert!(g.check(&p, "adv2", "tests failing", 61_000).is_ok());
368        assert!(g.check(&p, "adv", "tests failing", 700_000).is_ok());
369    }
370}
hatchery_server/routing.rs

hatchery_server/
routing.rs
