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/// Model ids may carry a vendor path or tag (`kimi-code/k3`, `name:tag`, `x@v`).
106pub fn ok_model_id(x: &str) -> bool {
107    !x.is_empty() && x.len() <= 128 && !x.starts_with(['/', ':']) && !x.contains("..") && x.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.[]/:@".contains(&b))
108}
109
110/// Validate a node's report (bounded, plain identifiers). A bad provider, model or effort
111/// entry is dropped on its own; the rest of the report is kept.
112pub fn clean_report(providers: Vec<ProviderInventory>) -> Result<Vec<ProviderInventory>, &'static str> {
113    if providers.len() > 16 {
114        return Err("too many providers");
115    }
116    let mut out = Vec::new();
117    for mut p in providers {
118        if !ok_name(&p.provider) || p.source.len() > 32 {
119            tracing::warn!(provider = %p.provider.chars().take(40).collect::<String>(), "inventory: provider entry dropped");
120            continue;
121        }
122        let before = p.models.len();
123        p.models.retain(|m| ok_model_id(&m.id));
124        p.models.truncate(256);
125        for m in &mut p.models {
126            m.efforts.retain(|e| ok_name(e));
127            m.efforts.truncate(16);
128        }
129        if p.models.len() < before {
130            tracing::warn!(provider = %p.provider, dropped = before - p.models.len(), "inventory: bad model entries dropped");
131        }
132        out.push(p);
133    }
134    Ok(out)
135}
136
137impl RoutingConfig {
138    /// Seats are the node-reported providers; roles are the user's stored choices.
139    pub fn from_store(inventory: &[NodeInventory], roles: BTreeMap<String, Role>) -> Self {
140        let seats = inventory
141            .iter()
142            .map(|i| Seat { id: format!("{}/{}", i.node, i.provider), node: i.node.clone(), provider: i.provider.clone(), account: None })
143            .collect();
144        Self { seats, roles, limit_cooldown_secs: None, advisor: AdvisorPolicy::default() }
145    }
146    fn cooldown_ms(&self) -> u64 {
147        self.limit_cooldown_secs.unwrap_or(3600) * 1000
148    }
149}
150
151/// `provider/model[:effort]`
152#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
153pub struct ModelRef {
154    pub provider: String,
155    pub model: String,
156    pub effort: Option<String>,
157}
158
159pub fn parse_model_ref(s: &str) -> Option<ModelRef> {
160    let (provider, rest) = s.split_once('/')?;
161    let ok = |x: &str| !x.is_empty() && x.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.".contains(&b));
162    // The effort is the part after the last ':' when that part is a plain name; a model id may
163    // itself hold '/' or ':' (`kimi/kimi-code/k3:high`).
164    let (model, effort) = match rest.rsplit_once(':') {
165        Some((m, e)) if ok(e) && !m.is_empty() => (m, Some(e.to_owned())),
166        Some((m, "")) => (m, None),
167        _ => (rest, None),
168    };
169    (ok(provider) && ok_model_id(model) && effort.as_deref().is_none_or(ok)).then(|| ModelRef { provider: provider.into(), model: model.into(), effort })
170}
171
172/// Latest limit knowledge per (node, provider), built from session signals.
173#[derive(Clone, Debug, Default, Serialize, PartialEq, Eq)]
174pub struct SeatState {
175    /// When the seat is expected to be free again (ms); `None` = free.
176    pub limited_until_ms: Option<u64>,
177    pub hint: Option<String>,
178}
179
180/// Fold signals (oldest first) into per-(node, provider) state.
181pub fn fold_signals<'a>(signals: impl IntoIterator<Item = &'a crate::tasks::SessionSignal>, cfg: &RoutingConfig) -> BTreeMap<(String, String), SeatState> {
182    let mut out: BTreeMap<(String, String), SeatState> = BTreeMap::new();
183    for s in signals {
184        let key = (s.node_id.clone(), s.provider.clone());
185        match s.kind.as_str() {
186            "limit" => {
187                let until = s.hint.as_deref().and_then(|h| parse_reset(h, s.at_ms)).unwrap_or(s.at_ms + cfg.cooldown_ms());
188                out.insert(key, SeatState { limited_until_ms: Some(until), hint: s.hint.clone() });
189            }
190            "ready" => {
191                out.insert(key, SeatState::default());
192            }
193            _ => {}
194        }
195    }
196    out
197}
198
199/// "try again in 2 hours", "in 45 minutes": relative reset, capped at 8 days.
200pub fn parse_reset(hint: &str, at_ms: u64) -> Option<u64> {
201    let lower = hint.to_lowercase();
202    let idx = lower.find(" in ")?;
203    let mut words = lower[idx + 4..].split_whitespace();
204    let n: f64 = words.next()?.parse().ok()?;
205    let unit = words.next()?;
206    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 };
207    let ms = (n * secs * 1000.0) as u64;
208    (ms > 0 && ms <= 8 * 86_400_000).then_some(at_ms + ms)
209}
210
211#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
212pub struct Candidate {
213    pub seat: Seat,
214    pub model: ModelRef,
215    pub available: bool,
216    pub limited_until_ms: Option<u64>,
217    pub hint: Option<String>,
218}
219
220#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
221pub struct Recommendation {
222    pub role: String,
223    /// First available candidate (or the soonest-free one when all are limited).
224    pub pick: Option<Candidate>,
225    /// Whether the router may act on `pick` without the operator (role allows auto).
226    pub auto: bool,
227    /// The preferred candidate when `pick` is a fallback from it.
228    pub fallback_from: Option<Candidate>,
229    pub candidates: Vec<Candidate>,
230}
231
232/// A role preset entry is valid only if some node reports that provider/model (and effort).
233pub fn offered(inventory: &[NodeInventory], m: &ModelRef) -> bool {
234    inventory.iter().any(|i| {
235        i.provider == m.provider
236            && i.models.iter().any(|x| x.id == m.model && m.effort.as_ref().is_none_or(|e| x.efforts.iter().any(|y| y == e)))
237    })
238}
239
240pub fn recommend(cfg: &RoutingConfig, state: &BTreeMap<(String, String), SeatState>, role: &str, now_ms: u64) -> Result<Recommendation, String> {
241    recommend_in(cfg, None, state, role, now_ms)
242}
243
244/// Like [`recommend`], but only seats whose node reports the model count (when `inventory` is given).
245pub fn recommend_in(cfg: &RoutingConfig, inventory: Option<&[NodeInventory]>, state: &BTreeMap<(String, String), SeatState>, role: &str, now_ms: u64) -> Result<Recommendation, String> {
246    let r = cfg.roles.get(role).ok_or("unknown role")?;
247    let mut candidates = Vec::new();
248    for entry in &r.preset {
249        let model = parse_model_ref(entry).ok_or_else(|| format!("bad preset entry: {entry}"))?;
250        let has = |seat: &Seat| {
251            inventory.is_none_or(|inv| {
252                inv.iter().any(|i| i.node == seat.node && i.provider == seat.provider && i.models.iter().any(|x| x.id == model.model))
253            })
254        };
255        for seat in cfg.seats.iter().filter(|s| s.provider == model.provider && has(s)) {
256            let st = state.get(&(seat.node.clone(), seat.provider.clone())).cloned().unwrap_or_default();
257            let limited = st.limited_until_ms.filter(|u| *u > now_ms);
258            candidates.push(Candidate { seat: seat.clone(), model: model.clone(), available: limited.is_none(), limited_until_ms: limited, hint: limited.and(st.hint) });
259        }
260    }
261    let pick = candidates
262        .iter()
263        .find(|c| c.available)
264        .or_else(|| candidates.iter().min_by_key(|c| c.limited_until_ms.unwrap_or(u64::MAX)))
265        .cloned();
266    let fallback_from = match (&pick, candidates.first()) {
267        (Some(p), Some(first)) if p != first => Some(first.clone()),
268        _ => None,
269    };
270    Ok(Recommendation { role: role.into(), pick, auto: r.auto, fallback_from, candidates })
271}
272
273/// Advisor anti-noise gate: one advice per `min_interval` per advisor, and the same topic
274/// not again within `dedup`. Memory only (a restart forgets; that is acceptable noise).
275#[derive(Default)]
276pub struct AdviceGate {
277    last: BTreeMap<String, u64>,
278    topics: BTreeMap<(String, String), u64>,
279}
280
281impl AdviceGate {
282    pub fn check(&mut self, policy: &AdvisorPolicy, advisor: &str, topic: &str, now_ms: u64) -> Result<(), &'static str> {
283        let topic = topic.trim().to_lowercase();
284        if topic.is_empty() {
285            return Err("topic required");
286        }
287        if self.last.get(advisor).is_some_and(|t| now_ms < t + policy.min_interval_secs * 1000) {
288            return Err("too soon");
289        }
290        let key = (advisor.to_owned(), topic);
291        if self.topics.get(&key).is_some_and(|t| now_ms < t + policy.dedup_secs * 1000) {
292            return Err("duplicate topic");
293        }
294        self.last.insert(advisor.to_owned(), now_ms);
295        self.topics.insert(key, now_ms);
296        let horizon = now_ms.saturating_sub(policy.dedup_secs * 1000);
297        self.topics.retain(|_, t| *t >= horizon);
298        Ok(())
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    #[test]
305    fn model_ids_with_paths_and_one_bad_entry() {
306        let r = super::parse_model_ref("kimi/kimi-code/k3:high").unwrap();
307        assert_eq!((r.model.as_str(), r.effort.as_deref()), ("kimi-code/k3", Some("high")));
308        let p: Vec<super::ProviderInventory> = serde_json::from_str(r#"[{"provider":"kimi","models":[{"id":"kimi-code/k3"},{"id":"bad id"}]},{"provider":"bad provider"}]"#).unwrap();
309        let c = super::clean_report(p).unwrap();
310        assert_eq!(c.len(), 1);
311        assert_eq!(c[0].models.len(), 1);
312    }
313    use super::*;
314    use crate::tasks::SessionSignal;
315
316    fn cfg() -> RoutingConfig {
317        serde_json::from_value(serde_json::json!({
318            "seats": [
319                { "id": "a-claude", "node": "a", "provider": "claude" },
320                { "id": "b-claude", "node": "b", "provider": "claude" },
321                { "id": "a-codex", "node": "a", "provider": "codex" }
322            ],
323            "roles": {
324                "dev": { "preset": ["claude/opus:high", "codex/gpt-5:high"], "auto": true },
325                "advisor": { "preset": ["codex/gpt-5:medium"] }
326            },
327            "limit_cooldown_secs": 600
328        }))
329        .unwrap()
330    }
331
332    fn sig(id: i64, node: &str, provider: &str, kind: &str, hint: Option<&str>, at: u64) -> SessionSignal {
333        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 }
334    }
335
336    #[test]
337    fn parses_model_refs() {
338        assert_eq!(parse_model_ref("claude/opus:high").unwrap().effort.as_deref(), Some("high"));
339        assert_eq!(parse_model_ref("codex/gpt-5").unwrap().effort, None);
340        assert!(parse_model_ref("bad").is_none());
341        assert!(parse_model_ref("a/b c").is_none());
342    }
343
344    #[test]
345    fn picks_first_free_seat_and_falls_back_on_limit() {
346        let c = cfg();
347        let r = recommend(&c, &BTreeMap::new(), "dev", 0).unwrap();
348        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
349        assert!(r.auto && r.fallback_from.is_none());
350        let sigs = [sig(1, "a", "claude", "limit", Some("usage limit, try again in 2 hours"), 1000)];
351        let st = fold_signals(&sigs, &c);
352        let r = recommend(&c, &st, "dev", 5000).unwrap();
353        assert_eq!(r.pick.as_ref().unwrap().seat.id, "b-claude");
354        assert_eq!(r.fallback_from.as_ref().unwrap().seat.id, "a-claude");
355        assert_eq!(r.candidates[0].limited_until_ms, Some(1000 + 7_200_000));
356        // Return after reset.
357        let r = recommend(&c, &st, "dev", 1000 + 7_200_001).unwrap();
358        assert_eq!(r.pick.unwrap().seat.id, "a-claude");
359    }
360
361    #[test]
362    fn ready_clears_and_cooldown_applies_without_hint() {
363        let c = cfg();
364        let sigs = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 0)];
365        let st = fold_signals(&sigs, &c);
366        let r = recommend(&c, &st, "dev", 1).unwrap();
367        assert_eq!(r.pick.as_ref().unwrap().seat.provider, "codex");
368        // All limited: soonest-free is still offered.
369        let sigs2 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "b", "claude", "limit", None, 10), sig(3, "a", "codex", "limit", None, 20)];
370        let r = recommend(&c, &fold_signals(&sigs2, &c), "dev", 30).unwrap();
371        assert_eq!(r.pick.as_ref().unwrap().seat.id, "a-claude");
372        assert!(!r.pick.unwrap().available);
373        // ready clears.
374        let sigs3 = [sig(1, "a", "claude", "limit", None, 0), sig(2, "a", "claude", "ready", None, 5)];
375        assert!(recommend(&c, &fold_signals(&sigs3, &c), "dev", 6).unwrap().pick.unwrap().available);
376        // After cooldown (600 s).
377        assert!(recommend(&c, &st, "dev", 600_001).unwrap().candidates.iter().all(|x| x.available));
378    }
379
380    #[test]
381    fn non_auto_role_and_unknown_role() {
382        let c = cfg();
383        assert!(!recommend(&c, &BTreeMap::new(), "advisor", 0).unwrap().auto);
384        assert!(recommend(&c, &BTreeMap::new(), "nope", 0).is_err());
385    }
386
387    #[test]
388    fn advice_gate_limits_noise() {
389        let p = AdvisorPolicy { min_interval_secs: 60, dedup_secs: 600 };
390        let mut g = AdviceGate::default();
391        assert!(g.check(&p, "adv", "tests failing", 0).is_ok());
392        assert_eq!(g.check(&p, "adv", "other", 1000), Err("too soon"));
393        assert_eq!(g.check(&p, "adv", "Tests failing ", 61_000), Err("duplicate topic"));
394        assert!(g.check(&p, "adv", "other", 61_000).is_ok());
395        assert!(g.check(&p, "adv2", "tests failing", 61_000).is_ok());
396        assert!(g.check(&p, "adv", "tests failing", 700_000).is_ok());
397    }
398}
hatchery_server/routing.rs

hatchery_server/
routing.rs
